[Spec Decode] Unified Parallel Drafting

[benchislett]

Purpose

This PR implements a single input preparation kernel for draft model support, and parallel drafting both with and without hidden states from the target model. As such we now have support for AMD's PARD, which proposes parallel drafting for fine-tuned external draft models, and AWS' P-EAGLE which implements parallel prediction for EAGLE3. Both of these are benchmarked as part of this PR effort.

Testing

E2E tests for parallel drafting and unit tests for the input preparation logic are all passing locally. Confirmed that E2E tests for draft models and EAGLE3 are also still passing locally.

Benchmarks

Benchmarks were conducted for AWS's 2-layer P-EAGLE on GPT-OSS 120B using Acceptance Lengths calculated by averaging AL over each of the MTBench categories, with 2048 max output tokens. The baseline is NVIDIA's EAGLE3 short-context. I also compare AMD's PARD Llama 3.2 1B for Llama 3.3 70B NVFP4, with the autoregressive drafter as a baseline. All benchmarks on 1xB200.

Best config for GPT-OSS at BS=1 is P-EAGLE with K=7, with ~560 output TPS, a speedup of 1.52x over baseline and 1.12x over EAGLE3 best config. At BS=8, P-EAGLE is optimal with K=3, a speedup of 1.34x over baseline and 1.07x over best EAGLE3.

Best config for Llama 3.3 70B-NVFP4 at BS=1 is PARD Llama-1B with K=11, with ~254 output TPS, a speedup of 3.10x over baseline and 1.61x over vanilla draft-model. At BS=8, PARD is optimal with K=7, a speedup of 2.87x over baseline and 1.61x over vanilla draft-model.

All data, with best-at-concurrency bolded for each model.

Checkpoint	K	BS	AL	Median Iter Time	Est. TPS
GPT-OSS 120B EAGLE3	0	1	1	0.0027	370
GPT-OSS 120B EAGLE3	3	1	2.4015427	0.0048	500.3213958
GPT-OSS 120B EAGLE3	4	1	2.5880982	0.0055	470.5633091
GPT-OSS 120B EAGLE3	5	1	2.685875283	0.0061	440.3074234
GPT-OSS 120B EAGLE3	7	1	2.776950087	0.007	396.7071553
GPT-OSS 120B P-EAGLE	0	1	1	0.0027	370
GPT-OSS 120B P-EAGLE	2	1	2.13727217	0.0042	508.8743262
GPT-OSS 120B P-EAGLE	3	1	2.3888	0.0043	555.5348837
GPT-OSS 120B P-EAGLE	4	1	2.545543	0.0046	553.378913
GPT-OSS 120B P-EAGLE	5	1	2.621487	0.0048	546.143125
GPT-OSS 120B P-EAGLE	7	1	2.6937	0.0048	561.1875
Llama 70B-NVFP4 Draft 1B	0	1	1	0.0122	81.96721311
Llama 70B-NVFP4 Draft 1B	3	1	2.964271208	0.0203	146.0232122
Llama 70B-NVFP4 Draft 1B	5	1	3.827371847	0.0243	157.5050143
Llama 70B-NVFP4 Draft 1B	7	1	4.408360614	0.0289	152.5384295
Llama 70B-NVFP4 PARD 1B	0	1	1	0.0122	81.96721311
Llama 70B-NVFP4 PARD 1B	3	1	2.72759473	0.0142	192.0841359
Llama 70B-NVFP4 PARD 1B	5	1	3.251206494	0.0141	230.5820208
Llama 70B-NVFP4 PARD 1B	7	1	3.568283773	0.0141	253.0697711
Llama 70B-NVFP4 PARD 1B	11	1	3.68784371	0.0145	254.334049
Llama 70B-NVFP4 PARD 1B	15	1	3.678693636	0.0145	253.7030094
GPT-OSS 120B EAGLE3	0	8	1	0.0045	1777.777778
GPT-OSS 120B EAGLE3	2	8	2.146052648	0.008	2146.052648
GPT-OSS 120B EAGLE3	3	8	2.4015427	0.0086	2233.993209
GPT-OSS 120B EAGLE3	4	8	2.5880982	0.01	2070.47856
GPT-OSS 120B EAGLE3	5	8	2.685875283	0.0108	1989.537247
GPT-OSS 120B EAGLE3	7	8	2.776950087	0.0117	1898.769291
GPT-OSS 120B P-EAGLE	0	8	1	0.0045	1777.777778
GPT-OSS 120B P-EAGLE	2	8	2.13727217	0.0080	2137.27217
GPT-OSS 120B P-EAGLE	3	8	2.3888	0.0080	2388.8
GPT-OSS 120B P-EAGLE	4	8	2.545543	0.009	2262.704889
GPT-OSS 120B P-EAGLE	5	8	2.621487	0.0094	2231.052766
GPT-OSS 120B P-EAGLE	7	8	2.6937	0.0095	2268.378947
Llama 70B-NVFP4 Draft 1B	0	8	1	0.0118	677.9661017
Llama 70B-NVFP4 Draft 1B	3	8	2.964271208	0.0208	1140.104311
Llama 70B-NVFP4 Draft 1B	5	8	3.827371847	0.0258	1186.781968
Llama 70B-NVFP4 Draft 1B	7	8	4.408360614	0.0299	1179.494479
Llama 70B-NVFP4 PARD 1B	0	8	1	0.012	666.6666667
Llama 70B-NVFP4 PARD 1B	3	8	2.72759473	0.0143	1525.927122
Llama 70B-NVFP4 PARD 1B	5	8	3.251206494	0.0147	1769.364078
Llama 70B-NVFP4 PARD 1B	7	8	3.568283773	0.0149	1915.857059
Llama 70B-NVFP4 PARD 1B	11	8	3.68784371	0.0156	1891.201903
Llama 70B-NVFP4 PARD 1B	15	8	3.678693636	0.0158	1862.629689

[mergify]

Documentation preview: https://vllm--32887.org.readthedocs.build/en/32887/

[gemini-code-assist]

Code Review

This pull request introduces a unified parallel drafting mechanism for speculative decoding, combining logic for EAGLE and other draft models. The changes are extensive, primarily refactoring the speculative decoding logic into a base proposer class and adding a new, complex Triton kernel for preparing inputs. While the overall refactoring appears sound, I've identified a potential critical issue in the new Triton kernel where a safeguard against out-of-bounds memory access is not being used, which could lead to memory corruption.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @benchislett.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @benchislett.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[zihaoanllm]

Hi @benchislett ,Thanks a lot for your great work!! I tested the PARD integration in vLLM and compared its performance with the PARD repo example. Under the same configuration, the acceptance length is well aligned between the two runs (3.56 vs 3.50).

Below are the full benchmark results and the exact script I used for vLLM testing.

Results

framework	target	draft method	bmk	device	bs	k	baseline tps	PARD tps	speedup	accept length
pard repo with transformers+	L3.1 8B	PARD	mt_bench	A100-40GB	1	8	76.55	197.77	2.58	3.50
vllm	L3.1 8B	PARD	mt_bench	A100-40GB	1	8	77.69	202.29	2.60	3.56

---

vLLM test script

# Start server
k=8
target=unsloth/Meta-Llama-3.1-8B-Instruct
draft=amd/PARD-Llama-3.2-1B

vllm serve $target \
  --tensor-parallel-size 1 \
  --max-model-len 4096 \
  --no-enable-prefix-caching \
  --port 8811 \
  --speculative-config '{"model": "'"$draft"'", "method": "draft_model", "num_speculative_tokens": '"$k"', "parallel_drafting": true}' \

# Benchmark
MAX_CONCURRENCY=1
NUM_PROMPTS=80
vllm bench serve --port 8811 \
    --temperature 0 \
    --backend openai-chat \
    --endpoint /v1/chat/completions \
    --dataset-name hf \
    --dataset-path philschmid/mt-bench \
    --num-prompts ${NUM_PROMPTS} \
    --max-concurrency ${MAX_CONCURRENCY}

[benchislett]

For future readers, here is a link to AWS's P-EAGLE arxiv paper:
https://arxiv.org/pdf/2602.01469

[AlecHenx]

why 1+1+ speculative_config.num_speculative_tokens without parallel_drafting here? Big issue here

[benchislett]

You are misunderstanding. It's 1 + 1 * num_spec_tokens. So it will be the same.