The populate_cuda_graph creates an unnecessary CUDA graph

[zcbenz]

In populate_cuda_graph, the cuDNN graph is first populated to a intermediate CUDA graph first, and then the intermediate CUDA graph is added as a child graph to the user's CUDA graph:

cudnn-frontend/include/cudnn_frontend/graph_interface.h

Lines 588 to 604 in be6c079

	// Finally get the backend cuda graph.
	cudaGraph_t backend_cuda_graph;
	// Initialize the cudnn cuda graph.
	// The responsibility to destroy is on the user.
	detail::cu_graph_create(&backend_cuda_graph, 0); // 0 is just what the API says to pass
	_CUDNN_CHECK_CUDNN_ERROR(detail::populate_cuda_graph(handle,
	plans.execution_plans[candidate]->get_raw_desc(),
	variant_pack_descriptor.get_ptr(),
	backend_cuda_graph));
	// Clone BE graph into a graph_node
	// This same call also places the newly created into FE's graph
	// TODO: BE graph is at the end, so put in appropriate dependencies
	cudaGraphNode_t backend_cuda_graph_node;
	detail::cuda_graph_add_child_graph_node(
	&backend_cuda_graph_node, cudnn_cuda_graph, &last_node, last_node != nullptr, backend_cuda_graph);

which I think is just adding overhead and complexity.

When using the populate_cuda_graph APIs, users are always using an intermediate CUDA graph to store the results from cuDNN, i.e.:

CudaGraph cuda_graph(device);
graph.populate_cuda_graph(..., cuda_graph);
cudaGraphNode_t node;
cudaGraphAddChildGraphNode(&node, root_cuda_graph, NULL, 0, cuda_graph));

I guess the current behavior of populate_cuda_graph assumed use case like below so it handles the work of creating an intermediate CUDA graph:

graph.populate_cuda_graph(..., root_cuda_graph);

However when updating the CUDA graph the assumption would break because you can't do things like below:

graph.update_cuda_graph(..., root_cuda_graph);

So to use the native CUDA graph APIs the users always need to create an intermediate CUDA graph themselves, and the final graph would look like this:

Root graph -> Intermediate child graph by user -> Intermediate child graph by cuDNN -> Actual kernel

and the intermediate child graph by cuDNN would be redundant.

Refs ml-explore/mlx#2857.

[egilliam-nv]

Thanks for reaching out!

As currently designed, cuDNN's native CUDA Graph API expects to populate an empty CUDA graph which the caller will typically add as a child node of a larger graph. In fact, since you're using the cuDNN frontend, there will be two levels deep of child graphs:

• the backend (BE) is given an empty graph, and populates it...
• which the frontend (FE) will copy into a child node of its own graph (this is the code snippet you pointed out above), and sometimes surround it with additional kernels/memsets/etc. of its own...
• which the user (framework) will typically copy into a child node of its own graph.

You're right that the idiom we used to embed the graph into a child node may be slightly inefficient because we create a temporary graph, only to copy all the nodes of that graph when we call cudaGraphNodeAddChildNode. It might be better to create the child node first with an empty graph, and then populate that child graph directly. Unfortunately, cudaGraphNodeAddChildNode still requires a temporary graph to be passed to it (even if empty), so we still would need to:

1. Create a temporary graph
2. Call cudaGraphNodeAddChildNode to add that empty graph as a child
3. Delete the temporary graph
4. Call cudaGraphChildGraphNodeGetGraph to get the handle to the cloned graph within the new child node
5. Call the backend to populate that new inner graph

The added complexity doesn't seem worth it (given that these graphs typically contain only a small number of nodes). Did you have a different solution in mind?

As a separate, longer-term change, we are wondering whether the whole use of child graphs was overkill and maybe cuDNN's native CUDA graph API should just populate nodes into an existing graph (which may already contain other nodes), resulting in a flat graph. This would require a breaking API change (or a "v2" API), so we are curious to hear from a framework developer whether this would work better for your needs.

[zcbenz]

I understand the necessity of having to create temporary graph to use the backend cudaGraphNodeAddChildNode API, but I think the work should be done by the user instead of the frontend API, i.e. I propose the change:

- // Finally get the backend cuda graph. 
- cudaGraph_t backend_cuda_graph; 
- // Initialize the cudnn cuda graph. 
- // The responsibility to destroy is on the user. 
- detail::cu_graph_create(&backend_cuda_graph, 0);  // 0 is just what the API says to pass 
-  
 _CUDNN_CHECK_CUDNN_ERROR(detail::populate_cuda_graph(handle, 
                                                      plans.execution_plans[candidate]->get_raw_desc(), 
                                                      variant_pack_descriptor.get_ptr(), 
-                                                     backend_cuda_graph)); 
+                                                     cudnn_cuda_graph)); 
  
- // Clone BE graph into a graph_node 
- // This same call also places the newly created into FE's graph 
- // TODO: BE graph is at the end, so put in appropriate dependencies 
- cudaGraphNode_t backend_cuda_graph_node; 
- detail::cuda_graph_add_child_graph_node( 
-     &backend_cuda_graph_node, cudnn_cuda_graph, &last_node, last_node != nullptr, backend_cuda_graph); 

And then users of frontend API would need to create the temporary graph and add it as child node themselves:

auto temp_graph = create_graph();
graph.populate_cuda_graph(..., temp_graph);
cudaGraphNode_t node;
cudaGraphAddChildGraphNode(&node, root_cuda_graph, NULL, 0, temp_graph));

This is because of the update_cuda_graph API — for performance we usually want to call populate_cuda_graph for the first time when creating the graph, and then only use update_cuda_graph for following calls. Which means that we need to cache the temporary graph passed to the populate API, and then pass it to the update API. (This is also JAX does.)

// First call.
auto temp_graph = create_graph();
graph.populate_cuda_graph(temp_graph);
cache[key] = temp_graph;
...
// Following calls.
auto temp_graph = cache[key];
graph.update_cuda_graph(temp_graph);

So as you can see for this scenario there is no necessarily for the frontend populate_cuda_graph API to create the temporary empty graph, because users will create one.

As a separate, longer-term change, we are wondering whether the whole use of child graphs was overkill and maybe cuDNN's native CUDA graph API should just populate nodes into an existing graph (which may already contain other nodes), resulting in a flat graph. This would require a breaking API change (or a "v2" API), so we are curious to hear from a framework developer whether this would work better for your needs.

From framework's side I think the thing matters most is the overhead when updating a graph. Generally we don't care if populate_cuda_graph is slow as it is usually only called once, we only need the update_cuda_graph API to be as fast as possible. A flat cuda graph would definitely help when we need to analyze the graph for possible optimizations, but we will probably not bother migrating to a new API if it does not bring performance gains.