Shaogang Wang

icon company nvidia icon location Changsha, China icon location Machine learning compiler engineer at NVidia.

Results 7 issues of Shaogang Wang

[XLA:GPU] selectively update command buffer opeators

[XLA:GPU] Optimize command buffer update by only operating on commands that have pointer change

3 comment

Head CI test

1 comment

[XLA:GPU] Use cuGraphInstantiateWithParams to instantiate cuda graphs for better debug.

This PR uses cuGraphInstantiateWithParms instead of cuGraphInstantiate to instantiate cuda graph executors, so current command_buffer_cmd_test and command_buffer_thunk_test should cover the changes in this PR.

Add cuda profiler hook

Explicit managing the output buffer of `jax.jit` function

1 comment

I found that for the common training pattern in Jax: `new_state, other_output = jitted_train_step_fn(old_state, other_input) ` Current XLA runtime may assign different backing device memory buffer for `old_state` and `new_state.`...

enhancement

Avoid calling nccl NVTX_PAYLOAD_EVTATTR_SET externally

nccl 2.23.4 has changed the API def of NVTX_PAYLOAD_EVTATTR_SET, this API is not stable, and it should be avoided to call it externally.

awaiting review
comp:xla
size:XS