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":
self.draft_logits = torch.zeros(
self.max_num_reqs,
self.num_speculative_steps,
self.vocab_size,
dtype=torch.float32,
device=device,
)
@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
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,
)
def set_attn(
self,
model_state: ModelState,
kv_cache_config: KVCacheConfig,
block_tables: BlockTables,
) -> None:
self.model_state = model_state
self.kv_cache_config = kv_cache_config
self.attn_groups, _, _ = init_attn_backend(
kv_cache_config,
self.vllm_config,
self.device,
active_layer_names=self.draft_attn_layer_names,
)
self.block_tables = block_tables
def _build_draft_attn_metadata(
self,
num_reqs: int,
num_reqs_padded: int,
num_tokens_padded: int,
num_query_per_req: int = 1,
causal: bool = True,
) -> dict[str, Any] | None:
# 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
)
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]
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=num_query_per_req,
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,
)
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,
output_processed_logits=draft_logits,
output_processed_logits_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)
if self.draft_logits is not None:
# 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)