init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

vllm_ascend/patch/worker/patch_deepseek_mtp.py

@@ -0,0 +1,80 @@
+import torch
+import torch.nn as nn
+import vllm
+from transformers import DeepseekV2Config, DeepseekV3Config
+from vllm.config import VllmConfig
+from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP, DeepSeekMultiTokenPredictorLayer
+from vllm.model_executor.models.deepseek_v2 import GlmMoeDsaForCausalLM
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+MTP_ROT_WEIGHT_NAME = "rot.weight"
+
+
+def get_spec_layer_idx_from_weight_name(config: DeepseekV2Config | DeepseekV3Config, weight_name: str) -> int | None:
+    if hasattr(config, "num_nextn_predict_layers") and config.num_nextn_predict_layers > 0:
+        layer_idx = config.num_hidden_layers
+        for i in range(config.num_nextn_predict_layers):
+            if (
+                weight_name.startswith(f"model.layers.{layer_idx + i}.")
+                or weight_name.startswith(MTP_ROT_WEIGHT_NAME)
+                or weight_name.startswith(f"layers.{layer_idx + i}.")
+            ):
+                return layer_idx + i
+    return None
+
+
+class AscendDeepSeekMultiTokenPredictorLayer(DeepSeekMultiTokenPredictorLayer):
+    def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
+        super().__init__(vllm_config, prefix)
+        quant_description = getattr(vllm_config.quant_config, "quant_description", None)
+        self.is_rot_used = quant_description.get("is_rot_used", False) if quant_description is not None else False
+        self.target_model_type = vllm_config.speculative_config.target_model_config.hf_text_config.model_type
+        if self.is_rot_used and self.target_model_type == "glm_moe_dsa":
+            self.rot = nn.Linear(self.config.hidden_size, self.config.hidden_size, bias=False)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        previous_hidden_states: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_index: int = 0,
+    ) -> torch.Tensor:
+        assert inputs_embeds is not None
+        # masking inputs at position 0, as not needed by MTP
+        inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
+        inputs_embeds = self.enorm(inputs_embeds)
+        if self.is_rot_used and self.target_model_type == "glm_moe_dsa":
+            previous_hidden_states = self.rot(previous_hidden_states)
+        previous_hidden_states = self.hnorm(previous_hidden_states)
+
+        hidden_states = self.eh_proj(torch.cat([inputs_embeds, previous_hidden_states], dim=-1))
+
+        hidden_states, residual = self.mtp_block(positions=positions, hidden_states=hidden_states, residual=None)
+        hidden_states = residual + hidden_states  # pre-final-norm (logits hidden)
+        # Recycle the post-final-norm hidden into the next draft step.
+        # compute_logits applies shared_head (== final norm) to the pre-norm
+        # element, so logits and the recycle each get exactly one final-norm.
+        # Matches SGLang's deepseek_nextn.
+        return hidden_states, self.shared_head(hidden_states)
+
+
+class AscendDeepSeekMTP(DeepSeekMTP):
+    def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
+        if name != MTP_ROT_WEIGHT_NAME:
+            return super()._rewrite_spec_layer_name(spec_layer, name)
+        else:
+            return f"model.layers.{spec_layer}.rot.weight"
+
+
+class AscendGlmMoeDsaForCausalLM(GlmMoeDsaForCausalLM):
+    def load_weights(self, weights):
+        loader = AutoWeightsLoader(self, skip_prefixes=[MTP_ROT_WEIGHT_NAME])
+        return loader.load_weights(weights)
+
+
+vllm.model_executor.models.deepseek_v2.get_spec_layer_idx_from_weight_name = get_spec_layer_idx_from_weight_name
+vllm.model_executor.models.deepseek_mtp.get_spec_layer_idx_from_weight_name = get_spec_layer_idx_from_weight_name
+vllm.model_executor.models.deepseek_mtp.DeepSeekMultiTokenPredictorLayer = AscendDeepSeekMultiTokenPredictorLayer
+vllm.model_executor.models.deepseek_mtp.DeepSeekMTP = AscendDeepSeekMTP
+vllm.model_executor.models.deepseek_v2.GlmMoeDsaForCausalLM = AscendGlmMoeDsaForCausalLM