wehub-resource-sync
312 lines
12 KiB
Python
312 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import Mapping
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from vllm.config import VllmConfig, get_layers_from_vllm_config
|
|
from vllm.config.compilation import CUDAGraphMode
|
|
from vllm.distributed.eplb.eplb_state import EplbState
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
|
from vllm.v1.kv_cache_interface import KVCacheConfig
|
|
from vllm.v1.worker.gpu.attn_utils import (
|
|
build_attn_metadata,
|
|
init_attn_backend,
|
|
)
|
|
from vllm.v1.worker.gpu.block_table import BlockTables
|
|
from vllm.v1.worker.gpu.cudagraph_utils import (
|
|
AttentionStatePair,
|
|
BatchExecutionDescriptor,
|
|
)
|
|
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.sample.gumbel import gumbel_sample
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class BaseSpeculator(ABC):
|
|
@abstractmethod
|
|
def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def capture(
|
|
self,
|
|
attn_states: dict[BatchExecutionDescriptor, AttentionStatePair],
|
|
) -> None:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def propose(
|
|
self,
|
|
input_batch: InputBatch,
|
|
attn_metadata: dict[str, Any],
|
|
slot_mappings: dict[str, torch.Tensor],
|
|
# [num_tokens, hidden_size]
|
|
last_hidden_states: torch.Tensor,
|
|
# num_layers x [num_tokens, hidden_size]
|
|
aux_hidden_states: list[torch.Tensor] | None,
|
|
# [num_reqs]
|
|
num_sampled: torch.Tensor,
|
|
# [num_reqs]
|
|
num_rejected: torch.Tensor,
|
|
# [max_num_reqs]
|
|
last_sampled: torch.Tensor,
|
|
# [max_num_reqs]
|
|
next_prefill_tokens: torch.Tensor,
|
|
# [max_num_reqs]
|
|
temperature: torch.Tensor,
|
|
# [max_num_reqs]
|
|
seeds: torch.Tensor,
|
|
num_tokens_across_dp: torch.Tensor | 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:
|
|
pass
|
|
|
|
|
|
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 | Mapping[int, 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)
|