Consider refactoring CTC greedy decoding
**Is your feature request related to a problem? Please describe.**

As part of my work on #8673 , I have been profiling to make sure we have low latency and high throughput. I have observed that this line can sometimes take https://github.com/NVIDIA/NeMo/blob/57444ae9b63afc5419a16948f7d032a1a9f7dc7f/nemo/collections/asr/parts/submodules/ctc_greedy_decoding.py#L218 several milliseconds. To be concrete, running at batch size 10, the model is taking 167 milliseconds, while decoding is taking 29 milliseconds, with almost all of the time spent on this max() operation. This is documented below.

<img width="1244" alt="image" src="https://github.com/NVIDIA/NeMo/assets/4767568/623c66d5-57b9-4d9d-8e35-62cfad51a702">

It is not acceptable to have model inference be bottlenecked on non-GPU work like this.

I uploaded a gzipped version of the nsight systems file if you would like to view it.

[report25.nsys-rep.gz](https://github.com/NVIDIA/NeMo/files/14950389/report25.nsys-rep.gz)

The most obvious way to handle this situation is to do the "max()" operation on the GPU. I previously did a refactor that sped up CTC decoding by improving copy from GPU to CPU https://github.com/NVIDIA/NeMo/pull/8521, but it seems that is better to copy only after the max() operation is done. I did not do this in my previous PR, because (1) I didn't think it would be a concern and (2) that change is much more intrusive. _greedy_decode_logprobs is built with running on only a single element at a time rather than a batch. Maybe vmap can save us...

BTW, why is this only sometimes slow? I'm not certain, but it probably has to do with the CPU cache. It's working directly on pinned memory. The cache line may be getting evicted only sometimes?
Consider refactoring CTC greedy decoding
**Is your feature request related to a problem? Please describe.**

As part of my work on #8673 , I have been profiling to make sure we have low latency and high throughput. I have observed that this line can sometimes take https://github.com/NVIDIA/NeMo/blob/57444ae9b63afc5419a16948f7d032a1a9f7dc7f/nemo/collections/asr/parts/submodules/ctc_greedy_decoding.py#L218 several milliseconds. To be concrete, running at batch size 10, the model is taking 167 milliseconds, while decoding is taking 29 milliseconds, with almost all of the time spent on this max() operation. This is documented below.

<img width="1244" alt="image" src="https://github.com/NVIDIA/NeMo/assets/4767568/623c66d5-57b9-4d9d-8e35-62cfad51a702">

It is not acceptable to have model inference be bottlenecked on non-GPU work like this.

I uploaded a gzipped version of the nsight systems file if you would like to view it.

[report25.nsys-rep.gz](https://github.com/NVIDIA/NeMo/files/14950389/report25.nsys-rep.gz)

The most obvious way to handle this situation is to do the "max()" operation on the GPU. I previously did a refactor that sped up CTC decoding by improving copy from GPU to CPU https://github.com/NVIDIA/NeMo/pull/8521, but it seems that is better to copy only after the max() operation is done. I did not do this in my previous PR, because (1) I didn't think it would be a concern and (2) that change is much more intrusive. _greedy_decode_logprobs is built with running on only a single element at a time rather than a batch. Maybe vmap can save us...

BTW, why is this only sometimes slow? I'm not certain, but it probably has to do with the CPU cache. It's working directly on pinned memory. The cache line may be getting evicted only sometimes?