pytorch 2.4 support

[yzh119]

Checklist

• add pytorch 2.4 and cuda 12.4 to release wheel matrix (ci: add torch 12.4 to the matrix configuration #398 )
• torch.compile compatibility, use new custom operator api

[abcdabcd987]

torch.compile compatibility: #554

[DefinitlyEvil]

About vLLM support:
Hello, there! I am trying to use FlashInfer with vLLM (Torch 2.4.0, CUDA 12.6) self-compiled. vLLM complains about .../torch/_library/infer_schema.py

ValueError: infer_schema(func): Parameter bitorder has an unsupported default value (we only support int, float, bool, None). Please file an issue on GitHub so we can prioritize this. Got func with signature (x: torch.Tensor, bitorder: str = 'big') -> torch.Tensor)

Does it relate to this issue? Thanks.

[abcdabcd987]

Hi @DefinitlyEvil

Thanks for raising this issue.

1. Can you try the main branch of FlashInfer?
2. After you installed the main branch, try FLASHINFER_TEST_TORCH_COMPILE=1 pytest -svx ./tests/test_quantization.py
3. If that doesn't work, can you try install PyTorch nightly?
4. If that doesn't work, can you please provide full logs of the following command: TORCHDYNAMO_REPRO_AFTER="dynamo" TORCH_COMPILE_DEBUG=1 TORCH_LOGS="+dynamo,+inductor" FLASHINFER_TEST_TORCH_COMPILE=1 pytest -svx ./tests/test_quantization.py

[DefinitlyEvil]

Hi @DefinitlyEvil

Thanks for raising this issue.

1. Can you try the main branch of FlashInfer?

2. After you installed the main branch, try `FLASHINFER_TEST_TORCH_COMPILE=1 pytest -svx ./tests/test_quantization.py`

3. If that doesn't work, can you try install PyTorch nightly?

4. If that doesn't work, can you please provide full logs of the following command: `TORCHDYNAMO_REPRO_AFTER="dynamo" TORCH_COMPILE_DEBUG=1 TORCH_LOGS="+dynamo,+inductor" FLASHINFER_TEST_TORCH_COMPILE=1 pytest -svx ./tests/test_quantization.py`

Hi, thanks for investivating it. My server was deployed on a production environment, I had to quickly revert to vLLM 0.6.2 and disable FlashInfer.
I couldn't answer all questions for now (I apologize), I did use main branch, all compilations are good.

As for question 2:

ImportError while loading conftest '/root/flashinfer/tests/conftest.py'.
tests/conftest.py:4: in <module>
    import flashinfer
flashinfer-aot/flashinfer/__init__.py:22: in <module>
    from .cascade import (
flashinfer-aot/flashinfer/cascade.py:21: in <module>
    from .decode import (
flashinfer-aot/flashinfer/decode.py:33: in <module>
    from .prefill import get_batch_prefill_module, get_single_prefill_module
flashinfer-aot/flashinfer/prefill.py:34: in <module>
    from .quantization import packbits, segment_packbits
flashinfer-aot/flashinfer/quantization.py:45: in <module>
    @register_custom_op("flashinfer::packbits", mutates_args=())
/opt/miniconda3/envs/vllm/lib/python3.12/site-packages/torch/_library/custom_ops.py:119: in inner
    schema_str = torch._custom_op.impl.infer_schema(fn, mutates_args)
/opt/miniconda3/envs/vllm/lib/python3.12/site-packages/torch/_library/infer_schema.py:61: in infer_schema
    error_fn(
/opt/miniconda3/envs/vllm/lib/python3.12/site-packages/torch/_library/infer_schema.py:21: in error_fn
    raise ValueError(
E   ValueError: infer_schema(func): Parameter bitorder has an unsupported default value (we only support int, float, bool, None). Please file an issue on GitHub so we can prioritize this. Got func with signature (x: torch.Tensor, bitorder: str = 'big') -> torch.Tensor)

As for 3, I just installed PyTorch 2.6.0 dev version, but is there a make clean command so I can recompile them?