[Feature] Reuse JIT code path for building AOT wheel

[abcdabcd987]

Currently we maintain AOT csrc and JIT csrc separately. Since JIT is the priority and AOT is tested less often, AOT falls out of sync with JIT from time to time. This causes lots of issues (e.g., compilation failure, linking failure, missing symbols at runtime, falling back to JIT) and duplicated work must be paid to bring the changes to AOT. This is clearly not sustainable.

On the other hand, we do want to have a wheel that have some common kernels precompiled.

My proposal is: Get rid of AOT csrc. Use JIT code path to generate and compile kernels during wheel build time.

[qiaoning]

Yes，lots of bugs when compiling&run unittest and benchmark.

For example:

1. generate_dispatch_inc.py config.cmake NOT match,
2. Possible Deadlock Detected when execute benchmark, block at pos_encoding_mode=1
3. ld returned 1 exit status
4. unknown file: Failure when run test_xxx
   C++ exception with description "Error in function 'single_mha' at /home/users/flashinfer/src/cpu_reference.h:90: std::vector<dtype_out> cpu_reference::single_mha(const std::vector&, const std::vector<dtype_kv>&, const std::vector<dtype_kv>&, size_t, size_t, size_t, size_t, size_t, bool, flashinfer::QKVLayout, flashinfer::PosEncodingMode, float, float) [with dtype_q = __half; dtype_kv = __half; dtype_out = __half; size_t = long unsigned int] failed to dispatch head_dim 64" thrown in the test body.

and so on...

[yzh119]

Hi @qiaoning yes the C++ side are terribly outdated and needs a thorough cleaning

[diptorupd]

@yzh119 I am taking a look. I will open a RFC once I have some more clarity on a reorganization. The work is part of my overall goal to add AMDGPU support for the kernels.

[diptorupd]

@abcdabcd987 have you already started on a redesign/refactor? I will collaborate with you in that case.

[abcdabcd987]

No I haven't started.

[ur4t]

#748 did part of the job. A translator from Python interface to C++ interface might be necessary.