NameError: name' flash_attn_with_kvcache 'is not defined
what is the output of `pip freeze`?

 I'm running `flash-attn==2.0.4` and `nemo_toolkit==1.23.0` and 

```
>>> from nemo.collections.nlp.modules.common.megatron.attention import CoreAttention
```

runs without issue This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days. This issue was closed because it has been inactive for 7 days since being marked as stale.
**Describe the bug**

When I use flash attn=2.0.4, running Nemo will result in an error `NameError: name' flash_attn_with_kvcache 'is not defined`

After checking the [code,](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/modules/common/megatron/attention.py) I found that when
```python
try:
  # Flash Attention 1.X
  from flash_attn.bert_padding import pad_input, unpad_input
  from flash_attn.flash_attn_interface import flash_attn_unpadded_func

  HAVE_FLASH_ATTENTION = True
  flash_attn_func = None
```
After success, below

```python
except (ImportError, ModuleNotFoundError):
  try:
    from flash_attn import flash_attn_with_kvcache
  except
    flash_attn_with_kvcache = None
```
I won't execute it
So below
```python
if (
  flash_attn_with_kvcache is not None
  and self.use_flash_attention
  and rotary_pos_emb is not None
  and inference_max_sequence_len
  and not set_inference_key_value_memory
  ）
```
The above error will be reported, I think this is a bug, right?

