From ef42093a36b7d7cc74ecfc0533099e01f2aa8286 Mon Sep 17 00:00:00 2001 From: Xiaoshuang Wang <1790571317@qq.com> Date: Tue, 29 Sep 2026 03:38:05 +0800 Subject: [PATCH] [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 | 2 - vllm/v1/worker/gpu/model_runner.py | 13 +- vllm/v1/worker/gpu/spec_decode/__init__.py | 6 + .../gpu/spec_decode/draft_model/__init__.py | 2 + .../gpu/spec_decode/draft_model/speculator.py | 527 ++++++++++++++++++ vllm/v1/worker/gpu/spec_decode/speculator.py | 6 +- 6 files changed, 552 insertions(+), 4 deletions(-) create mode 100644 vllm/v1/worker/gpu/spec_decode/draft_model/__init__.py create mode 100644 vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 220da0c72b..2af0018256 100644 --- a/vllm/config/vllm.py +++ b/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", diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 3d3dc9b162..fb95fc793a 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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) diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index ede733a755..0616f12dd5 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/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.") diff --git a/vllm/v1/worker/gpu/spec_decode/draft_model/__init__.py b/vllm/v1/worker/gpu/spec_decode/draft_model/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/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 diff --git a/vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py b/vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py new file mode 100644 index 0000000000..05bc997f7a --- /dev/null +++ b/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 diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 6f88febb66..c30b653196 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/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" )