chore: Update CI image to use pytorch 26.08 and vllm 0.29.0

.github/workflows/release.yaml

@@ -86,7 +86,7 @@ jobs:
       !cancelled() && !failure()
       && !(needs.pre-flight.outputs.docs_only == 'true'
            || needs.pre-flight.outputs.is_deployment_workflow == 'true')
-    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_release_library.yml@f7af92c4d1caea5a58577ce4ac5314f7de1bc5d4 # v1.4.3
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_release_library.yml@4fe957d4ea5816af1d67fae130bbab4b434c10b3 # v1.9.2
     with:
       release-ref: ${{ inputs.release-ref || github.sha }}
       python-package: nemo_export_deploy_common

docker/Dockerfile.pytorch

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-FROM nvcr.io/nvidian/nemo:export-deploy-pyt-26.06-vllm-v0.24.0
+FROM nvcr.io/nvidian/nemo:export-deploy-pyt-26.08-vllm-v0.29.0
 WORKDIR /workspace
 COPY . .
 ARG INFERENCE_FRAMEWORK

docker/common/install.sh

@@ -96,7 +96,7 @@ main() {
     if [[ "$USE_UV" == "true" ]]; then
 
         # Install uv
-        UV_VERSION="0.7.2"
+        UV_VERSION="0.8.22"
         curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh
         export PATH="$HOME/.local/bin:$PATH"
         export UV_PROJECT_ENVIRONMENT=/opt/venv
@@ -139,17 +139,18 @@ main() {
             UV_ARGS=()
         fi
 
-        UV_ARGS+=("--extra" "$INFERENCE_FRAMEWORK")
-
         # Create virtual environment and install dependencies
         uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages
 
         # Install dependencies
+        # Note: `--only-group` and `--extra` are mutually exclusive, so the extra is only
+        # passed to the second (`--all-groups`) sync.
         uv sync --locked --only-group build ${UV_ARGS[@]}
         uv sync \
             --link-mode copy \
             --locked \
-            --all-groups ${UV_ARGS[@]}
+            --all-groups ${UV_ARGS[@]} \
+            --extra "$INFERENCE_FRAMEWORK"
         # Install the package
         uv pip install --no-deps -e .
         # Remove pyelftools to address dependency/license issues.

nemo_export/vllm_exporter.py

@@ -112,7 +112,6 @@ def export(
         quantization: str = None,
         seed: int = 0,
         gpu_memory_utilization: float = 0.9,
-        swap_space: float = 4,
         cpu_offload_gb: float = 0,
         enforce_eager: bool = False,
         task: Literal["auto", "generate", "embedding"] = "auto",
@@ -132,7 +131,6 @@ def export(
             quantization (str, optional): Quantization type. Defaults to None.
             seed (int, optional): Random seed. Defaults to 0.
             gpu_memory_utilization (float, optional): Fraction of GPU memory to use. Defaults to 0.9.
-            swap_space (float, optional): Amount of swap space (in GB) to use. Defaults to 4.
             cpu_offload_gb (float, optional): Amount of CPU offload memory (in GB). Defaults to 0.
             enforce_eager (bool, optional): Whether to enforce eager execution. Defaults to False.
             task (Literal["auto", "generate", "embedding"], optional): Task type for vLLM. Defaults to "auto".
@@ -203,7 +201,6 @@ def export(
                     quantization=quantization,
                     seed=seed,
                     gpu_memory_utilization=gpu_memory_utilization,
-                    swap_space=swap_space,
                     cpu_offload_gb=cpu_offload_gb,
                     enforce_eager=enforce_eager,
                     runner=task,
@@ -219,7 +216,6 @@ def export(
                 quantization=quantization,
                 seed=seed,
                 gpu_memory_utilization=gpu_memory_utilization,
-                swap_space=swap_space,
                 cpu_offload_gb=cpu_offload_gb,
                 enforce_eager=enforce_eager,
                 runner=task,

pyproject.toml

@@ -67,7 +67,7 @@ dependencies = [
     "wandb==0.27.2",
     "webdataset>=0.2.86",
     "nvidia-pytriton ; platform_system != 'Darwin' ",
-    "flashinfer-python==0.6.8.post1 ; platform_system != 'Darwin'",
+    "flashinfer-python; platform_system != 'Darwin'",
     "Pillow ; platform_system != 'Darwin' and platform_machine != 'aarch64'",
     "decord ; platform_system != 'Darwin' and platform_machine != 'aarch64'",
     "pyparsing>2.0.2",
@@ -80,7 +80,7 @@ dependencies = [
 
 [project.optional-dependencies]
 inframework = []
-vllm = ["vllm~=0.24.0", "pandas", "timm"]
+vllm = ["vllm~=0.29.0", "pandas", "timm"]
 trt-onnx = [
     "tensorrt==10.16.1.11",
     "onnx==1.21.0",
@@ -106,11 +106,15 @@ nemo-run = ["nemo-run"]
 [tool.uv.sources]
 xformers = [{ index = "pytorch-cu130" }]
 torch = [{ index = "pytorch-cu130" }]
+# flashinfer-cubin isn't published to PyPI since 0.6.14; both packages are pulled from
+# flashinfer's own wheel index, same as vLLM's requirements/cuda.txt.
+flashinfer-python = [{ index = "flashinfer" }]
+flashinfer-cubin = [{ index = "flashinfer" }]
 vllm = [
     { index = "pytorch-cu130", marker = "python_version < '3.9' and platform_machine == 'x86_64'" },
     { index = "pypi", marker = "platform_machine == 'aarch64'" },
 ]
-megatron-bridge = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git", rev = "97e8dba4060ae74a110b9da6255eb0a66b016b2d" }
+megatron-bridge = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git", rev = "02fcce5a103ca0c88a945319075348729e0b666e" }
 # Use the released nvidia-resiliency-ext (>=0.6.0) rather than a git dev build:
 # megatron-core's runtime is_nvrx_min_version() asserts the installed package
 # version is >=0.6.0, and a git build reports 0.6.0.devN (< 0.6.0), which fails
@@ -150,7 +154,6 @@ override-dependencies = [
     "transformer-engine[pytorch]>=2.16.0a0,<2.19.0",
     "open-clip-torch>=3.2.0",
     "megatron-energon[av-decode]>=6.0,<7.dev0",
-    "flashinfer-python==0.6.8.post1",
     "datasets>=3.3.0",
     "flash-linear-attention>=0.3.0,<0.4.dev0",
     "patchelf; sys_platform=='never'",
@@ -164,15 +167,31 @@ override-dependencies = [
     "opencv-python-headless; sys_platform == 'never'",
     "cryptography>=43.0.0,<47",
     "onnx<1.20.0",
-    "flashinfer-python==0.6.8.post1",
-    "flashinfer-cubin==0.6.8.post1",
+    # flashinfer-python and flashinfer-cubin must be kept in lockstep (flashinfer-python does a
+    # runtime version-match check); flashinfer-cubin isn't published to PyPI since 0.6.14, so both
+    # are sourced from flashinfer's own wheel index below. Pin matches vLLM's own requirements/cuda.txt.
+    "flashinfer-python==0.6.18.post1",
+    "flashinfer-cubin==0.6.18.post1",
+    "fast-hadamard-transform @ git+https://github.com/Dao-AILab/fast-hadamard-transform.git@v1.1.0",  # pin 1.1.0; PyPI sdist is incomplete, install from git tag (see DSv4/GLM-5 READMEs)
+    "nvidia-cudnn-frontend==1.26.0",
+    "nvidia-cutlass-dsl[cu13]; sys_platform == 'never'",    # Rely on system site packages install
 ]
 prerelease = "allow"
 
 [[tool.uv.dependency-metadata]]
 name = "nvidia-resiliency-ext"
 version = "0.6.0"
 
+[[tool.uv.dependency-metadata]]
+name = "fast-hadamard-transform"
+version = "1.1.0"
+requires-dist = ["torch", "packaging", "ninja"]
+
+# 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" }
+
 [[tool.uv.index]]
 name = "pypi"
 url = "https://pypi.org/simple"
@@ -183,6 +202,19 @@ name = "pytorch-cu130"
 url = "https://download.pytorch.org/whl/cu130"
 explicit = true
 
+[[tool.uv.index]]
+name = "flashinfer"
+url = "https://flashinfer.ai/whl/"
+explicit = true
+
+# `[tool.uv.sources]`/`[[tool.uv.index]]` above are only honored by uv's project interface
+# (uv sync/lock/add). The release CI builds the wheel via the pip-compatible interface
+# (`uv pip install -e .`), which does NOT read that config and only reads `[tool.uv.pip]`.
+# flashinfer-cubin/flashinfer-python aren't published on PyPI since 0.6.14, so without this
+# extra index `uv pip install -e .` can't resolve megatron-bridge's flashinfer-cubin pin.
+[tool.uv.pip]
+extra-index-url = ["https://flashinfer.ai/whl/"]
+
 [tool.pytest.ini_options]
 addopts = "--durations=15 -s -rA -x"
 testpaths = ["tests"]

scripts/deploy/nlp/deploy_vllm_triton.py

@@ -109,13 +109,6 @@ def get_args(argv):
         type=float,
         help="GPU memory utilization percentage for vLLM.",
     )
-    parser.add_argument(
-        "-sp",
-        "--swap_space",
-        default=4,
-        type=float,
-        help="The size (GiB) of CPU memory per GPU to use as swap space.",
-    )
     parser.add_argument(
         "-cog",
         "--cpu_offload_gb",
@@ -203,7 +196,6 @@ def nemo_deploy(argv):
             quantization=args.quantization,
             seed=args.seed,
             gpu_memory_utilization=args.gpu_memory_utilization,
-            swap_space=args.swap_space,
             cpu_offload_gb=args.cpu_offload_gb,
             enforce_eager=args.enforce_eager,
             task="generate",

tests/unit_tests/export/test_vllm_exporter.py

@@ -67,7 +67,6 @@ def test_export(exporter, mock_llm):
         quantization=None,
         seed=0,
         gpu_memory_utilization=0.9,
-        swap_space=4,
         cpu_offload_gb=0,
         enforce_eager=False,
         runner="auto",
@@ -92,7 +91,6 @@ def test_export_with_lora(exporter, mock_llm):
         quantization=None,
         seed=0,
         gpu_memory_utilization=0.9,
-        swap_space=4,
         cpu_offload_gb=0,
         enforce_eager=False,
         runner="auto",
@@ -123,7 +121,6 @@ def test_export_with_custom_params(exporter, mock_llm):
         quantization=None,
         seed=0,
         gpu_memory_utilization=0.9,
-        swap_space=4,
         cpu_offload_gb=0,
         enforce_eager=False,
         runner="auto",
@@ -838,7 +835,6 @@ def test_export_megatron_bridge_with_all_vllm_params(exporter, mock_llm):
             quantization="awq",
             seed=42,
             gpu_memory_utilization=0.85,
-            swap_space=8,
             cpu_offload_gb=2,
             enforce_eager=False,
             task="generate",
@@ -856,7 +852,6 @@ def test_export_megatron_bridge_with_all_vllm_params(exporter, mock_llm):
         assert call_kwargs["quantization"] == "awq"
         assert call_kwargs["seed"] == 42
         assert call_kwargs["gpu_memory_utilization"] == 0.85
-        assert call_kwargs["swap_space"] == 8
         assert call_kwargs["cpu_offload_gb"] == 2
         assert call_kwargs["enforce_eager"] is False
         assert call_kwargs["runner"] == "generate"