From e1d9abc877048373b205922b0c9c1359574779df Mon Sep 17 00:00:00 2001 From: Kajal Jain Date: Fri, 28 Aug 2026 21:09:16 -0500 Subject: [PATCH] fix(ci): set TMS_CUDA_MAJOR via uv extra-build-variables torch-memory-saver's build backend needs TMS_CUDA_MAJOR when uv invokes its build backend during lockfile generation. Declare it once via uv's extra-build-variables so uv lock/sync pick it up automatically in every CI system, instead of deriving it via shell/sed in each caller (#3862, whose GitHub Actions workaround is reverted here as no longer needed). Signed-off-by: Kajal Jain --- .github/workflows/cicd-main.yml | 5 +---- pyproject.toml | 5 +++++ 2 files changed, 6 insertions(+), 4 deletions(-) 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"