MLA layer eliminates redundant index operators

Signed-off-by: huiying <chenhuiying4@huawei.com>

Author: huiyingCCCC

tests/multicard/test_offline_inference_distributed.py

@@ -63,3 +63,20 @@ def test_models_distributed_DeepSeek():
             distributed_executor_backend="mp",
     ) as vllm_model:
         vllm_model.generate_greedy(example_prompts, max_tokens)
+
+def test_models_eliminates_index_DeepSeek():
+    os.environ["VLLM_USE_V1"] = "1"
+    example_prompts = [
+        "vLLM is a high-throughput and memory-efficient inference and serving engine for LLMs.",
+        "Briefly describe the major milestones in the development of artificial intelligence from 1950 to 2020.",
+        "Compare and contrast artificial intelligence with human intelligence in terms of processing information.",
+    ]
+    dtype = "half"
+    max_tokens = 5
+    with VllmRunner(
+            "deepseek-ai/DeepSeek-V2-Lite",
+            dtype=dtype,
+            tensor_parallel_size=4,
+            distributed_executor_backend="mp",
+    ) as vllm_model:
+        vllm_model.generate_greedy(example_prompts, max_tokens)

vllm_ascend/attention/mla_v1.py

@@ -465,6 +465,10 @@ def __init__(
             self.enable_graph_mode = additional_config.get(
                 "enable_graph_mode", False)
 
+        self.cos = None
+        self.sin = None
+        self.debug_layer_idx = kwargs.get('debug_layer_idx', 0)
+
     def _v_up_proj_and_o_proj(self, x):
         # Convert from (B, N, L) to (N, B, L)
         x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
@@ -757,18 +761,20 @@ def forward(
             decode_ql_nope, decode_q_pe = \
                 self._q_proj_and_k_up_proj(decode_hs_or_q_c)
             if self.running_in_graph:
-                seq_len = self.rotary_emb.max_position_embeddings
-                cos = self.rotary_emb.cos_cached[:seq_len].to(
-                    dtype=decode_q_pe.dtype)
-                sin = self.rotary_emb.sin_cached[:seq_len].to(
-                    dtype=decode_q_pe.dtype)
-                cos = cos[attn_metadata.decode.input_positions]
-                sin = sin[attn_metadata.decode.input_positions]
-                cos = cos[:, None, None, :]
-                sin = sin[:, None, None, :]
-                decode_q_pe = self.rope_single(decode_q_pe, cos, sin)
+                # During the autoregressive decoding process, the cos and sin values are exactly the same for each layer
+                if self.debug_layer_idx == 0 or self.cos is None or self.sin is None:
+                    seq_len = self.rotary_emb.max_position_embeddings
+                    self.cos = self.rotary_emb.cos_cached[:seq_len].to(
+                        dtype=decode_q_pe.dtype)
+                    self.sin = self.rotary_emb.sin_cached[:seq_len].to(
+                        dtype=decode_q_pe.dtype)
+                    self.cos = self.cos[attn_metadata.decode.input_positions]
+                    self.sin = self.sin[attn_metadata.decode.input_positions]
+                    self.cos = self.cos[:, None, None, :]
+                    self.sin = self.sin[:, None, None, :]
+                decode_q_pe = self.rope_single(decode_q_pe, self.cos, self.sin)
                 decode_k_pe, decode_k_nope = self.exec_kv(
-                    hidden_states_or_kv_c_normed, cos, sin, kv_cache,
+                    hidden_states_or_kv_c_normed, self.cos, self.sin, kv_cache,
                     attn_metadata.slot_mapping)
             else:
                 decode_q_pe[...], decode_k_pe[...] = self.rotary_emb(

vllm_ascend/models/deepseek_v2.py

@@ -370,6 +370,9 @@ def __init__(
             mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
             self.scaling = self.scaling * mscale * mscale
 
+        self.prefix = prefix
+        self.debug_layer_idx = int(self.prefix.split(".")[-2])
+
         # In the MLA backend, kv_cache includes both k_c and
         # pe (i.e. decoupled position embeddings). In particular,
         # the concat_and_cache_mla op requires
@@ -398,10 +401,9 @@ def __init__(
             kv_a_layernorm=self.kv_a_layernorm,
             kv_b_proj=self.kv_b_proj,
             o_proj=self.o_proj,
+            debug_layer_idx=self.debug_layer_idx,
         )
 
-        self.prefix = prefix
-        self.debug_layer_idx = int(self.prefix.split(".")[-2])
         self.enable_graph_mode = False
         additional_config = get_current_vllm_config().additional_config
         if additional_config: