update

Author: Wanglongzhi2001

csrc/gpu/speculate_decoding_kernels/speculate_save_output.cc

@@ -19,7 +19,7 @@
 #include <sys/types.h>
 #include "paddle/extension.h"
 
-#define MAX_BSZ 512
+#define MAX_BSZ 256
 #define MAX_DRAFT_TOKENS 6
 
 struct msgdata {
@@ -31,7 +31,7 @@ void SpeculateSaveWithOutputMsg(const paddle::Tensor& accept_tokens,
                  const paddle::Tensor& accept_num,
                  const paddle::Tensor& not_need_stop,
                  int64_t rank_id,
-                 const int msg_queue_id) {          
+                 const int msg_queue_id) {        
     if (rank_id > 0) return;
 
     int max_draft_tokens = accept_tokens.shape()[1];
@@ -71,7 +71,7 @@ void SpeculateSaveWithOutputMsg(const paddle::Tensor& accept_tokens,
         }
     }
     if ((msgsnd(msgid, &msg_sed, (MAX_BSZ * MAX_DRAFT_TOKENS + MAX_BSZ + 2) * 4, 0)) == -1) {
-        printf("full msg buffer\n");
+      printf("full msg buffer\n");
     }
     return;
 }
@@ -98,10 +98,3 @@ PD_BUILD_OP(speculate_save_output)
     .Outputs({"x_out"})
     .SetInplaceMap({{"accept_tokens", "x_out"}})
     .SetKernelFn(PD_KERNEL(SpeculateSaveWithOutputMsgStatic));
-
-PD_BUILD_OP(speculate_save_output_dynamic)
-    .Inputs({"accept_tokens", "accept_num", "not_need_stop"})
-    .Attrs({"rank_id: int64_t", "msg_queue_id: int"})
-    .Outputs({"x_out"})
-    .SetInplaceMap({{"accept_tokens", "x_out"}})
-    .SetKernelFn(PD_KERNEL(SpeculateSaveWithOutputMsgDynamic));

csrc/gpu/speculate_decoding_kernels/speculate_step.cu

@@ -27,7 +27,7 @@
 from paddle.base.framework import in_cinn_mode, in_pir_executor_mode, use_pir_api
 from paddle.distributed import fleet
 
-from llm.speculate_decoding.proposer import InferenceWithReferenceProposer
+from llm.speculate_decoding.proposers import InferenceWithReferenceProposer
 from paddlenlp.generation import GenerationConfig, TextIteratorStreamer
 from paddlenlp.peft import LoRAConfig, LoRAModel, PrefixConfig, PrefixModelForCausalLM
 from paddlenlp.taskflow.utils import static_mode_guard
@@ -49,6 +49,8 @@
 
 # Note(@RochardWooSJTU): MAX_BSZ must be the same as definition in get_output / save_output
 MAX_BSZ = 512
+# Note(@Wanglongzhi2001): SPECULATE_MAX_BSZ must be the same as definition in speculate_get_output / speculate_save_output
+SPECULATE_MAX_BSZ = 256
 MAX_DRAFT_TOKENS = 6
 
 
@@ -106,7 +108,7 @@ class PredictorArgument:
         default="fp16",
         metadata={"help": "avx cachekv type. Supported values: fp16,int8"},
     )
-    batch_size: int = field(default=1, metadata={"help": "The batch size of data."})
+    batch_size: int = field(default=10, metadata={"help": "The batch size of data."})
     benchmark: bool = field(
         default=False,
         metadata={
@@ -142,7 +144,7 @@ class PredictorArgument:
             "help": "speculate method, it should be one of ['None', 'autoregressive', 'inference_with_reference']"
         },
     )
-    speculate_max_draft_tokens: int = field(
+    speculate_max_draft_token_num: int = field(
         default=1,
         metadata={"help": "the max length of draft tokens for speculate method."},
     )
@@ -1180,7 +1182,7 @@ def __init__(
         # init speculate components
         if config.speculate_method == "inference_with_reference":
             self.proposer = InferenceWithReferenceProposer(
-                config.speculate_max_draft_tokens,
+                config.speculate_max_draft_token_num,
                 config.speculate_max_ngram_size,
                 config.batch_size,
                 config.max_length,
@@ -1192,7 +1194,7 @@ def __init__(
     def predict(self, input_texts: list[str], return_tokens=False):
         self._preprocess(input_texts)
 
-        # Parameters such as seq_lens_encoder have been set in the preprocessor function, 
+        # Parameters such as seq_lens_encoder have been set in the preprocessor function,
         # then we use them to init the proposer's args
         self.init_proposer_args()
 
@@ -1206,7 +1208,7 @@ def predict(self, input_texts: list[str], return_tokens=False):
             read_res_process.start()
 
         output_tensor = paddle.full(
-            shape=[MAX_BSZ * MAX_DRAFT_TOKENS + MAX_BSZ + 2, 1], fill_value=2, dtype="int64"
+            shape=[SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2, 1], fill_value=2, dtype="int64"
         ).cpu()
         tensor_queue.put(output_tensor)
         if self.tensor_parallel_rank == 0:
@@ -1250,14 +1252,14 @@ def _preprocess(self, input_text: list[str]):
 
     def init_proposer_args(self):
         self.model_inputs["accept_tokens"] = paddle.full(
-            shape=[self.config.batch_size, self.config.speculate_max_draft_tokens + 1], fill_value=0, dtype="int64"
+            shape=[self.config.batch_size, self.config.speculate_max_draft_token_num + 1], fill_value=0, dtype="int64"
         )
         self.model_inputs["accept_num"] = paddle.full(shape=[self.config.batch_size], fill_value=0, dtype="int32")
         self.model_inputs["draft_tokens"] = paddle.full(
-            shape=[self.config.batch_size, self.config.speculate_max_draft_tokens + 1], fill_value=0, dtype="int64"
+            shape=[self.config.batch_size, self.config.speculate_max_draft_token_num + 1], fill_value=0, dtype="int64"
         )
         self.model_inputs["actual_draft_token_num"] = paddle.full(
-            shape=[self.config.batch_size], fill_value=self.config.speculate_max_draft_tokens, dtype="int32"
+            shape=[self.config.batch_size], fill_value=self.config.speculate_max_draft_token_num, dtype="int32"
         )
         if self.config.speculate_method == "inference_with_reference":
             self.proposer.input_ids_cpu = self.model_inputs["input_ids"].cpu()
@@ -1275,7 +1277,7 @@ def __init__(
         # init speculate components
         if config.speculate_method == "inference_with_reference":
             self.proposer = InferenceWithReferenceProposer(
-                config.speculate_max_draft_tokens,
+                config.speculate_max_draft_token_num,
                 config.speculate_max_ngram_size,
                 config.batch_size,
                 config.max_length,
@@ -1285,14 +1287,14 @@ def __init__(
 
     def init_proposer_args(self):
         self.model_inputs["accept_tokens"] = paddle.full(
-            shape=[self.config.batch_size, self.config.speculate_max_draft_tokens + 1], fill_value=0, dtype="int64"
+            shape=[self.config.batch_size, self.config.speculate_max_draft_token_num + 1], fill_value=0, dtype="int64"
         )
         self.model_inputs["accept_num"] = paddle.full(shape=[self.config.batch_size], fill_value=0, dtype="int32")
         self.model_inputs["draft_tokens"] = paddle.full(
-            shape=[self.config.batch_size, self.config.speculate_max_draft_tokens + 1], fill_value=0, dtype="int64"
+            shape=[self.config.batch_size, self.config.speculate_max_draft_token_num + 1], fill_value=0, dtype="int64"
         )
         self.model_inputs["actual_draft_token_num"] = paddle.full(
-            shape=[self.config.batch_size], fill_value=self.config.speculate_max_draft_tokens, dtype="int32"
+            shape=[self.config.batch_size], fill_value=self.config.speculate_max_draft_token_num, dtype="int32"
         )
         if self.config.speculate_method == "inference_with_reference":
             self.proposer.input_ids_cpu = self.model_inputs["input_ids"].cpu()
@@ -1309,7 +1311,7 @@ def predict(self, input_texts: list[str], return_tokens=False):
         s_time = time.time()
         self._preprocess(input_texts)
 
-        # Parameters such as seq_lens_encoder have been set in the preprocessor function, 
+        # Parameters such as seq_lens_encoder have been set in the preprocessor function,
         # then we use them to init the proposer's args
         self.init_proposer_args()
         logger.info(f"preprocess spend {time.time()  -  s_time}")
@@ -1332,7 +1334,7 @@ def predict(self, input_texts: list[str], return_tokens=False):
         if self.tensor_parallel_rank == 0:
             read_res_process.start()
         output_tensor = paddle.full(
-            shape=[MAX_BSZ * MAX_DRAFT_TOKENS + MAX_BSZ + 2, 1], fill_value=2, dtype="int64"
+            shape=[SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2, 1], fill_value=2, dtype="int64"
         ).cpu()
         tensor_queue.put(output_tensor)
         if self.tensor_parallel_rank == 0:
@@ -1505,7 +1507,7 @@ def create_predictor(
                 elif predictor_args.speculate_method is not None:
                     config.max_seq_len = predictor_args.total_max_length
                     config.block_size = predictor_args.block_size
-                    config.speculate_max_draft_tokens = predictor_args.speculate_max_draft_tokens
+                    config.speculate_max_draft_token_num = predictor_args.speculate_max_draft_token_num
                     config.speculate_max_ngram_size = predictor_args.speculate_max_ngram_size
                     config.speculate_verify_window = predictor_args.speculate_verify_window
                     config.speculate_max_candidate_len = predictor_args.speculate_max_candidate_len
@@ -1738,7 +1740,7 @@ def create_predictor(
                 elif predictor_args.speculate_method is not None:
                     config.max_seq_len = predictor_args.total_max_length
                     config.block_size = predictor_args.block_size
-                    config.speculate_max_draft_tokens = predictor_args.speculate_max_draft_tokens
+                    config.speculate_max_draft_token_num = predictor_args.speculate_max_draft_token_num
                     config.speculate_max_ngram_size = predictor_args.speculate_max_ngram_size
                     config.speculate_verify_window = predictor_args.speculate_verify_window
                     config.speculate_max_candidate_len = predictor_args.speculate_max_candidate_len

llm/predict/predictor.py

@@ -11,9 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from __future__ import annotations
 
 import paddle
 from paddlenlp_ops import ngram_match
@@ -43,10 +43,10 @@ class InferenceWithReferenceProposer(Proposer):
     It match tokens in the input and output as draft tokens.
     """
 
-    def __init__(self, max_draft_tokens: int, max_ngram_size: int, max_batch_size: int, max_seq_len: int, **kwargs):
+    def __init__(self, max_draft_token_num: int, max_ngram_size: int, max_batch_size: int, max_seq_len: int, **kwargs):
         """
         Args:
-        max_draft_tokens (int):
+        max_draft_token_num (int):
             Maximum number of tokens a proposer can generate at one time.
             The hyperparameter of k in the paper.
         max_ngram_size (int):
@@ -61,9 +61,15 @@ def __init__(self, max_draft_tokens: int, max_ngram_size: int, max_batch_size: i
         self.max_ngram_size = max_ngram_size
         self.input_ids_len = paddle.zeros(shape=[max_batch_size, 1], dtype="int64").cpu()
         self.max_batch_size = max_batch_size
-        self.max_draft_tokens = max_draft_tokens
+        self.max_draft_token_num = max_draft_token_num
         self.input_ids_cpu = paddle.full(shape=[max_batch_size, max_seq_len], fill_value=1, dtype="int64").cpu()
 
+    def update(self, bid: int, seq_len: int):
+        """
+        Used when inserting a new query to update the length of the input_ids.
+        """
+        self.input_ids_len[bid] = seq_len
+
     def run(self, model_inputs: dict[str, paddle.Tensor], **kargs):
         """
         Use ngram_match to get draft tokens from the input and output.
@@ -84,7 +90,7 @@ def run(self, model_inputs: dict[str, paddle.Tensor], **kargs):
             seq_lens_decoder,
             kargs["real_batch_size"],
             self.max_ngram_size,
-            self.max_draft_tokens,
+            self.max_draft_token_num,
         )
 
         model_inputs["draft_tokens"][:] = draft_tokens.cuda()

llm/speculate_decoding/proposer.py renamed to llm/speculate_decoding/proposers.py

@@ -649,6 +649,11 @@ def __init__(self, config: FusedMultiTransformerConfig):
 
         self.linear = fused_linear
 
+        # used in speculative decoding, if speculate_max_draft_token_num is 1
+        # and speculate_method is None, it will be autogressive decoding.
+        self.speculate_max_draft_token_num = 1
+        self.speculate_method = None
+
     def init_weight(self):
         self.qkv_weights = []
         self.linear_weights = []
@@ -1095,7 +1100,6 @@ def forward(
             kwargs["decoder_block_shape_q"] = 16
             kwargs["max_partition_size"] = 32768
             kwargs["encoder_max_partition_size"] = 32768
-            kwargs["speculate_max_draft_token_num"] = 5
 
             from paddlenlp_ops import get_block_shape_and_split_kv_block
 
@@ -1120,7 +1124,7 @@ def forward(
                 kwargs.get("decoder_block_shape_q", 16),
                 self.num_heads // self.kv_num_heads,
                 kwargs.get("block_size", 64),
-                kwargs["speculate_max_draft_token_num"],
+                self.speculate_max_draft_token_num,
             )
 
         residual_input = src
@@ -2259,9 +2263,9 @@ def compute_attn(
                 kwargs.get("decoder_block_shape_q", 16),
                 kwargs.get("max_partition_size", 32768),
                 kwargs.get("encoder_max_partition_size", 32768),
-                kwargs["speculate_max_draft_token_num"],  # speculate_max_draft_token_num
+                self.speculate_max_draft_token_num,  # speculate_max_draft_token_num
                 True,  # causal
-                False,  # speculate_decoder
+                self.speculate_method is not None,  # speculate_decoder
             )[0]
         else:
             if core.is_compiled_with_xpu():
@@ -2441,9 +2445,9 @@ def compute_attn(
                 kwargs.get("decoder_block_shape_q", 16),
                 kwargs.get("max_partition_size", 32768),
                 kwargs.get("encoder_max_partition_size", 32768),
-                kwargs["speculate_max_draft_token_num"],  # speculate_max_draft_token_num
+                self.speculate_max_draft_token_num,  # speculate_max_draft_token_num
                 True,  # causal
-                False,  # speculate_decoder
+                self.speculate_method is not None,  # speculate_decoder
             )[0]
         else:
             fmha_out = paddle.incubate.nn.functional.block_multihead_attention(
@@ -3258,7 +3262,17 @@ def forward(
         return out, caches
 
 
-class FusedSpeculateMultiTransformer(FusedAppendMultiTransformer):
+class FusedSpeculateMultiTransformer(FusedBlockMultiTransformer):
+    def __init__(
+        self,
+        speculate_max_draft_token_num: int,
+        speculate_method: str = None,
+        config: FusedMultiTransformerConfig = None,
+    ):
+        super().__init__(config)
+        self.speculate_max_draft_token_num = speculate_max_draft_token_num
+        self.speculate_method = speculate_method
+
     def post_process(self, **kwargs):
         embed_dim = self.config.embed_dim
         multi_block_output = kwargs.get("multi_block_output", None)
@@ -3279,12 +3293,18 @@ def post_process(self, **kwargs):
         return out
 
 
-class FusedSpeculateMultiTransformerA8W8(FusedAppendMultiTransformerA8W8):
-    def __init__(self, config: FusedMultiTransformerConfig):
+class FusedSpeculateMultiTransformerA8W8(FusedBlockMultiTransformerA8W8):
+    def __init__(
+        self,
+        speculate_max_draft_token_num: int,
+        speculate_method: str = None,
+        config: FusedMultiTransformerConfig = None,
+    ):
         super().__init__(config)
+        self.speculate_max_draft_token_num = speculate_max_draft_token_num
+        self.speculate_method = speculate_method
 
     def post_process(self, **kwargs):
-        logger.info("use FusedSpeculateMultiTransformerA8W8")
         embed_dim = self.config.embed_dim
         multi_block_output = kwargs.get("multi_block_output", None)
         cum_offsets = kwargs.get("cum_offsets", None)