[JIT] No module named 'flashinfer.jit.aot_config' when install from source

[ByronHsu]

1. A clean container from dockerfile
2. Follow Installation to install in JIT mode from source

pip install --no-build-isolation --verbose --editable .

3. When I import flashinfer, suffer from the following error

$ python                                                                                               04:51:33
Python 3.12.8 (main, Jan 14 2025, 22:49:14) [Clang 19.1.6 ] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import flashinfer
/workspaces/flashinfer/.venv/lib/python3.12/site-packages/torch/_subclasses/functional_tensor.py:295: UserWarning: Failed to initialize NumPy: No module named 'numpy' (Triggered internally at ../torch/csrc/utils/tensor_numpy.cpp:84.)
  cpu = _conversion_method_template(device=torch.device("cpu"))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/workspaces/flashinfer/flashinfer/__init__.py", line 17, in <module>
    from .activation import gelu_and_mul as gelu_and_mul
  File "/workspaces/flashinfer/flashinfer/activation.py", line 21, in <module>
    from .jit import gen_act_and_mul_module, has_prebuilt_ops, load_cuda_ops    # noqa: F401
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/flashinfer/flashinfer/jit/__init__.py", line 20, in <module>
    from .aot_config import prebuilt_ops_uri as prebuilt_ops_uri
ModuleNotFoundError: No module named 'flashinfer.jit.aot_config'
>>> 

I found aot_config is generated when aot is enabled. I tried to run the install command for AOT, and aot_config.py was indeed generated.

TORCH_CUDA_ARCH_LIST="9.0" FLASHINFER_ENABLE_AOT=1 pip install --no-build-isolation --verbose --editable .

I could import flashinfer afterwards. However, I think it should also work with JIT install from source?

[yzh119]

Should have been fixed in #757