From e57ca04a8a8a923709343b1d01212c85173bf9fe Mon Sep 17 00:00:00 2001 From: chenaoxuan Date: Wed, 11 Mar 2026 11:48:33 +0800 Subject: [PATCH] [Feature] Dflash on Ascend --- vllm/config/speculative.py | 13 +- vllm/model_executor/models/qwen3.py | 4 +- vllm/model_executor/models/qwen3_dflash.py | 503 +++++++++++++++++++++ vllm/model_executor/models/registry.py | 2 +- vllm/transformers_utils/configs/eagle.py | 12 +- vllm/v1/spec_decode/eagle.py | 14 +- 6 files changed, 536 insertions(+), 12 deletions(-) create mode 100644 vllm/model_executor/models/qwen3_dflash.py diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index bf533bf14e55..552827e125e2 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -38,7 +38,8 @@ "mtp", "pangu_ultra_moe_mtp", ] -EagleModelTypes = Literal["eagle", "eagle3", MTPModelTypes] +DFlashModelTypes = Literal["dflash"] +EagleModelTypes = Literal["eagle", "eagle3", MTPModelTypes, DFlashModelTypes] SpeculativeMethod = Literal[ "ngram", "medusa", @@ -161,7 +162,7 @@ def compute_hash(self) -> str: factors: list[Any] = [] # Eagle3 affects the computation graph because it returns intermediate # hidden states in addition to the final hidden state. - factors.append(self.method == "eagle3") + factors.append(self.method in ("eagle3", "dflash")) hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest() return hash_str @@ -341,7 +342,7 @@ def __post_init__(self): ) # Automatically detect the method - if self.method in ("eagle", "eagle3"): + if self.method in ("eagle", "eagle3", "dflash"): pass # examples: # yuhuili/EAGLE-LLaMA3-Instruct-8B @@ -351,6 +352,8 @@ def __post_init__(self): self.method = "eagle" elif "eagle3" in self.draft_model_config.model.lower(): self.method = "eagle3" + elif "dflash" in self.draft_model_config.model.lower(): + self.method = "dflash" elif self.draft_model_config.hf_config.model_type == "medusa": self.method = "medusa" elif self.draft_model_config.hf_config.model_type == "mlp_speculator": @@ -620,7 +623,7 @@ def _verify_args(self) -> Self: eagle3_target_supported = ["llama", "qwen", "minicpm", "gpt_oss"] if ( - self.method == "eagle3" + self.method in ("eagle3", "dflash") and self.target_model_config and not any( supported_model in self.target_model_config.hf_text_config.model_type @@ -635,7 +638,7 @@ def _verify_args(self) -> Self: return self def use_eagle(self) -> bool: - return self.method in ("eagle", "eagle3", "mtp") + return self.method in ("eagle", "eagle3", "mtp", "dflash") def __repr__(self) -> str: method = self.method diff --git a/vllm/model_executor/models/qwen3.py b/vllm/model_executor/models/qwen3.py index 0d0da52ed738..ff1b9ccc10df 100644 --- a/vllm/model_executor/models/qwen3.py +++ b/vllm/model_executor/models/qwen3.py @@ -297,7 +297,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: self.model.aux_hidden_state_layers = layers - def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]: + def get_eagle3_aux_hidden_state_layers(self, method: str | None = None) -> tuple[int, ...]: + if method is not None and method == "dflash": + return [1, 9, 17, 25, 33] num_layers = len(self.model.layers) return (2, num_layers // 2, num_layers - 3) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py new file mode 100644 index 000000000000..579f810db453 --- /dev/null +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -0,0 +1,503 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable +from typing import Any + +import torch +from torch import nn +from transformers import Qwen3Config + +from vllm.attention.layer import Attention +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.logger import init_logger +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.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.multimodal.inputs import NestedTensors +from vllm.transformers_utils.config import set_default_rope_theta +from vllm.attention.backends.abstract import AttentionType + +from .qwen2 import Qwen2MLP as Qwen3MLP +from .qwen3 import Qwen3ForCausalLM +from .utils import ( + AutoWeightsLoader, + extract_layer_index, + get_draft_quant_config, + maybe_prefix, + process_eagle_weight, +) + +logger = init_logger(__name__) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb_safe(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 + + +class Qwen3Attention(nn.Module): + 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, + qkv_bias: bool = False, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + attn_type: str = AttentionType.DECODER, + dual_chunk_attention_config: dict[str, Any] | None = None, + ) -> None: + super().__init__() + 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: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + 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.dual_chunk_attention_config = dual_chunk_attention_config + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=qkv_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position, + rope_parameters=rope_parameters, + dual_chunk_attention_config=dual_chunk_attention_config, + ) + 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, + prefix=f"{prefix}.attn", + attn_type=attn_type, + **{ + "layer_idx": extract_layer_index(prefix), + "dual_chunk_attention_config": dual_chunk_attention_config, + } + if dual_chunk_attention_config + else {}, + ) + 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, + context_states: torch.Tensor, + ) -> torch.Tensor: + assert context_states is not None + + num_context = context_states.shape[0] + + concat_states = torch.cat([context_states, hidden_states], dim=0) + qkv, _ = self.qkv_proj(concat_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim) + k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim) + q = self.q_norm(q_by_head).view(q.shape) + k = self.k_norm(k_by_head).view(k.shape) + + q, k = self.rotary_emb(positions, q, k) + + # Remove context states from Q + q = q[num_context:] + + attn_output = self.attn(q, k, v) + output, _ = self.o_proj(attn_output) + return output + + +class Qwen3DecoderLayer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + *, + config: Qwen3Config, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + layer_idx: int = 0, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + set_default_rope_theta(config, default_theta=1000000) + dual_chunk_attention_config = getattr( + config, "dual_chunk_attention_config", None + ) + attn_type = AttentionType.DECODER + + self.self_attn = Qwen3Attention( + 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, + qkv_bias=getattr(config, "attention_bias", False), + 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, + dual_chunk_attention_config=dual_chunk_attention_config, + ) + 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, + context_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, + context_states=context_states, + ) + + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class Qwen3Model(nn.Module): + 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 + + # Get drafter's quantization config + self.quant_config = get_draft_quant_config(vllm_config) + + drafter_config = getattr(self.config, "eagle_config", {}) + drafter_config.update(getattr(self.config, "dflash_config", {})) + + if drafter_config is not None and "use_aux_hidden_state" in drafter_config: + self.use_aux_hidden_state = drafter_config["use_aux_hidden_state"] + else: + self.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"), + ) + + self.layers = nn.ModuleList( + [ + Qwen3DecoderLayer( + current_vllm_config, + prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"), + config=self.config, + layer_idx=layer_idx, + ) + for layer_idx in range(self.config.num_hidden_layers) + ] + ) + if self.use_aux_hidden_state: + num_features_to_use = self.config.num_hidden_layers + if "layer_ids" in drafter_config: + num_features_to_use = len(drafter_config["layer_ids"]) + if hasattr(self.config, "target_hidden_size"): + fc_input_size = self.config.target_hidden_size * num_features_to_use + else: + fc_input_size = self.config.hidden_size * 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 embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + input_embeds: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if input_embeds is None: + input_embeds = self.embed_input_ids(input_ids) + assert hidden_states.shape[-1] == input_embeds.shape[-1] + + context_states = hidden_states + hidden_states = input_embeds + + residual = None + for layer in self.layers: + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + context_states=context_states, + residual=residual, + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "midlayer." in name: + name = name.replace("midlayer.", "layers.0.") + # Handle kv cache quantization scales + if self.quant_config is not None and ( + scale_name := self.quant_config.get_cache_scale(name) + ): + # Loading kv cache quantization scales + 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 + # Remapping the name FP8 kv-scale + if "scale" in name: + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + 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 DFlashQwen3ForCausalLM(Qwen3ForCausalLM): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + nn.Module.__init__(self) + self.config = vllm_config.speculative_config.draft_model_config.hf_config + # Ensure draft_vocab_size is set + # default to the base vocab size when absent + if getattr(self.config, "draft_vocab_size", None) is None: + base_vocab_size = getattr(self.config, "vocab_size", None) + self.config.draft_vocab_size = base_vocab_size + target_layer_num = vllm_config.model_config.get_num_layers( + vllm_config.parallel_config + ) + + # Store target layer count in draft config for + # proper layer_types indexing in draft models + self.config.target_layer_count = target_layer_num + self.model = Qwen3Model( + vllm_config=vllm_config, 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 + ) + # self.draft_id_to_target_id = nn.Parameter( + # torch.zeros(self.config.draft_vocab_size, dtype=torch.long), + # requires_grad=False, + # ) + 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, + hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.model(input_ids, positions, hidden_states, inputs_embeds) + + 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: + assert logits.shape[1] == self.config.vocab_size, ( + "Expected logits to have shape " + f"(*, {self.config.vocab_size}), but got {logits.shape}" + ) + 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 combine_hidden_states( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + if not self.model.use_aux_hidden_state: + return hidden_states + # combine multiple auxiliary hidden states returned by eagle3 + return self.model.hidden_norm(self.model.fc(hidden_states)) + + 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: + 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) + + skip_substrs = [] + if not includes_draft_id_mapping: + skip_substrs.append("draft_id_to_target_id") + if not includes_embed_tokens: + skip_substrs.append("embed_tokens") + if not self.model.use_aux_hidden_state: + skip_substrs.append("fc.") + loader = AutoWeightsLoader( + self, + skip_prefixes=None, + skip_substrs=skip_substrs, + ) + loader.load_weights(model_weights.items()) \ No newline at end of file diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index d332f5115248..352148d9e6ba 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -426,7 +426,7 @@ "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"), "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"), "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"), - "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), + "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), "LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"), "Eagle3Qwen2_5vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), "Eagle3Qwen3vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), diff --git a/vllm/transformers_utils/configs/eagle.py b/vllm/transformers_utils/configs/eagle.py index ce428e567c84..fb882f56a2c8 100644 --- a/vllm/transformers_utils/configs/eagle.py +++ b/vllm/transformers_utils/configs/eagle.py @@ -60,9 +60,19 @@ def __init__( else f"Eagle3{arch}" for arch in self.model.architectures ] + elif method == "dflash": + assert self.model is not None, ( + "model should not be None when method is dflash" + ) + kwargs["architectures"] = [ + arch + if arch.startswith("DFlash") or arch.endswith("DFlash") + else f"DFlash{arch}" + for arch in self.model.architectures + ] else: raise ValueError( - f"Invalid method {method}. Supported methods are eagle and eagle3." + f"Invalid method {method}. Supported methods are eagle, eagle3 and dflash." ) super().__init__(**kwargs) diff --git a/vllm/v1/spec_decode/eagle.py b/vllm/v1/spec_decode/eagle.py index 65a0a88ec0f5..644601262572 100644 --- a/vllm/v1/spec_decode/eagle.py +++ b/vllm/v1/spec_decode/eagle.py @@ -23,6 +23,7 @@ from vllm.model_executor.models import supports_multimodal from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM +from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.platforms import current_platform from vllm.triton_utils import triton @@ -136,9 +137,14 @@ def __init__( ) else: # RoPE need (max_num_tokens,) - self.positions = torch.zeros( - self.max_num_tokens, dtype=torch.int64, device=device - ) + if self.method == "dflash": + self.positions = torch.zeros( + self.max_num_tokens, dtype=torch.int64, device=device + ) + else: + self.positions = torch.zeros( + 2 * self.max_num_tokens, dtype=torch.int64, device=device + ) self.hidden_states = torch.zeros( (self.max_num_tokens, self.hidden_size), dtype=self.dtype, device=device ) @@ -1226,7 +1232,7 @@ def _get_eagle3_use_aux_hidden_state_from_config(self) -> bool: They might indicate this by setting "use_aux_hidden_state" to False inside the "eagle_config" dict of their hf_config. """ - if self.method != "eagle3": + if self.method != "eagle3" and self.method != "dflash": return False # Assume that eagle3 heads use aux hidden states by default use_aux_hidden_state = True