[RFC] "Multi-Streaming" in CUDA Graphs

[ezhulenev]

Hi everyone,

We are working on adding multi-streaming to CUDA graphs in XLA. The high level plan is capture in this document.

[MoFHeka]

Any progress?

[ezhulenev]

It's behind the flag today: https://github.com/openxla/xla/blob/main/xla/xla.proto#L615-L616, we are working on more explicit multi-streaming + cuda graphs support on XLA, it's here: https://github.com/openxla/xla/blob/main/xla/service/gpu/execution_stream_assignment.h

[MoFHeka]

Thanks for your reply. May I ask how can I confirm the concurrence regions have been assigned? The after optimized hlo graph (html) seems no different between the flag on and off.

[ezhulenev]

you can run with TF_CPP_VMODULE=gpu_command_buffer=100, it will dump all CUDA graphs to logs (and temp files), they are dot files and you can plot them to see if kernels form a DAG