NameError: name' flash_attn_with_kvcache 'is not defined

[1074224619]

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, I found that when

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

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

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?

[csn1011]

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