@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e1e4646b06

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 07:30 UTC

[Model Runner V2] Rebuild attn metadata between draft decode steps (#41162)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>

Giancarlo Delfin committed May 4, 2026 at 17:44 UTC e1e4646b06f289475ee57f31e3df817e06351321

3 files changed +198 -121

vllm/v1/worker/gpu/sample/gumbel.py

+18 -4

@@ -76,6 +76,8 @@ def gumbel_block_argmax(

76 pos_ptr,

77 processed_logits_ptr,

78 processed_logits_stride,

79 + processed_logits_col_ptr,

80 + vocab_size,

81 APPLY_TEMPERATURE: tl.constexpr,

82 ):

83 req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)

@@ -88,8 +90,15 @@ def gumbel_block_argmax(

91 if processed_logits_ptr is not None:

92 # Store the temperature-applied logits.

93 + if processed_logits_col_ptr is not None:

94 + col = tl.load(processed_logits_col_ptr)

95 + else:

96 + col = 0

97 tl.store(

92 - processed_logits_ptr + req_state_idx * processed_logits_stride + block,

98 + processed_logits_ptr

99 + + req_state_idx * processed_logits_stride

100 + + col * vocab_size

101 + + block,

102 logits,

103 mask=mask,

104 )

@@ -121,6 +130,7 @@ def _gumbel_sample_kernel(

130 local_max_stride,

131 processed_logits_ptr,

132 processed_logits_stride,

133 + processed_logits_col_ptr,

134 logits_ptr,

135 logits_stride,

136 expanded_idx_mapping_ptr,

@@ -153,6 +163,8 @@ def _gumbel_sample_kernel(

163 pos_ptr,

164 processed_logits_ptr,

165 processed_logits_stride,

166 + processed_logits_col_ptr,

167 + vocab_size,

168 APPLY_TEMPERATURE=APPLY_TEMPERATURE,

169 )

170 token_id = block_idx * BLOCK_SIZE + idx

@@ -167,7 +179,8 @@ def gumbel_sample(

179 seed: torch.Tensor, # [max_num_reqs]

180 pos: torch.Tensor, # [num_tokens]

181 apply_temperature: bool,

170 - processed_logits_out: torch.Tensor | None = None, # [num_reqs, vocab_size]

182 + output_processed_logits: torch.Tensor | None = None,

183 + output_processed_logits_col: torch.Tensor | None = None,

184 ) -> torch.Tensor:

185 num_tokens, vocab_size = logits.shape

186 BLOCK_SIZE = 1024

@@ -179,8 +192,9 @@ def gumbel_sample(

192 local_argmax.stride(0),

193 local_max,

194 local_max.stride(0),

182 - processed_logits_out,

183 - processed_logits_out.stride(0) if processed_logits_out is not None else 0,

195 + output_processed_logits,

196 + output_processed_logits.stride(0) if output_processed_logits is not None else 0,

197 + output_processed_logits_col,

198 logits,

199 logits.stride(0),

200 expanded_idx_mapping,

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

+176 -115

@@ -89,9 +89,13 @@ class EagleSpeculator:

89 dtype=torch.int64,

90 device=device,

91 )

92 + self.current_draft_step = torch.tensor(0, dtype=torch.int64, device=device)

93 self.last_token_indices = torch.zeros(

94 self.max_num_reqs, dtype=torch.int64, device=device

95 )

96 + self.arange = torch.arange(

97 + self.max_num_reqs + 1, dtype=torch.int32, device="cpu"

98 + )

100 self.supports_mm_inputs = MULTIMODAL_REGISTRY.supports_multimodal_inputs(

101 self.draft_model_config

@@ -228,9 +232,10 @@ class EagleSpeculator:

232 logits: torch.Tensor,

233 idx_mapping: torch.Tensor,

234 pos: torch.Tensor,

231 - step: int,

235 + draft_step: torch.Tensor,

236 + draft_logits: torch.Tensor | None,

237 ) -> torch.Tensor:

233 - if self.draft_logits is not None:

238 + if draft_logits is not None:

239 # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise

240 # used for draft and target sampling.

241 return gumbel_sample(

@@ -240,7 +245,8 @@ class EagleSpeculator:

245 self.seeds,

246 pos + 1,

247 apply_temperature=True,

243 - processed_logits_out=self.draft_logits[:, step],

248 + output_processed_logits=draft_logits,

249 + output_processed_logits_col=draft_step,

250 )

251 else:

252 return logits.argmax(dim=-1)

@@ -274,11 +280,63 @@ class EagleSpeculator:

280 logits,

281 idx_mapping,

282 pos,

277 - step=0,

283 + self.current_draft_step,

284 + self.draft_logits,

285 )

286 self.hidden_states[:num_reqs] = hidden_states[last_token_indices]

287 self.input_buffers.positions[:num_reqs] = pos

288

289 + def multi_step_decode(

290 + self,

291 + num_reqs: int,

292 + skip_attn: bool,

293 + batch_desc: BatchExecutionDescriptor,

294 + num_tokens_across_dp: torch.Tensor | None,

295 + ) -> None:

296 + positions = self.input_buffers.positions[:num_reqs]

297 + query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]

298 + idx_mapping = self.idx_mapping[:num_reqs]

299 +

300 + for step in range(1, self.num_speculative_steps):

301 + attn_metadata = None

302 + slot_mappings_by_layer = None

303 + if not skip_attn:

304 + # Build attention metadata and slot mappings for each draft

305 + # decode step. It is necessary to rebuild the attention

306 + # metadata even when replaying the FULL graph so that any

307 + # attention metadata builder state is updated.

308 + slot_mappings = self.block_tables.compute_slot_mappings(

309 + idx_mapping,

310 + query_start_loc,

311 + positions,

312 + batch_desc.num_tokens,

313 + )

314 + slot_mappings_by_layer = build_slot_mappings_by_layer(

315 + slot_mappings, self.kv_cache_config

316 + )

317 + attn_metadata = self._build_draft_attn_metadata(

318 + num_reqs=num_reqs,

319 + num_reqs_padded=batch_desc.num_reqs or num_reqs,

320 + num_tokens_padded=batch_desc.num_tokens,

321 + )

322 +

323 + # Update the current draft step.

324 + self.current_draft_step.fill_(step)

325 +

326 + # Generate draft tokens for the current step.

327 + if batch_desc.cg_mode == CUDAGraphMode.FULL:

328 + assert self.decode_cudagraph_manager is not None

329 + self.decode_cudagraph_manager.run_fullgraph(batch_desc)

330 + else:

331 + self.generate_draft(

332 + num_reqs,

333 + batch_desc.num_tokens,

334 + attn_metadata,

335 + slot_mappings_by_layer,

336 + num_tokens_across_dp=num_tokens_across_dp,

337 + cudagraph_runtime_mode=batch_desc.cg_mode,

338 + )

339 +

340 def generate_draft(

341 self,

342 num_reqs: int,

@@ -288,59 +346,52 @@ class EagleSpeculator:

346 num_tokens_across_dp: torch.Tensor | None,

347 cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,

348 ) -> None:

291 - pos = self.input_buffers.positions[:num_reqs]

292 - query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]

349 idx_mapping = self.idx_mapping[:num_reqs]

294 - for step in range(1, self.num_speculative_steps):

295 - # Run the eagle model.

296 - last_hidden_states, hidden_states = self.run_model(

297 - num_tokens_padded,

298 - attn_metadata,

299 - slot_mappings,

300 - num_tokens_across_dp,

301 - cudagraph_runtime_mode,

302 - )

303 - last_hidden_states = last_hidden_states[:num_reqs]

304 - hidden_states = hidden_states[:num_reqs]

305 - logits = self.model.compute_logits(last_hidden_states)

350 + positions = self.input_buffers.positions[:num_reqs]

351 + # Run the eagle model forward pass.

352 + last_hidden_states, hidden_states = self.run_model(

353 + num_tokens_padded,

354 + attn_metadata,

355 + slot_mappings,

356 + num_tokens_across_dp,

357 + cudagraph_runtime_mode,

358 + )

359 + last_hidden_states = last_hidden_states[:num_reqs]

360

307 - draft_tokens = self._sample_draft(

308 - logits,

309 - idx_mapping,

310 - pos,

311 - step=step,

312 - )

313 - self.draft_tokens[:num_reqs, step] = draft_tokens

314 -

315 - if step < self.num_speculative_steps - 1:

316 - # Update the inputs for the next step.

317 - update_eagle_inputs(

318 - draft_tokens,

319 - hidden_states,

320 - self.input_buffers,

321 - self.hidden_states,

322 - self.max_model_len,

323 - )

324 - if attn_metadata is not None:

325 - self.block_tables.compute_slot_mappings(

326 - idx_mapping, query_start_loc, pos, num_tokens_padded

327 - )

361 + # Sample the draft tokens.

362 + logits = self.model.compute_logits(last_hidden_states)

363 + draft_tokens = self._sample_draft(

364 + logits,

365 + idx_mapping,

366 + positions,

367 + self.current_draft_step,

368 + self.draft_logits,

369 + )

370 +

371 + # Update the inputs for the next step.

372 + update_eagle_draft_inputs(

373 + draft_tokens,

374 + self.current_draft_step,

375 + hidden_states,

376 + self.draft_tokens,

377 + self.hidden_states,

378 + self.input_buffers,

379 + num_reqs,

380 + self.max_model_len,

381 + self.num_speculative_steps,

382 + )

383

384 def _build_draft_attn_metadata(

385 self,

386 num_reqs: int,

387 num_reqs_padded: int,

388 num_tokens_padded: int,

334 - max_query_len: int,

389 ) -> dict[str, Any] | None:

390 if not self.draft_attn_layer_names:

391 return None

392

339 - query_start_loc_cpu = (

340 - torch.arange(num_reqs_padded + 1, dtype=torch.int32, device="cpu").clamp_(

341 - max=num_reqs

342 - )

343 - * max_query_len

393 + query_start_loc_cpu = torch.clamp(

394 + self.arange[: num_reqs_padded + 1], max=num_reqs

395 )

396 block_tables = [

397 x[:num_reqs_padded] for x in self.block_tables.input_block_tables

@@ -354,7 +405,7 @@ class EagleSpeculator:

405 : num_reqs_padded + 1

406 ],

407 query_start_loc_cpu=query_start_loc_cpu,

357 - max_query_len=max_query_len,

408 + max_query_len=1,

409 seq_lens=self.input_buffers.seq_lens[:num_reqs_padded],

410 max_seq_len=self.max_model_len,

411 block_tables=block_tables,

@@ -373,7 +424,7 @@ class EagleSpeculator:

424 self.last_token_indices.zero_()

425

426 # Capture the prefill routine (model forward + compute_logits +

376 - # gumbel_sample).

427 + # sample).

428 # For FULL graphs, the entire routine is recorded as one graph.

429 # For PIECEWISE, only the model's compiled regions are captured

430 # and the rest (compute_logits, gumbel_sample) runs eagerly.

@@ -387,10 +438,9 @@ class EagleSpeculator:

438 if self.num_speculative_steps == 1:

439 return

440

390 - # Capture the decode draft generation loop (model forward +

391 - # compute_logits + gumbel_sample + update_eagle_inputs, for

392 - # each step). For FULL graphs, the entire multi-step loop is

393 - # recorded as one graph.

441 + # Capture the decode draft generation routine (model forward +

442 + # compute_logits + sample + update_eagle_inputs) for a single

443 + # step.

444 assert self.decode_cudagraph_manager is not None

445 self.decode_cudagraph_manager.capture(

446 self.generate_draft,

@@ -461,9 +511,10 @@ class EagleSpeculator:

511

512 # Get the input ids and last token indices for the speculator.

513 prepare_eagle_inputs(

514 + self.last_token_indices,

515 + self.current_draft_step,

516 self.input_buffers,

517 input_batch,

466 - self.last_token_indices,

518 num_sampled,

519 num_rejected,

520 last_sampled,

@@ -473,12 +524,18 @@ class EagleSpeculator:

524

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

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

476 - # Mixed or prefill-only batches fall back to PIECEWISE.

527 + uniform_token_count = get_uniform_token_count(

528 + num_reqs,

529 + # Use the actual number of tokens without padding added by

530 + # the target model during FULL cudagraph.

531 + input_batch.num_tokens,

532 + max_query_len,

533 + )

534 prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp(

535 self.prefill_cudagraph_manager,

536 num_reqs,

537 num_tokens,

481 - get_uniform_token_count(num_reqs, num_tokens, max_query_len),

538 + uniform_token_count,

539 dp_size=self.dp_size,

540 dp_rank=self.dp_rank,

541 need_eager=is_profile,

@@ -528,48 +585,21 @@ class EagleSpeculator:

585 need_eager=is_profile,

586 )

587

531 - attn_metadata_updated = None

532 - slot_mappings_updated = None

533 - if not (dummy_run and skip_attn_for_dummy_run):

534 - # Build attention metadata and slot mappings for the draft

535 - # decode steps. It is necessary to rebuild the attention

536 - # metadata even when replaying the FULL graph so that any

537 - # attention metadata builder state is updated.

538 - slot_mappings = self.block_tables.compute_slot_mappings(

539 - self.idx_mapping[:num_reqs],

540 - self.input_buffers.query_start_loc[: num_reqs + 1],

541 - self.input_buffers.positions[:num_reqs],

542 - decode_batch_desc.num_tokens,

543 - )

544 - slot_mappings_updated = build_slot_mappings_by_layer(

545 - slot_mappings, self.kv_cache_config

546 - )

547 - attn_metadata_updated = self._build_draft_attn_metadata(

548 - num_reqs=num_reqs,

549 - num_reqs_padded=decode_batch_desc.num_reqs or num_reqs,

550 - num_tokens_padded=decode_batch_desc.num_tokens,

551 - max_query_len=1,

552 - )

588 + # Generate the remaining num_speculative_steps - 1 draft tokens.

589 + self.multi_step_decode(

590 + num_reqs,

591 + dummy_run and skip_attn_for_dummy_run,

592 + decode_batch_desc,

593 + num_tokens_across_dp,

594 + )

595

554 - if decode_batch_desc.cg_mode == CUDAGraphMode.FULL:

555 - # Replay the full graph for draft generation.

556 - assert self.decode_cudagraph_manager is not None

557 - self.decode_cudagraph_manager.run_fullgraph(decode_batch_desc)

558 - else:

559 - self.generate_draft(

560 - num_reqs,

561 - decode_batch_desc.num_tokens,

562 - attn_metadata_updated,

563 - slot_mappings_updated,

564 - num_tokens_across_dp=num_tokens_across_dp,

565 - cudagraph_runtime_mode=decode_batch_desc.cg_mode,

566 - )

596 return self.draft_tokens[:num_reqs]

597

598

599 @triton.jit

600 def _prepare_eagle_inputs_kernel(

601 last_token_indices_ptr,

602 + eagle_current_draft_step_ptr,

603 eagle_input_ids_ptr,

604 eagle_positions_ptr,

605 eagle_query_start_loc_ptr,

@@ -630,6 +660,8 @@ def _prepare_eagle_inputs_kernel(

660 # Copy sequence lengths.

661 tl.store(eagle_seq_lens_ptr + req_idx, seq_len)

662 if req_idx == (num_reqs - 1):

663 + # Reset the current draft step to 0.

664 + tl.store(eagle_current_draft_step_ptr, 0)

665 # Pad query_start_loc for CUDA graphs.

666 for i in range(num_reqs, max_num_reqs + 1, BLOCK_SIZE):

667 block = i + tl.arange(0, BLOCK_SIZE)

@@ -648,10 +680,11 @@ def _prepare_eagle_inputs_kernel(

680

681

682 def prepare_eagle_inputs(

651 - input_buffers: InputBuffers,

652 - input_batch: InputBatch,

683 # [num_reqs]

684 last_token_indices: torch.Tensor,

685 + current_draft_step: torch.Tensor,

686 + input_buffers: InputBuffers,

687 + input_batch: InputBatch,

688 # [num_reqs]

689 num_sampled: torch.Tensor,

690 # [num_reqs]

@@ -665,6 +698,7 @@ def prepare_eagle_inputs(

698 num_reqs = input_batch.num_reqs

699 _prepare_eagle_inputs_kernel[(num_reqs,)](

700 last_token_indices,

701 + current_draft_step,

702 input_buffers.input_ids,

703 input_buffers.positions,

704 input_buffers.query_start_loc,

@@ -685,7 +719,7 @@ def prepare_eagle_inputs(

719

720

721 @triton.jit

688 -def _prepare_eagle_docode_kernel(

722 +def _prepare_eagle_decode_kernel(

723 draft_tokens_ptr,

724 draft_tokens_stride,

725 target_seq_lens_ptr,

@@ -742,7 +776,7 @@ def prepare_eagle_decode(

776 max_num_reqs: int,

777 ):

778 num_reqs = draft_tokens.shape[0]

745 - _prepare_eagle_docode_kernel[(num_reqs + 1,)](

779 + _prepare_eagle_decode_kernel[(num_reqs + 1,)](

780 draft_tokens,

781 draft_tokens.stride(0),

782 target_seq_lens,

@@ -758,36 +792,55 @@ def prepare_eagle_decode(

792

793

794 @triton.jit

761 -def _update_eagle_inputs_kernel(

795 +def _update_eagle_draft_inputs_kernel(

796 + output_draft_tokens_ptr,

797 + output_draft_tokens_stride,

798 + next_input_hidden_states_ptr,

799 + next_input_hidden_states_stride,

800 input_ids_ptr,

801 positions_ptr,

764 - input_hidden_states_ptr,

765 - input_hidden_states_stride,

802 seq_lens_ptr,

767 - max_model_len,

803 draft_tokens_ptr,

769 - output_hidden_states_ptr,

770 - output_hidden_states_stride,

804 + current_draft_step_ptr,

805 + hidden_states_ptr,

806 + hidden_states_stride,

807 hidden_size,

808 + max_model_len,

809 + num_speculative_steps,

810 BLOCK_SIZE: tl.constexpr,

811 ):

812 req_idx = tl.program_id(0)

813

776 - # Draft token -> Input ID.

814 + # Write the sampled draft token into self.draft_tokens[req_idx, step].

815 draft_token = tl.load(draft_tokens_ptr + req_idx)

816 + step = tl.load(current_draft_step_ptr)

817 + tl.store(

818 + output_draft_tokens_ptr + req_idx * output_draft_tokens_stride + step,

819 + draft_token,

820 + )

821 +

822 + if step >= num_speculative_steps - 1:

823 + # This is the final step. Skip updating draft forward inputs.

824 + return

825 +

826 + # Write the sampled draft token into the input ids tensor for the next

827 + # forward pass.

828 tl.store(input_ids_ptr + req_idx, draft_token)

829

780 - # Output hidden states -> Input hidden states.

830 + # Copy hidden states into the input hidden states tensor for the next

831 + # forward pass.

832 for i in range(0, hidden_size, BLOCK_SIZE):

833 block = i + tl.arange(0, BLOCK_SIZE)

834 mask = block < hidden_size

784 - output_hidden_states = tl.load(

785 - output_hidden_states_ptr + req_idx * output_hidden_states_stride + block,

835 + hidden_states = tl.load(

836 + hidden_states_ptr + req_idx * hidden_states_stride + block,

837 mask=mask,

838 )

839 tl.store(

789 - input_hidden_states_ptr + req_idx * input_hidden_states_stride + block,

790 - output_hidden_states,

840 + next_input_hidden_states_ptr

841 + + req_idx * next_input_hidden_states_stride

842 + + block,

843 + hidden_states,

844 mask=mask,

845 )

846

@@ -803,24 +856,32 @@ def _update_eagle_inputs_kernel(

856 tl.store(seq_lens_ptr + req_idx, seq_len)

857

858

806 -def update_eagle_inputs(

859 +def update_eagle_draft_inputs(

860 draft_tokens: torch.Tensor,

808 - output_hidden_states: torch.Tensor,

809 - input_buffers: InputBuffers,

861 + current_draft_step: torch.Tensor,

862 hidden_states: torch.Tensor,

863 + output_draft_tokens: torch.Tensor,

864 + next_input_hidden_states: torch.Tensor,

865 + input_buffers: InputBuffers,

866 + num_reqs: int,

867 max_model_len: int,

868 + num_speculative_steps: int,

869 ):

813 - num_reqs, hidden_size = output_hidden_states.shape

814 - _update_eagle_inputs_kernel[(num_reqs,)](

870 + _, hidden_size = hidden_states.shape

871 + _update_eagle_draft_inputs_kernel[(num_reqs,)](

872 + output_draft_tokens,

873 + output_draft_tokens.stride(0),

874 + next_input_hidden_states,

875 + next_input_hidden_states.stride(0),

876 input_buffers.input_ids,

877 input_buffers.positions,

817 - hidden_states,

818 - hidden_states.stride(0),

878 input_buffers.seq_lens,

820 - max_model_len,

879 draft_tokens,

822 - output_hidden_states,

823 - output_hidden_states.stride(0),

880 + current_draft_step,

881 + hidden_states,

882 + hidden_states.stride(0),

883 hidden_size,

884 + max_model_len,

885 + num_speculative_steps,

886 BLOCK_SIZE=1024,

887 )

vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py

+4 -2

@@ -392,8 +392,10 @@ def _resample_kernel(

392 temp_ptr,

393 seed_ptr,

394 pos_ptr,

395 - None,

396 - 0,

395 + None, # processed_logits_ptr

396 + 0, # processed_logits_stride

397 + None, # processed_logits_col_ptr

398 + vocab_size,

399 APPLY_TEMPERATURE=False,

400 )

401 token_id = block_idx * BLOCK_SIZE + idx