[Model Runner V2][Spec Decode] Support spec decode with draft model (#43091)

Signed-off-by: Icey <1790571317@qq.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

vllm/config/vllm.py

@@ -3067,8 +3067,6 @@ class VllmConfig:
                 # https://github.com/vllm-project/vllm/pull/40704
                 "ngram",
                 "ngram_gpu",
-                # https://github.com/vllm-project/vllm/pull/43091
-                "draft_model",
                 "suffix",
                 "medusa",
                 "mlp_speculator",

vllm/v1/worker/gpu/model_runner.py

@@ -799,7 +799,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 "skip_attn must only be True for initial memory profiling."
             )
 
-        # Create a dummy scheduler output.
+        # Create a dummy scheduler output. Plain draft-model speculation adds
+        # one correction slot per request during prefill. The scheduler
+        # accounts for these slots, while dummy runs bypass the scheduler.
+        # Adjust the dummy token count only when the expanded draft batch
+        # would exceed the token budget.
         num_reqs = min(num_tokens, self.max_num_reqs)
         if uniform_decode:
             # HACK(lucas): for now since the worker is shared between MRV1 and MRV2,
@@ -813,8 +817,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             num_reqs = min(
                 num_reqs, self.max_num_tokens // self.speculator.num_query_per_req
             )
+            num_tokens = min(
+                num_tokens,
+                self.max_num_tokens
+                - num_reqs * self.speculator.num_extra_query_per_req,
+            )
+
         # Distribute the remainder evenly so no dummy request exceeds
         # ceil(num_tokens / num_reqs) <= max_model_len tokens.
+
         num_tokens_per_request = [
             num_tokens // num_reqs + (i >= num_reqs - num_tokens % num_reqs)
             for i in range(num_reqs)

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

@@ -60,5 +60,11 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device):
         )
 
         return EagleSpeculator(vllm_config, device)
+    elif speculative_config.uses_draft_model():
+        from vllm.v1.worker.gpu.spec_decode.draft_model.speculator import (
+            PlainDraftModelSpeculator,
+        )
+
+        return PlainDraftModelSpeculator(vllm_config, device)
     else:
         raise NotImplementedError(f"{speculative_config.method} is not supported yet.")

vllm/v1/worker/gpu/spec_decode/draft_model/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

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

@@ -0,0 +1,527 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from vllm.compilation.backends import set_model_tag
+from vllm.config import VllmConfig, replace
+from vllm.config.compilation import CUDAGraphMode
+from vllm.forward_context import BatchDescriptor, set_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.model_loader import get_model
+from vllm.triton_utils import tl, triton
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.spec_decode.utils import next_power_of_2
+from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
+from vllm.v1.worker.gpu.block_table import BlockTables
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
+from vllm.v1.worker.gpu.dp_utils import DPSyncState
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.gpu.model_states.interface import ModelState
+from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
+from vllm.v1.worker.utils import AttentionGroup
+
+logger = init_logger(__name__)
+
+
+class PlainDraftModelSpeculator(DraftModelSpeculator):
+    """Speculative decoding using a separate smaller draft LM.
+
+    Unlike Eagle, the draft model runs fully independently of the target model.
+    Step 0 builds an expanded buffer (accepted + correction token + rejected
+    slots) via a Triton kernel; steps 1..k-1 are single-token decode steps.
+    """
+
+    # Plain draft model needs one extra slot per request for correction.
+    num_extra_query_per_req = 1
+
+    def __init__(self, vllm_config: VllmConfig, device: torch.device):
+        super().__init__(vllm_config, device)
+
+        # draft_max_seq_len is read by the parent's attention metadata builder.
+        # Plain draft model doesn't do per-batch adjustment; cap at max.
+        self.draft_max_seq_len = self.max_model_len
+
+        self.last_token_indices = torch.zeros(
+            self.max_num_reqs, dtype=torch.int64, device=device
+        )
+        # GPU-side arange for decode query_start_loc initialisation.
+        self.arange_gpu = torch.arange(
+            self.max_num_reqs + 1, dtype=torch.int32, device=device
+        )
+        # Scalar step counter used as column index into draft_logits.
+        self.current_draft_step = torch.tensor(0, dtype=torch.int64, device=device)
+
+        _expanded_max = self.max_num_tokens + self.max_num_reqs
+        self.expanded_input_ids = torch.zeros(
+            _expanded_max, dtype=torch.int32, device=device
+        )
+        self.expanded_positions = torch.zeros(
+            _expanded_max, dtype=torch.int64, device=device
+        )
+        self.expanded_slot_mappings: torch.Tensor
+        self.supports_mm_inputs = False
+
+    def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
+        pass  # CUDA graph not yet supported for plain draft model speculator.
+
+    def capture(self) -> None:
+        pass  # CUDA graph not yet supported for plain draft model speculator.
+
+    def set_attn(
+        self,
+        model_state: ModelState,
+        kv_cache_config: KVCacheConfig,
+        block_tables: BlockTables,
+        target_input_buffers: InputBuffers,
+        target_attn_groups: list[list[AttentionGroup]],
+    ) -> None:
+        super().set_attn(
+            model_state,
+            kv_cache_config,
+            block_tables,
+            target_input_buffers,
+            target_attn_groups,
+        )
+        self.expanded_slot_mappings = torch.empty(
+            block_tables.num_kv_cache_groups,
+            self.max_num_tokens + self.max_num_reqs,
+            dtype=torch.int64,
+            device=self.device,
+        )
+
+    def load_draft_model(
+        self,
+        target_model: nn.Module,
+        target_attn_layer_names: set[str],
+    ) -> nn.Module:
+        spec = self.speculative_config
+        assert spec is not None
+        draft_vllm_config = replace(
+            self.vllm_config,
+            model_config=self.draft_model_config,
+            quant_config=None,
+            parallel_config=replace(
+                spec.draft_parallel_config,
+                rank=self.vllm_config.parallel_config.rank,
+            ),
+        )
+        with set_model_tag("draft_model"):
+            return get_model(
+                vllm_config=draft_vllm_config,
+                prefix="draft_model",
+            )
+
+    @torch.inference_mode()
+    def _run_model(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        attn_metadata: dict[str, Any] | None,
+        slot_mappings: dict[str, torch.Tensor] | None,
+        num_tokens_across_dp: torch.Tensor | None,
+    ) -> torch.Tensor:
+        num_tokens = input_ids.shape[0]
+        batch_descriptor = BatchDescriptor(num_tokens=num_tokens)
+        with set_forward_context(
+            attn_metadata,
+            self.vllm_config,
+            num_tokens=num_tokens,
+            cudagraph_runtime_mode=CUDAGraphMode.NONE,
+            num_tokens_across_dp=num_tokens_across_dp,
+            slot_mapping=slot_mappings,
+            batch_descriptor=batch_descriptor,
+        ):
+            hidden_states = self.model(  # type: ignore[misc]
+                input_ids=input_ids,
+                positions=positions,
+            )
+        if isinstance(hidden_states, tuple):
+            hidden_states = hidden_states[0]
+        return hidden_states
+
+    def _accepted_last_indices(
+        self,
+        input_batch: InputBatch,
+        num_rejected: torch.Tensor,
+        num_reqs: int,
+    ) -> torch.Tensor:
+        qsl = input_batch.query_start_loc
+        adjusted_lens = qsl[1 : num_reqs + 1] - qsl[:num_reqs] - num_rejected[:num_reqs]
+        return qsl[:num_reqs] + adjusted_lens - 1
+
+    def _prefill(
+        self,
+        input_batch: InputBatch,
+        num_rejected: torch.Tensor,
+        last_sampled: torch.Tensor,
+        num_reqs: int,
+        skip_attn: bool,
+        num_tokens_across_dp: torch.Tensor | None,
+    ) -> None:
+        if skip_attn:
+            src_tokens = input_batch.num_tokens
+            self.last_token_indices[:num_reqs] = self._accepted_last_indices(
+                input_batch, num_rejected, num_reqs
+            )
+            hidden_states = self._run_model(
+                input_batch.input_ids[:src_tokens],
+                input_batch.positions[:src_tokens],
+                None,
+                None,
+                num_tokens_across_dp,
+            )
+            positions_buf = input_batch.positions
+        else:
+            block_tables = self.block_tables
+            kv_cache_config = self.kv_cache_config
+            assert block_tables is not None
+            assert kv_cache_config is not None
+
+            total_expanded = prepare_prefill_inputs(
+                input_buffers=self.input_buffers,
+                input_batch=input_batch,
+                last_sampled=last_sampled,
+                num_rejected=num_rejected,
+                expanded_input_ids=self.expanded_input_ids,
+                expanded_positions=self.expanded_positions,
+                last_token_indices=self.last_token_indices,
+                max_num_reqs=self.max_num_reqs,
+                max_model_len=self.max_model_len,
+            )
+
+            prefill_slot_mappings = block_tables.compute_slot_mappings(
+                input_batch.idx_mapping,
+                self.input_buffers.query_start_loc[: num_reqs + 1],
+                self.expanded_positions[:total_expanded],
+                total_expanded,
+                out=self.expanded_slot_mappings,
+            )
+            qsl_np = input_batch.query_start_loc_np
+            query_start_loc_cpu_expanded = (
+                torch.from_numpy(qsl_np[: num_reqs + 1]).int()
+                + torch.from_numpy(self.arange_np[: num_reqs + 1]).int()
+            )
+            prefill_attn_md = self._build_attn_metadata(
+                batch_desc=BatchExecutionDescriptor(
+                    cg_mode=CUDAGraphMode.NONE,
+                    num_tokens=total_expanded,
+                    num_reqs=num_reqs,
+                ),
+                num_reqs=num_reqs,
+                query_start_loc_np=query_start_loc_cpu_expanded.numpy(),
+                seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound,
+                step=0,
+                slot_mappings=prefill_slot_mappings,
+            )
+            prefill_slot_maps_by_layer = build_slot_mappings_by_layer(
+                prefill_slot_mappings, kv_cache_config
+            )
+            hidden_states = self._run_model(
+                self.expanded_input_ids[:total_expanded],
+                self.expanded_positions[:total_expanded],
+                prefill_attn_md,
+                prefill_slot_maps_by_layer,
+                num_tokens_across_dp,
+            )
+            positions_buf = self.expanded_positions
+
+        last_indices = self.last_token_indices[:num_reqs]
+        self.current_draft_step.fill_(0)
+        self.draft_tokens[:num_reqs, 0] = self.sample_draft(
+            hidden_states[last_indices],
+            positions_buf[last_indices],
+            self.idx_mapping[:num_reqs],
+            self.temperature,
+            self.seeds,
+            self.current_draft_step,
+            self.draft_logits,
+        )
+
+    def _multi_step_decode(
+        self,
+        input_batch: InputBatch,
+        num_rejected: torch.Tensor,
+        num_reqs: int,
+        skip_attn: bool,
+        num_tokens_across_dp: torch.Tensor | None,
+    ) -> None:
+        if skip_attn:
+            last_positions = input_batch.positions[
+                self._accepted_last_indices(input_batch, num_rejected, num_reqs)
+            ]
+        else:
+            last_positions = self.expanded_positions[self.last_token_indices[:num_reqs]]
+
+        self.input_buffers.positions[:num_reqs].copy_(last_positions)
+        # The decode loop increments seq_lens BEFORE each forward.
+        # Initial seq_lens = target_seq_lens - num_rejected + 1.
+        # Step 1 forward uses      target_seq_lens - num_rejected + 2.
+        # Step 2 forward uses      target_seq_lens - num_rejected + 3.
+        self.input_buffers.seq_lens[:num_reqs].copy_(
+            torch.clamp(
+                input_batch.seq_lens[:num_reqs] - num_rejected[:num_reqs].int() + 1,
+                max=self.max_model_len,
+            )
+        )
+        self.input_buffers.query_start_loc[: num_reqs + 1].copy_(
+            self.arange_gpu[: num_reqs + 1]
+        )
+
+        idx_mapping = input_batch.idx_mapping
+
+        for step in range(1, self.num_speculative_steps):
+            self.input_buffers.input_ids[:num_reqs].copy_(
+                self.draft_tokens[:num_reqs, step - 1].int()
+            )
+            torch.clamp(
+                self.input_buffers.positions[:num_reqs] + 1,
+                max=self.max_model_len - 1,
+                out=self.input_buffers.positions[:num_reqs],
+            )
+            # Increment seq_lens BEFORE the forward so the attention reads from
+            # the KV slot written by the previous step.
+            torch.clamp(
+                self.input_buffers.seq_lens[:num_reqs] + 1,
+                max=self.max_model_len,
+                out=self.input_buffers.seq_lens[:num_reqs],
+            )
+
+            decode_attn_md = None
+            step_slot_maps_by_layer = None
+            if not skip_attn:
+                block_tables = self.block_tables
+                kv_cache_config = self.kv_cache_config
+                assert block_tables is not None
+                assert kv_cache_config is not None
+                q_start = self.input_buffers.query_start_loc[: num_reqs + 1]
+                positions = self.input_buffers.positions[:num_reqs]
+                step_slot_mappings = block_tables.compute_slot_mappings(
+                    idx_mapping, q_start, positions, num_reqs
+                )
+                step_slot_maps_by_layer = build_slot_mappings_by_layer(
+                    step_slot_mappings, kv_cache_config
+                )
+                decode_attn_md = self._build_uniform_attn_metadata(
+                    batch_desc=BatchExecutionDescriptor(
+                        cg_mode=CUDAGraphMode.NONE,
+                        num_tokens=num_reqs,
+                        num_reqs=num_reqs,
+                        uniform_token_count=1,
+                    ),
+                    num_reqs=num_reqs,
+                    num_query_per_req=1,
+                    seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound,
+                    # Include the correction token inserted during prefill.
+                    step=step + 1,
+                )
+            hidden_states = self._run_model(
+                self.input_buffers.input_ids[:num_reqs],
+                self.input_buffers.positions[:num_reqs],
+                decode_attn_md,
+                step_slot_maps_by_layer,
+                num_tokens_across_dp,
+            )
+
+            self.current_draft_step.fill_(step)
+            self.draft_tokens[:num_reqs, step] = self.sample_draft(
+                hidden_states[:num_reqs],
+                self.input_buffers.positions[:num_reqs],
+                self.idx_mapping[:num_reqs],
+                self.temperature,
+                self.seeds,
+                self.current_draft_step,
+                self.draft_logits,
+            )
+
+    @torch.inference_mode()
+    def propose(
+        self,
+        input_batch: InputBatch,
+        attn_metadata: dict[str, Any],
+        slot_mappings: dict[str, torch.Tensor],
+        last_hidden_states: torch.Tensor,  # unused
+        aux_hidden_states: list[torch.Tensor] | None,  # unused
+        num_sampled: torch.Tensor,
+        num_rejected: torch.Tensor,
+        last_sampled: torch.Tensor,
+        next_prefill_tokens: torch.Tensor,
+        temperature: torch.Tensor,
+        seeds: torch.Tensor,
+        dp_sync: DPSyncState | None = None,
+        dummy_run: bool = False,
+        skip_attn_for_dummy_run: bool = False,
+        mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
+        is_profile: bool = False,
+    ) -> torch.Tensor:
+        assert self.model is not None
+
+        num_reqs = input_batch.num_reqs
+        skip_attn = dummy_run and skip_attn_for_dummy_run
+        num_tokens_across_dp = (
+            dp_sync.num_tokens_across_dp if dp_sync is not None else None
+        )
+
+        # Copy per-request temperature/seeds/idx_mapping into pre-allocated
+        # buffers so sample_draft can read them without extra slicing.
+        self._copy_request_inputs(
+            num_reqs,
+            input_batch.idx_mapping,
+            temperature,
+            seeds,
+        )
+
+        self._prefill(
+            input_batch,
+            num_rejected,
+            last_sampled,
+            num_reqs,
+            skip_attn,
+            num_tokens_across_dp,
+        )
+
+        if self.num_speculative_steps == 1:
+            return self.draft_tokens[:num_reqs, :1]
+
+        self._multi_step_decode(
+            input_batch,
+            num_rejected,
+            num_reqs,
+            skip_attn,
+            num_tokens_across_dp,
+        )
+
+        return self.draft_tokens[:num_reqs]
+
+
+@triton.jit
+def _prepare_prefill_inputs_kernel(
+    target_input_ids_ptr,  # [src_tokens] int32
+    target_positions_ptr,  # [src_tokens] int64
+    last_sampled_ptr,  # [max_num_reqs] int32
+    idx_mapping_ptr,  # [num_reqs] int32
+    out_input_ids_ptr,  # [src_tokens + num_reqs] int32
+    out_positions_ptr,  # [src_tokens + num_reqs] int64
+    last_token_indices_ptr,  # [max_num_reqs] int64
+    out_query_start_loc_ptr,  # [max_num_reqs + 1] int32
+    out_seq_lens_ptr,  # [max_num_reqs] int32
+    query_start_loc_ptr,  # [num_reqs + 1] int32
+    seq_lens_ptr,  # [num_reqs] int32
+    num_rejected_ptr,  # [num_reqs] int32
+    max_num_reqs,
+    max_model_len,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Per-request step-0 input preparation for the plain draft-model speculator.
+
+    Output layout for request i (out_start = query_start_loc[i] + i):
+        [out_start,              out_start + num_valid)      accepted tokens
+        [out_start + num_valid]                              correction token
+        (out_start + num_valid, out_start + total_out)      rejected slots
+
+    where num_valid  = query_lens[i] - num_rejected[i]
+          total_out  = query_lens[i] + 1
+    """
+    req_idx = tl.program_id(0)
+    num_reqs = tl.num_programs(0)
+
+    req_state_idx = tl.load(idx_mapping_ptr + req_idx)
+
+    q_start = tl.load(query_start_loc_ptr + req_idx)
+    q_next = tl.load(query_start_loc_ptr + req_idx + 1)
+    seq_len = tl.load(seq_lens_ptr + req_idx)
+    num_rejected = tl.load(num_rejected_ptr + req_idx)
+
+    num_valid = q_next - q_start - num_rejected
+    correction_token = tl.load(last_sampled_ptr + req_state_idx).to(tl.int32)
+    start_pos = tl.load(target_positions_ptr + q_start)
+    out_start = q_start + req_idx
+    total_out = q_next - q_start + 1
+
+    for i in range(0, total_out, BLOCK_SIZE):
+        j = i + tl.arange(0, BLOCK_SIZE)
+        in_bounds = j < total_out
+
+        is_valid = j < num_valid
+        is_correction = j == num_valid
+        src_idx = q_start + j
+        token_ids = tl.load(target_input_ids_ptr + src_idx, mask=is_valid, other=0)
+        positions = tl.minimum(start_pos + j, max_model_len - 1)
+        token_ids = tl.where(is_correction, correction_token, token_ids)
+
+        out_idx = out_start + j
+        tl.store(out_input_ids_ptr + out_idx, token_ids, mask=in_bounds)
+        tl.store(out_positions_ptr + out_idx, positions, mask=in_bounds)
+
+    tl.store(last_token_indices_ptr + req_idx, out_start + num_valid)
+    tl.store(out_query_start_loc_ptr + req_idx, out_start)
+    # seqlen_k for the expanded prefill attention:
+    #   pre_existing = seq_len - query_len (draft KV before this round)
+    #   expanded_query_len = query_len + 1 (adds correction token)
+    #   seqlen_k = pre_existing + expanded_query_len = seq_len + 1
+    # num_rejected does NOT change seqlen_k — the expanded query always has
+    # query_len+1 slots. Rejected positions write temporary KV entries that
+    # are overwritten with valid tokens in a later round.
+    new_seq_len = tl.minimum(seq_len + 1, max_model_len)
+    tl.store(out_seq_lens_ptr + req_idx, new_seq_len)
+
+    if req_idx == num_reqs - 1:
+        total_expanded = out_start + total_out
+        tl.store(out_query_start_loc_ptr + num_reqs, total_expanded)
+        for i in range(num_reqs + 1, max_num_reqs + 1, BLOCK_SIZE):
+            block = i + tl.arange(0, BLOCK_SIZE)
+            mask = block <= max_num_reqs
+            tl.store(out_query_start_loc_ptr + block, total_expanded, mask=mask)
+        for i in range(num_reqs, max_num_reqs, BLOCK_SIZE):
+            block = i + tl.arange(0, BLOCK_SIZE)
+            mask = block < max_num_reqs
+            tl.store(out_seq_lens_ptr + block, 0, mask=mask)
+            tl.store(last_token_indices_ptr + block, 0, mask=mask)
+
+
+def prepare_prefill_inputs(
+    input_buffers: InputBuffers,
+    input_batch: InputBatch,
+    last_sampled: torch.Tensor,  # [max_num_reqs] int64
+    num_rejected: torch.Tensor,  # [num_reqs]     int64
+    expanded_input_ids: torch.Tensor,  # [max_tokens + max_reqs] int32
+    expanded_positions: torch.Tensor,  # [max_tokens + max_reqs] int64
+    last_token_indices: torch.Tensor,  # [max_num_reqs] int64
+    max_num_reqs: int,
+    max_model_len: int,
+) -> int:
+    """Call _prepare_prefill_inputs_kernel and return total_expanded tokens.
+
+    Side-effects (kernel writes):
+      - expanded_input_ids, expanded_positions
+      - last_token_indices
+      - input_buffers.query_start_loc  (expanded: original[i] + i)
+      - input_buffers.seq_lens         (= target_seq_lens - num_rejected + 1)
+    """
+    num_reqs = input_batch.num_reqs
+    src_tokens = input_batch.num_tokens
+    qsl_np = input_batch.query_start_loc_np
+    query_lens = qsl_np[1 : num_reqs + 1] - qsl_np[:num_reqs]
+    max_total_out = int(query_lens.max()) + 2  # +1 correction, +1 alignment
+    BLOCK_SIZE = min(512, next_power_of_2(max_total_out))
+
+    _prepare_prefill_inputs_kernel[(num_reqs,)](
+        input_batch.input_ids,
+        input_batch.positions,
+        last_sampled.int(),
+        input_batch.idx_mapping,
+        expanded_input_ids,
+        expanded_positions,
+        last_token_indices,
+        input_buffers.query_start_loc,
+        input_buffers.seq_lens,
+        input_batch.query_start_loc,
+        input_batch.seq_lens,
+        num_rejected.int(),
+        max_num_reqs,
+        max_model_len,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+    return src_tokens + num_reqs

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

@@ -58,6 +58,8 @@ def _target_feeds_hc_residual(vllm_config: VllmConfig) -> bool:
 
 class BaseSpeculator(ABC):
     num_query_per_req: int = 1
+    # Extra query slots reserved per request outside the regular queries.
+    num_extra_query_per_req: int = 0
 
     @abstractmethod
     def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
@@ -302,6 +304,7 @@ class DraftModelSpeculator(BaseSpeculator):
         step: int,
         causal: bool | Mapping[int, bool] = True,
         dcp_local_seq_lens: torch.Tensor | None = None,
+        slot_mappings: torch.Tensor | None = None,
     ) -> dict[str, Any] | None:
         num_reqs_padded = batch_desc.num_reqs or num_reqs
         # A FULL graph replays a captured shape whose padded requests each hold
@@ -323,7 +326,8 @@ class DraftModelSpeculator(BaseSpeculator):
         block_tables = [
             x[:num_reqs_padded] for x in self.block_tables.input_block_tables
         ]
-        slot_mappings = self.block_tables.slot_mappings[:, :num_tokens]
+        if slot_mappings is None:
+            slot_mappings = self.block_tables.slot_mappings[:, :num_tokens]
         draft_seq_lens_cpu_upper_bound = torch.zeros(
             num_reqs_padded, dtype=torch.int32, device="cpu"
         )