[SpecDecode] Reduce TP communication for large-vocab draft models speculative decoding

[EanWang211123]

Purpose

Use local argmax + gather reduction instead of all-gathering full-vocabulary logits to reduce TP communication overhead in speculative decoding modes such as DFlash and PARD.

Enabling use_local_argmax_reduction improves performance in all tested scenarios.

In terms of Total token throughput (tok/s), the gains are:

• +9.38% for Qwen3-8B + Qwen3-8B-DFlash + tp2 at 1 concurrency
• +10.21% for Qwen3-8B + Qwen3-8B-DFlash + tp2 at 32 concurrency
• +11.54% for Qwen3-8B + Qwen3-0.6B-PARD + tp2 at 32 concurrency
• +30.63% for Qwen3-8B + Qwen3-8B-DFlash + tp4 at 32 concurrency

Current Issue

Speculative decoding methods such as DFlash and PARD often require the draft model and target model to have the same vocabulary size. Under the current implementation, this requires all-gathering the full logits tensor across TP ranks, which incurs O(batch_size × vocab_size) communication cost per step. For example, PARD-Qwen3-0.6B uses a vocabulary size of 40960.

We can reuse the approach introduced in #34049 to reduce this communication overhead.

This PR adds a use_local_argmax_reduction option to speculative_config that:

• Computes the local argmax on each TP rank's vocab shard
• Gathers only (max_value, global_index) pairs — O(batch_size × 2 × tp_size) instead of O(batch_size × vocab_size)
• Reduces across ranks to find the global argmax

Changes

Add get_top_tokens() to the corresponding model files.

Test Plan

Environment: 4090D
Target model: Qwen3-8B
Draft models: Qwen3-8B-DFlash-b16 and PARD-Qwen3-0.6B
Dataset: gsm8k

Test Commands

CUDA_VISIBLE_DEVICES=0,1,2,3 \
python -m vllm.entrypoints.openai.api_server \
  --model Qwen3-8B \
  --trust-remote-code \
  --tensor-parallel-size 4 \
  --max-model-len 8000 \
  --max-num-batched-tokens 4096 \
  --gpu-memory-utilization 0.85 \
  --speculative-config '{
    "method": "dflash",
    "model": "Qwen3-8B-DFlash-b16",
    "num_speculative_tokens": 15,
    "use_local_argmax_reduction": true
  }'

CUDA_VISIBLE_DEVICES=6,7 \
vllm serve Qwen3-8B \
  --host 0.0.0.0 \
  --port 8000 \
  --tensor-parallel-size 2 \
  --max-model-len 4096 \
  --max-num-batched-tokens 4096 \
  --gpu-memory-utilization 0.8 \
  --speculative-config '{
    "method": "draft_model",
    "model": "PARD-Qwen3-0.6B",
    "num_speculative_tokens": 8,
    "parallel_drafting": true,
    "use_local_argmax_reduction": true
  }'

Test Results

Accuracy

Enabling use_local_argmax_reduction does not change the acceptance rate.

Qwen3-8B + Qwen3-8B-DFlash-b16
Concurrency: 1

use_local_argmax_reduction=true

Per-position acceptance (%):
  Position 0:                            69.05
  Position 1:                            43.76
  Position 2:                            28.46
  Position 3:                            18.59
  Position 4:                            12.24
  Position 5:                            7.82
  Position 6:                            4.76
  Position 7:                            3.17
  Position 8:                            1.25
  Position 9:                            0.68
  Position 10:                           0.23
  Position 11:                           0.11
  Position 12:                           0.00
  Position 13:                           0.00
  Position 14:                           0.00

use_local_argmax_reduction=false

Per-position acceptance (%):
  Position 0:                            69.05
  Position 1:                            43.76
  Position 2:                            28.46
  Position 3:                            18.59
  Position 4:                            12.24
  Position 5:                            7.82
  Position 6:                            4.76
  Position 7:                            3.17
  Position 8:                            1.25
  Position 9:                            0.68
  Position 10:                           0.23
  Position 11:                           0.11
  Position 12:                           0.00
  Position 13:                           0.00
  Position 14:                           0.00

Performance

Scenario	Concurrency	use_local_argmax_reduction	Total token throughput (tok/s)	Mean TTFT (ms)	Mean TPOT (ms)	Mean ITL (ms)	Total token throughput improvement
Qwen3-8B + Qwen3-8B-DFlash + tp2	1	false	199.25	179.88	5.59	15.97	Baseline
Qwen3-8B + Qwen3-8B-DFlash + tp2	1	true	217.94	98.47	5.37	15.34	+9.38%
Qwen3-8B + Qwen3-8B-DFlash + tp2	32	false	1407.37	885.07	24.77	74.44	Baseline
Qwen3-8B + Qwen3-8B-DFlash + tp2	32	true	1551.11	558.23	22.82	68.69	+10.21%
Qwen3-8B + Qwen3-0.6B-PARD + tp2	32	false	2747.77	323.18	12.84	57.88	Baseline
Qwen3-8B + Qwen3-0.6B-PARD + tp2	32	true	3064.77	197.74	11.89	53.80	+11.54%
Qwen3-8B + Qwen3-8B-DFlash + tp4	32	false	1338.82	1366.44	24.00	71.91	Baseline
Qwen3-8B + Qwen3-8B-DFlash + tp4	32	true	1748.92	479.53	20.71	62.02	+30.63%

Open Question

I noticed that almost all speculative decoding setups using a draft model, whether parallelized or not, seem to benefit from this optimization. Is there a more general way to integrate this feature?

---

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 the get_top_tokens method to the Qwen3 and Qwen3-DFlash models, enabling vocab-parallel argmax operations that avoid all-gathering full logits. Feedback for qwen3_dflash.py suggests optimizing the fallback logic when draft_id_to_target_id is present; instead of falling back to full logits computation, the remapping can be applied directly to the result of the optimized get_top_tokens call to maintain performance benefits.

[benchislett]

Since this is an opt-in feature, can we also add it to llama.py (for llama_eagle3.py) and deepseek_v2.py (for deepseek_mtp.py and deepseek_v4_mtp.py?

[EanWang211123]

I have resubmitted the adaptation for llama_eagle3.py and llama.py.

Previously, the EAGLE3 code always created self.draft_id_to_target_id unconditionally:

self.draft_id_to_target_id = nn.Parameter(
    torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
    requires_grad=False,
)

Following the approach in qwen3_dflash, I added a check to determine whether the vocabularies are different before creating draft_id_to_target_id.

I also ran P-EAGLE experiments on 4x 4090 with /models/Qwen3-Coder-30B-A3B-Instruct and amazon/Qwen3-Coder-30B-A3B-Instruct-P-EAGLE.

Dataset: gsm8k
Concurrency: 32

Setting	Total E2E Output Throughput (tok/s)	Mean TPOT (ms)	Mean ITL (ms)
Before enabling	722.63	23.82	76.00
After enabling	925.42	19.88	64.16

The results show a clear performance improvement after enabling this optimization.

[benchislett]

I'm not sure this is the right way to handle it. If there's a D2T in the checkpoint, it probably means it should be used and we messed up by setting draft_id_to_target_id to None

[benchislett]

I don't love that we have to duplicate this code on all the EAGLE models. Maybe we can have a mixin similar to EagleModelMixin or some similar pattern?

[benchislett]

It sems like compute_logits is already expecting this to be set to None. Are you sure there isn't some other place in the code or weight loading pathway that's handling this already?

[benchislett]

Is this fallback necessary? I feel like we could still just do get_top_tokens and then map those tokens into D2T, right? What's the problem with that?

[EanWang211123]

Fixed by changing the logic to perform local argmax first, then apply D2T mapping

[mergify]

Hi @EanWang211123, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Hi @EanWang211123, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[EanWang211123]

@benchislett

---

Refactor local argmax code

Add LocalArgmaxMixin to interfaces.py and unify the D2T-aware get_top_tokens implementation.

Core logic: perform local argmax first, then apply D2T mapping as needed:

target_id = k + draft_id_to_target_id[k]

Add qwen3.5-mtp support

Tested with 32 concurrent requests on 4090 × 4, TP4.

Dataset: GSM8K

enable

============ Serving Benchmark Result ============
Successful requests:                     128
Failed requests:                         0
Maximum request concurrency:             32
Benchmark duration (s):                  9.46
Total input tokens:                      10121
Total generated tokens:                  32768
Request throughput (req/s):              13.52
Output token throughput (tok/s):         3462.25
Peak output token throughput (tok/s):    1150.00
Peak concurrent requests:                60.00
Total token throughput (tok/s):          4531.62
---------------Time to First Token----------------
Mean TTFT (ms):                          412.57
Median TTFT (ms):                        178.91
P99 TTFT (ms):                           1182.18
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          7.41
Median TPOT (ms):                        7.27
P99 TPOT (ms):                           9.99
---------------Inter-token Latency----------------
Mean ITL (ms):                           33.55
Median ITL (ms):                         26.50
P99 ITL (ms):                            82.74
---------------Speculative Decoding---------------
Acceptance rate (%):                     72.95
Acceptance length:                       4.65
Drafts:                                  7085
Draft tokens:                            35425
Accepted tokens:                         25841
Per-position acceptance (%):
  Position 0:                            92.51
  Position 1:                            81.72
  Position 2:                            72.21
  Position 3:                            63.35
  Position 4:                            54.95
==================================================

disable

============ Serving Benchmark Result ============
Successful requests:                     128
Failed requests:                         0
Maximum request concurrency:             32
Benchmark duration (s):                  10.41
Total input tokens:                      10121
Total generated tokens:                  32768
Request throughput (req/s):              12.29
Output token throughput (tok/s):         3146.56
Peak output token throughput (tok/s):    1023.00
Peak concurrent requests:                61.00
Total token throughput (tok/s):          4118.43
---------------Time to First Token----------------
Mean TTFT (ms):                          436.33
Median TTFT (ms):                        182.45
P99 TTFT (ms):                           1255.88
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          8.08
Median TPOT (ms):                        8.07
P99 TPOT (ms):                           10.76
---------------Inter-token Latency----------------
Mean ITL (ms):                           36.43
Median ITL (ms):                         30.71
P99 ITL (ms):                            83.22
---------------Speculative Decoding---------------
Acceptance rate (%):                     72.48
Acceptance length:                       4.62
Drafts:                                  7115
Draft tokens:                            35575
Accepted tokens:                         25786
Per-position acceptance (%):
  Position 0:                            92.24
  Position 1:                            81.48
  Position 2:                            71.83
  Position 3:                            62.64
  Position 4:                            54.22
==================================================

[EanWang211123]

Since this is an opt-in feature, can we also add it to llama.py (for llama_eagle3.py) and deepseek_v2.py (for deepseek_mtp.py and deepseek_v4_mtp.py?

For the local-argmax integration for dpsk-v4-mtp, as I understand it, the logits of DeepSeek V4 MTP are not computed as “a fixed lm_head + raw hidden_states”. Instead, they are computed by “selecting a layer according to spec_step_idx → passing it through hc_head/rmsnorm first → then using that layer’s shared_head.head to compute the vocabulary logits”. Therefore, the LocalArgmaxMixin used by PARD/EAGLE, which assumes a fixed self.lm_head, no step index, and no preprocessing, cannot be copied over directly.

I also think this optimization mainly benefits devices like the 4090 that do not have NVLink. Perhaps most devices capable of deploying dpsk-v4 already have NVLink, so the improvement may not be very large. That said, I have not tested it yet. If the adaptation works and proves effective, I will implement a dedicated local-argmax adaptation for dpsk-v4.

[benchislett]

@EanWang211123 I think we could also add it for deepseek in a follow-up if needed. I'll move forward for now

[benchislett]

Thanks!