Cache-aware Conformer Encoder has high latency with ONNX runtime
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

Is ONNX runtime fully supported for the cache-aware Conformer Encoder?
With a side-by-side comparison of calling `asr_model.encoder.forward()` and ONNX encoder `session.run()`, the latency for PyTorch and ONNX version is different, with ONNX 2-3 times slower than the PyTorch version.
Initially, the latencies are similar, but the ONNX version's latency keeps increasing linearly with time. (I was expecting the latency to plateau after reaching a maximum size of `last_channel_cache_size`)

I saw a similar issue reported on https://github.com/microsoft/onnxruntime/issues/11293, wondering if it could be related.
Also, if this is a known issue, does https://github.com/NVIDIA/NeMo/pull/5837 provides a solution?

Model is exported to ONNX as follows

```python
asr_model.encoder.export_cache_support = True
asr_model.encoder.export(opts)
```

PyTorch version

```python
(encoded, encoded_len, cache_last_channel, cache_last_time) = asr_model.encoder.forward(
                                                                                                    audio_signal=feats,
                                                                                                    length=feats_len,
                                                                                                    cache_last_channel=cache_last_channel,
                                                                                                    cache_last_time=cache_last_time,
                                                                                                )
cache_last_channel = cache_last_channel[
                    :, :, -asr_model.encoder.streaming_cfg.last_channel_cache_size :, :
                ]
```
![image](https://user-images.githubusercontent.com/3745694/214763960-c0d6741e-dfd4-4428-b8f8-ab6c97997dee.png)


ONNX version

```python
(encoded, encoded_len, cache_last_channel, cache_last_time) = enc_session.run(
                    encoder_outputs,
                    {'audio_signal': feats.numpy(),
                     'length': feats_len.numpy(),
                     'cache_last_channel': cache_last_channel,
                     'cache_last_time': cache_last_time,
                    }
                )
cache_last_channel = cache_last_channel[
                    :, :, -asr_model.encoder.streaming_cfg.last_channel_cache_size :, :
                ]
```

![image](https://user-images.githubusercontent.com/3745694/214763981-526be54b-413b-4f4f-bbd5-c4428fc305d3.png)


**Steps/Code to reproduce bug**

- Export the cache-aware Conformer Encoder to ONNX using
```python
asr_model.encoder.export_cache_support = True
asr_model.encoder.export(opts)
```
- Pass in random feats and caches and measure latency
- Compare PyTorch and ONNX latencies

**Expected behavior**

ONNX runtime is supposed to be at least as fast (if not faster) than the PyTorch version.
Also, since latency increases linearly with signal length, it suggests that maybe ONNX version is not resetting the cache internally?

**Environment overview (please complete the following information)**

 - Environment location: Bare metal on a GCP VM with Intel Skylake CPU
 - Method of NeMo install: From source

**Environment details**
- OS version: Debian GNU/Linux 10 (buster) x86_64
- PyTorch version: 1.12.1
- Python version: 3.8.15
- onnxruneime version: 1.13.1

**Additional context**
Decoding is done on the CPU.

Please let me know if you need any other details.
Thanks for the help!