wehub-resource-sync
343 lines
13 KiB
Python
343 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""DFlash speculator for Laguna target models.
|
|
|
|
Laguna DFlash uses a uniform drafter layer flavor (`layer_types` all full
|
|
or all sliding). The draft checkpoint shares token embedding and lm_head
|
|
weights with the target model through the generic spec-decode proposer.
|
|
"""
|
|
|
|
from collections.abc import Iterable
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from vllm import _custom_ops as ops
|
|
from vllm.compilation.decorators import support_torch_compile
|
|
from vllm.config import VllmConfig
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.layernorm import RMSNorm
|
|
from vllm.model_executor.layers.linear import ReplicatedLinear
|
|
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
|
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
|
ParallelLMHead,
|
|
VocabParallelEmbedding,
|
|
)
|
|
from vllm.model_executor.model_loader.weight_utils import (
|
|
default_weight_loader,
|
|
maybe_remap_kv_scale_name,
|
|
)
|
|
from vllm.model_executor.models.interfaces import EagleModelMixin, SupportsEagle3
|
|
from vllm.multimodal.inputs import NestedTensors
|
|
|
|
from .laguna import LagunaDecoderLayer
|
|
from .qwen3_dflash import DFlashQwen3Model
|
|
from .utils import (
|
|
AutoWeightsLoader,
|
|
get_draft_quant_config,
|
|
maybe_prefix,
|
|
process_eagle_weight,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _get_dflash_layer_types(config) -> tuple[str, ...]:
|
|
layer_types = getattr(config, "layer_types", None)
|
|
if layer_types is None:
|
|
raise ValueError("Laguna DFlash config requires `layer_types`.")
|
|
if len(layer_types) != config.num_hidden_layers:
|
|
raise ValueError(
|
|
f"DFlash layer_types length {len(layer_types)} does not match "
|
|
f"num_hidden_layers {config.num_hidden_layers}."
|
|
)
|
|
# Laguna DFlash checkpoints use a uniform drafter attention flavor.
|
|
if len(set(layer_types)) > 1:
|
|
raise NotImplementedError(
|
|
"Laguna DFlash drafter requires a uniform `layer_types` "
|
|
f"(got {sorted(set(layer_types))})."
|
|
)
|
|
return tuple(layer_types)
|
|
|
|
|
|
@support_torch_compile
|
|
class DFlashLagunaModel(DFlashQwen3Model, EagleModelMixin):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
vllm_config: VllmConfig,
|
|
prefix: str = "",
|
|
) -> None:
|
|
nn.Module.__init__(self)
|
|
self.config = vllm_config.speculative_config.draft_model_config.hf_config
|
|
self.vocab_size = self.config.vocab_size
|
|
self.quant_config = get_draft_quant_config(vllm_config)
|
|
|
|
target_layer_ids = self.config.dflash_config["target_layer_ids"]
|
|
if not target_layer_ids:
|
|
raise ValueError(
|
|
"Laguna DFlash config requires non-empty "
|
|
"`dflash_config.target_layer_ids`."
|
|
)
|
|
|
|
self.embed_tokens = VocabParallelEmbedding(
|
|
self.config.vocab_size,
|
|
self.config.hidden_size,
|
|
prefix=maybe_prefix(prefix, "embed_tokens"),
|
|
)
|
|
|
|
self.mask_token_id = self.config.dflash_config.get("mask_token_id")
|
|
self.register_buffer(
|
|
"mask_embedding",
|
|
torch.zeros(
|
|
self.config.hidden_size,
|
|
dtype=vllm_config.model_config.dtype,
|
|
),
|
|
persistent=False,
|
|
)
|
|
self.has_separate_mask_embedding = False
|
|
|
|
self.layer_types = _get_dflash_layer_types(self.config)
|
|
target_layer_count = self.config.target_layer_count
|
|
self.layers = nn.ModuleList(
|
|
[
|
|
LagunaDecoderLayer(
|
|
prefix=maybe_prefix(prefix, f"layers.{layer_idx}"),
|
|
config=self.config,
|
|
cache_config=vllm_config.cache_config,
|
|
quant_config=self.quant_config,
|
|
layer_idx=layer_idx,
|
|
attention_prefix=maybe_prefix(
|
|
prefix, f"layers.{layer_idx + target_layer_count}"
|
|
),
|
|
)
|
|
for layer_idx in range(self.config.num_hidden_layers)
|
|
]
|
|
)
|
|
for layer in self.layers:
|
|
if getattr(layer.self_attn, "sliding_window", None) is not None:
|
|
# DFlash inserts verifier-context K/V at absolute cache slots.
|
|
# Keep full KV allocation; SWA remains a compute-time limit.
|
|
layer.self_attn.attn.sliding_window = None
|
|
num_features_to_use = len(target_layer_ids)
|
|
target_hidden_size = vllm_config.model_config.get_hidden_size()
|
|
fc_input_size = target_hidden_size * num_features_to_use
|
|
self.num_aux_slices = num_features_to_use
|
|
self.aux_hidden_norms = nn.ModuleList(
|
|
[
|
|
RMSNorm(
|
|
fc_input_size // num_features_to_use,
|
|
eps=self.config.rms_norm_eps,
|
|
)
|
|
for _ in range(num_features_to_use)
|
|
]
|
|
)
|
|
self.fc = ReplicatedLinear(
|
|
input_size=fc_input_size,
|
|
output_size=self.config.hidden_size,
|
|
bias=False,
|
|
params_dtype=vllm_config.model_config.dtype,
|
|
quant_config=self.quant_config,
|
|
prefix=maybe_prefix(prefix, "fc"),
|
|
return_bias=False,
|
|
)
|
|
self.hidden_norm = RMSNorm(
|
|
self.config.hidden_size,
|
|
eps=self.config.rms_norm_eps,
|
|
)
|
|
self.norm = RMSNorm(
|
|
self.config.hidden_size,
|
|
eps=self.config.rms_norm_eps,
|
|
)
|
|
|
|
def _build_context_kv_buffers(
|
|
self,
|
|
layers_attn: list[nn.Module],
|
|
has_bias: bool,
|
|
) -> None:
|
|
self._kv_weights = torch.stack(
|
|
[a.qkv_proj.weight[a.q_size :] for a in layers_attn], dim=0
|
|
).contiguous()
|
|
if has_bias:
|
|
self._kv_biases: torch.Tensor | None = torch.stack(
|
|
[a.qkv_proj.bias[a.q_size :] for a in layers_attn], dim=0
|
|
).contiguous()
|
|
else:
|
|
self._kv_biases = None
|
|
self._input_layernorm_weights = torch.stack(
|
|
[layer.input_layernorm.weight.data for layer in self.layers], dim=0
|
|
).contiguous()
|
|
self._k_norm_weights = torch.stack(
|
|
[a.k_norm.weight.data for a in layers_attn], dim=0
|
|
).contiguous()
|
|
|
|
def _project_context_kv(
|
|
self,
|
|
context_states: torch.Tensor,
|
|
num_ctx: int,
|
|
num_layers: int,
|
|
num_kv_heads: int,
|
|
head_dim: int,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
normed_context_states = torch.empty(
|
|
(num_layers, num_ctx, context_states.shape[-1]),
|
|
dtype=context_states.dtype,
|
|
device=context_states.device,
|
|
)
|
|
ops.rms_norm(
|
|
normed_context_states,
|
|
context_states.unsqueeze(0).expand(num_layers, -1, -1),
|
|
self._input_layernorm_weights,
|
|
self._rms_norm_eps,
|
|
)
|
|
all_kv_flat = torch.bmm(
|
|
normed_context_states,
|
|
self._kv_weights.transpose(1, 2),
|
|
)
|
|
if self._kv_biases is not None:
|
|
all_kv_flat += self._kv_biases[:, None, :]
|
|
all_kv = (
|
|
all_kv_flat.view(num_layers, num_ctx, 2, num_kv_heads, head_dim)
|
|
.permute(2, 0, 1, 3, 4)
|
|
.contiguous()
|
|
)
|
|
all_k = all_kv[0]
|
|
all_v = all_kv[1]
|
|
return all_k, all_v
|
|
|
|
def _normalize_context_k(self, all_k: torch.Tensor) -> torch.Tensor:
|
|
all_k_normed = torch.empty_like(all_k)
|
|
ops.rms_norm(
|
|
all_k_normed,
|
|
all_k,
|
|
self._k_norm_weights,
|
|
self._rms_norm_eps,
|
|
)
|
|
return all_k_normed
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
params_dict = dict(self.named_parameters())
|
|
loaded_params: set[str] = set()
|
|
for name, loaded_weight in weights:
|
|
if self.quant_config is not None and (
|
|
scale_name := self.quant_config.get_cache_scale(name)
|
|
):
|
|
param = params_dict[scale_name]
|
|
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
|
loaded_weight = (
|
|
loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
|
|
)
|
|
weight_loader(param, loaded_weight)
|
|
loaded_params.add(scale_name)
|
|
continue
|
|
if "scale" in name:
|
|
name = maybe_remap_kv_scale_name(name, params_dict)
|
|
if name is None:
|
|
continue
|
|
param = params_dict[name]
|
|
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
|
weight_loader(param, loaded_weight)
|
|
loaded_params.add(name)
|
|
return loaded_params
|
|
|
|
|
|
class DFlashLagunaForCausalLM(nn.Module, SupportsEagle3):
|
|
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
|
nn.Module.__init__(self)
|
|
self.config = vllm_config.speculative_config.draft_model_config.hf_config
|
|
if getattr(self.config, "draft_vocab_size", None) is None:
|
|
raise ValueError("Laguna DFlash config requires `draft_vocab_size`.")
|
|
self.has_own_embed_tokens = False
|
|
self.has_own_lm_head = False
|
|
target_layer_num = vllm_config.model_config.get_num_layers(
|
|
vllm_config.parallel_config
|
|
)
|
|
self.config.target_layer_count = target_layer_num
|
|
target_vocab_size = vllm_config.model_config.get_vocab_size()
|
|
if self.config.draft_vocab_size != target_vocab_size:
|
|
raise ValueError(
|
|
"Laguna DFlash shares the target lm_head and requires "
|
|
"`draft_vocab_size` to match the target vocabulary size "
|
|
f"({self.config.draft_vocab_size} != {target_vocab_size})."
|
|
)
|
|
self.model = DFlashLagunaModel(
|
|
vllm_config=vllm_config,
|
|
prefix="model",
|
|
)
|
|
|
|
self.lm_head = ParallelLMHead(
|
|
self.config.draft_vocab_size,
|
|
self.config.hidden_size,
|
|
prefix=maybe_prefix(prefix, "lm_head"),
|
|
)
|
|
self.logits_processor = LogitsProcessor(self.config.draft_vocab_size)
|
|
|
|
def embed_input_ids(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
multimodal_embeddings: NestedTensors | None = None,
|
|
is_multimodal: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
return self.model.embed_input_ids(input_ids)
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
inputs_embeds: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
return self.model(input_ids, positions, inputs_embeds)
|
|
|
|
def compute_logits(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
) -> torch.Tensor | None:
|
|
return self.logits_processor(self.lm_head, hidden_states)
|
|
|
|
def precompute_and_store_context_kv(
|
|
self,
|
|
context_states: torch.Tensor,
|
|
context_positions: torch.Tensor,
|
|
context_slot_mapping: torch.Tensor | None = None,
|
|
) -> None:
|
|
self.model.precompute_and_store_context_kv(
|
|
context_states, context_positions, context_slot_mapping
|
|
)
|
|
|
|
def combine_hidden_states(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
# Normalize each verifier hidden-state slice, concatenate them, then
|
|
# project into the drafter hidden size used as DFlash context.
|
|
needs_squeeze = hidden_states.dim() == 1
|
|
if needs_squeeze:
|
|
hidden_states = hidden_states.unsqueeze(0)
|
|
num_slices = self.model.num_aux_slices
|
|
slice_size = hidden_states.shape[-1] // num_slices
|
|
slices = hidden_states.view(hidden_states.shape[0], num_slices, slice_size)
|
|
normed = torch.empty_like(slices)
|
|
for i, norm in enumerate(self.model.aux_hidden_norms):
|
|
normed[:, i, :] = norm(slices[:, i, :])
|
|
hidden_states = normed.reshape(hidden_states.shape[0], -1)
|
|
result = self.model.fc(hidden_states)
|
|
result = self.model.hidden_norm(result)
|
|
if needs_squeeze:
|
|
result = result.squeeze(0)
|
|
return result
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
|
|
model_weights = {}
|
|
for name, loaded_weight in weights:
|
|
if "lm_head" not in name:
|
|
name = "model." + name
|
|
model_weights[name] = loaded_weight
|
|
process_eagle_weight(self, name)
|
|
|
|
loader = AutoWeightsLoader(self)
|
|
loaded_weight_names = loader.load_weights(model_weights.items())
|
|
loaded_weight_names.add("lm_head.weight")
|
|
loaded_weight_names.add("model.embed_tokens.weight")
|
|
self.model._build_fused_kv_buffers()
|
|
return loaded_weight_names
|