[Core][PD] Allow FULL decode graphs for one-token prompt continuations on MRv2 by NickLucche · Pull Request #58374 · vllm-project/vllm

NickLucche · 2026-09-23T14:24:53Z

On MRV2, the decode node in P/D leaves FULL CUDA graphs every time a new request arrives.
This affects every KV connector. It was first reported with HiSparse (#57049 / #57939), but HiSparse doesn't cause it.

How it happens

Suppose D has speculation disabled and two requests already decoding.
A third request arrives with a 129-token prompt whose KV has been imported from P.

D replays the last prompt token to obtain logits for the first generated token.
The scheduler therefore sets num_computed_tokens = 128 and schedules one token:

   Request    Work scheduled on D                                Logical status
  ━━━━━━━━━  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━  ━━━━━━━━━━━━━━━━━━━━━━━━━━━━
   A          One decode token over existing context             Decode
  ─────────  ─────────────────────────────────────────────────  ────────────────────────────
   B          One decode token over existing context             Decode
  ─────────  ─────────────────────────────────────────────────  ────────────────────────────
   C          Last prompt token over the preceding 128 tokens    Prefill, because 128 < 129

Before: C makes has_prefill=True. MRV2 rejects the FULL decode graph for the entire batch, so A and B also fall back to PIECEWISE or eager execution. In P/D serving, every new arrival can trigger this.

After: MRV2 recognizes that every request schedules one token and already has context, making the batch eligible for a FULL decode graph. C remains logically prefilling, so input preparation still loads its token from the prompt.

The fix

This changes only graph classification in gather_batch_req_state.
On D nodes with decode_query_len == 1, max_query_len == 1 and no PCP, a prefilling request with prior state (num_computed_prefill_tokens > 0) no longer blocks FULL graphs.
Fresh one-token prompts (num_computed == 0) are still excluded; that is the #50532 bug #51865 fixed.

This is safe because a one-token query over existing context is what the decode kernels already expect. Sparse MLA and the indexer split by query length, and Mamba, GDN and KDA already treat one-token chunks with prior state as decodes
(mamba_attn.py:474, gdn_attn.py:250).

Testing

Ignore ITL I will post bigger benchs soon

B200, P/D TP4, Inferact/GLM-5.3-NVFP4 + HiSparse on D
160 requests arriving at 4 per second (2000 tokens in, 400 out), so new requests keep landing while others decode

rate(req/s)  arm       FULL steps  PIECEWISE steps  p99 ITL(ms)  median TTFT(ms)
4           baseline  16451       120              1.90         30.6
4           fix       17193       0                1.90         30.5
8           baseline  19525       238              4.21         31.0
8           fix       20217       0                1.97         28.1
16          baseline  19179       476              4.29         31.1
16          fix       20406       0                2.45         30.5
32          baseline  15276       937              5.38         39.6
32          fix       18335       0                3.74         33.1

Qwen3-0.6B, PD TP1
1000 tokens in and 256 out. Graph-mode counts are from D's --cudagraph-metrics. 

rate(req/s)  arm       FULL steps  PIECEWISE steps  p99 ITL(ms)  median TTFT(ms)
4           baseline  16451       120              1.90         30.6
4           fix       17193       0                1.90         30.5
8           baseline  19525       238              4.21         31.0
8           fix       20217       0                1.97         28.1
16          baseline  19179       476              4.29         31.1
16          fix       20406       0                2.45         30.5
32          baseline  15276       937              5.38         39.6
32          fix       18335       0                3.74         33.1

On PD checking

Yeah I know we're missing a cleaner way to answer the "am I P on a PD setup" question.
I will follow up with something better (it doesn't really belong to this PR), but that's why I am basically enumerating all PD connectors for this fix.

PS: Spec decode on D isn't covered: the scheduler pads the replay to K+1 tokens, which still trips has_prefill.
That will be a follow-up.

cc @njhill @S1ro1 @LucasWilkinson

…V consumers On the decode side of P/D, a full-prompt KV import resumes with a single-token replay of the last prompt token. That request is still logically prefilling, so vllm-project#51865's strict has_prefill check pushes the whole batch off the FULL decode graph on every new arrival. With decode_query_len == 1 the replay is decode-shaped, and linear attention backends already treat one-token chunks with prior state as decodes. Relax only the graph classification for KV consumers; logical prefill state is unchanged. Co-authored-by: Claude Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>

Co-authored-by: Claude Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>

Add KVTransferConfig.is_pd_consumer: kv_consumer role with a known P/D connector (NIXL, Mooncake, MoRIIO), directly or inside MultiConnector. Co-authored-by: Claude Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>

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.

Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>

NickLucche · 2026-09-23T16:08:06Z

Pre-PR vs Post-PR

Baseline = 955bd6abef (before), fix = this PR (after), both with VLLM_USE_V2_MODEL_RUNNER=1. non-FULL steps = share of decode-side steps not dispatched as FULL CUDA graph (--cudagraph-metrics). A → B (Δ) is before → after.

Qwen3-0.6B, P/D TP1/TP1

NixlConnector 1P+1D on B200, random dataset with --ignore-eos, 60s per point, mean of 2 reps.

ISL/OSL	req/s	non-FULL steps	mean TPOT (ms)	p99 TPOT (ms)	p99 ITL (ms)	median TTFT (ms)
1024/256	8	1.31% → 0.00%	1.49 → 1.47 (-1.7%)	1.67 → 1.63 (-2.3%)	4.28 → 3.23 (-24.5%)	37.12 → 35.16 (-5.3%)
1024/256	16	2.60% → 0.00%	1.68 → 1.60 (-4.6%)	2.14 → 1.97 (-8.0%)	4.94 → 4.86 (-1.7%)	42.52 → 39.78 (-6.4%)
1024/256	32	6.80% → 0.00%	2.27 → 1.99 (-12.1%)	2.93 → 2.51 (-14.4%)	8.61 → 7.61 (-11.7%)	52.86 → 47.08 (-10.9%)
4096/1024	4	0.94% → 0.00%	2.66 → 2.53 (-4.7%)	3.65 → 3.42 (-6.1%)	14.33 → 10.28 (-28.2%)	98.01 → 88.79 (-9.4%)
4096/1024	8	4.31% → 0.00%	7.38 → 6.31 (-14.5%)	10.22 → 8.59 (-15.9%)	36.94 → 31.29 (-15.3%)	139.15 → 122.73 (-11.8%)
4096/1024	16	12.35% → 0.04%	26.03 → 24.62 (-5.4%)	30.90 → 29.31 (-5.1%)	81.62 → 69.65 (-14.7%)	15.48s → 13.49s (-12.9%)

Llama-3-8B-Instruct, P/D TP1/TP1

NixlConnector 1P+1D on B200, random dataset with --ignore-eos, 60s per point, mean of 3 reps.

ISL/OSL	req/s	non-FULL steps	mean TPOT (ms)	p99 TPOT (ms)	p99 ITL (ms)	median TTFT (ms)
1024/256	8	3.08% → 0.00%	3.78 → 3.78 (-0.1%)	3.91 → 3.88 (-0.7%)	4.96 → 4.76 (-4.1%)	51.32 → 51.11 (-0.4%)
1024/256	16	6.29% → 0.00%	4.07 → 4.06 (-0.3%)	4.59 → 4.52 (-1.5%)	8.07 → 8.05 (-0.3%)	59.13 → 60.07 (+1.6%)
1024/256	32	14.99% → 0.00%	5.04 → 4.86 (-3.6%)	6.02 → 5.75 (-4.5%)	15.06 → 13.94 (-7.4%)	70.29 → 68.56 (-2.5%)
4096/1024	4	2.40% → 0.00%	6.92 → 6.98 (+0.8%)	8.22 → 8.31 (+1.1%)	38.06 → 39.76 (+4.5%)	157.81 → 160.89 (+2.0%)
4096/1024	8	8.93% → 0.00%	16.52 → 16.91 (+2.4%)	20.44 → 21.04 (+3.0%)	47.57 → 48.61 (+2.2%)	228.86 → 232.82 (+1.7%)
4096/1024	16	16.00% → 2.34%	29.49 → 29.83 (+1.1%)	40.19 → 41.63 (+3.6%)	134.19 → 132.36 (-1.4%)	18.34s → 18.61s (+1.5%)

MRV1 vs MRV2, Qwen3-0.6B, P/D TP1/TP1

NixlConnector 1P+1D on B200, random dataset with --ignore-eos, 60s per point, mean of 2 reps. Δ is relative to MRV2 before. MRV1 = VLLM_USE_V2_MODEL_RUNNER=0 on this PR (no MRV1 code touched).

ISL/OSL	req/s	arm	non-FULL steps	mean TPOT (ms)	p99 TPOT (ms)	p99 ITL (ms)	median TTFT (ms)
1024/256	8	MRV1	0.00%	1.45 (-2.5%)	1.61 (-3.1%)	3.13 (-26.9%)	34.41 (-7.7%)
1024/256	8	MRV2 before	1.36%	1.49	1.66	4.28	37.29
1024/256	8	MRV2 after	0.00%	1.46 (-2.2%)	1.59 (-4.1%)	3.20 (-25.2%)	34.36 (-7.9%)

gsm8k (5-shot, strict / flexible)

model	MRV1	MRV2 before	MRV2 after
Qwen3-0.6B, P/D TP1/TP1	0.4223 / 0.4177	0.4230 / 0.4185	0.4215 / 0.4170
Llama-3-8B-Instruct, P/D TP1/TP1	–	0.7627 / 0.7582	0.7635 / 0.7589
Qwen3-Next-80B-A3B-Instruct (hybrid GDN), P/D TP2/TP2	–	0.8059 / 0.8469	0.8082 / 0.8529

mergify · 2026-09-29T04:59:01Z

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

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

NickLucche · 2026-09-29T08:11:41Z

we merged Nick's #58400

NickLucche added 4 commits September 23, 2026 10:38

[Perf][MRV2] Require max_query_len == 1 for the prompt-tail exception

5c54d6c

Co-authored-by: Claude Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>

[Perf][MRV2] Limit prompt-tail FULL graph exception to kv_consumer

799a3a3

Co-authored-by: Claude Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>

NickLucche requested review from ProExpertProg, WoosukKwon, houseroad, mgoin, njhill, robertgshaw2-redhat, tlrmchlsmth, yewentao256 and youkaichao as code owners September 23, 2026 14:24

claude Bot reviewed Sep 23, 2026

View reviewed changes

mergify Bot added mrv2 Model Runner V2 specific kv-connector labels Sep 23, 2026

comment

1e0e999

Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>

njhill mentioned this pull request Sep 23, 2026

[Perf][MRV2] Allow FULL decode graphs for one-token prompt tails #58400

Merged

mergify Bot added the needs-rebase label Sep 29, 2026

NickLucche closed this Sep 29, 2026

Navigation Menu

Repository navigation

[Core][PD] Allow FULL decode graphs for one-token prompt continuations on MRv2 - #58374

[Core][PD] Allow FULL decode graphs for one-token prompt continuations on MRv2#58374
NickLucche wants to merge 5 commits into
vllm-project:mainfrom
NickLucche:mrv2-full-cg-one-token-prefill

NickLucche commented Sep 23, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

NickLucche commented Sep 23, 2026 •

edited

Loading

Uh oh!

mergify Bot commented Sep 29, 2026

Uh oh!

NickLucche commented Sep 29, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

NickLucche commented Sep 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

How it happens

The fix

Testing

On PD checking

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

NickLucche commented Sep 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pre-PR vs Post-PR

Qwen3-0.6B, P/D TP1/TP1

Llama-3-8B-Instruct, P/D TP1/TP1

MRV1 vs MRV2, Qwen3-0.6B, P/D TP1/TP1

gsm8k (5-shot, strict / flexible)

Uh oh!

mergify Bot commented Sep 29, 2026

Uh oh!

NickLucche commented Sep 29, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

NickLucche commented Sep 23, 2026 •

edited

Loading

NickLucche commented Sep 23, 2026 •

edited

Loading