[CudaGraph] support cudagraph use shared pool (#4199)
Some checks failed
CE Compile Job / ce_job_pre_check (push) Has been cancelled
CE Compile Job / print_ce_job_pre_check_outputs (push) Has been cancelled
CE Compile Job / FD-Clone-Linux (push) Has been cancelled
CE Compile Job / Show Code Archive Output (push) Has been cancelled
CE Compile Job / BUILD_SM8090 (push) Has been cancelled
CE Compile Job / BUILD_SM8689 (push) Has been cancelled
CE Compile Job / CE_UPLOAD (push) Has been cancelled

* support cudagraph use shared pool

* add envs

* change CUDAGRAPH_POOL_ID to int

* change CUDAGRAPH_POOL_ID to use_memory_pool

* unify use_unique_memory_pool

* fix use_unique_memory_pool
This commit is contained in:
lizhenyun01
2025-09-24 21:32:04 +08:00
committed by GitHub
parent e2b68b33c9
commit bab779011c
2 changed files with 12 additions and 7 deletions

View File

@@ -588,6 +588,9 @@ class GraphOptimizationConfig:
Thus this flag cannot be used together with splitting_ops."""
self.full_cuda_graph: bool = True
""" Whether to use shared memory pool for multi capture_size """
self.use_unique_memory_pool: bool = False
self.max_capture_size: int = None
self.real_shape_to_captured_size: dict[int, int] = None
# CINN Config ...