diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml
index cf800629872..84dfc425cb3 100644
--- a/.github/workflows/cicd-main.yml
+++ b/.github/workflows/cicd-main.yml
@@ -477,10 +477,7 @@ jobs:
           enable-cache: true
 
       - name: Generate lockfile
-        run: |
-          export TMS_CUDA_MAJOR="$(sed -nE 's/^ARG[[:space:]]+BASE_IMAGE=.*cuda([0-9]+)\..*/\1/p' docker/Dockerfile | head -1)"
-          test -n "$TMS_CUDA_MAJOR"
-          uv lock
+        run: uv lock
 
       - name: Package lockfile
         id: package
diff --git a/pyproject.toml b/pyproject.toml
index 4a4a6516df6..6b20966223d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -594,6 +594,11 @@ fast-hadamard-transform = [{ requirement = "torch", match-runtime = true }]
 nv-grouped-gemm = [{ requirement = "torch", match-runtime = true }]
 flash_mla = [{ requirement = "torch", match-runtime = true }]
 
+# torch-memory-saver's build backend reads TMS_CUDA_MAJOR at build time. Keep this in
+# sync with docker/Dockerfile's BASE_IMAGE CUDA major version (currently cuda13.x).
+[tool.uv.extra-build-variables]
+torch-memory-saver = { TMS_CUDA_MAJOR = "13" }
+
 # Needed when building from source
 [[tool.uv.dependency-metadata]]
 name = "flash-attn"