add mc2 mask

Signed-off-by: weiguihua2 <weiguihua2@huawei.com>

Author: weiguihua2

vllm_ascend/attention/mla_v1.py

@@ -11,6 +11,7 @@
 from vllm.config import get_current_vllm_config
 from vllm.model_executor.layers.linear import (LinearBase,
                                                UnquantizedLinearMethod)
+from vllm.platforms import current_platform
 from vllm.utils import cdiv, round_down
 
 from vllm_ascend import envs
@@ -98,6 +99,7 @@ class AscendMLADecodeMetadata:
     attn_mask: Optional[torch.Tensor] = None
     sin: torch.Tensor = None
     cos: torch.Tensor = None
+    mc2_mask: Optional[torch.Tensor] = None
 
 
 @dataclass
@@ -213,6 +215,13 @@ def __init__(self,
         self.cos_cache = None
         self.sin_cache = None
 
+    def generate_activate_mask(self, actual_seqs_num, batch_size):
+        mc2_mask = torch.zeros(batch_size,
+                               dtype=torch.bool,
+                               device=current_platform.device_type)
+        mc2_mask[:actual_seqs_num].fill_(True)
+        return mc2_mask
+
     def reorder_batch(self, input_batch: "InputBatch",
                       scheduler_output: "SchedulerOutput") -> bool:
         # We now want to reorder the batch so that the "decode" requests are at
@@ -355,6 +364,7 @@ def build_torchair_graph_dummy(
                          self.rope_dim,
                          dtype=self.runner.dtype,
                          device=device)
+        mc2_mask = self.generate_activate_mask(num_actual_tokens, num_reqs)
         decode_metadata = AscendMLADecodeMetadata(
             input_positions=input_positions,
             block_table=block_table,
@@ -364,7 +374,8 @@ def build_torchair_graph_dummy(
             attn_mask=self.runner.spec_attn_mask,
             actual_seq_q_lens=self.runner.actual_seq_q_lens[:num_reqs],
             sin=sin,
-            cos=cos)
+            cos=cos,
+            mc2_mask=mc2_mask)
         return self.metadata_cls(  # type: ignore
             num_input_tokens=num_actual_tokens,
             num_actual_tokens=num_actual_tokens,
@@ -545,6 +556,8 @@ def build(
             else:
                 seq_lens_list = seq_lens.tolist()
                 cos, sin = None, None
+            mc2_mask = self.generate_activate_mask(
+                num_actual_tokens, num_reqs + num_reqs_pad_size)
 
             decode_metadata = AscendMLADecodeMetadata(
                 input_positions=input_positions,
@@ -555,7 +568,8 @@ def build(
                 attn_mask=self.runner.spec_attn_mask,
                 actual_seq_q_lens=actual_seq_q_lens,
                 sin=sin,
-                cos=cos)
+                cos=cos,
+                mc2_mask=mc2_mask)
 
         return self.metadata_cls(  # type: ignore
             num_actual_tokens=num_actual_tokens,

vllm_ascend/worker/model_runner_v1.py

@@ -1668,6 +1668,8 @@ def _dummy_run(
                             attn_metadata.decode.input_positions)
                         torch._dynamo.mark_static(attn_metadata.decode.sin)
                         torch._dynamo.mark_static(attn_metadata.decode.cos)
+                        torch._dynamo.mark_static(
+                            attn_metadata.decode.mc2_mask)
                         torch._dynamo.mark_static(attn_metadata.slot_mapping)
                         for kv in self.kv_caches:
                             assert isinstance(