From 1769204ade8c9f738518fa8faac9975ed15a2ce3 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 24 May 2026 11:35:07 +0000
Subject: [PATCH] Dockerfile: arm64 DGX Spark NVRTC + ptxas fix (cu13 alongside
 cu128)

Empirically (cu128 wheel SASS list `sm_80;90;90a;100;100a;120;120a` on
aarch64) the cu128 wheel covers DGX Spark sm_121 via sm_120 binary
forward-compat. BUT two CPU-side compilers shipped at cu12.8 do not know
sm_121 and need a cu13 swap:

  (1) torch's bundled libnvrtc.so.12 from CUDA 12.8 rejects sm_121 as a
      --gpu-architecture. Symlinks libnvrtc.so.13 over it.

  (2) Triton's nvidia backend runs ptxas. Wheels older than 3.6.0 bundled
      cu12.8 ptxas which silently downgrades sm_121 to sm_80 (see
      triton-lang/triton#8335). Bump pin triton>=3.6.0 (3.6 bundles cu13
      ptxas) AND install cuda-nvcc-13-0 so the entrypoint can point
      TRITON_PTXAS_PATH at it as defense in depth.

Both fixes are arm64-only (gated on TARGETARCH, ~400 MB on the arm64
image; amd64 is untouched, no sm_121 hardware exists on x86_64). Neither
component talks to libcuda, so this does NOT bump the toolkit driver
floor away from cu128's 570+.

TRITON_PTXAS_PATH is set from the entrypoint (only when the cu13 ptxas
actually exists in the image) rather than via a Dockerfile ENV, because
ENV is unconditional and Triton errors out if TRITON_PTXAS_PATH points
at a nonexistent file.

Sources: martimramos/dgx-spark-ml-guide Challenge 14; triton-lang/triton
issue #8335; ptrblck PyTorch forum thread on sm_121 fwd-compat from
sm_120.
---
 docker/Dockerfile    | 47 +++++++++++++++++++++++++++++++++++++++++++-
 docker/entrypoint.sh |  8 ++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index af081f08fd..f74cde4560 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -159,7 +159,7 @@ RUN set -eux \
         --index-strategy unsafe-best-match \
         --extra-index-url https://download.pytorch.org/whl/cu128 \
         "torch==2.10.0" "torchvision==0.25.0" "torchaudio==2.11.0" \
-        "triton>=3.3.1" \
+        "triton>=3.6.0" \
         "bitsandbytes>=0.49.2,!=0.46.0,!=0.48.0" \
         "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo@${UNSLOTH_ZOO_REF}" \
         "unsloth[${UNSLOTH_EXTRA}] @ git+https://github.com/unslothai/unsloth@${UNSLOTH_REF}"
@@ -279,6 +279,51 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 
 COPY --from=builder /opt/unsloth-venv /opt/unsloth-venv
 
+# DGX Spark / GB10 (sm_121) fix, arm64 ONLY.
+#
+# Two cu13 components need to override what the cu128 stack ships, because
+# nothing in CUDA 12.8 -- toolkit or wheel -- knows about sm_121:
+#
+#   (1) torch's bundled libnvrtc.so.12 (from CUDA 12.8) does not accept
+#       sm_121 as --gpu-architecture. The jiterator C++ side queries the
+#       device cap directly, so any path that JIT-compiles a kernel (e.g.
+#       torch.fft.rfft(complex).abs(), used inside mel-spectrogram code)
+#       errors out. Fix: symlink libnvrtc.so.13 over the bundled .so.12.
+#
+#   (2) Triton's nvidia backend invokes ptxas. Triton wheels older than
+#       3.6.0 bundled cu12.8 ptxas which tops out at sm_120 and refuses
+#       sm_121, silently downgrading to sm_80 per triton-lang/triton#8335.
+#       Triton 3.6.0 (which we pin above) bundles cu13 ptxas, but for
+#       defense in depth we ALSO install cuda-nvcc-13-0 and point Triton
+#       at it via TRITON_PTXAS_PATH.
+#
+# NVRTC and ptxas are CPU-side compilers; they do NOT call into libcuda,
+# so we can install cu13 alongside the cu128 runtime without any driver
+# requirement bump (toolkit driver floor stays 570+).
+#
+# amd64 image is untouched: no sm_121 hardware exists on amd64, and the
+# extra ~400 MB would be dead weight.
+RUN if [ "${TARGETARCH:-amd64}" = "arm64" ]; then \
+        set -eux; \
+        # SBSA = Server Base System Architecture; the NVIDIA repo path for
+        # Grace / GH200 / GB200 / DGX Spark aarch64 hosts.
+        curl -fsSL "https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/sbsa/cuda-keyring_1.1-1_all.deb" \
+             -o /tmp/cuda-keyring.deb; \
+        dpkg -i /tmp/cuda-keyring.deb; \
+        rm /tmp/cuda-keyring.deb; \
+        apt-get update; \
+        apt-get install -y --no-install-recommends \
+            cuda-nvrtc-13-0 \
+            cuda-nvcc-13-0; \
+        rm -rf /var/lib/apt/lists/*; \
+        # (1) NVRTC swap. torch's wheel-bundled cu128 NVRTC -> cu13 NVRTC.
+        NVRTC_DIR=/opt/unsloth-venv/lib/python${PYTHON_VERSION}/site-packages/nvidia/cuda_nvrtc/lib; \
+        if [ -f "${NVRTC_DIR}/libnvrtc.so.12" ]; then \
+            mv "${NVRTC_DIR}/libnvrtc.so.12" "${NVRTC_DIR}/libnvrtc.so.12.cu128.orig"; \
+            ln -s /usr/local/cuda-13.0/lib64/libnvrtc.so.13 "${NVRTC_DIR}/libnvrtc.so.12"; \
+        fi; \
+    fi
+
 WORKDIR /workspace
 RUN mkdir -p ${HF_HOME} ${TRITON_CACHE_DIR}
 
diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh
index ab05547978..9874ca2c89 100755
--- a/docker/entrypoint.sh
+++ b/docker/entrypoint.sh
@@ -17,6 +17,14 @@
 #   docker run -e UNSLOTH_SKIP_GPU_CHECK=1 ...
 set -euo pipefail
 
+# DGX Spark fix, arm64 image only: prefer the cu13 ptxas we baked into the
+# image at /usr/local/cuda-13.0/bin/ptxas over Triton's bundled tools. The
+# file only exists on the arm64 variant; amd64 images skip this and use
+# Triton's own ptxas (cu13 in triton>=3.6.0).
+if [[ -x /usr/local/cuda-13.0/bin/ptxas ]] && [[ -z "${TRITON_PTXAS_PATH:-}" ]]; then
+    export TRITON_PTXAS_PATH=/usr/local/cuda-13.0/bin/ptxas
+fi
+
 if [[ "${UNSLOTH_SKIP_GPU_CHECK:-0}" == "1" ]]; then
     exec "$@"
 fi