[Excutor] Experiment-Support Prefill in cudagraph #3459

littledgg · 2025-08-18T13:48:07Z

目前支持Prefill-Only的batch进cudagraph。在确定graph可以共用之前，只能选择要么capture decode-only的，要么capture prefill-only。
1.想要开启，需要使用以下参数启动，重点是use_cudagraph和cudagraph_only_prefill都设为True

python -m fastdeploy.entrypoints.openai.api_server --model ${model_path} \
    --max-num-seqs 64 --max-model-len 32768 \
    --port 8988 --engine-worker-queue-port 7732 \
    --metrics-port 7733 --tensor-parallel-size 1 \
    --graph-optimization-config ' {"use_cudagraph":true,"cudagraph_only_prefill":true}' \

2.在当前动态插入的背景下，假设发送4个80 tokens的prompt，那么seq_lens_this_time第一轮是[80]，第二轮是[1, 80, 80, 80]，很明显只有第一轮是纯P，可以进cudagraph，第二轮就是MIX了，进不了cudagraph，可以通过修改fastdeploy/engine/engine.py中的函数_insert_task_to_worker中

tasks = self.scheduler.get_requests(
                    available_blocks=self.resource_manager.available_block_num(),
                    block_size=self.cfg.cache_config.block_size,
                    reserved_output_blocks=self.cfg.cache_config.enc_dec_block_num,
                    max_num_batched_tokens=self.cfg.max_num_batched_tokens,
                    batch=num_prefill_batch,
                )

改为

                tasks = list()
                while (len(tasks) < 8):
                    print("===RyanDebug, Begin to collect tasks ===")
                    print("====The self.resource_manager.available_block_num is:", self.resource_manager.available_block_num())
                    print("====The self.cfg.cache_config.block_size is:", self.cfg.cache_config.block_size)
                    print("====The self.cfg.cache_config.enc_dec_block_num is:", self.cfg.cache_config.enc_dec_block_num)
                    print("====The self.cfg.max_num_batched_tokens is:", self.cfg.max_num_batched_tokens)
                    print("===RyanDebug, num_prefill_batch is: ",8)

                    tmp_task = self.scheduler.get_requests(
                        available_blocks=5000,
                        block_size=self.cfg.cache_config.block_size,
                        reserved_output_blocks=self.cfg.cache_config.
                        enc_dec_block_num,
                        max_num_batched_tokens=self.cfg.max_num_batched_tokens,
                        batch=8)
                    print("===RyanDebug, the tmp_task is :", tmp_task)

                    if isinstance(tmp_task, list):
                        tasks.extend(tmp_task)
                    elif tmp_task is not None:
                        tasks.append(tmp_task)

                print("===RyanDebug, Finish Fix task, the len of tasks is {} ===", len(tasks))
                print("===RyanDebug, Finish Fix task, the tasks is {} ===", tasks)

这样就是不开启动态插入的逻辑，需要等待8个prompt来（数字可更改），这8个prompt才会一起进入prefill（多个prompt纯P加速），一起进入decode。
3.在fastdeploy/config.py的init_with_cudagrpah_size中，512为capture prefill时最大capture size，可以手动更改。

if self.graph_opt_config.cudagraph_only_prefill:
            self.graph_opt_config.init_with_cudagrpah_size(max_capture_size=512)

TODO：buffer_size的大小需要进一步确认。

…into draft

paddle-bot · 2025-08-18T13:48:14Z

Thanks for your contribution!

gongshaotian · 2025-08-18T13:58:30Z

fastdeploy/model_executor/layers/attention/append_attn_backend.py

+        self.share_inputs["encoder_batch_ids"] = paddle.full(
+            shape=[self.max_seq_len], fill_value=0, dtype="int32"
+        )  # gpu
+        self.share_inputs["encoder_tile_ids_per_batch"] = paddle.full(
+            shape=[self.max_seq_len], fill_value=0, dtype="int32"
+        )  # gpu
+        self.share_inputs["encoder_num_blocks"] = paddle.full(shape=[1], fill_value=0, dtype="int32").cpu()  # cpu


这些 buffer 改为在 gpu_model_runner 里面管理，同时改造下 get_block_shape_and_split_kv 这个Kernel，把 encoder 相关的tensor 改为 Inplace 的实现，不然前处理的 copy_ 开销太高了

暂时提交了一版还没验证过的草稿版，怕服务器挂代码没了

gongshaotian · 2025-08-18T13:59:11Z

fastdeploy/model_executor/layers/attention/append_attn_backend.py

+        self.share_inputs["encoder_batch_ids"].copy_(temp_encoder_batch_ids, False)
+        metadata.encoder_batch_ids = self.share_inputs["encoder_batch_ids"]
+
+        self.share_inputs["encoder_tile_ids_per_batch"].copy_(temp_encoder_tile_ids_per_batch, False)
+        metadata.encoder_tile_ids_per_batch = self.share_inputs["encoder_tile_ids_per_batch"]
+
+        self.share_inputs["encoder_num_blocks"].copy_(temp_encoder_num_blocks, False)
+        metadata.encoder_num_blocks = self.share_inputs["encoder_num_blocks"]
+
+        self.share_inputs["kv_batch_ids"].copy_(temp_kv_batch_ids, False)
+        metadata.kv_batch_ids = self.share_inputs["kv_batch_ids"]
+
+        self.share_inputs["kv_tile_ids_per_batch"].copy_(temp_kv_tile_ids_per_batch, False)
+        metadata.kv_tile_ids_per_batch = self.share_inputs["kv_tile_ids_per_batch"]
+
+        self.share_inputs["kv_num_blocks"].copy_(temp_kv_num_blocks, False)
+        metadata.kv_num_blocks = self.share_inputs["kv_num_blocks"]
+
+        self.share_inputs["max_len_kv"].copy_(temp_max_len_kv, False)
+        metadata.max_len_kv = self.share_inputs["max_len_kv"]
+


copy_ 开销太高了

gongshaotian · 2025-08-18T14:00:22Z