link_cuda_headers.sh 2.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  1. #!/usr/bin/env bash
  2. # Put the CUDA headers where nvcc looks for them.
  3. #
  4. # This base ships a trimmed /usr/local/cuda/include: the libraries are all
  5. # present (libnvrtc.so.13 included) but the matching headers are not. The pip
  6. # nvidia-cu13 package carries the full set, so the fix is to link what is
  7. # missing rather than install anything.
  8. #
  9. # The same trimming applies to the development symlinks: libnvrtc.so.13 is
  10. # present but libnvrtc.so is not, and -lnvrtc resolves against the unversioned
  11. # name. (libcudart.so is there, which is why only some links fail.)
  12. #
  13. # Without the headers, FlashInfer's DeepGEMM JIT dies compiling the CUTLASS MoE
  14. # kernels at step 17 of 97:
  15. #
  16. # jit_utils.cuh:21:10: fatal error: nvrtc.h: No such file or directory
  17. #
  18. # and without the symlink it gets to 78 of 78 and dies linking:
  19. #
  20. # /usr/bin/ld: cannot find -lnvrtc
  21. #
  22. # Either way the worker takes SIGKILL and the engine never starts, which reads
  23. # as a hang rather than a build failure.
  24. set -euo pipefail
  25. SRC=/usr/local/lib/python3.12/dist-packages/nvidia/cu13/include
  26. DST=/usr/local/cuda/include
  27. [[ -d "$SRC" ]] || { echo "no $SRC; base image changed" >&2; exit 1; }
  28. linked=0
  29. for h in "$SRC"/*; do
  30. name=$(basename "$h")
  31. [[ -e "$DST/$name" ]] && continue
  32. ln -s "$h" "$DST/$name"
  33. linked=$(( linked + 1 ))
  34. done
  35. # The header that motivated this must now resolve on the default include path.
  36. echo '#include <nvrtc.h>' > /tmp/_nvrtc_probe.c
  37. gcc -fsyntax-only -I"$DST" /tmp/_nvrtc_probe.c
  38. rm -f /tmp/_nvrtc_probe.c
  39. # -l<name> needs the unversioned symlink, which this base ships for some
  40. # libraries and not others.
  41. LIB=/usr/local/cuda/lib64
  42. solinks=0
  43. for so in "$LIB"/lib*.so.[0-9]*; do
  44. base=$(basename "$so"); stem=${base%%.so.*}
  45. [[ -e "$LIB/$stem.so" ]] && continue
  46. ln -s "$so" "$LIB/$stem.so"
  47. solinks=$(( solinks + 1 ))
  48. done
  49. ldconfig 2>/dev/null || true
  50. [[ -e "$LIB/libnvrtc.so" ]] || { echo "libnvrtc.so still missing" >&2; exit 1; }
  51. echo "[cuda-headers] linked $linked headers and $solinks dev symlinks; nvrtc resolves"