support deepseek quant & mix-parallel with graphmode (#585)

### What this PR does / why we need it?
1. support deepseek with w8a8 quant;
2. support deepseek with mix-parallel(multi-DP, EP+TP);
3. support deepseek with graphmode.
---------

Signed-off-by: wen-jie666 <wenjie39@huawei.com>
Signed-off-by: Yizhou Liu <liuyizhou5@h-partners.com>
Signed-off-by: libaokui <libaokui@huawei.com>
Signed-off-by: linfeng-yuan <1102311262@qq.com>
Co-authored-by: wen-jie666 <wenjie39@huawei.com>

Author: zzzzwwjj

examples/dp_offline/data_parallel.py

@@ -11,17 +11,15 @@
 import gc
 import os
 
-VLLM_ENABLE_GRAPGH_MODE = os.environ.get("VLLM_ENABLE_GRAPH_MODE") == "1"
-
 
 def main():
     dp_rank = int(os.environ['RANK'])
     local_rank = int(os.environ['LOCAL_RANK'])
     dp_size = int(os.environ['WORLD_SIZE'])
     master_addr = os.environ['MASTER_ADDR']
     master_port = os.environ['MASTER_PORT']
-    tp_size = 4
-    etp_size = 2
+    tp_size = 1
+    etp_size = 1
 
     os.environ["VLLM_DP_RANK"] = str(dp_rank)
     os.environ["VLLM_DP_SIZE"] = str(dp_size)
@@ -58,15 +56,15 @@ def main():
                                      max_tokens=4,
                                      min_tokens=4)
     # Create an LLM.
-    llm = LLM(
-        model="deepseek-ai/DeepSeek-V2-Lite-Chat",
-        tensor_parallel_size=tp_size,
-        trust_remote_code=True,
-        expert_tensor_parallel_size=etp_size,
-        max_model_len=4096,
-        max_num_seqs=num_seqs,
-        compilation_config=1 if VLLM_ENABLE_GRAPGH_MODE else 0,
-    )
+    llm = LLM(model="deepseek-ai/DeepSeek-V2-Lite-Chat",
+              tensor_parallel_size=tp_size,
+              trust_remote_code=True,
+              max_model_len=4096,
+              max_num_seqs=num_seqs,
+              additional_config={
+                  'expert_tensor_parallel_size': etp_size,
+                  'enable_graph_mode': False,
+              })
 
     outputs = llm.generate(prompts, sampling_params)
     for output in outputs:

examples/dp_offline/run_dp.sh

@@ -6,15 +6,13 @@ export HCCL_SOCKET_IFNAME=${ifname}
 # dp_size = node_size * dp_per_node
 node_size=1
 node_rank=0
-dp_per_node=2
+dp_per_node=4
 master_addr=127.0.0.1
 master_port=12345
 
 rm -rf ./.torchair_cache/
 rm -rf ./dynamo_*
 rm -rf /root/ascend/log/debug/plog/*
-export VLLM_ENABLE_GRAPH_MODE=0
-export VLLM_ENABLE_MC2=0
 
 torchrun --nproc_per_node ${dp_per_node} --nnodes ${node_size} \
     --node_rank ${node_rank} --master_addr ${master_addr} --master_port ${master_port} \

vllm_ascend/attention/attention.py

@@ -27,6 +27,7 @@
 except ImportError:
     print("Failed to import torch_npu.")
 
+import torchair._contrib.custom_torch_ops  # type: ignore  # noqa: F401
 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
                                               AttentionLayer,
                                               AttentionMetadata, AttentionType,
@@ -36,9 +37,9 @@
                                            compute_slot_mapping,
                                            compute_slot_mapping_start_idx,
                                            is_block_tables_empty)
+from vllm.config import get_current_vllm_config
 from vllm.utils import async_tensor_h2d, make_tensor_with_pad
 
-from vllm_ascend.utils import VLLM_ENABLE_GRAPH_MODE
 from vllm_ascend.worker.model_runner import (
     ModelInputForNPUBuilder, ModelInputForNPUWithSamplingMetadata)
 
@@ -913,6 +914,12 @@ def __init__(
         self.w_kc = None
         self.w_vc = None
 
+        self.enable_graph_mode = False
+        additional_config = get_current_vllm_config().additional_config
+        if additional_config:
+            self.enable_graph_mode = additional_config.get(
+                "enable_graph_mode", False)
+
     def exec_kv(
         self,
         hidden_states: torch.Tensor,
@@ -1084,7 +1091,7 @@ def forward(
                                                        self.num_heads, -1)
 
         # TODO: Replace the env with more flexible expressions
-        if VLLM_ENABLE_GRAPH_MODE == '1':
+        if self.enable_graph_mode:
             if len(kv_cache) > 0 and kv_cache[0].numel(
             ) > 0 and attn_metadata.num_prefills > 0:
                 slots = attn_metadata.slot_mapping
@@ -1141,7 +1148,7 @@ def forward(
                 )
         elif attn_metadata.decode_metadata:
             assert kv_cache is not None
-            if VLLM_ENABLE_GRAPH_MODE == '1':
+            if self.enable_graph_mode:
                 # TorchAir's shape is [bs, num_heads_per_rank, seq_len, dim]
                 q_nope = q_nope.view(num_tokens, self.num_heads, 1, -1)
                 q_pe = q_pe.view(num_tokens, self.num_heads, 1, -1)

vllm_ascend/models/deepseek_v2.py

@@ -26,13 +26,13 @@
 # """Inference-only DeepseekV2/DeepseekV3 model."""
 
 import os
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import torch
 import torch.distributed as dist
 from torch import nn
 from transformers import PretrainedConfig
-from vllm.attention import Attention
+from vllm.attention import Attention, AttentionMetadata
 from vllm.config import (CacheConfig, ModelConfig, VllmConfig,
                          get_current_vllm_config)
 from vllm.distributed import (get_dp_group, get_pp_group,
@@ -64,7 +64,6 @@
 from vllm.sequence import IntermediateTensors
 
 from vllm_ascend.ops.fused_moe import AscendFusedMoE
-from vllm_ascend.utils import VLLM_ENABLE_GRAPH_MODE
 
 
 class CustomDeepseekV2MoE(nn.Module):
@@ -133,7 +132,7 @@ def __init__(
         vllm_config = get_current_vllm_config()
         self.dp_size = get_dp_group().world_size
         batch_size = vllm_config.scheduler_config.max_num_seqs
-        self.enable_mc2 = int(os.environ.get("VLLM_ENABLE_MC2", 0)) == 1
+        self.enable_mc2 = int(os.environ.get("VLLM_ENABLE_MC2", '0')) == 1
 
         params_dtype = torch.get_default_dtype()
         self.final_hidden_states = torch.zeros(
@@ -309,38 +308,36 @@ def __init__(
 
         self.prefix = prefix
         self.debug_layer_idx = int(self.prefix.split(".")[-2])
-        if VLLM_ENABLE_GRAPH_MODE == "1":
-            self.forward = self.forward_torchair
-        else:
-            self.forward = self.forward_eager  # type: ignore
+        self.enable_graph_mode = False
+        additional_config = get_current_vllm_config().additional_config
+        if additional_config:
+            self.enable_graph_mode = additional_config.get(
+                "enable_graph_mode", False)
 
-    def forward_torchair(self,
-                         positions: torch.Tensor,
-                         hidden_states: torch.Tensor,
-                         kv_cache: torch.Tensor = None,
-                         attn_metadata=None):
+    def forward(
+            self,
+            positions: torch.Tensor,
+            hidden_states: torch.Tensor,
+            kv_cache: Optional[torch.Tensor] = None,
+            attn_metadata: Optional[AttentionMetadata] = None) -> torch.Tensor:
         if self.q_lora_rank is not None:
             ckq = self.q_a_proj(hidden_states)[0]
             hidden_states_or_q_c = self.q_a_layernorm(ckq)
         else:
             hidden_states_or_q_c = hidden_states
-        return self.mla_attn(hidden_states_or_q_c, hidden_states, None,
-                             kv_cache, attn_metadata)
-
-    def forward_eager(self, positions: torch.Tensor,
-                      hidden_states: torch.Tensor):
-        if self.q_lora_rank is not None:
-            ckq = self.q_a_proj(hidden_states)[0]
-            hidden_states_or_q_c = self.q_a_layernorm(ckq)
+        if self.enable_graph_mode:
+            return self.mla_attn.impl.forward(self.mla_attn,
+                                              hidden_states_or_q_c,
+                                              hidden_states, None, kv_cache,
+                                              attn_metadata)
         else:
-            hidden_states_or_q_c = hidden_states
-        kv_c, k_pe = self.kv_a_proj_with_mqa(hidden_states)[0].split(
-            [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
-        kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())
-        return self.mla_attn(hidden_states_or_q_c,
-                             kv_c_normed,
-                             k_pe,
-                             output_shape=hidden_states.shape)
+            kv_c, k_pe = self.kv_a_proj_with_mqa(hidden_states)[0].split(
+                [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+            kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())
+            return self.mla_attn(hidden_states_or_q_c,
+                                 kv_c_normed,
+                                 k_pe,
+                                 output_shape=hidden_states.shape)
 
 
 class CustomDeepseekV2DecoderLayer(DeepseekV2DecoderLayer):
@@ -408,6 +405,54 @@ def __init__(
                                                 eps=config.rms_norm_eps)
         self.routed_scaling_factor = config.routed_scaling_factor
 
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: Optional[torch.Tensor],
+        kv_cache: Optional[torch.Tensor] = None,
+        attn_metadata: Optional[AttentionMetadata] = None,
+    ) -> torch.Tensor:
+        # Self Attention
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(
+                hidden_states, residual)
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+            kv_cache=kv_cache,
+            attn_metadata=attn_metadata,
+        )
+
+        if hidden_states.dtype == torch.float16:
+            # Fix FP16 overflow
+            # We scale both hidden_states and residual before
+            # rmsnorm, and rmsnorm result would not affect by scale.
+            hidden_states *= 1. / self.routed_scaling_factor
+            if self.layer_idx == 0:
+                # The residual is shared by all layers, we only scale it on
+                # first layer.
+                residual *= 1. / self.routed_scaling_factor
+
+        # Fully Connected
+        hidden_states, residual = self.post_attention_layernorm(
+            hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+
+        if isinstance(self.mlp,
+                      DeepseekV2MLP) and hidden_states.dtype == torch.float16:
+            # Fix FP16 overflow
+            # Scaling the DeepseekV2MLP output, it is the input of
+            # input_layernorm of next decoder layer.
+            # The scaling of DeepseekV2MOE output would be done in the forward
+            # of DeepseekV2MOE
+            hidden_states *= 1. / self.routed_scaling_factor
+
+        return hidden_states, residual
+
 
 class CustomDeepseekV2Model(nn.Module):
 
@@ -459,7 +504,9 @@ def forward(
         self,
         input_ids: torch.Tensor,
         positions: torch.Tensor,
-        intermediate_tensors: Optional[IntermediateTensors],
+        kv_caches: Optional[List[torch.Tensor]] = None,
+        attn_metadata: Optional[AttentionMetadata] = None,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
         inputs_embeds: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, IntermediateTensors]:
         if get_pp_group().is_first_rank:
@@ -473,8 +520,13 @@ def forward(
             hidden_states = intermediate_tensors["hidden_states"]
             residual = intermediate_tensors["residual"]
 
-        for layer in self.layers[self.start_layer:self.end_layer]:
-            hidden_states, residual = layer(positions, hidden_states, residual)
+        for i in range(self.start_layer, self.end_layer):
+            layer = self.layers[i]
+            hidden_states, residual = layer(
+                positions, hidden_states, residual,
+                kv_caches[i -
+                          self.start_layer] if kv_caches is not None else None,
+                attn_metadata)
 
         if not get_pp_group().is_last_rank:
             return IntermediateTensors({
@@ -514,6 +566,20 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.make_empty_intermediate_tensors = (
             self.model.make_empty_intermediate_tensors)
 
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        kv_caches: Optional[List[torch.Tensor]] = None,
+        attn_metadata: Optional[AttentionMetadata] = None,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, IntermediateTensors]:
+        hidden_states = self.model(input_ids, positions, kv_caches,
+                                   attn_metadata, intermediate_tensors,
+                                   inputs_embeds)
+        return hidden_states
+
 
 class CustomDeepseekV3ForCausalLM(CustomDeepseekV2ForCausalLM):
     pass