No module named 'fast_self_multihead_attn'

Hi,
Iwas trying run the fast-MHA example here:
https://github.com/NVIDIA/apex/tree/master/apex/contrib/multihead_attn

I got the following errors when running the performance:

Traceback (most recent call last):
  File "perf_test_multihead_attn.py", line 5, in <module>
    from apex.contrib.multihead_attn import SelfMultiheadAttn
  File "/home/apex/apex/contrib/multihead_attn/__init__.py", line 1, in <module>
    from .self_multihead_attn import SelfMultiheadAttn
  File "/home/apex/apex/contrib/multihead_attn/self_multihead_attn.py", line 9, in <module>
    from .fast_self_multihead_attn_func          import fast_self_attn_func
  File "/home/apex/apex/contrib/multihead_attn/fast_self_multihead_attn_func.py", line 2, in <module>
    import fast_self_multihead_attn
ModuleNotFoundError: No module named 'fast_self_multihead_attn'

The Apex packaged is installed with CPP and CUDA extension turned on. And I was not able to find the definition of this missing function anywhere.

Start with apex/contrib/multihead_attn/fast_self_multihead_attn_func.py and the perf_test_multihead_attn.py example, then inspect how the native extension is built and installed. Re-run the performance example after investigating the missing fast_self_multihead_attn import and confirm that the module imports successfully.