[Bug] flashinfer 0.2.11+0.2.12 - No module named `nvidia.nvshmem`

[mgoin]

If I flashinfer-python from pypi or try to aot build the package, I get an error from the jit module that nvshem is not found.

During AOT build in vLLM docker build:

[2025-08-18T14:00:14Z] #33 14.41 Traceback (most recent call last):
[2025-08-18T14:00:14Z] #33 14.41   File "<frozen runpy>", line 198, in _run_module_as_main
[2025-08-18T14:00:14Z] #33 14.41   File "<frozen runpy>", line 88, in _run_code
[2025-08-18T14:00:14Z] #33 14.41   File "/vllm-workspace/flashinfer/flashinfer/aot.py", line 669, in <module>
[2025-08-18T14:00:14Z] #33 14.41     main()
[2025-08-18T14:00:14Z] #33 14.41   File "/vllm-workspace/flashinfer/flashinfer/aot.py", line 642, in main
[2025-08-18T14:00:14Z] #33 14.41     jit_specs += gen_all_modules(
[2025-08-18T14:00:14Z] #33 14.41                  ^^^^^^^^^^^^^^^^
[2025-08-18T14:00:14Z] #33 14.41   File "/vllm-workspace/flashinfer/flashinfer/aot.py", line 384, in gen_all_modules
[2025-08-18T14:00:14Z] #33 14.41     gen_nvshmem_module(),
[2025-08-18T14:00:14Z] #33 14.41     ^^^^^^^^^^^^^^^^^^^^
[2025-08-18T14:00:14Z] #33 14.41   File "/vllm-workspace/flashinfer/flashinfer/comm/nvshmem.py", line 15, in gen_nvshmem_module
[2025-08-18T14:00:14Z] #33 14.41     lib_dirs = jit_env.get_nvshmem_lib_dirs()
[2025-08-18T14:00:14Z] #33 14.41                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[2025-08-18T14:00:14Z] #33 14.41   File "/vllm-workspace/flashinfer/flashinfer/jit/env.py", line 84, in get_nvshmem_lib_dirs
[2025-08-18T14:00:14Z] #33 14.41     import nvidia.nvshmem
[2025-08-18T14:00:14Z] #33 14.41 ModuleNotFoundError: No module named 'nvidia.nvshmem'

Locally installing flashinfer and seeing that nvshmem is not installed to my venv

> uv pip install flashinfer-python==0.2.12 --torch-backend=cu128
Resolved 40 packages in 380ms
      Built flashinfer-python==0.2.12
Prepared 7 packages in 32.82s
Installed 7 packages in 14ms
 + cuda-bindings==12.9.2
 + cuda-pathfinder==1.1.0
 + cuda-python==12.9.0
 + flashinfer-python==0.2.12
 + nvidia-cudnn-frontend==1.14.0
 + nvidia-ml-py==12.575.51
 + pynvml==12.0.0
> python                                                        
Python 3.12.3 (main, Feb  4 2025, 14:48:35) [GCC 13.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import nvidia.nvshmem
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'nvidia.nvshmem'

[yzh119]

It's part of nvidia-nvshmem-cu12 library, we didn't explicitly specify this as dependency because it's part of torch's dependency (at least on x86_64).

pip install nvidia-nvshmem-cu12

would resolve the issue.

It might be worthwhile to check why nvidia-nvshmem-cu12 is not installed together with torch.