class DraftModelSpeculator(BaseSpeculator):
def __init__(self, vllm_config: VllmConfig, device: torch.device):
self.vllm_config = vllm_config
self.device = device
assert vllm_config.speculative_config is not None
self.speculative_config = vllm_config.speculative_config
self.method = self.speculative_config.method
self.num_speculative_steps = self.speculative_config.num_speculative_tokens
self.draft_model_config = self.speculative_config.draft_model_config
self.scheduler_config = vllm_config.scheduler_config
self.max_num_reqs = self.scheduler_config.max_num_seqs
self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
self.max_model_len = vllm_config.model_config.max_model_len
self.draft_max_seq_len = self.max_model_len
# We need to get the hidden size from the draft model config because
# the draft model's hidden size can be different from the target model's
# hidden size (e.g., Llama 3.3 70B).
self.hidden_size = self.draft_model_config.get_hidden_size()
# Widen for HC-multiplexed residuals (e.g. DeepSeek V4 feeds the MTP
# draft the target's pre-hc_head (T, hc_mult * hidden_size) residual).
# Non-HC models default to hc_mult=1 and are unaffected.
hc_mult = getattr(self.draft_model_config.hf_config, "hc_mult", 1)
self.hidden_size = self.hidden_size * hc_mult
self.vocab_size = self.draft_model_config.get_vocab_size()
self.dtype = vllm_config.model_config.dtype
self.use_fp64_gumbel = vllm_config.model_config.use_fp64_gumbel
self.use_local_argmax_reduction = (
self.speculative_config.use_local_argmax_reduction
)
# DP configuration
self.dp_size = vllm_config.parallel_config.data_parallel_size
self.dp_rank = vllm_config.parallel_config.data_parallel_rank
self.eplb_state: EplbState | None = None
self.input_buffers = InputBuffers(
max_num_reqs=self.max_num_reqs,
max_num_tokens=self.max_num_tokens,
device=device,
)
self.idx_mapping = torch.zeros(
self.max_num_reqs, dtype=torch.int32, device=device
)
self.temperature = torch.zeros(
self.max_num_reqs, dtype=torch.float32, device=device
)
self.seeds = torch.zeros(self.max_num_reqs, dtype=torch.int64, device=device)
self.draft_tokens = torch.zeros(
self.max_num_reqs,
self.num_speculative_steps,
dtype=torch.int64,
device=device,
)
self.arange = torch.arange(
self.max_num_reqs + 1, dtype=torch.int32, device="cpu"
)
self.draft_logits: torch.Tensor | None = None
if self.speculative_config.draft_sample_method == "probabilistic":
# Pre-temperature logits, cached from the previous decode step.
self.draft_logits = torch.zeros(
self.max_num_reqs,
self.num_speculative_steps,
self.vocab_size,
dtype=vllm_config.model_config.head_dtype,
device=device,
)
self.supports_mm_inputs = False
@abstractmethod
def load_draft_model(
self,
target_model: nn.Module,
target_attn_layer_names: set[str],
) -> nn.Module:
pass
def load_model(self, target_model: nn.Module) -> None:
target_attn_layer_names = set(
get_layers_from_vllm_config(
self.vllm_config,
AttentionLayerBase, # type: ignore[type-abstract]
).keys()
)
self.model = self.load_draft_model(target_model, target_attn_layer_names)
self._validate_local_argmax_reduction()
all_attn_layers = set[str](
get_layers_from_vllm_config(
self.vllm_config,
AttentionLayerBase, # type: ignore[type-abstract]
).keys()
)
self.draft_attn_layer_names = all_attn_layers - target_attn_layer_names
target_supports_mm = MULTIMODAL_REGISTRY.supports_multimodal_inputs(
self.vllm_config.model_config
)
draft_supports_mm = supports_multimodal_embeddings(self.model)
self.supports_mm_inputs = target_supports_mm and draft_supports_mm
if target_supports_mm and not draft_supports_mm:
logger.warning_once(
"Draft model %s does not support external multimodal embeddings. "
"Embeddings from the target model will not be passed to the "
"drafter; using text-only draft inputs instead.",
type(self.model).__name__,
)
def set_eplb_state(self, eplb_state: EplbState) -> None:
"""Inject EPLB state after construction."""
self.eplb_state = eplb_state
def _prepare_eplb_forward(self, num_unpadded_tokens: int) -> None:
"""Call EPLB prepare_forward if EPLB is active for the draft model."""
if self.eplb_state is not None:
self.eplb_state.prepare_forward(
self.speculative_config.draft_model_config,
num_unpadded_tokens,
)
@property
def attn_vllm_config(self) -> VllmConfig:
"""Config for the draft's attention metadata builders. Overridden by
speculators whose attention mode differs from the target's."""
return self.vllm_config
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:
self.model_state = model_state
self.kv_cache_config = kv_cache_config
self.attn_groups, self.attn_cg_support, _ = init_attn_backend(
kv_cache_config,
self.attn_vllm_config,
self.device,
active_layer_names=self.draft_attn_layer_names,
)
self.block_tables = block_tables
# The target model runner's buffers and attention groups. Draft
# prefill reuses the target model's attention metadata, so its
# cudagraph capture must build dummy metadata through the same
# builders and buffers.
self.target_input_buffers = target_input_buffers
self.target_attn_groups = target_attn_groups
def _build_draft_attn_metadata(
self,
num_reqs: int,
num_reqs_padded: int,
num_tokens_padded: int,
seq_lens_cpu_upper_bound: torch.Tensor,
step: int,
num_query_per_req: int = 1,
causal: bool | Mapping[int, bool] = True,
query_start_loc_np: np.ndarray | None = None,
) -> dict[str, Any] | None:
if query_start_loc_np is not None:
# Non-uniform query layout (e.g. multi-module MTP's mixed
# prefill/decode queries); num_query_per_req is ignored.
query_start_loc_cpu = torch.empty(num_reqs_padded + 1, dtype=torch.int32)
query_start_loc_cpu[: num_reqs + 1] = torch.from_numpy(
query_start_loc_np[: num_reqs + 1]
)
query_start_loc_cpu[num_reqs:] = query_start_loc_cpu[num_reqs]
max_query_len = int(
(query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).max()
)
else:
# Uniform query: query_start_loc[i] = min(i, num_reqs) * num_query_per_req.
# Clamp keeps the series non-decreasing past num_reqs, which some
# attention backends require.
query_start_loc_cpu = (
torch.clamp(self.arange[: num_reqs_padded + 1], max=num_reqs)
* num_query_per_req
)
max_query_len = num_query_per_req
block_tables = [
x[:num_reqs_padded] for x in self.block_tables.input_block_tables
]
slot_mappings = self.block_tables.slot_mappings[:, :num_tokens_padded]
draft_seq_lens_cpu_upper_bound = torch.zeros(
num_reqs_padded, dtype=torch.int32, device="cpu"
)
torch.add(
seq_lens_cpu_upper_bound[:num_reqs],
step,
out=draft_seq_lens_cpu_upper_bound[:num_reqs],
)
draft_seq_lens_cpu_upper_bound[:num_reqs].clamp_(max=self.max_model_len)
attn_metadata = build_attn_metadata(
attn_groups=self.attn_groups,
num_reqs=num_reqs_padded,
num_tokens=num_tokens_padded,
query_start_loc_gpu=self.input_buffers.query_start_loc[
: num_reqs_padded + 1
],
query_start_loc_cpu=query_start_loc_cpu,
max_query_len=max_query_len,
seq_lens=self.input_buffers.seq_lens[:num_reqs_padded],
max_seq_len=self.draft_max_seq_len,
block_tables=block_tables,
slot_mappings=slot_mappings,
kv_cache_config=self.kv_cache_config,
causal=causal,
seq_lens_cpu_upper_bound=draft_seq_lens_cpu_upper_bound,
)
return attn_metadata
def _validate_local_argmax_reduction(self) -> None:
if not self.use_local_argmax_reduction:
return
if self.speculative_config.draft_sample_method == "probabilistic":
raise ValueError(
"use_local_argmax_reduction is not compatible with "
"draft_sample_method='probabilistic'."
)
if not hasattr(self.model, "get_top_tokens"):
raise ValueError(
"use_local_argmax_reduction is enabled but draft model "
f"{self.model.__class__.__name__} does not implement "
"get_top_tokens()."
)
logger.info(
"Using local argmax reduction for draft token generation "
"(communication: O(2*tp_size) vs O(vocab_size))."
)
def _greedy_sample_draft(self, hidden_states: torch.Tensor) -> torch.Tensor:
if self.use_local_argmax_reduction:
return self.model.get_top_tokens(hidden_states)
logits = self.model.compute_logits(hidden_states)
return logits.argmax(dim=-1)
def sample_draft(
self,
hidden_states: torch.Tensor,
positions: torch.Tensor,
idx_mapping: torch.Tensor,
temperature: torch.Tensor,
seeds: torch.Tensor,
draft_step: torch.Tensor,
draft_logits: torch.Tensor | None,
) -> torch.Tensor:
if draft_logits is not None:
logits = self.model.compute_logits(hidden_states)
# NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
# used for draft and target sampling.
return gumbel_sample(
logits,
idx_mapping,
temperature,
seeds,
positions + 1,
apply_temperature=True,
logits_cache=draft_logits,
logits_cache_col=draft_step,
use_fp64=self.use_fp64_gumbel,
)
return self._greedy_sample_draft(hidden_states)
def _copy_request_inputs(
self,
num_reqs: int,
# [num_reqs]
idx_mapping: torch.Tensor,
# [max_num_reqs]
temperature: torch.Tensor,
# [max_num_reqs]
seeds: torch.Tensor,
) -> None:
# Copy temperature, seeds, and idx mapping to the pre-allocated buffers.
# NOTE(woosuk): For draft sampling, we only consider the temperature
# and ignore the other sampling parameters such as top_k and top_p,
# for simplicity and performance.
# While this may slightly degrade the acceptance rate, it does not
# affect the output distribution after rejection sampling.
self.temperature.copy_(temperature)
self.seeds.copy_(seeds)
self.idx_mapping[:num_reqs].copy_(idx_mapping)
# idx_mapping for CG padded requests points to -1, which is ignored
# during sampling to prevent writing stale values to draft logits.
self.idx_mapping[num_reqs:].fill_(-1)