去掉内部关于tgt_attention_mask的一些修改 (#7696)

* commit

* 删掉一些东西

Author: zhoutianzi666

csrc/generation/set_alibi_mask_value.cu

@@ -55,7 +55,6 @@ def get_gencode_flags():
     ext_modules=CUDAExtension(
         sources=[
             "./generation/save_with_output.cc",
-            "./generation/set_mask_value.cu",
             "./generation/set_value_by_flags.cu",
             "./generation/token_penalty_multi_scores.cu",
             "./generation/stop_generation_multi_ends.cu",
@@ -66,7 +65,6 @@ def get_gencode_flags():
             "./generation/transpose_removing_padding.cu",
             "./generation/write_cache_kv.cu",
             "./generation/encode_rotary_qk.cu",
-            "./generation/set_alibi_mask_value.cu",
             "./generation/quant_int8.cu",
             "./generation/dequant_int8.cu",
         ],

csrc/generation/set_mask_value.cu

@@ -375,13 +375,11 @@ def __init__(self, config: PredictorArgument, tokenizer: PretrainedTokenizer):
                 dtype=self.dtype,
             )
 
-        self.tgt_generation_mask = paddle.zeros(
+        self.tgt_generation_mask = paddle.ones(
             shape=[config.batch_size, 1, 1, config.total_max_length],
             dtype=self.dtype,
         )
-        self.arange_tensor_encoder = paddle.zeros(
-            shape=(config.batch_size, 1, config.total_max_length), dtype=self.dtype
-        )
+        self.arange_tensor_encoder = paddle.arange(config.total_max_length, dtype=self.dtype)
 
         if config.export_precache:
             if config.prefix_path:
@@ -427,7 +425,7 @@ def _postprocess(self, predictions):
 
     def _preprocess(self, source):
         self.attention_mask[:] = 0
-        self.tgt_generation_mask[:] = 0
+        self.tgt_generation_mask[:] = 1
         pre_caches_length = 0 if not self.config.export_precache else self.pre_caches[0].shape[-2]
 
         if self.tokenizer.chat_template is not None:
@@ -468,15 +466,6 @@ def _preprocess(self, source):
                         [prefix_attention_mask, post_attention_mask], axis=2
                     )
 
-                if self.config.prefix_path is None:
-                    self.tgt_generation_mask[i, 0, 0, pre_caches_length : length + pre_caches_length] = paddle.ones(
-                        shape=[1, length], dtype=self.config.dtype
-                    )
-                else:
-                    self.tgt_generation_mask[i, 0, 0, : length + pre_caches_length] = paddle.ones(
-                        shape=[1, length + pre_caches_length], dtype=self.config.dtype
-                    )
-
             inputs["tgt_pos"] = self.tgt_pos
         elif "bloom" in self.architectures:
             for i in range(inputs["input_ids"].shape[0]):
@@ -496,20 +485,13 @@ def _preprocess(self, source):
                     self.attention_mask[i, :, :length, : length + pre_caches_length] = paddle.concat(
                         [prefix_attention_mask, post_attention_mask], axis=2
                     )
-                self.arange_tensor_encoder[i, :, : length + pre_caches_length] = paddle.arange(
-                    length + pre_caches_length
-                ).astype(self.config.dtype)
 
-                self.tgt_generation_mask[i, :, 0, : length + pre_caches_length] = paddle.ones(
-                    shape=[1, length + pre_caches_length], dtype=self.config.dtype
-                )
             inputs["tgt_pos"] = inputs["tgt_pos"] + pre_caches_length
             # alibi encoder
             alibi_slopes = get_alibi_slopes(self.model_config.n_head)
             inputs["position_ids"] = paddle.to_tensor(alibi_slopes, dtype="float32")
 
-            alibi = alibi_slopes[..., None] * self.arange_tensor_encoder
-            alibi = alibi[:, :, None, :]
+            alibi = alibi_slopes[None, :, None, None] * self.arange_tensor_encoder
 
             if self.model_config.tensor_parallel_degree > 1:
                 block_size = self.model_config.n_head // self.model_config.tensor_parallel_degree
@@ -534,6 +516,9 @@ def _preprocess(self, source):
                     self.config.total_max_length,
                 ]
             )
+            # only generate valid encoder attention mask, other place set 0.
+            alibi_encoder[i, :, length:, length:] = 0
+
             alibi_decoder = alibi.expand(
                 [
                     self.config.batch_size,
@@ -572,15 +557,6 @@ def _preprocess(self, source):
                         [prefix_attention_mask, post_attention_mask], axis=2
                     )
 
-                if self.config.prefix_path is None:
-                    self.tgt_generation_mask[i, 0, 0, pre_caches_length : length + pre_caches_length] = paddle.ones(
-                        shape=[1, length], dtype="float16"
-                    )
-                else:
-                    self.tgt_generation_mask[i, 0, 0, : length + pre_caches_length] = paddle.ones(
-                        shape=[1, length + pre_caches_length], dtype=self.config.dtype
-                    )
-
         inputs["pre_ids"] = self.pre_ids
         inputs["attention_mask"] = self.attention_mask
         inputs["tgt_generation_mask"] = self.tgt_generation_mask