fix: update nixl setup for arm builds (#1061) (#1062)

Cargo.lock

@@ -4081,9 +4081,9 @@ dependencies = [
 
 [[package]]
 name = "nixl-sys"
-version = "0.2.1-rc.1"
+version = "0.2.1-rc.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e4698155ec791ae888482b0c1c5d19792a1b457167fa8ec0ffa487ffef8c8e5a"
+checksum = "bfeec33e0229249e8688309a70c677f522446f9ac92105a85aad4a5ceef2dd2f"
 dependencies = [
  "bindgen 0.71.1",
  "cc",

container/Dockerfile.vllm

@@ -108,12 +108,21 @@ WORKDIR /workspace
 # Copy nixl source, and use commit hash as cache hint
 COPY --from=nixl_base /opt/nixl /opt/nixl
 COPY --from=nixl_base /opt/nixl/commit.txt /opt/nixl/commit.txt
-RUN cd /opt/nixl && \
+RUN if [ "$ARCH" = "arm64" ]; then \
+        cd /opt/nixl && \
+        mkdir build && \
+        meson setup build/ --prefix=/usr/local/nixl -Dgds_path=/usr/local/cuda/targets/sbsa-linux && \
+        cd build/ && \
+        ninja && \
+        ninja install; \
+    else \
+        cd /opt/nixl && \
         mkdir build && \
         meson setup build/ --prefix=/usr/local/nixl && \
         cd build/ && \
         ninja && \
-    ninja install
+        ninja install; \
+    fi
 
 ### NATS & ETCD SETUP ###
 # nats
@@ -161,7 +170,8 @@ RUN --mount=type=bind,source=./container/deps/,target=/tmp/deps \
     # NOTE: vLLM build from source on ARM can take several hours, see VLLM_MAX_JOBS details.
     if [ "$ARCH" = "arm64" ]; then \
         # PyTorch 2.7 supports CUDA 12.8 and aarch64 installs
-        uv pip install torch==2.7.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 && \
+        # NIXL has a torch dependency, so need to force-reinstall to install the correct version
+        uv pip install torch==2.7.0 torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/cu128 && \
         # Download vLLM source with version matching patch
         git clone --branch v${VLLM_REF} --depth 1 https://github.com/vllm-project/vllm.git /tmp/vllm/vllm-${VLLM_REF} && \
         cd /tmp/vllm/vllm-${VLLM_REF}/ && \

lib/llm/Cargo.toml

@@ -77,7 +77,7 @@ regex = "1"
 rayon = "1"
 
 # block_manager
-nixl-sys = { version = "0.2.1-rc.1", optional = true }
+nixl-sys = { version = "0.2.1-rc.3", optional = true }
 cudarc = { version = "0.16.2", features = ["cuda-12020"], optional = true }
 ndarray = { version = "0.16", optional = true }