then run the following script (you'll need Nsight Systems CLI installed), which will show whether each kernel was launched as part of a CUDA graph or not:
Hi @ezhulenev,
I confirm that cudaGraph showed up in the nsys reports with the fixes introduced in #14921 and #15021.
We can close this issue.
Many thanks.
This seems to happen because this piece of logic only looks at registrations for generic platform
gpu
: https://github.com/openxla/xla/blob/5e9ce9715ca3c43f8bf81c7ebdcf1c8642fa5a5d/xla/service/gpu/command_buffer_scheduling.cc#L162-L165 Hence custom calls registered forCUDA
are not taken into account.@ezhulenev has suggested offline that the fix might be to do platform name canonicalization more thoroughly.
A quick way to repro is to modify the JAX
cuda_custom_call
test as follows:then run the following script (you'll need Nsight Systems CLI installed), which will show whether each kernel was launched as part of a CUDA graph or not: