AttributeError: module 'flashinfer._kernels' has no attribute 'apply_rope_pos_ids_cos_sin_cache'

[fergusfinn]

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

[yzh119]

Hi @fergusfinn , that bug should have been fixed in #697 and already in the post1 release. But let me double check.