Role of save_cuda_graph and the uncommented native capture

[fengyuanyu1]

Hello, @ShaoxunZeng I read the code and have two questions:

1. I find the save_cuda_graph at https://github.com/ShaoxunZeng/PyTorch-Medusa/blob/de68b8092d45893e45489b39a219d66b7897c73d/aten/src/ATen/cuda/CUDAGraph.cpp#L216. This function reads the information of CUDAGraph from files generated in offline execution. Right?
2. During online execution, the native capture still exists at https://github.com/ShaoxunZeng/PyTorch-Medusa/blob/de68b8092d45893e45489b39a219d66b7897c73d/aten/src/ATen/cuda/CUDAGraph.cpp#L128. Why Medusa can boot faster with additional reading operations?

Maybe I understand it wrong, can you explain it in more detail?

[ShaoxunZeng]

1. The save_cuda_graph function would save the captured CUDAGraph. In terms of reading the information of CUDAGraph from files generated in offline execution, it is the function: https://github.com/ShaoxunZeng/PyTorch-Medusa/blob/de68b8092d45893e45489b39a219d66b7897c73d/aten/src/ATen/cuda/CUDAGraph.cpp#L1696
2. Medusa would only capture the first layer of the original LLM, in order to get the right kernel address. The overhead of capturing a single layer is much less than capturing a whole model (e.g, 1/40 for a 40-layers model).

[fengyuanyu1]

1. The save_cuda_graph function would save the captured CUDAGraph. In terms of reading the information of CUDAGraph from files generated in offline execution, it is the function: https://github.com/ShaoxunZeng/PyTorch-Medusa/blob/de68b8092d45893e45489b39a219d66b7897c73d/aten/src/ATen/cuda/CUDAGraph.cpp#L1696

Yes, our opinions mean the same thing. Because there is the only call chain of save_cuda_graph->_save_cuda_graph->load_kernel_graph

2. Medusa would only capture the first layer of the original LLM, in order to get the right kernel address. The overhead of capturing a single layer is much less than capturing a whole model (e.g, 1/40 for a 40-layers model).

New information. Thanks for your reply.