mirror of
https://github.com/PaddlePaddle/FastDeploy.git
synced 2025-10-04 00:06:38 +08:00
[CUDAGraph]CUDA Graph support unique memory pool (#4230)
* cuda graph use unique memory pool * fix custom device import bug * refine code * refine code * refine code
This commit is contained in:
@@ -841,8 +841,13 @@ class GraphOptimizationConfig:
|
|||||||
Now don't support capture both decode-only and prefill-only"""
|
Now don't support capture both decode-only and prefill-only"""
|
||||||
self.full_cuda_graph: bool = True
|
self.full_cuda_graph: bool = True
|
||||||
|
|
||||||
|
""" Maximum CUDA Graph capture size """
|
||||||
self.max_capture_size: int = None
|
self.max_capture_size: int = None
|
||||||
|
""" Record maps mapped from real shape to captured size to reduce runtime overhead """
|
||||||
self.real_shape_to_captured_size: dict[int, int] = None
|
self.real_shape_to_captured_size: dict[int, int] = None
|
||||||
|
""" Whether to use shared memory pool for multi capture_size """
|
||||||
|
self.use_unique_memory_pool: bool = False
|
||||||
|
|
||||||
# CINN Config ...
|
# CINN Config ...
|
||||||
if args is not None:
|
if args is not None:
|
||||||
for key, value in args.items():
|
for key, value in args.items():
|
||||||
|
@@ -96,6 +96,13 @@ class CudaGraphPiecewiseBackend:
|
|||||||
self.cudagraph_capture_sizes = fd_config.graph_opt_config.cudagraph_capture_sizes
|
self.cudagraph_capture_sizes = fd_config.graph_opt_config.cudagraph_capture_sizes
|
||||||
self.warm_up_size = fd_config.graph_opt_config.cudagraph_num_of_warmups
|
self.warm_up_size = fd_config.graph_opt_config.cudagraph_num_of_warmups
|
||||||
self.real_shape_to_captured_size = fd_config.graph_opt_config.real_shape_to_captured_size
|
self.real_shape_to_captured_size = fd_config.graph_opt_config.real_shape_to_captured_size
|
||||||
|
self.unique_memory_pool_id = None
|
||||||
|
if self.fd_config.graph_opt_config.use_unique_memory_pool:
|
||||||
|
# TODO(gongshaotian): Optimize code
|
||||||
|
if paddle.is_compiled_with_cuda():
|
||||||
|
from paddle.base.core import CUDAGraph
|
||||||
|
|
||||||
|
self.unique_memory_pool_id = CUDAGraph.gen_new_memory_pool_id()
|
||||||
|
|
||||||
self._create_entry_dict()
|
self._create_entry_dict()
|
||||||
|
|
||||||
@@ -169,7 +176,7 @@ class CudaGraphPiecewiseBackend:
|
|||||||
input_addresses = [x.data_ptr() for (_, x) in kwargs.items() if isinstance(x, paddle.Tensor)]
|
input_addresses = [x.data_ptr() for (_, x) in kwargs.items() if isinstance(x, paddle.Tensor)]
|
||||||
entry.input_addresses = input_addresses
|
entry.input_addresses = input_addresses
|
||||||
|
|
||||||
new_grpah = graphs.CUDAGraph()
|
new_grpah = graphs.CUDAGraph(pool_id=self.unique_memory_pool_id)
|
||||||
paddle.device.synchronize()
|
paddle.device.synchronize()
|
||||||
|
|
||||||
# Capture
|
# Capture
|
||||||
|
Reference in New Issue
Block a user