[Others] add assert and only count the actual load in cuda_graph (#5445)
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
Deploy GitHub Pages / deploy (push) Has been cancelled

This commit is contained in:
周周周
2025-12-10 11:22:54 +08:00
committed by GitHub
parent e38709b499
commit 83a9ef51d7
2 changed files with 15 additions and 2 deletions

View File

@@ -130,7 +130,6 @@ class GraphOptBackend:
)
def __call__(self, **kwargs):
self._debug_count_total_step += 1
if not self.fd_config.graph_opt_config.use_cudagraph:
return self.runnable(**kwargs)
if self.cudagraph_piecewise_backend is None:
@@ -140,6 +139,9 @@ class GraphOptBackend:
assert kwargs["forward_meta"].ids_remove_padding is not None
real_shape = kwargs["forward_meta"].ids_remove_padding.shape[0]
if real_shape > 0:
# only count the actual load.
self._debug_count_total_step += 1
if (not kwargs["forward_meta"].step_use_cudagraph) or (real_shape > self.cudagraph_switch_threshold):
return self.dy_runnable(**kwargs)