class AutoRegressiveSpeculator(DraftModelSpeculator):
"""Shared CUDA graph and decode loop for autoregressive speculators."""
def __init__(self, vllm_config: VllmConfig, device: torch.device):
super().__init__(vllm_config, device)
self.current_draft_step = torch.tensor(0, dtype=torch.int64, device=device)
self.last_token_indices = torch.zeros(
self.max_num_reqs, dtype=torch.int64, device=device
)
self.prefill_cudagraph_manager: SpeculatorCudaGraphManager | None = None
self.decode_cudagraph_manager: SpeculatorCudaGraphManager | None = None
def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
# Subclasses create prefill_cudagraph_manager before calling this.
# PIECEWISE cudagraphs are not supported for draft decodes.
decode_mode = (
CUDAGraphMode.FULL_DECODE_ONLY
if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
else CUDAGraphMode.NONE
)
self.decode_cudagraph_manager = SpeculatorCudaGraphManager(
self.vllm_config,
self.device,
decode_mode,
decode_query_lens={1: self.max_num_reqs},
decode_request_sizes=(
self.vllm_config.get_draft_decode_cudagraph_request_sizes()
),
)
@abstractmethod
def _prefill(
self,
num_reqs: int,
num_tokens: int,
attn_metadata: dict[str, Any] | None,
slot_mappings: dict[str, torch.Tensor] | None,
num_tokens_across_dp: torch.Tensor | None,
cudagraph_runtime_mode: CUDAGraphMode,
mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
) -> None:
"""Run the draft prefill and sample draft token zero."""
@abstractmethod
def _decode(
self,
num_reqs: int,
num_tokens_padded: int,
attn_metadata: dict[str, Any] | None,
slot_mappings: dict[str, torch.Tensor] | None,
num_tokens_across_dp: torch.Tensor | None,
cudagraph_runtime_mode: CUDAGraphMode,
num_speculative_steps: int,
) -> None:
"""Generate the current draft token and prepare the next decode step."""
@abstractmethod
def _prepare_decode_inputs(
self,
input_batch: InputBatch,
num_rejected: torch.Tensor,
num_reqs: int,
skip_attn: bool,
batch_desc: BatchExecutionDescriptor,
) -> None:
"""Prepare initial decode inputs for this speculator."""
def _prefill_capture_inputs(
self,
) -> tuple[InputBuffers, list[list[AttentionGroup]]]:
return self.input_buffers, self.attn_groups
def _prefill_capture_slot_mappings(self) -> torch.Tensor | None:
return None
def _configure_decode_metadata_update(self) -> None:
if self.num_speculative_steps == 1:
self.supports_decode_metadata_update = False
return
super()._configure_decode_metadata_update()
def on_prefill_begin(self, num_reqs: int) -> None: ...
def on_prefill_end(self, num_reqs: int) -> None: ...
def on_multi_step_decode_begin(self, num_reqs: int) -> None: ...
def on_multi_step_decode_end(self, num_reqs: int) -> None: ...
@property
def advance_draft_positions(self) -> bool:
"""Whether each draft step writes a new KV position."""
return True
def capture(self) -> None:
logger.info("Capturing autoregressive speculator...")
self.last_token_indices.zero_()
self.idx_mapping.zero_()
self.current_draft_step.zero_()
assert self.prefill_cudagraph_manager is not None
if self.prefill_cudagraph_manager.use_breakable_cg:
self.prefill_cudagraph_manager.init_breakable_cg_runner(self.model)
input_buffers, attn_groups = self._prefill_capture_inputs()
self.on_prefill_begin(self.max_num_reqs)
self.prefill_cudagraph_manager.capture(
self._prefill,
self.model_state,
input_buffers,
self.block_tables,
attn_groups,
self.kv_cache_config,
progress_bar_desc="Capturing prefill CUDA graphs",
slot_mappings_buffer=self._prefill_capture_slot_mappings(),
)
self.on_prefill_end(self.max_num_reqs)
if self.num_speculative_steps == 1:
return
assert self.decode_cudagraph_manager is not None
self.on_multi_step_decode_begin(self.max_num_reqs)
decode_fn = (
self._fused_multi_step_decode
if self.supports_decode_metadata_update
else self._decode
)
self.decode_cudagraph_manager.capture(
decode_fn,
self.model_state,
self.input_buffers,
self.block_tables,
self.attn_groups,
self.kv_cache_config,
progress_bar_desc="Capturing decode CUDA graphs",
specialize_spec_tokens=self.supports_decode_metadata_update
and self.speculative_config.uses_dynamic_speculative_decoding(),
)
self.on_multi_step_decode_end(self.max_num_reqs)
def _run_multi_step_decode(
self,
input_batch: InputBatch,
num_rejected: torch.Tensor,
num_reqs: int,
skip_attn: bool,
dp_sync_state: DPSyncState | None,
seq_lens_cpu_upper_bound: torch.Tensor,
is_profile: bool,
num_speculative_steps: int,
) -> None:
decode_batch_sync, num_batch_tokens = None, num_reqs
if dp_sync_state is not None:
# Padded dummy batches can exceed the requests the DP sync covers.
num_reqs = min(num_reqs, dp_sync_state.num_reqs)
decode_batch_sync, num_batch_tokens = self._build_uniform_batch_dp_sync(
dp_sync_state, num_reqs
)
decode_batch_desc, decode_batch_sync = dispatch_cg_and_sync_dp(
self.decode_cudagraph_manager,
num_reqs,
num_batch_tokens,
uniform_token_count=1,
dp_size=self.dp_size,
dp_rank=self.dp_rank,
need_eager=is_profile,
dp_sync_state=decode_batch_sync,
parallel_config=self.vllm_config.parallel_config,
)
if self.decode_cudagraph_manager is not None:
decode_batch_desc = self.decode_cudagraph_manager.specialize_spec_tokens(
decode_batch_desc, num_speculative_steps
)
self._prepare_decode_inputs(
input_batch, num_rejected, num_reqs, skip_attn, decode_batch_desc
)
num_tokens_across_dp = (
decode_batch_sync.num_tokens_across_dp
if decode_batch_sync is not None
else None
)
self.on_multi_step_decode_begin(num_reqs)
multi_step_decode_fn = (
self._run_fused_multi_step_decode
if self.supports_decode_metadata_update
else self._multi_step_decode
)
multi_step_decode_fn(
num_reqs,
skip_attn,
decode_batch_desc,
num_tokens_across_dp,
seq_lens_cpu_upper_bound,
num_speculative_steps,
)
self.on_multi_step_decode_end(num_reqs)
def _multi_step_decode(
self,
num_reqs: int,
skip_attn: bool,
batch_desc: BatchExecutionDescriptor,
num_tokens_across_dp: torch.Tensor | None,
seq_lens_cpu_upper_bound: torch.Tensor,
num_speculative_steps: int,
) -> None:
positions = self.input_buffers.positions[:num_reqs]
query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
idx_mapping = self.idx_mapping[:num_reqs]
attn_metadata = None
slot_mappings_by_layer = None
for step in range(1, num_speculative_steps):
if not skip_attn and (self.advance_draft_positions or step == 1):
slot_mappings = self.block_tables.compute_slot_mappings(
idx_mapping,
query_start_loc,
positions,
batch_desc.num_tokens,
)
slot_mappings_by_layer = build_slot_mappings_by_layer(
slot_mappings, self.kv_cache_config
)
attn_metadata = self._build_uniform_attn_metadata(
num_reqs=num_reqs,
batch_desc=batch_desc,
num_query_per_req=1,
seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
step=step,
)
self.current_draft_step.fill_(step)
if batch_desc.cg_mode == CUDAGraphMode.FULL:
assert self.decode_cudagraph_manager is not None
self.decode_cudagraph_manager.run_fullgraph(batch_desc)
else:
self._decode(
num_reqs,
batch_desc.num_tokens,
attn_metadata,
slot_mappings_by_layer,
num_tokens_across_dp,
batch_desc.cg_mode,
num_speculative_steps,
)
def _run_fused_multi_step_decode(
self,
num_reqs: int,
skip_attn: bool,
batch_desc: BatchExecutionDescriptor,
num_tokens_across_dp: torch.Tensor | None,
seq_lens_cpu_upper_bound: torch.Tensor,
num_speculative_steps: int,
) -> None:
if batch_desc.cg_mode == CUDAGraphMode.FULL:
self._map_kernel_block_tables()
assert self.decode_cudagraph_manager is not None
self.decode_cudagraph_manager.run_fullgraph(batch_desc)
return
attn_metadata = None
slot_mappings_by_layer = None
if not skip_attn:
slot_mappings = self.block_tables.compute_slot_mappings(
self.idx_mapping[:num_reqs],
self.input_buffers.query_start_loc[: num_reqs + 1],
self.input_buffers.positions[:num_reqs],
batch_desc.num_tokens,
)
slot_mappings_by_layer = build_slot_mappings_by_layer(
slot_mappings, self.kv_cache_config
)
attn_metadata = self._build_uniform_attn_metadata(
batch_desc=batch_desc,
num_reqs=num_reqs,
num_query_per_req=1,
seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
step=1,
)
self._fused_multi_step_decode(
num_reqs,
batch_desc.num_tokens,
attn_metadata,
slot_mappings_by_layer,
num_tokens_across_dp,
batch_desc.cg_mode,
num_speculative_steps,
)
def _fused_multi_step_decode(
self,
num_reqs: int,
num_tokens_padded: int,
attn_metadata: dict[str, Any] | None,
slot_mappings: dict[str, torch.Tensor] | None,
num_tokens_across_dp: torch.Tensor | None,
cudagraph_runtime_mode: CUDAGraphMode,
num_speculative_steps: int,
) -> None:
idx_mapping = self.idx_mapping[:num_reqs]
positions = self.input_buffers.positions[:num_reqs]
query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
for step in range(1, num_speculative_steps):
if attn_metadata is not None and (
self.advance_draft_positions or step == 1
):
self.block_tables.compute_slot_mappings(
idx_mapping,
query_start_loc,
positions,
num_tokens_padded,
)
self._update_draft_decode_metadata(attn_metadata, num_reqs)
self.current_draft_step.fill_(step)
self._decode(
num_reqs,
num_tokens_padded,
attn_metadata,
slot_mappings,
num_tokens_across_dp,
cudagraph_runtime_mode,
num_speculative_steps,
)