|
|
@@ -0,0 +1,974 @@
|
|
|
+# SPDX-License-Identifier: Apache-2.0
|
|
|
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
+
|
|
|
+import io
|
|
|
+from collections.abc import Iterable
|
|
|
+
|
|
|
+import torch
|
|
|
+import torch.nn.functional as F
|
|
|
+from torch import nn
|
|
|
+from transformers import Qwen3Config
|
|
|
+
|
|
|
+from vllm import _custom_ops as ops
|
|
|
+from vllm.compilation.decorators import support_torch_compile
|
|
|
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
|
|
|
+from vllm.distributed import (
|
|
|
+ get_tensor_model_parallel_rank,
|
|
|
+ get_tensor_model_parallel_world_size,
|
|
|
+)
|
|
|
+from vllm.logger import init_logger
|
|
|
+from vllm.model_executor.layers.attention import Attention
|
|
|
+from vllm.model_executor.layers.layernorm import RMSNorm
|
|
|
+from vllm.model_executor.layers.linear import (
|
|
|
+ QKVParallelLinear,
|
|
|
+ ReplicatedLinear,
|
|
|
+ RowParallelLinear,
|
|
|
+)
|
|
|
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
|
|
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
|
|
+from vllm.model_executor.layers.rotary_embedding import get_rope
|
|
|
+from vllm.model_executor.layers.vocab_parallel_embedding import (
|
|
|
+ ParallelLMHead,
|
|
|
+ VocabParallelEmbedding,
|
|
|
+)
|
|
|
+from vllm.multimodal.inputs import NestedTensors
|
|
|
+from vllm.transformers_utils.config import set_default_rope_theta
|
|
|
+from vllm.transformers_utils.repo_utils import get_hf_file_bytes
|
|
|
+from vllm.v1.attention.backend import AttentionType
|
|
|
+from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
|
|
|
+ get_eagle3_aux_layers_from_config,
|
|
|
+)
|
|
|
+
|
|
|
+from .qwen2 import Qwen2MLP as Qwen3MLP
|
|
|
+from .qwen3 import Qwen3ForCausalLM
|
|
|
+from .utils import (
|
|
|
+ AutoWeightsLoader,
|
|
|
+ WeightsMapper,
|
|
|
+ get_draft_quant_config,
|
|
|
+ maybe_prefix,
|
|
|
+ process_eagle_weight,
|
|
|
+)
|
|
|
+
|
|
|
+logger = init_logger(__name__)
|
|
|
+
|
|
|
+
|
|
|
+_SLIDING_ATTENTION = "sliding_attention"
|
|
|
+
|
|
|
+
|
|
|
+def _dflash_layer_causal(config: Qwen3Config, layer_idx: int) -> bool:
|
|
|
+ """Resolve explicit causality before falling back to legacy layer defaults."""
|
|
|
+ is_causal = getattr(config, "is_causal", None)
|
|
|
+ if is_causal is not None:
|
|
|
+ return bool(is_causal)
|
|
|
+ override = (getattr(config, "dflash_config", None) or {}).get("causal")
|
|
|
+ if override is not None:
|
|
|
+ return bool(override)
|
|
|
+ layer_types = getattr(config, "layer_types", None)
|
|
|
+ return bool(layer_types) and layer_types[layer_idx] == _SLIDING_ATTENTION
|
|
|
+
|
|
|
+
|
|
|
+def dflash_has_any_non_causal(config: Qwen3Config) -> bool:
|
|
|
+ """Whether the draft needs a non-causal-capable backend, resolved from config
|
|
|
+ (config mirror of the model's ``get_draft_attn_causal``, usable pre-build)."""
|
|
|
+ return not all(
|
|
|
+ _dflash_layer_causal(config, i) for i in range(config.num_hidden_layers)
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def dflash_target_rope_is_neox_style(target_model: nn.Module) -> bool | None:
|
|
|
+ """The target's RoPE layout, from its first attention layer.
|
|
|
+
|
|
|
+ A DFlash head must rotate Q/K the way the target it was distilled against
|
|
|
+ does, and a mismatch is silent — acceptance collapses but nothing errors and
|
|
|
+ the output stays correct. Draft checkpoints do not carry this, so take it
|
|
|
+ from the target. None if the target uses no RoPE.
|
|
|
+ """
|
|
|
+ language_model = (
|
|
|
+ target_model.get_language_model()
|
|
|
+ if hasattr(target_model, "get_language_model")
|
|
|
+ else target_model
|
|
|
+ )
|
|
|
+ for module in language_model.modules():
|
|
|
+ style = getattr(module, "is_neox_style", None)
|
|
|
+ if isinstance(style, bool):
|
|
|
+ return style
|
|
|
+ return None
|
|
|
+
|
|
|
+
|
|
|
+def _get_dflash_fc_input_size(vllm_config: VllmConfig) -> int:
|
|
|
+ spec_config = vllm_config.speculative_config
|
|
|
+ config = spec_config.draft_model_config.hf_config
|
|
|
+ aux_layers = get_eagle3_aux_layers_from_config(spec_config)
|
|
|
+ num_features_to_use = len(aux_layers) if aux_layers else config.num_hidden_layers
|
|
|
+ target_hidden_size = (
|
|
|
+ getattr(config, "target_hidden_size", None) or config.hidden_size
|
|
|
+ )
|
|
|
+ return target_hidden_size * num_features_to_use
|
|
|
+
|
|
|
+
|
|
|
+def _resolve_layer_attention(
|
|
|
+ config: Qwen3Config, layer_idx: int
|
|
|
+) -> tuple[int | None, bool]:
|
|
|
+ """Resolve ``(sliding_window, causal)`` for one DFlash draft layer.
|
|
|
+
|
|
|
+ +----------------------+-------------------------+--------------------------------+
|
|
|
+ | Config | ``layer_type`` | *``causal`` |
|
|
|
+ +======================+=========================+================================+
|
|
|
+ | ``layer_types`` | SWA if ``use_swa`` | True if ``layer_types[i]=SWA`` |
|
|
|
+ | | else ``layer_types[i]`` | else False |
|
|
|
+ +----------------------+-------------------------+--------------------------------+
|
|
|
+ | ``layer_types=None`` | SWA | False |
|
|
|
+ | + ``use_swa=True`` | | |
|
|
|
+ +----------------------+-------------------------+--------------------------------+
|
|
|
+ | ``layer_types=None`` | Full | False |
|
|
|
+ | + ``use_swa=False`` | | |
|
|
|
+ +----------------------+-------------------------+--------------------------------+
|
|
|
+ * If ``dflash_config.causal`` is set, its value overrides ``causal`` for all layers.
|
|
|
+
|
|
|
+ This is to support a varied ecosystem of checkpoints, including:
|
|
|
+ - XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash (sets "use_swa", assumes non-causal)
|
|
|
+ - z-lab/gemma-4-31B-it-DFlash (has mixed layer types, assumes causal only for SWA)
|
|
|
+ - z-lab/Qwen3.5-9B-DFlash ("standard" DFlash, all full attn, assumes non-causal)
|
|
|
+ """
|
|
|
+ dflash_config = getattr(config, "dflash_config", None) or {}
|
|
|
+ layer_types = getattr(config, "layer_types", None)
|
|
|
+ use_swa = dflash_config.get("use_swa", False)
|
|
|
+
|
|
|
+ any_sliding = False
|
|
|
+ if layer_types is not None:
|
|
|
+ num_sliding = sum(lt == _SLIDING_ATTENTION for lt in layer_types)
|
|
|
+ any_sliding = num_sliding > 0
|
|
|
+ # Mixed sliding/full attention needs multiple KV groups (V2 runner only).
|
|
|
+ if (
|
|
|
+ 0 < num_sliding < len(layer_types)
|
|
|
+ and not get_current_vllm_config().use_v2_model_runner
|
|
|
+ ):
|
|
|
+ raise NotImplementedError(
|
|
|
+ "DFlash drafters with mixed sliding/full attention require "
|
|
|
+ "the V2 model runner; relaunch with "
|
|
|
+ "VLLM_USE_V2_MODEL_RUNNER=1."
|
|
|
+ )
|
|
|
+
|
|
|
+ # ``use_swa`` forces SWA on every layer, even an all-full ``layer_types``.
|
|
|
+ if layer_types is None or (use_swa and not any_sliding):
|
|
|
+ is_sliding = use_swa
|
|
|
+ else:
|
|
|
+ is_sliding = layer_types[layer_idx] == _SLIDING_ATTENTION
|
|
|
+
|
|
|
+ sliding_window = None
|
|
|
+ if is_sliding:
|
|
|
+ sliding_window = dflash_config.get(
|
|
|
+ "swa_window_size", getattr(config, "sliding_window", None)
|
|
|
+ )
|
|
|
+ if sliding_window is None:
|
|
|
+ raise ValueError(
|
|
|
+ "DFlash sliding attention requires a window size configured in "
|
|
|
+ "dflash_config.swa_window_size or the top-level sliding_window."
|
|
|
+ )
|
|
|
+
|
|
|
+ return sliding_window, _dflash_layer_causal(config, layer_idx)
|
|
|
+
|
|
|
+
|
|
|
+class DFlashQwen3Attention(nn.Module):
|
|
|
+ """Attention for DFlash speculative decoding.
|
|
|
+
|
|
|
+ Context KVs are pre-inserted into the KV cache before the forward pass.
|
|
|
+ This layer handles only query tokens via standard attention.
|
|
|
+ Adapted from Qwen3Attention."""
|
|
|
+
|
|
|
+ def __init__(
|
|
|
+ self,
|
|
|
+ hidden_size: int,
|
|
|
+ num_heads: int,
|
|
|
+ num_kv_heads: int,
|
|
|
+ rope_parameters: dict,
|
|
|
+ max_position: int = 4096 * 32,
|
|
|
+ head_dim: int | None = None,
|
|
|
+ rms_norm_eps: float = 1e-06,
|
|
|
+ attention_bias: bool = False,
|
|
|
+ add_swa_attention_sink_bias: bool = False,
|
|
|
+ sliding_window: int | None = None,
|
|
|
+ causal: bool = False,
|
|
|
+ is_neox_style: bool = True,
|
|
|
+ cache_config: CacheConfig | None = None,
|
|
|
+ quant_config: QuantizationConfig | None = None,
|
|
|
+ prefix: str = "",
|
|
|
+ attn_type: str = AttentionType.DECODER,
|
|
|
+ ) -> None:
|
|
|
+ super().__init__()
|
|
|
+ self.layer_name = prefix
|
|
|
+ self.hidden_size = hidden_size
|
|
|
+ tp_size = get_tensor_model_parallel_world_size()
|
|
|
+ self.total_num_heads = num_heads
|
|
|
+ assert self.total_num_heads % tp_size == 0
|
|
|
+ self.num_heads = self.total_num_heads // tp_size
|
|
|
+ self.total_num_kv_heads = num_kv_heads
|
|
|
+ if self.total_num_kv_heads >= tp_size:
|
|
|
+ assert self.total_num_kv_heads % tp_size == 0
|
|
|
+ else:
|
|
|
+ assert tp_size % self.total_num_kv_heads == 0
|
|
|
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
|
|
|
+ self.head_dim = head_dim or hidden_size // self.total_num_heads
|
|
|
+ self.q_size = self.num_heads * self.head_dim
|
|
|
+ self.kv_size = self.num_kv_heads * self.head_dim
|
|
|
+ self.scaling = self.head_dim**-0.5
|
|
|
+
|
|
|
+ self.qkv_proj = QKVParallelLinear(
|
|
|
+ hidden_size,
|
|
|
+ self.head_dim,
|
|
|
+ self.total_num_heads,
|
|
|
+ self.total_num_kv_heads,
|
|
|
+ bias=attention_bias,
|
|
|
+ quant_config=quant_config,
|
|
|
+ prefix=f"{prefix}.qkv_proj",
|
|
|
+ )
|
|
|
+ self.o_proj = RowParallelLinear(
|
|
|
+ self.total_num_heads * self.head_dim,
|
|
|
+ hidden_size,
|
|
|
+ bias=attention_bias, # DFlash has o_proj bias when using attention bias
|
|
|
+ quant_config=quant_config,
|
|
|
+ prefix=f"{prefix}.o_proj",
|
|
|
+ )
|
|
|
+
|
|
|
+ self.rotary_emb = get_rope(
|
|
|
+ self.head_dim,
|
|
|
+ max_position=max_position,
|
|
|
+ is_neox_style=is_neox_style,
|
|
|
+ rope_parameters=rope_parameters,
|
|
|
+ )
|
|
|
+
|
|
|
+ self.attention_sink_bias = (
|
|
|
+ torch.nn.Parameter(torch.empty(self.num_heads), requires_grad=False)
|
|
|
+ if add_swa_attention_sink_bias
|
|
|
+ else None
|
|
|
+ )
|
|
|
+
|
|
|
+ self.sliding_window = sliding_window
|
|
|
+ self.attn = Attention(
|
|
|
+ self.num_heads,
|
|
|
+ self.head_dim,
|
|
|
+ self.scaling,
|
|
|
+ num_kv_heads=self.num_kv_heads,
|
|
|
+ cache_config=cache_config,
|
|
|
+ quant_config=quant_config,
|
|
|
+ per_layer_sliding_window=sliding_window,
|
|
|
+ prefix=f"{prefix}.attn",
|
|
|
+ attn_type=attn_type,
|
|
|
+ sinks=self.attention_sink_bias,
|
|
|
+ )
|
|
|
+ self.causal = causal
|
|
|
+ self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
|
|
|
+ self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
|
|
|
+
|
|
|
+ def forward(
|
|
|
+ self,
|
|
|
+ positions: torch.Tensor,
|
|
|
+ hidden_states: torch.Tensor,
|
|
|
+ ) -> torch.Tensor:
|
|
|
+ """DFlash attention assumes that the KV cache is already populated
|
|
|
+ with the context K/V from the target model's hidden states. This forward op
|
|
|
+ computes attention for the query tokens only.
|
|
|
+ See also: precompute_and_store_context_kv"""
|
|
|
+ qkv, _ = self.qkv_proj(hidden_states)
|
|
|
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
|
|
+
|
|
|
+ # Per-head RMSNorm
|
|
|
+ q_shape, k_shape = q.shape, k.shape
|
|
|
+ q = self.q_norm(
|
|
|
+ q.view(*q_shape[:-1], q_shape[-1] // self.head_dim, self.head_dim)
|
|
|
+ ).view(q_shape)
|
|
|
+ k = self.k_norm(
|
|
|
+ k.view(*k_shape[:-1], k_shape[-1] // self.head_dim, self.head_dim)
|
|
|
+ ).view(k_shape)
|
|
|
+
|
|
|
+ q, k = self.rotary_emb(positions, q, k)
|
|
|
+
|
|
|
+ attn_output = self.attn(q, k, v)
|
|
|
+ output, _ = self.o_proj(attn_output)
|
|
|
+ return output
|
|
|
+
|
|
|
+
|
|
|
+class DFlashQwen3DecoderLayer(nn.Module):
|
|
|
+ def __init__(
|
|
|
+ self,
|
|
|
+ vllm_config: VllmConfig,
|
|
|
+ *,
|
|
|
+ config: Qwen3Config,
|
|
|
+ layer_idx: int,
|
|
|
+ cache_config: CacheConfig | None = None,
|
|
|
+ quant_config: QuantizationConfig | None = None,
|
|
|
+ prefix: str = "",
|
|
|
+ ) -> None:
|
|
|
+ super().__init__()
|
|
|
+ self.hidden_size = config.hidden_size
|
|
|
+ set_default_rope_theta(config, default_theta=1000000)
|
|
|
+ attn_type = AttentionType.DECODER
|
|
|
+
|
|
|
+ # DFlash drafts store the sink-bias flag inside dflash_config; fall back
|
|
|
+ # to the top-level attribute used by other (e.g. MiMo) configs.
|
|
|
+ dflash_config = getattr(config, "dflash_config", None) or {}
|
|
|
+ add_swa_attention_sink_bias = dflash_config.get(
|
|
|
+ "attention_sink_bias",
|
|
|
+ getattr(config, "add_swa_attention_sink_bias", False),
|
|
|
+ )
|
|
|
+
|
|
|
+ # Resolve this layer's attention mode (full vs sliding window, causal vs
|
|
|
+ # non-causal) from the draft config.
|
|
|
+ sliding_window, causal = _resolve_layer_attention(config, layer_idx)
|
|
|
+
|
|
|
+ # RoPE layout, copied off the target at load time by the draft loader
|
|
|
+ # (see `dflash_target_rope_is_neox_style`). Checkpoints do not carry it:
|
|
|
+ # a head distilled from an interleaved-RoPE target must rotate the way
|
|
|
+ # that target does, or every drafted Q/K is wrong and acceptance
|
|
|
+ # collapses with no error raised.
|
|
|
+ is_neox_style = getattr(config, "is_neox_style", True)
|
|
|
+
|
|
|
+ self.self_attn = DFlashQwen3Attention(
|
|
|
+ hidden_size=self.hidden_size,
|
|
|
+ num_heads=config.num_attention_heads,
|
|
|
+ max_position=config.max_position_embeddings,
|
|
|
+ num_kv_heads=config.num_key_value_heads,
|
|
|
+ rms_norm_eps=config.rms_norm_eps,
|
|
|
+ attention_bias=getattr(config, "attention_bias", False),
|
|
|
+ add_swa_attention_sink_bias=add_swa_attention_sink_bias,
|
|
|
+ sliding_window=sliding_window,
|
|
|
+ causal=causal,
|
|
|
+ is_neox_style=is_neox_style,
|
|
|
+ head_dim=getattr(config, "head_dim", None),
|
|
|
+ cache_config=cache_config,
|
|
|
+ quant_config=quant_config,
|
|
|
+ rope_parameters=config.rope_parameters,
|
|
|
+ prefix=f"{prefix}.self_attn",
|
|
|
+ attn_type=attn_type,
|
|
|
+ )
|
|
|
+ self.mlp = Qwen3MLP(
|
|
|
+ hidden_size=self.hidden_size,
|
|
|
+ intermediate_size=config.intermediate_size,
|
|
|
+ hidden_act=config.hidden_act,
|
|
|
+ quant_config=quant_config,
|
|
|
+ prefix=f"{prefix}.mlp",
|
|
|
+ )
|
|
|
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
|
+ self.post_attention_layernorm = RMSNorm(
|
|
|
+ config.hidden_size, eps=config.rms_norm_eps
|
|
|
+ )
|
|
|
+
|
|
|
+ def forward(
|
|
|
+ self,
|
|
|
+ positions: torch.Tensor,
|
|
|
+ hidden_states: torch.Tensor,
|
|
|
+ residual: torch.Tensor | None,
|
|
|
+ ) -> tuple[torch.Tensor, torch.Tensor]:
|
|
|
+ if residual is not None:
|
|
|
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
|
|
|
+ else:
|
|
|
+ residual = hidden_states
|
|
|
+ hidden_states = self.input_layernorm(hidden_states)
|
|
|
+
|
|
|
+ hidden_states = self.self_attn(
|
|
|
+ positions=positions,
|
|
|
+ hidden_states=hidden_states,
|
|
|
+ )
|
|
|
+
|
|
|
+ hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
|
|
|
+ hidden_states = self.mlp(hidden_states)
|
|
|
+ return hidden_states, residual
|
|
|
+
|
|
|
+
|
|
|
+@support_torch_compile
|
|
|
+class DFlashQwen3Model(nn.Module):
|
|
|
+ decoder_layer_cls = DFlashQwen3DecoderLayer
|
|
|
+
|
|
|
+ hf_to_vllm_mapper = WeightsMapper(
|
|
|
+ orig_to_new_substr={
|
|
|
+ "midlayer.": "layers.0.",
|
|
|
+ # Muse-Glimmer-30B-assistant names the aux-hidden-state encoder
|
|
|
+ # `encoder.fc` / `encoder.output_norm_enc`; this head calls them
|
|
|
+ # `fc` / `hidden_norm`. Same tensors and shapes, different names.
|
|
|
+ "encoder.output_norm_enc.": "hidden_norm.",
|
|
|
+ "encoder.fc.": "fc.",
|
|
|
+ },
|
|
|
+ orig_to_new_stacked={
|
|
|
+ ".q_proj": (".qkv_proj", "q"),
|
|
|
+ ".k_proj": (".qkv_proj", "k"),
|
|
|
+ ".v_proj": (".qkv_proj", "v"),
|
|
|
+ ".gate_proj": (".gate_up_proj", 0),
|
|
|
+ ".up_proj": (".gate_up_proj", 1),
|
|
|
+ },
|
|
|
+ )
|
|
|
+
|
|
|
+ def __init__(
|
|
|
+ self,
|
|
|
+ *,
|
|
|
+ vllm_config: VllmConfig,
|
|
|
+ start_layer_id: int = 0,
|
|
|
+ prefix: str = "",
|
|
|
+ ) -> None:
|
|
|
+ super().__init__()
|
|
|
+ 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)
|
|
|
+
|
|
|
+ drafter_config = getattr(self.config, "eagle_config", {})
|
|
|
+ drafter_config.update(getattr(self.config, "dflash_config", {}))
|
|
|
+
|
|
|
+ self.use_aux_hidden_state = drafter_config.get(
|
|
|
+ "use_aux_hidden_state",
|
|
|
+ getattr(self.config, "use_aux_hidden_state", True),
|
|
|
+ )
|
|
|
+
|
|
|
+ current_vllm_config = get_current_vllm_config()
|
|
|
+
|
|
|
+ self.embed_tokens = VocabParallelEmbedding(
|
|
|
+ self.config.vocab_size,
|
|
|
+ self.config.hidden_size,
|
|
|
+ prefix=maybe_prefix(prefix, "embed_tokens"),
|
|
|
+ )
|
|
|
+
|
|
|
+ # Masked query slots are fed to the draft as `mask_token_id`. Most DFlash
|
|
|
+ # checkpoints will have the mask embedding in the vocabulary embedding table
|
|
|
+ # at that slot id. Some checkpoints (XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash) ship
|
|
|
+ # with a separate mask embedding tensor to use instead. When present, we load it
|
|
|
+ # and substitute it for embed_tokens[mask_token_id] when computing embeddings.
|
|
|
+ self.mask_token_id = drafter_config.get(
|
|
|
+ "mask_token_id", getattr(self.config, "mask_token_id", None)
|
|
|
+ )
|
|
|
+ self.mask_embedding = nn.Parameter(
|
|
|
+ torch.zeros(self.config.hidden_size, dtype=vllm_config.model_config.dtype),
|
|
|
+ requires_grad=False,
|
|
|
+ )
|
|
|
+ self.has_separate_mask_embedding = False
|
|
|
+
|
|
|
+ self.layers = nn.ModuleList(
|
|
|
+ [
|
|
|
+ self.decoder_layer_cls(
|
|
|
+ current_vllm_config,
|
|
|
+ config=self.config,
|
|
|
+ layer_idx=layer_idx,
|
|
|
+ cache_config=current_vllm_config.cache_config,
|
|
|
+ quant_config=self.quant_config,
|
|
|
+ prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"),
|
|
|
+ )
|
|
|
+ for layer_idx in range(self.config.num_hidden_layers)
|
|
|
+ ]
|
|
|
+ )
|
|
|
+ if self.use_aux_hidden_state:
|
|
|
+ self.fc = ReplicatedLinear(
|
|
|
+ input_size=_get_dflash_fc_input_size(
|
|
|
+ vllm_config,
|
|
|
+ ),
|
|
|
+ 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,
|
|
|
+ )
|
|
|
+ # The context projection concatenates K/V weights from every draft
|
|
|
+ # layer. It is not a LinearBase module because its output layout is a
|
|
|
+ # DFlash-specific fusion. Serialized MXFP8 checkpoints retain an
|
|
|
+ # independently packed copy in this parameter container.
|
|
|
+ self._fused_kv_linear = nn.Module()
|
|
|
+ self._fused_kv_quant_method = None
|
|
|
+ self._fused_kv_weight: torch.Tensor | None = None
|
|
|
+ self._fused_kv_weight_scale: torch.Tensor | None = None
|
|
|
+
|
|
|
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
|
|
+ embeds = self.embed_tokens(input_ids)
|
|
|
+ if self.has_separate_mask_embedding and self.mask_token_id is not None:
|
|
|
+ # Replace masked slots with the dedicated mask embedding.
|
|
|
+ is_mask = (input_ids == self.mask_token_id).unsqueeze(-1)
|
|
|
+ embeds = torch.where(is_mask, self.mask_embedding.to(embeds.dtype), embeds)
|
|
|
+ return embeds
|
|
|
+
|
|
|
+ def _build_context_kv_buffers(
|
|
|
+ self,
|
|
|
+ layers_attn: list[nn.Module],
|
|
|
+ has_bias: bool,
|
|
|
+ ) -> None:
|
|
|
+ from vllm.model_executor.layers.quantization.modelopt import (
|
|
|
+ ModelOptMxFp8LinearMethod,
|
|
|
+ )
|
|
|
+
|
|
|
+ quant_methods = [a.qkv_proj.quant_method for a in layers_attn]
|
|
|
+ uses_mxfp8 = [
|
|
|
+ isinstance(method, ModelOptMxFp8LinearMethod) for method in quant_methods
|
|
|
+ ]
|
|
|
+ if any(uses_mxfp8) and not all(uses_mxfp8):
|
|
|
+ raise ValueError(
|
|
|
+ "Every DFlash attention layer must use the same MXFP8 "
|
|
|
+ "linear format for the fused context K/V projection."
|
|
|
+ )
|
|
|
+
|
|
|
+ self._hidden_norm_weight = self.hidden_norm.weight.data
|
|
|
+
|
|
|
+ # KV projection weights: [num_layers * 2 * kv_size, hidden_size]
|
|
|
+ kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn]
|
|
|
+ self._fused_kv_weight = torch.cat(kv_weights, dim=0)
|
|
|
+
|
|
|
+ if all(uses_mxfp8):
|
|
|
+ kv_scales = [a.qkv_proj.weight_scale[a.q_size :] for a in layers_attn]
|
|
|
+ self._fused_kv_weight_scale = torch.cat(kv_scales, dim=0)
|
|
|
+ else:
|
|
|
+ self._fused_kv_weight_scale = None
|
|
|
+ if has_bias:
|
|
|
+ kv_biases = [a.qkv_proj.bias[a.q_size :] for a in layers_attn]
|
|
|
+ self._fused_kv_bias: torch.Tensor | None = torch.cat(kv_biases, dim=0)
|
|
|
+ else:
|
|
|
+ self._fused_kv_bias = None
|
|
|
+
|
|
|
+ # K-norm weights stacked into one contiguous [num_layers, head_dim]
|
|
|
+ # tensor so the per-layer K-norm runs as a single grouped kernel.
|
|
|
+ self._k_norm_weights = torch.stack(
|
|
|
+ [a.k_norm.weight.data for a in layers_attn], dim=0
|
|
|
+ ).contiguous()
|
|
|
+
|
|
|
+ def _build_fused_kv_buffers(self) -> None:
|
|
|
+ """Build fused weight buffers for precompute_and_store_context_kv.
|
|
|
+
|
|
|
+ Must be called after weights are loaded. Stacks the KV-projection
|
|
|
+ weights, K-norm weights, and RoPE parameters from every attention
|
|
|
+ layer so that precompute_and_store_context_kv can run one fused
|
|
|
+ GEMM for all layers at once. Also aliases the weight of the hidden_norm.
|
|
|
+ """
|
|
|
+ layers_attn = [layer.self_attn for layer in self.layers]
|
|
|
+ attn0 = layers_attn[0]
|
|
|
+ has_bias = attn0.qkv_proj.bias is not None
|
|
|
+
|
|
|
+ self._build_context_kv_buffers(layers_attn, has_bias)
|
|
|
+
|
|
|
+ # RoPE parameters
|
|
|
+ self._rope_head_size = attn0.rotary_emb.head_size
|
|
|
+ self._rope_cos_sin_cache = attn0.rotary_emb.cos_sin_cache
|
|
|
+ self._rope_is_neox = attn0.rotary_emb.is_neox_style
|
|
|
+ # Validation that RoPE params are the same across all layers
|
|
|
+ for attn in layers_attn[1:]:
|
|
|
+ assert (
|
|
|
+ attn.rotary_emb.head_size == self._rope_head_size
|
|
|
+ and attn.rotary_emb.is_neox_style == self._rope_is_neox
|
|
|
+ ), "All layers must have the same RoPE parameters for DFlash precomputation"
|
|
|
+
|
|
|
+ # Layer metadata
|
|
|
+ self._num_attn_layers = len(layers_attn)
|
|
|
+ self._kv_size = attn0.kv_size
|
|
|
+ self._head_dim = attn0.head_dim
|
|
|
+ self._num_kv_heads = attn0.num_kv_heads
|
|
|
+ self._rms_norm_eps = attn0.q_norm.variance_epsilon
|
|
|
+ # Validation that all layers have the same attention config
|
|
|
+ for attn in layers_attn[1:]:
|
|
|
+ assert (
|
|
|
+ attn.kv_size == self._kv_size
|
|
|
+ and attn.head_dim == self._head_dim
|
|
|
+ and attn.num_kv_heads == self._num_kv_heads
|
|
|
+ and attn.q_norm.variance_epsilon == self._rms_norm_eps
|
|
|
+ ), "All layers must have the same attn config for DFlash precomputation"
|
|
|
+
|
|
|
+ # References to inner Attention layers for direct cache writes
|
|
|
+ self._attn_layers = [layer.self_attn.attn for layer in self.layers]
|
|
|
+
|
|
|
+ def process_weights_after_loading(self) -> None:
|
|
|
+ """Pack the serialized MXFP8 context projection for its GEMM backend."""
|
|
|
+ from vllm.model_executor.layers.quantization.modelopt import (
|
|
|
+ ModelOptMxFp8LinearMethod,
|
|
|
+ )
|
|
|
+
|
|
|
+ quant_method = self.layers[0].self_attn.qkv_proj.quant_method
|
|
|
+ if not isinstance(quant_method, ModelOptMxFp8LinearMethod):
|
|
|
+ return
|
|
|
+ if self._fused_kv_weight is None or self._fused_kv_weight_scale is None:
|
|
|
+ raise RuntimeError(
|
|
|
+ "The DFlash MXFP8 context projection requires serialized K/V "
|
|
|
+ "weights and block scales to be fused during checkpoint loading."
|
|
|
+ )
|
|
|
+
|
|
|
+ output_size, input_size = self._fused_kv_weight.shape
|
|
|
+ self._fused_kv_linear.input_size_per_partition = input_size
|
|
|
+ self._fused_kv_linear.output_size_per_partition = output_size
|
|
|
+ self._fused_kv_linear.logical_widths = [output_size]
|
|
|
+ self._fused_kv_linear.register_parameter(
|
|
|
+ "weight", nn.Parameter(self._fused_kv_weight, requires_grad=False)
|
|
|
+ )
|
|
|
+ self._fused_kv_linear.register_parameter(
|
|
|
+ "weight_scale",
|
|
|
+ nn.Parameter(self._fused_kv_weight_scale, requires_grad=False),
|
|
|
+ )
|
|
|
+ quant_method.process_weights_after_loading(self._fused_kv_linear)
|
|
|
+ self._fused_kv_quant_method = quant_method
|
|
|
+ self._fused_kv_weight = None
|
|
|
+ self._fused_kv_weight_scale = None
|
|
|
+ logger.info_once(
|
|
|
+ "Using %s for the fused DFlash context K/V projection.",
|
|
|
+ type(quant_method.kernel).__name__,
|
|
|
+ )
|
|
|
+
|
|
|
+ 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]:
|
|
|
+ # --- Fused KV projection (one GEMM for all layers) ---
|
|
|
+ normed_context_states = torch.empty_like(context_states)
|
|
|
+ ops.rms_norm(
|
|
|
+ normed_context_states,
|
|
|
+ context_states,
|
|
|
+ self._hidden_norm_weight,
|
|
|
+ self._rms_norm_eps,
|
|
|
+ )
|
|
|
+ if self._fused_kv_quant_method is None:
|
|
|
+ if self._fused_kv_weight is None:
|
|
|
+ raise RuntimeError("DFlash context K/V projection is not initialized.")
|
|
|
+ all_kv_flat = F.linear(
|
|
|
+ normed_context_states, self._fused_kv_weight, self._fused_kv_bias
|
|
|
+ )
|
|
|
+ else:
|
|
|
+ all_kv_flat = self._fused_kv_quant_method.apply(
|
|
|
+ self._fused_kv_linear,
|
|
|
+ normed_context_states,
|
|
|
+ self._fused_kv_bias,
|
|
|
+ )
|
|
|
+ # Single contiguous copy that separates K/V and transposes to
|
|
|
+ # layer-major layout. Result: [2, L, num_ctx, nkv, hd] contiguous.
|
|
|
+ # Indexing dim-0 gives contiguous [L, num_ctx, nkv, hd] for K and V.
|
|
|
+ all_kv = (
|
|
|
+ all_kv_flat.view(num_ctx, num_layers, 2, num_kv_heads, head_dim)
|
|
|
+ .permute(2, 1, 0, 3, 4)
|
|
|
+ .contiguous()
|
|
|
+ )
|
|
|
+ all_k = all_kv[0] # [L, num_ctx, nkv, hd], contiguous
|
|
|
+ all_v = all_kv[1] # [L, num_ctx, nkv, hd], contiguous
|
|
|
+ return all_k, all_v
|
|
|
+
|
|
|
+ def _normalize_context_k(self, all_k: torch.Tensor) -> torch.Tensor:
|
|
|
+ # --- Grouped RMSNorm K across all layers ([L, num_ctx, nkv, hd]) ---
|
|
|
+ # The weight is selected per layer by the outermost (layer) index.
|
|
|
+ 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 precompute_and_store_context_kv(
|
|
|
+ self,
|
|
|
+ context_states: torch.Tensor,
|
|
|
+ context_positions: torch.Tensor,
|
|
|
+ context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
|
|
|
+ ) -> None:
|
|
|
+ """Precompute K/V for context states write them into each layer's KV cache.
|
|
|
+
|
|
|
+ Input context states are projected to K/V, normed, and have RoPE applied.
|
|
|
+ Since the context shape is different than the query shape, we can't rely on the
|
|
|
+ regular forward pass to apply torch.compile and CUDA graphs to this section.
|
|
|
+ As such, this function is optimized to minimize the number of torch ops present:
|
|
|
+ we use fused vLLM kernels for RMSNorm and RoPE, fuse the GEMM into one
|
|
|
+ large projection, and avoid cloning buffers (with .contiguous()) where possible.
|
|
|
+
|
|
|
+ When context_slot_mapping is None (e.g. during dummy_run) only
|
|
|
+ the computation runs, and no K/V is written to cache.
|
|
|
+ """
|
|
|
+ if not hasattr(self, "_num_attn_layers"):
|
|
|
+ logger.warning_once(
|
|
|
+ "DFlash buffer initialization was skipped. If dummy weights are not "
|
|
|
+ "in use, this may indicate an error in weight loading."
|
|
|
+ )
|
|
|
+ self._build_fused_kv_buffers()
|
|
|
+
|
|
|
+ num_ctx = context_states.shape[0]
|
|
|
+ L = self._num_attn_layers
|
|
|
+ kv = self._kv_size
|
|
|
+ hd = self._head_dim
|
|
|
+ nkv = self._num_kv_heads
|
|
|
+
|
|
|
+ all_k, all_v = self._project_context_kv(context_states, num_ctx, L, nkv, hd)
|
|
|
+ all_k_normed = self._normalize_context_k(all_k)
|
|
|
+
|
|
|
+ # --- Fused RoPE across all layers ---
|
|
|
+ # View as [L * num_ctx, kv] so RoPE sees one big batch (no copy).
|
|
|
+ # In-place RoPE: pass K as the "query" arg with key=None.
|
|
|
+ all_k_flat = all_k_normed.view(L * num_ctx, kv)
|
|
|
+ positions_repeated = context_positions.repeat(L)
|
|
|
+ cos_sin_cache = self._rope_cos_sin_cache
|
|
|
+ if cos_sin_cache.dtype != all_k_flat.dtype:
|
|
|
+ cos_sin_cache = cos_sin_cache.to(dtype=all_k_flat.dtype)
|
|
|
+ ops.rotary_embedding(
|
|
|
+ positions_repeated,
|
|
|
+ all_k_flat,
|
|
|
+ None,
|
|
|
+ self._rope_head_size,
|
|
|
+ cos_sin_cache,
|
|
|
+ self._rope_is_neox,
|
|
|
+ )
|
|
|
+
|
|
|
+ if context_slot_mapping is None:
|
|
|
+ return
|
|
|
+
|
|
|
+ # --- Per-layer cache insert ---
|
|
|
+ all_k_final = all_k_flat.view(L, num_ctx, nkv, hd)
|
|
|
+ per_layer = isinstance(context_slot_mapping, (list, tuple))
|
|
|
+ for i in range(L):
|
|
|
+ slot_mapping = (
|
|
|
+ context_slot_mapping[i] if per_layer else context_slot_mapping
|
|
|
+ )
|
|
|
+ if slot_mapping is None:
|
|
|
+ continue # dummy run: skip cache ops
|
|
|
+ attn = self._attn_layers[i]
|
|
|
+ kv_cache = attn.kv_cache
|
|
|
+ attn.impl.do_kv_cache_update(
|
|
|
+ attn,
|
|
|
+ all_k_final[i],
|
|
|
+ all_v[i],
|
|
|
+ kv_cache,
|
|
|
+ slot_mapping,
|
|
|
+ )
|
|
|
+
|
|
|
+ def forward(
|
|
|
+ self,
|
|
|
+ input_ids: torch.Tensor,
|
|
|
+ positions: torch.Tensor,
|
|
|
+ input_embeds: torch.Tensor | None = None,
|
|
|
+ ) -> torch.Tensor:
|
|
|
+ if input_embeds is None:
|
|
|
+ input_embeds = self.embed_input_ids(input_ids)
|
|
|
+
|
|
|
+ hidden_states = input_embeds
|
|
|
+
|
|
|
+ residual = None
|
|
|
+ for layer in self.layers:
|
|
|
+ hidden_states, residual = layer(
|
|
|
+ positions=positions,
|
|
|
+ hidden_states=hidden_states,
|
|
|
+ residual=residual,
|
|
|
+ )
|
|
|
+ hidden_states, _ = self.norm(hidden_states, residual)
|
|
|
+ return hidden_states
|
|
|
+
|
|
|
+ def _preprocess(
|
|
|
+ self, weights: Iterable[tuple[str, torch.Tensor]]
|
|
|
+ ) -> Iterable[tuple[str, torch.Tensor]]:
|
|
|
+ tp_size = get_tensor_model_parallel_world_size()
|
|
|
+ tp_rank = get_tensor_model_parallel_rank()
|
|
|
+ for name, loaded_weight in weights:
|
|
|
+ if "attention_sink_bias" in name:
|
|
|
+ # Sink bias is per-head; shard it across TP ranks like the
|
|
|
+ # attention heads themselves.
|
|
|
+ heads_per_rank = loaded_weight.shape[0] // tp_size
|
|
|
+ loaded_weight = loaded_weight.narrow(
|
|
|
+ 0, tp_rank * heads_per_rank, heads_per_rank
|
|
|
+ )
|
|
|
+ yield name, loaded_weight
|
|
|
+
|
|
|
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
|
+ loader = AutoWeightsLoader(self)
|
|
|
+ return loader.load_weights(
|
|
|
+ self._preprocess(weights), mapper=self.hf_to_vllm_mapper
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
|
|
|
+ model_cls = DFlashQwen3Model
|
|
|
+
|
|
|
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
|
|
+ nn.Module.__init__(self)
|
|
|
+ self.draft_model_config = vllm_config.speculative_config.draft_model_config
|
|
|
+ self.config = self.draft_model_config.hf_config
|
|
|
+ if getattr(self.config, "draft_vocab_size", None) is None:
|
|
|
+ self.config.draft_vocab_size = getattr(self.config, "vocab_size", None)
|
|
|
+ target_layer_num = vllm_config.model_config.get_num_layers(
|
|
|
+ vllm_config.parallel_config
|
|
|
+ )
|
|
|
+ self.model = self.model_cls(
|
|
|
+ vllm_config=vllm_config,
|
|
|
+ prefix=maybe_prefix(prefix, "model"),
|
|
|
+ start_layer_id=target_layer_num,
|
|
|
+ )
|
|
|
+
|
|
|
+ logit_scale = getattr(self.config, "logit_scale", 1.0)
|
|
|
+ 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, scale=logit_scale
|
|
|
+ )
|
|
|
+ target_vocab_size = vllm_config.model_config.get_vocab_size()
|
|
|
+ if self.config.draft_vocab_size != target_vocab_size:
|
|
|
+ self.draft_id_to_target_id = nn.Parameter(
|
|
|
+ torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
|
|
|
+ requires_grad=False,
|
|
|
+ )
|
|
|
+ else:
|
|
|
+ self.draft_id_to_target_id = None
|
|
|
+
|
|
|
+ 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 get_draft_kv_cache_layer_names(self) -> list[str]:
|
|
|
+ return [layer.self_attn.attn.layer_name for layer in self.model.layers]
|
|
|
+
|
|
|
+ def get_draft_attn_causal(self) -> list[bool]:
|
|
|
+ """Per-layer attention causality, aligned with
|
|
|
+ get_draft_kv_cache_layer_names."""
|
|
|
+ return [layer.self_attn.causal for layer in self.model.layers]
|
|
|
+
|
|
|
+ def compute_logits(
|
|
|
+ self,
|
|
|
+ hidden_states: torch.Tensor,
|
|
|
+ ) -> torch.Tensor | None:
|
|
|
+ logits = self.logits_processor(self.lm_head, hidden_states)
|
|
|
+ if self.draft_id_to_target_id is None:
|
|
|
+ return logits
|
|
|
+
|
|
|
+ base = torch.arange(self.config.draft_vocab_size, device=logits.device)
|
|
|
+ targets = base + self.draft_id_to_target_id
|
|
|
+ logits_new = logits.new_full(
|
|
|
+ (logits.shape[0], self.config.vocab_size),
|
|
|
+ float("-inf"),
|
|
|
+ )
|
|
|
+ logits_new[:, targets] = logits
|
|
|
+ return logits_new
|
|
|
+
|
|
|
+ def precompute_and_store_context_kv(
|
|
|
+ self,
|
|
|
+ context_states: torch.Tensor,
|
|
|
+ context_positions: torch.Tensor,
|
|
|
+ context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
|
|
|
+ ) -> None:
|
|
|
+ """Precompute projected + RoPE'd K/V and write to cache."""
|
|
|
+ self.model.precompute_and_store_context_kv(
|
|
|
+ context_states, context_positions, context_slot_mapping
|
|
|
+ )
|
|
|
+
|
|
|
+ def process_weights_after_loading(self) -> None:
|
|
|
+ """Finalize the DFlash fused context projection after linear packing."""
|
|
|
+ self.model.process_weights_after_loading()
|
|
|
+
|
|
|
+ def combine_hidden_states(
|
|
|
+ self,
|
|
|
+ hidden_states: torch.Tensor,
|
|
|
+ ) -> torch.Tensor:
|
|
|
+ if not self.model.use_aux_hidden_state:
|
|
|
+ return hidden_states
|
|
|
+ needs_squeeze = hidden_states.dim() == 1
|
|
|
+ if needs_squeeze:
|
|
|
+ hidden_states = hidden_states.unsqueeze(0)
|
|
|
+ expected = self.model.fc.input_size
|
|
|
+ if hidden_states.shape[-1] != expected:
|
|
|
+ raise ValueError(
|
|
|
+ f"DFlash drafter expects {expected} concatenated aux hidden "
|
|
|
+ f"features but received {hidden_states.shape[-1]}. This usually "
|
|
|
+ "means the draft model's target_layer_ids reference layers that "
|
|
|
+ "do not exist in the target model (incompatible draft/target pair)."
|
|
|
+ )
|
|
|
+ result = self.model.fc(hidden_states)
|
|
|
+ if needs_squeeze:
|
|
|
+ result = result.squeeze(0)
|
|
|
+ return result
|
|
|
+
|
|
|
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
|
|
|
+ model_weights = {}
|
|
|
+ includes_draft_id_mapping = False
|
|
|
+ includes_embed_tokens = False
|
|
|
+ for name, loaded_weight in weights:
|
|
|
+ assert "mask_hidden" not in name, (
|
|
|
+ "DFlash embeds masked slots via mask_token_id (optionally "
|
|
|
+ "overridden by a mask_embedding.pt file); it should not ship a "
|
|
|
+ "mask_hidden weight."
|
|
|
+ )
|
|
|
+ if "t2d" in name:
|
|
|
+ continue
|
|
|
+ if "d2t" in name:
|
|
|
+ name = name.replace("d2t", "draft_id_to_target_id")
|
|
|
+ includes_draft_id_mapping = True
|
|
|
+ elif "lm_head" not in name:
|
|
|
+ name = "model." + name
|
|
|
+ if "embed_tokens" in name:
|
|
|
+ includes_embed_tokens = True
|
|
|
+ model_weights[name] = loaded_weight
|
|
|
+ process_eagle_weight(self, name)
|
|
|
+
|
|
|
+ # Route the separately-trained mask embedding (if shipped) through the
|
|
|
+ # standard weight loader alongside the rest of the draft weights.
|
|
|
+ mask_embedding = self._read_mask_embedding()
|
|
|
+ if mask_embedding is not None:
|
|
|
+ model_weights["model.mask_embedding"] = mask_embedding
|
|
|
+ self.model.has_separate_mask_embedding = True
|
|
|
+
|
|
|
+ orig_to_new_substr = {}
|
|
|
+ if not includes_draft_id_mapping:
|
|
|
+ orig_to_new_substr["draft_id_to_target_id"] = None
|
|
|
+ if not includes_embed_tokens:
|
|
|
+ orig_to_new_substr["embed_tokens"] = None
|
|
|
+ if not self.model.use_aux_hidden_state:
|
|
|
+ orig_to_new_substr["fc."] = None
|
|
|
+ if not self.model.has_separate_mask_embedding:
|
|
|
+ orig_to_new_substr["mask_embedding"] = None
|
|
|
+ mapper = WeightsMapper(orig_to_new_substr=orig_to_new_substr)
|
|
|
+ loader = AutoWeightsLoader(self)
|
|
|
+ loader.load_weights(model_weights.items(), mapper=mapper)
|
|
|
+ self.model._build_fused_kv_buffers()
|
|
|
+
|
|
|
+ def _read_mask_embedding(self) -> torch.Tensor | None:
|
|
|
+ """Checks for an override mask embedding in `mask_embedding.pt` and returns it.
|
|
|
+
|
|
|
+ Some checkpoints ship a separately-trained mask embedding for the mask token,
|
|
|
+ which we use to overwrite the embedding for `mask_token_id`. This helper
|
|
|
+ checks for the file, loads the pytorch tensor, and returns the embedding to use.
|
|
|
+
|
|
|
+ Returns None if the override file is not present.
|
|
|
+ """
|
|
|
+ mask_token_id = self.model.mask_token_id
|
|
|
+ if mask_token_id is None:
|
|
|
+ return None
|
|
|
+
|
|
|
+ MASK_EMBEDDING_FILENAME = "mask_embedding.pt"
|
|
|
+ data = get_hf_file_bytes(
|
|
|
+ MASK_EMBEDDING_FILENAME,
|
|
|
+ self.draft_model_config.model,
|
|
|
+ self.draft_model_config.revision,
|
|
|
+ )
|
|
|
+ if data is None:
|
|
|
+ return None
|
|
|
+
|
|
|
+ state = torch.load(io.BytesIO(data), weights_only=True)
|
|
|
+ if isinstance(state, dict):
|
|
|
+ if state.get("mask_token_id", mask_token_id) != mask_token_id:
|
|
|
+ raise ValueError(
|
|
|
+ f"{MASK_EMBEDDING_FILENAME} mask_token_id does not match "
|
|
|
+ f"dflash_config.mask_token_id ({mask_token_id}). "
|
|
|
+ f"Got {state.get('mask_token_id')}."
|
|
|
+ )
|
|
|
+ state = state["embedding"]
|
|
|
+
|
|
|
+ logger.info(
|
|
|
+ "Loaded DFlash mask embedding for mask_token_id %s from %s",
|
|
|
+ mask_token_id,
|
|
|
+ MASK_EMBEDDING_FILENAME,
|
|
|
+ )
|
|
|
+ return state.reshape(-1)
|
