python3Packages.vllm: fix rocm build (#495680)

pkgs/development/python-modules/vllm/0006-drop-rocm-extra-reqs.patch

@@ -0,0 +1,20 @@
+diff --git a/requirements/rocm.txt b/requirements/rocm.txt
+index 375f0a019..04a59dc3b 100644
+--- a/requirements/rocm.txt
++++ b/requirements/rocm.txt
+@@ -9,4 +9,2 @@ ray[cgraph]>=2.48.0 # Ray Compiled Graph, required for pipeline parallelism in V1.
+ peft
+-pytest-asyncio
+-tensorizer==2.10.1
+ packaging>=24.2
+@@ -12,6 +10,2 @@ tensorizer==2.10.1
+ packaging>=24.2
+-setuptools>=77.0.3,<80.0.0
+-setuptools-scm>=8
+-runai-model-streamer[s3,gcs]==0.15.3
+-conch-triton-kernels==1.2.1
+ timm>=1.0.17
+ grpcio-tools>=1.76.0
+@@ -17,2 +13,1 @@ conch-triton-kernels==1.2.1
+ timm>=1.0.17
+-grpcio-tools>=1.76.0

pkgs/development/python-modules/vllm/default.nix

@@ -27,12 +27,14 @@
 
   # dependencies
   aioprometheus,
+  amdsmi,
   anthropic,
   bitsandbytes,
   blake3,
   cachetools,
   cbor2,
   compressed-tensors,
+  datasets,
   depyf,
   einops,
   fastapi,
@@ -58,6 +60,7 @@
   outlines,
   pandas,
   partial-json-parser,
+  peft,
   prometheus-fastapi-instrumentator,
   py-cpuinfo,
   pyarrow,
@@ -70,6 +73,7 @@
   sentencepiece,
   setproctitle,
   tiktoken,
+  timm,
   tokenizers,
   torch,
   torchaudio,
@@ -294,7 +298,7 @@ let
     else if cudaSupport then
       gpuArchWarner supportedCudaCapabilities unsupportedCudaCapabilities
     else if rocmSupport then
-      rocmPackages.clr.gpuTargets
+      rocmPackages.clr.localGpuTargets or rocmPackages.clr.gpuTargets
     else
       throw "No GPU targets specified"
   );
@@ -311,6 +315,13 @@ let
     libcublas
   ];
 
+  # header path ends up missing rocthrust & its deps
+  rocmExtraIncludeFlags = lib.concatMapStringsSep " " (pkg: "-I${lib.getInclude pkg}/include") [
+    rocmPackages.rocthrust
+    rocmPackages.rocprim
+    rocmPackages.hipcub
+  ];
+
   # Some packages are not available on all platforms
   nccl = shouldUsePkg (cudaPackages.nccl or null);
 
@@ -338,6 +349,7 @@ buildPythonPackage.override { stdenv = torch.stdenv; } (finalAttrs: {
     ./0002-setup.py-nix-support-respect-cmakeFlags.patch
     ./0003-propagate-pythonpath.patch
     ./0005-drop-intel-reqs.patch
+    ./0006-drop-rocm-extra-reqs.patch
   ];
 
   postPatch = ''
@@ -413,6 +425,16 @@ buildPythonPackage.override { stdenv = torch.stdenv; } (finalAttrs: {
         rocprim
         hipsparse
         hipblas
+        rocrand
+        hiprand
+        rocblas
+        miopen-hip
+        hipfft
+        hipcub
+        hipsolver
+        rocsolver
+        hipblaslt
+        rocm-runtime
       ]
     )
     ++ lib.optionals stdenv.cc.isClang [
@@ -485,6 +507,13 @@ buildPythonPackage.override { stdenv = torch.stdenv; } (finalAttrs: {
     cupy
     flashinfer
     nvidia-ml-py
+  ]
+  ++ lib.optionals rocmSupport [
+    rocmPackages.rocminfo
+    amdsmi
+    datasets
+    peft
+    timm
   ];
 
   optional-dependencies = {
@@ -523,9 +552,12 @@ buildPythonPackage.override { stdenv = torch.stdenv; } (finalAttrs: {
     }
     // lib.optionalAttrs rocmSupport {
       VLLM_TARGET_DEVICE = "rocm";
-      # Otherwise it tries to enumerate host supported ROCM gfx archs, and that is not possible due to sandboxing.
-      PYTORCH_ROCM_ARCH = lib.strings.concatStringsSep ";" rocmPackages.clr.gpuTargets;
-      ROCM_HOME = "${rocmPackages.clr}";
+      PYTORCH_ROCM_ARCH = gpuTargetString;
+      # vLLM's CMake logic checks `ROCM_PATH` to decide whether HIP/ROCm is available.
+      ROCM_PATH = "${rocmPackages.clr}";
+      TRITON_KERNELS_SRC_DIR = "${lib.getDev triton-kernels}/python/triton_kernels/triton_kernels";
+      HIPFLAGS = rocmExtraIncludeFlags;
+      CXXFLAGS = rocmExtraIncludeFlags;
     }
     // lib.optionalAttrs cpuSupport {
       VLLM_TARGET_DEVICE = "cpu";
@@ -558,6 +590,7 @@ buildPythonPackage.override { stdenv = torch.stdenv; } (finalAttrs: {
       happysalada
       lach
       daniel-fahey
+      LunNova # esp. for ROCm
     ];
     badPlatforms = [
       # CMake Error at cmake/cpu_extension.cmake:188 (message):