@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e7c903609f

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

Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@mistral.ai> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>

Nick Hill committed Sep 28, 2026 at 12:48 UTC e7c903609f6ca10f8d6c8e9a5d09a39ff7f4c2e4

12 files changed +144 -42

tests/v1/spec_decode/test_dynamic_sd_cug.py

+6 -6

@@ -126,7 +126,7 @@ def test_dynamic_sd_full_cudagraph_covers_all_uniform_decode_shapes(monkeypatch)

126 num_reqs,

127 num_tokens,

128 max_query_len,

129 - has_prefill=False,

129 + decode_graph_eligible=True,

130 )

131

132 # The scheduler should mark every one of these shapes as a uniform

@@ -262,14 +262,14 @@ def test_prompt_chunks_shaped_like_spec_decode_miss_the_full_graph(monkeypatch):

262 # The batch shape is indistinguishable from a full batch of spec decodes.

263 assert (

264 get_uniform_decode_token_count(

265 - num_reqs, num_tokens, decode_query_len, has_prefill=False

265 + num_reqs, num_tokens, decode_query_len, decode_graph_eligible=True

266 )

267 == decode_query_len

268 )

269

270 # Two of the requests are decode_query_len tokens into a longer prompt.

271 uniform_tok_count = get_uniform_decode_token_count(

272 - num_reqs, num_tokens, decode_query_len, has_prefill=True

272 + num_reqs, num_tokens, decode_query_len, decode_graph_eligible=False

273 )

274 assert uniform_tok_count is None

275 desc = manager.dispatch(

@@ -283,7 +283,7 @@ def test_prompt_chunks_shaped_like_spec_decode_miss_the_full_graph(monkeypatch):

283

284 # The same shape with every request decoding still gets its FULL graph.

285 uniform_tok_count = get_uniform_decode_token_count(

286 - num_reqs, num_tokens, decode_query_len, has_prefill=False

286 + num_reqs, num_tokens, decode_query_len, decode_graph_eligible=True

287 )

288 assert uniform_tok_count == decode_query_len

289 desc = manager.dispatch(

@@ -336,7 +336,7 @@ def test_basic_sd_does_not_capture_shorter_full_decode_shapes(monkeypatch):

336 num_reqs,

337 num_tokens,

338 max_query_len,

339 - has_prefill=False,

339 + decode_graph_eligible=True,

340 )

341 assert uniform_tok_count == max_query_len

342

@@ -401,7 +401,7 @@ def test_dynamic_sd_only_captures_scheduled_query_lengths(monkeypatch):

401 num_reqs,

402 num_tokens,

403 max_query_len,

404 - has_prefill=False,

404 + decode_graph_eligible=True,

405 )

406 assert uniform_tok_count == max_query_len

407

tests/v1/worker/test_gpu_batch_ordering.py

+65 -4

@@ -56,6 +56,7 @@ def _make_runner(

56 runner: Any = GPUModelRunner.__new__(GPUModelRunner)

57 runner.decode_query_len = decode_query_len

58 runner.adaptive_verification = None

59 + runner.pcp_manager = None

60 runner.req_states = SimpleNamespace(

61 req_id_to_index={req_id: i for i, req_id in enumerate(req_states)},

62 # The runner keeps this as min(num_computed_tokens, prefill_len).

@@ -128,6 +129,65 @@ def test_prompt_chunk_of_decode_query_len_is_not_uniform_decode():

129 assert _uniform_token_count({"p0": (0, 8)}, 8) is None

130

131

132 +def test_one_new_prompt_token_over_context_runs_as_decode():

133 + # A prompt tail padded with placeholder drafts and a one-token mid-prompt

134 + # chunk both compute like decodes. A fresh one-token prompt and a real

135 + # multi-token chunk don't.

136 + batch = {"d0": (16, 16), "tail": (128, 129), "mid": (64, 129)}

137 + batch.update({"fresh": (0, 1), "chunk": (120, 129)})

138 + runner = _make_runner(batch, decode_query_len=4)

139 + scheduler_output = SimpleNamespace(

140 + num_scheduled_tokens={"d0": 4, "tail": 4, "mid": 1, "fresh": 1, "chunk": 4},

141 + total_num_scheduled_tokens=14,

142 + scheduled_spec_decode_tokens={"d0": [1, 2, 3], "tail": [-1] * 3},

143 + )

144 + state, _ = runner.gather_batch_req_state(scheduler_output, False)

145 + assert state is not None

146 + runs_as_decode = dict(zip(state.req_ids, state.prefill_runs_as_decode_np))

147 + assert runs_as_decode == {

148 + "d0": False,

149 + "tail": True,

150 + "mid": True,

151 + "fresh": False,

152 + "chunk": False,

153 + }

154 +

155 +

156 +def test_one_token_prompt_tail_with_prior_context_is_uniform_decode():

157 + # Replaying the last prompt token over existing context (full prefix-cache

158 + # hit or P/D KV import) computes exactly like a decode.

159 + batch = {"d0": (16, 16), "d1": (16, 16), "p0": (128, 129)}

160 + assert _uniform_token_count(batch, 1) == 1

161 +

162 + # A fresh 1-token prompt alongside it has no prior state.

163 + batch["p1"] = (0, 1)

164 + assert _uniform_token_count(batch, 1) is None

165 +

166 + # Only one-token tails qualify, not K+1 chunks with prior context.

167 + assert _uniform_token_count({"d0": (16, 16), "p0": (120, 129)}, 8) is None

168 +

169 +

170 +def test_padded_prompt_tail_with_prior_context_is_uniform_decode():

171 + # The scheduler pads a one-token prompt tail with K placeholder drafts to

172 + # keep the K+1 spec-decode shape; it verifies like any other decode.

173 + batch = {"d0": (16, 16), "d1": (16, 16), "p0": (128, 129)}

174 + runner = _make_runner(batch, decode_query_len=4)

175 + scheduler_output = SimpleNamespace(

176 + num_scheduled_tokens={req_id: 4 for req_id in batch},

177 + total_num_scheduled_tokens=12,

178 + scheduled_spec_decode_tokens={req_id: [-1] * 3 for req_id in batch},

179 + )

180 + state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False)

181 + assert state.has_prefill and state.decode_graph_eligible

182 + assert uniform_tok_count == 4

183 +

184 + # Without the placeholder drafts, it is a real 4-token prompt chunk.

185 + del scheduler_output.scheduled_spec_decode_tokens["p0"]

186 + state, uniform_tok_count = runner.gather_batch_req_state(scheduler_output, False)

187 + assert not state.decode_graph_eligible

188 + assert uniform_tok_count is None

189 +

190 +

191 def test_dummy_batches_stay_uniform_decode():

192 # Dummy batches have no request state to consult and must stay classified

193 # by shape alone; rejecting them would make capture and replay disagree.

@@ -185,6 +245,7 @@ def test_uniform_decode_uses_state_index_not_batch_position():

245 runner: Any = GPUModelRunner.__new__(GPUModelRunner)

246 runner.decode_query_len = 8

247 runner.adaptive_verification = None

248 + runner.pcp_manager = None

249 # State arrays in state-index order: a prefilling request, then two decodes.

250 runner.req_states = SimpleNamespace(

251 req_id_to_index={"prefilling": 0, "decode_a": 1, "decode_b": 2},

@@ -214,11 +275,11 @@ def test_uniform_decode_uses_state_index_not_batch_position():

275

276

277 def test_uniform_decode_predicate():

217 - # Shape and prefill state must both pass.

218 - assert get_uniform_decode_token_count(2, 16, 8, False) == 8

219 - assert get_uniform_decode_token_count(2, 16, 8, True) is None

278 + # Shape and decode-graph eligibility must both pass.

279 + assert get_uniform_decode_token_count(2, 16, 8, True) == 8

280 + assert get_uniform_decode_token_count(2, 16, 8, False) is None

281 # 12 tokens over 2 requests is no shared query length.

221 - assert get_uniform_decode_token_count(2, 12, 8, False) is None

282 + assert get_uniform_decode_token_count(2, 12, 8, True) is None

283

284

285 def test_no_speculator_dispatches_on_query_length_alone():

tests/v1/worker/test_gpu_batch_shard.py

@@ -96,6 +96,7 @@ def _make_batch(

96 num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32),

97 is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_),

98 has_prefill=False,

99 + decode_graph_eligible=True,

100 input_ids=torch.zeros(num_tokens, dtype=torch.int32, device=DEVICE),

101 positions=torch.arange(num_tokens, dtype=torch.int64, device=DEVICE),

102 is_padding=torch.zeros(num_tokens, dtype=torch.bool, device=DEVICE),

tests/v1/worker/test_mamba_hybrid_model_state.py

+4 -3

@@ -45,6 +45,7 @@ def test_prepare_attn_forwards_positions(monkeypatch: pytest.MonkeyPatch) -> Non

45 seq_lens_cpu_upper_bound=torch.tensor([1537], dtype=torch.int32),

46 seq_lens=torch.tensor([1537], dtype=torch.int32),

47 is_prefilling_np=torch.tensor([False]).numpy(),

48 + prefill_runs_as_decode_np=None,

49 dcp_local_seq_lens=None,

50 positions=positions,

51 prompt_lens=torch.tensor([1024], dtype=torch.int32),

@@ -102,8 +103,7 @@ def test_padded_prompt_tail_builds_as_spec_decode(

103 seq_lens_cpu_upper_bound=torch.tensor(seq_lens, dtype=torch.int32),

104 seq_lens=torch.tensor(seq_lens, dtype=torch.int32),

105 is_prefilling_np=np.array(is_prefilling),

105 - prefill_len_np=np.array([40, 129, 100], dtype=np.int32),

106 - num_computed_prefill_tokens_np=np.array([40, 128, 0], dtype=np.int32),

106 + prefill_runs_as_decode_np=np.array([False, True, False]),

107 dcp_local_seq_lens=None,

108 positions=torch.zeros(12, dtype=torch.int64),

109 prompt_lens=None,

@@ -121,11 +121,12 @@ def test_padded_prompt_tail_builds_as_spec_decode(

121 mamba_metadata = build_attn_metadata.call_args.kwargs[

122 "model_specific_attn_metadata"

123 ]

124 + assert mamba_metadata.is_prefilling.tolist() == [False, False, True]

125

126 builder = _create_gdn_builder(num_speculative_tokens=k)

127 common = create_common_attn_metadata(

128 BatchSpec(seq_lens=seq_lens, query_lens=query_lens), BLOCK_SIZE, DEVICE

128 - ).replace(is_prefilling=torch.tensor(is_prefilling))

129 + ).replace(**mamba_metadata.get_extra_common_attn_kwargs(0, 3))

130 meta = builder.build(

131 common_prefix_len=0,

132 common_attn_metadata=common,

vllm/v1/worker/gpu/input_batch.py

@@ -81,6 +81,8 @@ class InputBatch:

81 is_prefilling_np: np.ndarray

82 # == np.any(is_prefilling_np)

83 has_prefill: bool

84 + # No prefills, or only prefill rows that run as decodes.

85 + decode_graph_eligible: bool

87 # [num_tokens_after_padding]

88 input_ids: torch.Tensor

@@ -113,6 +115,11 @@ class InputBatch:

115 # [num_reqs] set only under PCP+DCP (see CommonAttentionMetadata).

116 dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None

117

118 + # [num_reqs] prefilling rows that schedule one new prompt token (excluding

119 + # drafts) over existing context and so compute exactly like decodes.

120 + # None if there are no prefills.

121 + prefill_runs_as_decode_np: np.ndarray | None = None

122 +

123 @classmethod

124 def make_dummy(

125 cls,

@@ -196,6 +203,7 @@ class InputBatch:

203 num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32),

204 is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_),

205 has_prefill=False,

206 + decode_graph_eligible=True,

207 input_ids=input_ids,

208 positions=positions,

209 is_padding=is_padding,

vllm/v1/worker/gpu/model_runner.py

+38 -11

@@ -1242,7 +1242,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1242 if dummy_run:

1243 # Dummy batches are uniform by construction.

1244 return None, get_uniform_decode_token_count(

1245 - num_reqs, num_toks, max_query_len, has_prefill=False

1245 + num_reqs, num_toks, max_query_len, decode_graph_eligible=True

1246 )

1247

1248 draft_tokens = scheduler_output.scheduled_spec_decode_tokens

@@ -1261,24 +1261,48 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1261 idx_mapping_np

1262 ]

1263 is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np

1264 + has_prefill = bool(is_prefilling_np.any())

1265

1265 - if self.adaptive_verification is not None and draft_tokens:

1266 - num_toks = self.adaptive_verification.get_num_tokens(

1267 - num_tokens_per_req, draft_tokens

1266 + num_draft_tokens_np = None

1267 + if draft_tokens:

1268 + num_drafts_iter = (len(draft_tokens.get(req_id, ())) for req_id in req_ids)

1269 + num_draft_tokens_np = np.fromiter(

1270 + num_drafts_iter, dtype=np.int32, count=num_reqs

1271 + )

1272 + if self.adaptive_verification is not None:

1273 + num_toks = self.adaptive_verification.get_num_tokens(

1274 + num_tokens_per_req, draft_tokens

1275 + )

1276 +

1277 + prefill_runs_as_decode_np = None

1278 + decode_graph_eligible = not has_prefill

1279 + if has_prefill and self.pcp_manager is None:

1280 + # One new prompt token over existing context (possibly padded with

1281 + # placeholder drafts) runs as a decode. PCP partitions prefill rows.

1282 + num_new_tokens = num_scheduled_tokens

1283 + if num_draft_tokens_np is not None:

1284 + num_new_tokens = num_new_tokens - num_draft_tokens_np

1285 + one_token_ext = (num_new_tokens == 1) & (num_computed_prefill_tokens_np > 0)

1286 + prefill_runs_as_decode_np = is_prefilling_np & one_token_ext

1287 + decode_graph_eligible = bool(

1288 + (prefill_runs_as_decode_np | ~is_prefilling_np).all()

1289 )

1290

1291 batch_state = BatchReqState(

1292 req_ids=req_ids,

1293 num_scheduled_tokens=num_scheduled_tokens,

1294 num_tokens=num_toks,

1295 + num_draft_tokens_np=num_draft_tokens_np,

1296 idx_mapping_np=idx_mapping_np,

1297 prefill_len_np=prefill_len_np,

1298 num_computed_prefill_tokens_np=num_computed_prefill_tokens_np,

1299 is_prefilling_np=is_prefilling_np,

1278 - has_prefill=bool(is_prefilling_np.any()),

1300 + has_prefill=has_prefill,

1301 + prefill_runs_as_decode_np=prefill_runs_as_decode_np,

1302 + decode_graph_eligible=decode_graph_eligible,

1303 )

1304 return batch_state, get_uniform_decode_token_count(

1281 - num_reqs, num_toks, max_query_len, batch_state.has_prefill

1305 + num_reqs, num_toks, max_query_len, decode_graph_eligible

1306 )

1307

1308 def prepare_inputs(

@@ -1320,11 +1344,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1344 num_reqs, dtype=torch.int32, device=self.device

1345 )

1346 else:

1323 - num_draft_tokens_per_req = np.fromiter(

1324 - (len(draft_tokens.get(req_id, ())) for req_id in req_ids),

1325 - dtype=np.int32,

1326 - count=num_reqs,

1327 - )

1347 + num_draft_tokens_per_req = batch_req_state.num_draft_tokens_np

1348 + if num_draft_tokens_per_req is None:

1349 + num_draft_tokens_per_req = np.zeros(num_reqs, dtype=np.int32)

1350 num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step

1351 total_num_draft_tokens = int(num_draft_tokens_per_req.sum())

1352 total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens

@@ -1463,6 +1485,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1485 num_computed_prefill_tokens_np=batch_req_state.num_computed_prefill_tokens_np,

1486 is_prefilling_np=batch_req_state.is_prefilling_np,

1487 has_prefill=batch_req_state.has_prefill,

1488 + prefill_runs_as_decode_np=batch_req_state.prefill_runs_as_decode_np,

1489 + decode_graph_eligible=batch_req_state.decode_graph_eligible,

1490 input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],

1491 positions=self.input_buffers.positions[:num_tokens_after_padding],

1492 is_padding=is_padding,

@@ -2345,11 +2369,14 @@ class BatchReqState(NamedTuple):

2369 # May be less than scheduler_output.total_num_scheduled_tokens:

2370 # adaptive verification trims the draft budget before running.

2371 num_tokens: int

2372 + num_draft_tokens_np: np.ndarray | None # [num_reqs], None if no drafts

2373 idx_mapping_np: np.ndarray # [num_reqs]

2374 prefill_len_np: np.ndarray # [num_reqs]

2375 num_computed_prefill_tokens_np: np.ndarray # [num_reqs]

2376 is_prefilling_np: np.ndarray # [num_reqs]

2377 has_prefill: bool

2378 + prefill_runs_as_decode_np: np.ndarray | None # [num_reqs]

2379 + decode_graph_eligible: bool

2380

2381

2382 def sort_batch_req_ids(

vllm/v1/worker/gpu/model_states/mamba_hybrid.py

+7 -13

@@ -258,10 +258,13 @@ class MambaHybridModelState(DefaultModelState):

258 else:

259 max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item()

260

261 + is_prefilling_np = input_batch.is_prefilling_np

262 + if input_batch.prefill_runs_as_decode_np is not None:

263 + # A prompt tail the scheduler padded with placeholder drafts must run

264 + # as a spec-decode row: the prefill kernels can't roll them back.

265 + is_prefilling_np = is_prefilling_np & ~input_batch.prefill_runs_as_decode_np

266 is_prefilling = torch.zeros(num_reqs, dtype=torch.bool, device="cpu")

262 - is_prefilling[: input_batch.num_reqs] = torch.from_numpy(

263 - input_batch.is_prefilling_np

264 - )

267 + is_prefilling[: input_batch.num_reqs] = torch.from_numpy(is_prefilling_np)

268 # During CUDAGraph capture, num_decode_draft_tokens_cpu and num_accepted_tokens

269 # are created by attn_metadata_builder.build_for_cudagraph_capture, so we only

270 # compute them during actual (non-capture) forward execution.

@@ -281,16 +284,7 @@ class MambaHybridModelState(DefaultModelState):

284 # Test request state, not num_scheduled_tokens == draft_count+1:

285 # adaptive rewrites num_scheduled_tokens to an even split, so that

286 # equality rarely holds and would demote every verify row to decode.

284 - # A one-token prompt tail over prior state that the scheduler padded

285 - # with placeholder drafts is also a spec-decode row: the prefill

286 - # kernels can't roll the placeholders back.

287 - num_computed = input_batch.num_computed_prefill_tokens_np

288 - is_prompt_tail = (num_computed > 0) & (

289 - input_batch.prefill_len_np - num_computed == 1

290 - )

291 - is_decode = (~input_batch.is_prefilling_np | is_prompt_tail) & (

292 - input_batch.num_scheduled_tokens > 0

293 - )

287 + is_decode = ~is_prefilling_np & (input_batch.num_scheduled_tokens > 0)

288 spec_decode_mask = (num_draft_tokens_per_req > 0) & is_decode

289 num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where(

290 spec_decode_mask, num_draft_tokens_per_req, -1

vllm/v1/worker/gpu/pcp_manager.py

+4 -1

@@ -595,6 +595,7 @@ class PCPManager:

595 )

596 local_is_prefilling_np = np.zeros(num_reqs_after_padding, dtype=np.bool_)

597 local_is_prefilling_np[:num_local_reqs] = real_local_is_prefilling_np

598 + local_has_prefill = bool(local_is_prefilling_np.any())

599 seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_after_padding, dtype=np.int32)

600 seq_lens_cpu_upper_bound_np[:num_local_reqs] = (

601 local_start_pos_np + local_num_scheduled_tokens

@@ -645,7 +646,9 @@ class PCPManager:

646 prefill_len_np=local_prefill_len_np,

647 num_computed_prefill_tokens_np=local_num_computed_prefill_tokens_np,

648 is_prefilling_np=local_is_prefilling_np,

648 - has_prefill=bool(local_is_prefilling_np.any()),

649 + has_prefill=local_has_prefill,

650 + decode_graph_eligible=not local_has_prefill,

651 + prefill_runs_as_decode_np=None,

652 input_ids=input_buffers.input_ids[:num_local_tokens_padded],

653 positions=input_buffers.positions[:num_local_tokens_padded],

654 is_padding=is_padding,

vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py

+1 -1

@@ -282,7 +282,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

282 # the target model during FULL cudagraph.

283 num_tokens,

284 max_query_len,

285 - input_batch.has_prefill,

285 + input_batch.decode_graph_eligible,

286 )

287 prefill_batch_desc, prefill_batch_sync = dispatch_cg_and_sync_dp(

288 self.prefill_cudagraph_manager,

vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py

+1 -1

@@ -190,7 +190,7 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):

190 # When all requests are decoding (no true prefills), each has

191 # num_speculative_steps + 1 tokens, enabling FULL graph replay.

192 uniform_token_count = get_uniform_decode_token_count(

193 - num_reqs, num_tokens, max_query_len, input_batch.has_prefill

193 + num_reqs, num_tokens, max_query_len, input_batch.decode_graph_eligible

194 )

195 batch_desc, batch_sync = dispatch_cg_and_sync_dp(

196 self.cudagraph_manager,

vllm/v1/worker/gpu/ubatch_utils.py

@@ -181,6 +181,11 @@ def _slice_input_batch(

181 req_start:req_stop

182 ],

183 is_prefilling_np=input_batch.is_prefilling_np[req_start:req_stop],

184 + prefill_runs_as_decode_np=(

185 + None

186 + if input_batch.prefill_runs_as_decode_np is None

187 + else input_batch.prefill_runs_as_decode_np[req_start:req_stop]

188 + ),

189 input_ids=input_batch.input_ids[tok_start:tok_stop],

190 positions=input_batch.positions[tok_start:tok_stop],

191 is_padding=input_batch.is_padding[tok_start:tok_stop],

vllm/v1/worker/utils.py

+4 -2

@@ -747,10 +747,12 @@ def is_uniform_query_len(num_reqs: int, num_tokens: int, max_query_len: int) ->

747

748

749 def get_uniform_decode_token_count(

750 - num_reqs: int, num_tokens: int, max_query_len: int, has_prefill: bool

750 + num_reqs: int, num_tokens: int, max_query_len: int, decode_graph_eligible: bool

751 ) -> int | None:

752 """Per-request token count of a uniform decode batch, or None."""

753 - if not has_prefill and is_uniform_query_len(num_reqs, num_tokens, max_query_len):

753 + if decode_graph_eligible and is_uniform_query_len(

754 + num_reqs, num_tokens, max_query_len

755 + ):

756 return max_query_len

757 return None

758