| import logging |
| from typing import Callable, Optional |
|
|
| import torch |
| from torch import nn |
| from transformers import DynamicCache |
| from transformers.cache_utils import Cache |
| from transformers.modeling_outputs import CausalLMOutputWithPast |
| from transformers.models.qwen3.modeling_qwen3 import ( |
| ALL_ATTENTION_FUNCTIONS, |
| FlashAttentionKwargs, |
| GradientCheckpointingLayer, |
| Qwen3Config, |
| Qwen3MLP, |
| Qwen3PreTrainedModel, |
| Qwen3RMSNorm, |
| Qwen3RotaryEmbedding, |
| eager_attention_forward, |
| rotate_half, |
| ) |
| from typing_extensions import Tuple, Unpack |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): |
| cos = cos.unsqueeze(unsqueeze_dim) |
| sin = sin.unsqueeze(unsqueeze_dim) |
| q_len = q.size(-2) |
| q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :]) |
| k_embed = (k * cos) + (rotate_half(k) * sin) |
| return q_embed, k_embed |
|
|
|
|
| def _to_additive_attention_mask( |
| attention_mask: torch.Tensor, |
| *, |
| query_dtype: torch.dtype, |
| device: torch.device, |
| key_len: int, |
| ) -> torch.Tensor: |
| if attention_mask.ndim == 4: |
| attention_mask = attention_mask[:, :, :, :key_len] |
| if attention_mask.dtype == torch.bool: |
| additive_mask = torch.zeros_like(attention_mask, dtype=query_dtype, device=device) |
| return additive_mask.masked_fill( |
| attention_mask.logical_not().to(device=device), |
| torch.finfo(query_dtype).min, |
| ) |
| return attention_mask.to(device=device, dtype=query_dtype) |
|
|
|
|
| def _build_dflash_causal_attention_mask( |
| *, |
| query: torch.Tensor, |
| key: torch.Tensor, |
| cached_kv_len: int, |
| ctx_len: int, |
| ) -> torch.Tensor: |
| q_len = query.shape[-2] |
| kv_len = key.shape[-2] |
| key_positions = torch.arange(kv_len, device=query.device) |
| query_positions = cached_kv_len + ctx_len + torch.arange(q_len, device=query.device) |
| can_attend = key_positions.unsqueeze(0) <= query_positions.unsqueeze(1) |
| mask = torch.zeros((1, 1, q_len, kv_len), dtype=query.dtype, device=query.device) |
| return mask.masked_fill(can_attend.logical_not().unsqueeze(0).unsqueeze(0), torch.finfo(query.dtype).min) |
|
|
|
|
| class Qwen3DFlashAttention(nn.Module): |
| """Multi-headed attention from 'Attention Is All You Need' paper""" |
|
|
| def __init__(self, config: Qwen3Config, layer_idx: int): |
| super().__init__() |
| self.config = config |
| self.layer_idx = layer_idx |
| self.head_dim = getattr( |
| config, "head_dim", config.hidden_size // config.num_attention_heads |
| ) |
| self.num_key_value_groups = ( |
| config.num_attention_heads // config.num_key_value_heads |
| ) |
| self.scaling = self.head_dim**-0.5 |
| self.attention_dropout = config.attention_dropout |
| dflash_config = getattr(config, "dflash_config", {}) or {} |
| self.is_causal = bool(dflash_config.get("causal_head", False)) |
| self.q_proj = nn.Linear( |
| config.hidden_size, |
| config.num_attention_heads * self.head_dim, |
| bias=config.attention_bias, |
| ) |
| self.k_proj = nn.Linear( |
| config.hidden_size, |
| config.num_key_value_heads * self.head_dim, |
| bias=config.attention_bias, |
| ) |
| self.v_proj = nn.Linear( |
| config.hidden_size, |
| config.num_key_value_heads * self.head_dim, |
| bias=config.attention_bias, |
| ) |
| self.o_proj = nn.Linear( |
| config.num_attention_heads * self.head_dim, |
| config.hidden_size, |
| bias=config.attention_bias, |
| ) |
| self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) |
| self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) |
| self.sliding_window = ( |
| config.sliding_window |
| if config.layer_types[layer_idx] == "sliding_attention" |
| else None |
| ) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| target_hidden: torch.Tensor, |
| position_embeddings: tuple[torch.Tensor, torch.Tensor], |
| attention_mask: Optional[torch.Tensor], |
| inject_target_context: bool = True, |
| past_key_values: Optional[Cache] = None, |
| cache_position: Optional[torch.LongTensor] = None, |
| **kwargs: Unpack[FlashAttentionKwargs], |
| ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: |
| bsz, q_len = hidden_states.shape[:-1] |
| ctx_len = target_hidden.shape[1] if inject_target_context else 0 |
| is_causal = kwargs.pop("is_causal", None) |
| if is_causal is None: |
| is_causal = self.is_causal |
| q = self.q_proj(hidden_states) |
| q = q.view(bsz, q_len, -1, self.head_dim) |
| q = self.q_norm(q).transpose(1, 2) |
| if inject_target_context: |
| k_ctx = self.k_proj(target_hidden) |
| k_noise = self.k_proj(hidden_states) |
| v_ctx = self.v_proj(target_hidden) |
| v_noise = self.v_proj(hidden_states) |
| k = torch.cat([k_ctx, k_noise], dim=1) |
| v = torch.cat([v_ctx, v_noise], dim=1) |
| else: |
| k = self.k_proj(hidden_states) |
| v = self.v_proj(hidden_states) |
| k = k.view(bsz, ctx_len + q_len, -1, self.head_dim) |
| v = v.view(bsz, ctx_len + q_len, -1, self.head_dim) |
| k = self.k_norm(k).transpose(1, 2) |
| v = v.transpose(1, 2) |
| cos, sin = position_embeddings |
| if not inject_target_context: |
| cos = cos[..., -q_len:, :] |
| sin = sin[..., -q_len:, :] |
| q, k = apply_rotary_pos_emb(q, k, cos, sin) |
| cached_kv_len = past_key_values.get_seq_length() if past_key_values is not None else 0 |
| if past_key_values is not None: |
| cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} |
| k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs) |
|
|
| attn_backend = kwargs.pop( |
| "attention_backend", |
| self.config._attn_implementation, |
| ) |
| use_explicit_dflash_causal_mask = bool(is_causal) and attn_backend in {"eager", "sdpa"} |
| if not inject_target_context and attn_backend in {"eager", "sdpa"}: |
| block_local_mask = _build_block_local_causal_attention_mask( |
| query=q, |
| key=k, |
| block_size=self.config.block_size, |
| ) |
| if attention_mask is not None: |
| if attention_mask.shape[-1] != k.shape[-2]: |
| attention_mask = attention_mask[..., -k.shape[-2] :] |
| block_local_mask = block_local_mask + _to_additive_attention_mask( |
| attention_mask, |
| query_dtype=q.dtype, |
| device=q.device, |
| key_len=k.shape[-2], |
| ) |
| attention_mask = block_local_mask |
| is_causal = False |
| elif use_explicit_dflash_causal_mask: |
| dflash_causal_mask = _build_dflash_causal_attention_mask( |
| query=q, key=k, cached_kv_len=cached_kv_len, ctx_len=ctx_len, |
| ) |
| if attention_mask is not None: |
| dflash_causal_mask = dflash_causal_mask + _to_additive_attention_mask( |
| attention_mask, query_dtype=q.dtype, device=q.device, key_len=k.shape[-2], |
| ) |
| attention_mask = dflash_causal_mask |
| is_causal = False |
|
|
| kwargs["is_causal"] = is_causal |
|
|
| attn_fn: Callable = eager_attention_forward |
| if attn_backend != "eager": |
| attn_fn = ALL_ATTENTION_FUNCTIONS[attn_backend] |
| attn_output, attn_weights = attn_fn( |
| self, |
| q, |
| k, |
| v, |
| attention_mask, |
| dropout=0.0 if not self.training else self.attention_dropout, |
| scaling=self.scaling, |
| sliding_window=self.sliding_window, |
| **kwargs, |
| ) |
| attn_output = attn_output.reshape(bsz, q_len, -1) |
| attn_output = self.o_proj(attn_output) |
| return attn_output, attn_weights |
|
|
|
|
| def _resolve_drafter_mlp(config: Qwen3Config) -> bool: |
| dflash_config = getattr(config, "dflash_config", {}) or {} |
| raw_value = dflash_config.get("drafter_mlp", True) |
| if not isinstance(raw_value, bool): |
| raise ValueError( |
| f"dflash_config.drafter_mlp must be a boolean, got {raw_value!r}." |
| ) |
| return raw_value |
|
|
|
|
| def _resolve_inject_layer_ids(config: Qwen3Config) -> Optional[list[int]]: |
| dflash_config = getattr(config, "dflash_config", {}) or {} |
| if "inject_layer_ids" not in dflash_config: |
| return None |
| inject_layer_ids = dflash_config["inject_layer_ids"] |
| if not isinstance(inject_layer_ids, list) or any( |
| not isinstance(layer_idx, int) or isinstance(layer_idx, bool) |
| for layer_idx in inject_layer_ids |
| ): |
| raise ValueError( |
| "dflash_config.inject_layer_ids must be a list of integer layer indices, " |
| f"got {inject_layer_ids!r}." |
| ) |
| invalid_ids = [ |
| layer_idx |
| for layer_idx in inject_layer_ids |
| if layer_idx < 0 or layer_idx >= config.num_hidden_layers |
| ] |
| if invalid_ids: |
| raise ValueError( |
| "dflash_config.inject_layer_ids contains indices outside " |
| f"[0, {config.num_hidden_layers}): {invalid_ids}." |
| ) |
| return inject_layer_ids |
|
|
|
|
| class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer): |
| def __init__(self, config: Qwen3Config, layer_idx: int): |
| super().__init__() |
| self.hidden_size = config.hidden_size |
| self.drafter_mlp = _resolve_drafter_mlp(config) |
| inject_layer_ids = _resolve_inject_layer_ids(config) |
| self.inject_target_context = ( |
| inject_layer_ids is None or layer_idx in inject_layer_ids |
| ) |
| self.self_attn = Qwen3DFlashAttention(config=config, layer_idx=layer_idx) |
| self.mlp = Qwen3MLP(config) |
| self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| self.post_attention_layernorm = Qwen3RMSNorm( |
| config.hidden_size, eps=config.rms_norm_eps |
| ) |
|
|
| def forward( |
| self, |
| target_hidden: Optional[torch.Tensor] = None, |
| hidden_states: Optional[torch.Tensor] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| position_ids: Optional[torch.LongTensor] = None, |
| past_key_value: Optional[Cache] = None, |
| output_attentions: Optional[bool] = False, |
| use_cache: Optional[bool] = False, |
| cache_position: Optional[torch.LongTensor] = None, |
| position_embeddings: Optional[ |
| Tuple[torch.Tensor, torch.Tensor] |
| ] = None, |
| **kwargs: Unpack[FlashAttentionKwargs], |
| ) -> Tuple[ |
| torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] |
| ]: |
| residual = hidden_states |
| hidden_states = self.input_layernorm(hidden_states) |
| hidden_states = self.self_attn( |
| hidden_states=hidden_states, |
| target_hidden=target_hidden, |
| attention_mask=attention_mask, |
| inject_target_context=self.inject_target_context, |
| position_ids=position_ids, |
| past_key_values=past_key_value, |
| output_attentions=output_attentions, |
| use_cache=use_cache, |
| cache_position=cache_position, |
| position_embeddings=position_embeddings, |
| **kwargs, |
| )[0] |
| hidden_states = residual + hidden_states |
| if self.drafter_mlp: |
| residual = hidden_states |
| hidden_states = self.post_attention_layernorm(hidden_states) |
| hidden_states = self.mlp(hidden_states) |
| hidden_states = residual + hidden_states |
| return hidden_states |
|
|
|
|
| def build_target_layer_ids(num_target_layers: int, num_draft_layers: int): |
| if num_draft_layers == 1: |
| return [(num_target_layers // 2)] |
| start = 1 |
| end = num_target_layers - 3 |
| span = end - start |
| target_layer_ids = [ |
| int(round(start + (i * span) / (num_draft_layers - 1))) |
| for i in range(num_draft_layers) |
| ] |
| return target_layer_ids |
|
|
|
|
| def extract_context_feature( |
| hidden_states: list[torch.Tensor], |
| layer_ids: Optional[list[int]], |
| ) -> torch.Tensor: |
| offset = 1 |
| selected_states = [] |
| for layer_id in layer_ids: |
| selected_states.append(hidden_states[layer_id + offset]) |
| target_hidden = torch.cat(selected_states, dim=-1) |
| return target_hidden |
|
|
|
|
| class DFlashDraftModel(Qwen3PreTrainedModel): |
| config_class = Qwen3Config |
| _no_split_modules = ["Qwen3DFlashDecoderLayer"] |
|
|
| def __init__(self, config) -> None: |
| super().__init__(config) |
| self.config = config |
| if not hasattr(self.config, "dflash_config") or self.config.dflash_config is None: |
| self.config.dflash_config = {} |
| self.causal_head = bool(self.config.dflash_config.get("causal_head", False)) |
| self.drafter_mlp = _resolve_drafter_mlp(self.config) |
| self.config.dflash_config["drafter_mlp"] = self.drafter_mlp |
| configured_inject_layer_ids = _resolve_inject_layer_ids(self.config) |
| self.inject_layer_ids = ( |
| list(range(config.num_hidden_layers)) |
| if configured_inject_layer_ids is None |
| else list(configured_inject_layer_ids) |
| ) |
| self.has_masked_inject_layers = set(self.inject_layer_ids) != set( |
| range(config.num_hidden_layers) |
| ) |
| self.layers = nn.ModuleList( |
| [ |
| Qwen3DFlashDecoderLayer(config, layer_idx) |
| for layer_idx in range(config.num_hidden_layers) |
| ] |
| ) |
| dflash_config = getattr(config, "dflash_config", {}) or {} |
| self.target_layer_ids = dflash_config.get( |
| "target_layer_ids", |
| build_target_layer_ids(config.num_target_layers, config.num_hidden_layers), |
| ) |
| self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| self.rotary_emb = Qwen3RotaryEmbedding(config) |
| target_hidden_size = getattr(config, "target_hidden_size", config.hidden_size) |
| self.fc = nn.Linear( |
| len(self.target_layer_ids) * target_hidden_size, |
| config.hidden_size, |
| bias=False, |
| ) |
| self.hidden_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| jacobi_forcing = dflash_config.get("jacobi_forcing", False) |
| if not isinstance(jacobi_forcing, bool): |
| raise ValueError("dflash_config.jacobi_forcing must be a boolean.") |
| self.jacobi_forcing = jacobi_forcing |
| self.stock_drafter = dflash_config.get("stock_drafter", False) |
| if not isinstance(self.stock_drafter, bool): |
| raise ValueError("dflash_config.stock_drafter must be a boolean.") |
| if self.stock_drafter and dflash_config.get("fresh_fraction") != 1.0: |
| raise ValueError("stock_drafter requires an explicit fresh_fraction=1.0.") |
| self.config.dflash_config["stock_drafter"] = self.stock_drafter |
| |
| |
| if self.jacobi_forcing and not self.stock_drafter: |
| self.adapted_position_embedding = nn.Embedding( |
| 2, |
| len(self.target_layer_ids) * target_hidden_size, |
| ) |
| self.block_size = config.block_size |
| self.mask_token_id = dflash_config.get("mask_token_id", None) |
| self.post_init() |
| if self.jacobi_forcing and not self.stock_drafter: |
| nn.init.zeros_(self.adapted_position_embedding.weight) |
| if not self.drafter_mlp: |
| for layer in self.layers: |
| del layer.mlp |
|
|
|
|
| def resolve_causal_head(self, head_type: str = "auto") -> bool: |
| if head_type == "auto": |
| return bool(self.causal_head) |
| if head_type == "bidirectional": |
| return False |
| if head_type == "causal": |
| return True |
| raise ValueError( |
| f"Unsupported head_type={head_type!r}. Expected one of: auto, bidirectional, causal." |
| ) |
|
|
| def forward( |
| self, |
| position_ids: torch.LongTensor, |
| attention_mask: Optional[torch.Tensor] = None, |
| noise_embedding: Optional[torch.Tensor] = None, |
| target_hidden: Optional[torch.Tensor] = None, |
| past_key_values: Optional[Cache] = None, |
| use_cache: bool = False, |
| **kwargs, |
| ) -> CausalLMOutputWithPast: |
| no_inject_attention_mask = kwargs.pop("no_inject_attention_mask", None) |
| attention_backend = kwargs.get( |
| "attention_backend", |
| self.config._attn_implementation, |
| ) |
| hidden_states = noise_embedding |
| target_hidden = self.hidden_norm(self.fc(target_hidden)) |
| position_embeddings = self.rotary_emb(hidden_states, position_ids) |
| for layer_idx, layer in enumerate(self.layers): |
| layer_attention_mask = attention_mask |
| if not layer.inject_target_context and no_inject_attention_mask is not None: |
| layer_attention_mask = no_inject_attention_mask |
| hidden_states = layer( |
| hidden_states=hidden_states, |
| target_hidden=target_hidden, |
| attention_mask=layer_attention_mask, |
| position_ids=position_ids, |
| past_key_value=past_key_values, |
| use_cache=use_cache, |
| position_embeddings=position_embeddings, |
| **kwargs, |
| ) |
| hidden_states = self.norm(hidden_states) |
| return hidden_states |
|
|
|
|
