Faster compilation times

[skrider]

Thank you for the awesome project! I am interested in doing some experimentation using this kernel as a base, however the compilation times are quite long. What temporary changes can be made to the codebase to speed it up, selecting only a particular set of kernel template arguments?

I attempted to modify https://github.com/flashinfer-ai/flashinfer/blob/main/python/setup.py#L51-L58 and comment out branches in https://github.com/flashinfer-ai/flashinfer/blob/main/include/flashinfer/utils.cuh, unfortunately still got "symbol not found" errors upon trying to load the C++ extension. Any help would be welcome. Thanks again.

[zhyncs]

Thank you for the awesome project! I am interested in doing some experimentation using this kernel as a base, however the compilation times are quite long. What temporary changes can be made to the codebase to speed it up, selecting only a particular set of kernel template arguments?

I attempted to modify https://github.com/flashinfer-ai/flashinfer/blob/main/python/setup.py#L51-L58 and comment out branches in https://github.com/flashinfer-ai/flashinfer/blob/main/include/flashinfer/utils.cuh, unfortunately still got "symbol not found" errors upon trying to load the C++ extension. Any help would be welcome. Thanks again.

You can use cmake -G Ninja to speedup. And you can decide which kernels you need to compile by enabling or disabling the parameters FLASHINFER_* in https://github.com/flashinfer-ai/flashinfer/blob/main/CMakeLists.txt. Perhaps you don't need all types of kernels when using them.

[yzh119]

Hi @skrider , you can remove some auto-generated files by:

rm -rf csrc/generated/

You can use environment variables to specify which shapes/configurations to compile:

rm -rf build/ # ninja fails to track file changes in `*.cuh` files
FLASHINFER_HEAD_DIMS=64,128 pip install -e .

Hi @zhyncs , thanks for your suggestions, the cmake files are used to compile benchmarks and tests and they are not related to python wheel building.

[esmeetu]

@yzh119 I can't use flashinfer(on commit b3fef8a) when following below commands:

rm -rf csrc/generated/
rm -rf build/
FLASHINFER_HEAD_DIMS=128 pip install -e .

After running above commands, i can succeed to install flashinfer. But when run example script, it will raise undefined symbol error:

Traceback (most recent call last):
  File "/home/roy/flashinfer/python/test.py", line 2, in <module>
    import flashinfer
  File "/home/roy/flashinfer/python/flashinfer/__init__.py", line 17, in <module>
    from .decode import (
  File "/home/roy/flashinfer/python/flashinfer/decode.py", line 31, in <module>
    raise e
  File "/home/roy/flashinfer/python/flashinfer/decode.py", line 22, in <module>
    from . import _kernels
ImportError: /home/roy/flashinfer/python/flashinfer/_kernels.cpython-310-x86_64-linux-gnu.so: undefined symbol: _ZN10flashinfer33BatchDecodeWithPagedKVCacheKernelILb1ELNS_15PosEncodingModeE1ELj2ELj4ELj8ELj8ELj1ELj16ELNS_11PageStorageE0ELNS_9QKVLayoutE1E6__halfS4_iEEvPT9_PT11_NS_10paged_kv_tIXT7_EXT8_ES5_S7_EENS_19kv_partition_info_tIS7_EEPT10_SE_Pffff

If not specific FLASHINFER_HEAD_DIMS, the error is gone and everything is ok.

[ZSL98]

Hi @esmeetu @yzh119 , how long does it take on average to build from source? My installation seems to get stuck:

Looking in indexes: https://pypi.org/simple, https://pypi.ngc.nvidia.com
Obtaining file:///app/LLMs/repos/flashinfer/python
  Preparing metadata (setup.py) ... done
Installing collected packages: flashinfer
  Running setup.py develop for flashinfer

Yes. And I face the same undefined symbol error as @esmeetu

[esmeetu]

@ZSL98 It's up to your cpu performance and MAX_JOBS config. You can see the progress by adding verbose args, like pip install -e . -v.