Customizable SM90 prefill kernels.

[hyhieu]

Added customizable SM90 prefill kernels.

AOT command:

MAX_JOBS=224 \
FLASHINFER_ENABLE_AOT=1 \
TORCH_CUDA_ARCH_LIST=9.0a \
pip install -e . \
  --no-cache-dir --ignore-installed --verbose --force-reinstall

JIT command:

MAX_JOBS=224 \
FLASHINFER_ENABLE_AOT=0 \
TORCH_CUDA_ARCH_LIST=9.0a \
pip install -e . \
  --no-cache-dir --ignore-installed --verbose --force-reinstall

After either command, you can make sure the SM90 kernels work by:

pytest -x tests/test_hopper.py

An example for a customizable_sm90_prefill_kernel has also been added to `tests/test_jit_example.py". You can follow that example to add custom kernels.

[yzh119]

@hyhieu Thanks so much for the contribution! It works great in my environment, left some suggestions.

[yzh119]

LGTM, thank you!