update torch config

Signed-off-by: David9857 <985700846@qq.com>

Author: David9857

vllm_ascend/ascend_config.py

@@ -53,6 +53,8 @@ def __init__(self, torchair_graph_config):
             "graph_batch_sizes", [])
         self.graph_batch_sizes_init = torchair_graph_config.get(
             "graph_batch_sizes_init", False)
+        self.enable_multistream_shared_expert = torchair_graph_config.get(
+            "enable_multistream_shared_expert", False)
 
         if not isinstance(self.graph_batch_sizes, list):
             raise TypeError("graph_batch_sizes must be list[int]")

vllm_ascend/models/deepseek_v2.py

@@ -30,7 +30,6 @@
 import torch
 import torch.distributed as dist
 import torch_npu
-import torchair as tng  # type: ignore
 import vllm.envs as envs
 from torch import nn
 from transformers import PretrainedConfig
@@ -278,10 +277,10 @@ def forward(
             top_k=CustomDeepseekV2MoE.top_k,
             enable_force_load_balance=enable_force_load_balance,
             **kwargs)
-        
+
         if multistream:
             hidden_states, shared_output = hidden_states
-        
+
         hidden_states = hidden_states * self.routed_scaling_factor
 
         if self.tp_size > 1: