When installing from PyPI, the functions apply_rope_pos_ids_cos_sin_cache and apply_rope_pos_ids_cos_sin_cache_inplace don't seem to be available. They do work when installing from source, both with with AOT or JIT compilation.
To reproduce:
pip install flashinfer==0.2.0.post1 -i https://flashinfer.ai/whl/cu124/torch2.4/
import torch
from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
if __name__ == "__main__":
apply_rope_with_cos_sin_cache_inplace(
torch.randn(10, 2, 64, device="cuda", dtype=torch.bfloat16),
torch.randn(10, 2, 64, device="cuda", dtype=torch.bfloat16),
torch.randn(10, 64, device="cuda", dtype=torch.float32),
torch.randn(10, 64, device="cuda", dtype=torch.float32),
torch.arange(10, device="cuda"),
)The result is:
Traceback (most recent call last):
File "/home/fergus/Documents/testing/hello.py", line 6, in <module>
apply_rope_with_cos_sin_cache_inplace(
File "/home/fergus/Documents/testing/.venv/lib/python3.11/site-packages/flashinfer/rope.py", line 1092, in apply_rope_with_cos_sin_cache_inplace
_apply_rope_pos_ids_cos_sin_cache(
File "/home/fergus/Documents/testing/.venv/lib/python3.11/site-packages/torch/_library/custom_ops.py", line 506, in __call__
return self._opoverload(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/fergus/Documents/testing/.venv/lib/python3.11/site-packages/torch/_ops.py", line 667, in __call__
return self_._op(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/fergus/Documents/testing/.venv/lib/python3.11/site-packages/torch/_library/autograd.py", line 98, in autograd_impl
result = Generated.apply(*args, Metadata(keyset, keyword_only_args)) # type: ignore[attr-defined]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/fergus/Documents/testing/.venv/lib/python3.11/site-packages/torch/autograd/function.py", line 574, in apply
return super().apply(*args, **kwargs) # type: ignore[misc]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/fergus/Documents/testing/.venv/lib/python3.11/site-packages/torch/_library/autograd.py", line 40, in forward
result = op.redispatch(keyset & _C._after_autograd_keyset, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/fergus/Documents/testing/.venv/lib/python3.11/site-packages/torch/_ops.py", line 672, in redispatch
return self_._handle.redispatch_boxed(keyset, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/fergus/Documents/testing/.venv/lib/python3.11/site-packages/torch/_library/custom_ops.py", line 494, in adinplaceorview_impl
return self._opoverload.redispatch(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/fergus/Documents/testing/.venv/lib/python3.11/site-packages/torch/_ops.py", line 672, in redispatch
return self_._handle.redispatch_boxed(keyset, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/fergus/Documents/testing/.venv/lib/python3.11/site-packages/torch/_library/custom_ops.py", line 236, in backend_impl
result = self._backend_fns[device_type](*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/fergus/Documents/testing/.venv/lib/python3.11/site-packages/flashinfer/rope.py", line 200, in _apply_rope_pos_ids_cos_sin_cache
get_rope_module().apply_rope_pos_ids_cos_sin_cache(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: module 'flashinfer._kernels' has no attribute 'apply_rope_pos_ids_cos_sin_cache'
python=3.11
flashinfer==0.2.0.post1
torch=2.4.0
When installing from PyPI, the functions
apply_rope_pos_ids_cos_sin_cacheandapply_rope_pos_ids_cos_sin_cache_inplacedon't seem to be available. They do work when installing from source, both with with AOT or JIT compilation.To reproduce:
The result is:
python=3.11
flashinfer==0.2.0.post1
torch=2.4.0