[Draft, not for merge] GPU CI speed-up series (runtime measurement)

.github/actions/cache-extensions/action.yml

@@ -1,24 +1,52 @@
 name: Cache extensions
-description: Cache and reuse JIT-compiled extensions (e.g. CUDA extensions) across runs.
+description: Restore JIT-compiled CUDA extensions and Triton kernels cached by earlier runs. Pair with cache-extensions/save after the tests.
 
 inputs:
   cache-key:
-    description: Environment discriminator for the cache key (e.g. GPU arch + container image).
+    description: Environment discriminator for the cache keys (e.g. GPU arch + lane + container image).
     required: true
 
+outputs:
+  ext-key:
+    description: Key under which cache-extensions/save stores the torch extensions cache.
+    value: ${{ steps.ext.outputs.cache-primary-key }}
+  ext-hit:
+    description: "'true' if the torch extensions cache was restored with an exact key match."
+    value: ${{ steps.ext.outputs.cache-hit }}
+  triton-key:
+    description: Key under which cache-extensions/save stores the Triton cache.
+    value: ${{ steps.triton.outputs.cache-primary-key }}
+  triton-hit:
+    description: "'true' if the Triton cache was restored with an exact key match."
+    value: ${{ steps.triton.outputs.cache-hit }}
+
 runs:
   using: composite
   steps:
     - shell: bash
-      run: echo "TORCH_EXTENSIONS_DIR=/root/.cache/torch_extensions" >> "$GITHUB_ENV"
-    - id: cache
-      uses: actions/cache@v6
+      run: |
+        echo "TORCH_EXTENSIONS_DIR=/root/.cache/torch_extensions" >> "$GITHUB_ENV"
+        echo "TRITON_CACHE_DIR=/root/.cache/triton" >> "$GITHUB_ENV"
+    # CUDA extensions: exact key only, because ninja decides what to rebuild from file timestamps.
+    # Only CUDA/C++ sources and the loaders are hashed, so Triton-only edits keep the entry valid.
+    - id: ext
+      uses: actions/cache/restore@v6
       with:
         path: /root/.cache/torch_extensions
-        key: torch-ext-${{ inputs.cache-key }}-${{ hashFiles('modelopt/torch/kernels/quantization/**', 'modelopt/torch/quantization/extensions.py', 'modelopt/torch/utils/cpp_extension.py')
-          }}
+        key: torch-ext-${{ inputs.cache-key }}-${{ hashFiles('modelopt/torch/kernels/quantization/**/*.cu', 'modelopt/torch/kernels/quantization/**/*.cuh',
+          'modelopt/torch/kernels/quantization/**/*.cpp', 'modelopt/torch/kernels/quantization/**/*.h', 'modelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.py',
+          'modelopt/torch/quantization/extensions.py', 'modelopt/torch/utils/cpp_extension.py') }}
     # On a cache hit, backdate sources below the cached objects so ninja reuses them (touching
     # the objects instead would desync ninja's .ninja_deps and force a rebuild).
-    - if: steps.cache.outputs.cache-hit == 'true'
+    - if: steps.ext.outputs.cache-hit == 'true'
       shell: bash
       run: find modelopt/torch/kernels/quantization -exec touch -d '2000-01-01' {} +
+    # Triton: entries are content-addressed, so restoring a cache from an older commit is safe
+    # and only costs the kernels that changed.
+    - id: triton
+      uses: actions/cache/restore@v6
+      with:
+        path: /root/.cache/triton
+        key: triton-${{ inputs.cache-key }}-${{ hashFiles('modelopt/torch/kernels/**/*.py') }}
+        restore-keys: |
+          triton-${{ inputs.cache-key }}-

.github/actions/cache-extensions/save/action.yml

@@ -0,0 +1,31 @@
+name: Save extension caches
+description: Save the JIT caches that cache-extensions did not restore with an exact key match. Run it under always() so a failed or cancelled run still
+  seeds the next one.
+
+inputs:
+  ext-key:
+    description: The ext-key output of cache-extensions.
+    required: true
+  ext-hit:
+    description: The ext-hit output of cache-extensions.
+    required: true
+  triton-key:
+    description: The triton-key output of cache-extensions.
+    required: true
+  triton-hit:
+    description: The triton-hit output of cache-extensions.
+    required: true
+
+runs:
+  using: composite
+  steps:
+    - if: inputs.ext-hit != 'true'
+      uses: actions/cache/save@v6
+      with:
+        path: /root/.cache/torch_extensions
+        key: ${{ inputs.ext-key }}
+    - if: inputs.triton-hit != 'true'
+      uses: actions/cache/save@v6
+      with:
+        path: /root/.cache/triton
+        key: ${{ inputs.triton-key }}

.github/workflows/_example_tests_runner.yml

@@ -58,9 +58,12 @@ jobs:
         run: apt-get update && apt-get install -y git
       - uses: actions/checkout@v6
       - uses: nv-gha-runners/setup-proxy-cache@main
-      - uses: ./.github/actions/cache-extensions
+      - id: jit-cache
+        uses: ./.github/actions/cache-extensions
         with:
-          cache-key: rtxpro6000-${{ inputs.docker_image }}
+          # One entry per example: the GPU-tests lanes and the examples that share an image would
+          # otherwise share one immutable, partially filled entry.
+          cache-key: rtxpro6000-example-${{ inputs.example }}-${{ inputs.docker_image }}
       - name: Setup environment variables
         run: |
           echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu:/usr/local/tensorrt/targets/x86_64-linux-gnu/lib" >> $GITHUB_ENV
@@ -113,6 +116,15 @@ jobs:
             pytest_args+=(-n 4)
           fi
           python3 -m pytest tests/examples/${{ inputs.example }} --cov "${pytest_args[@]}"
+      - name: Save JIT caches
+        # Also after a failed or cancelled run, so the next run starts warm
+        if: always()
+        uses: ./.github/actions/cache-extensions/save
+        with:
+          ext-key: ${{ steps.jit-cache.outputs.ext-key }}
+          ext-hit: ${{ steps.jit-cache.outputs.ext-hit }}
+          triton-key: ${{ steps.jit-cache.outputs.triton-key }}
+          triton-hit: ${{ steps.jit-cache.outputs.triton-hit }}
       - name: Flag allowed failure
         if: ${{ inputs.allow_failure && steps.run_tests.outcome == 'failure' }}
         run: |

.github/workflows/gpu_tests.yml

@@ -70,6 +70,8 @@ jobs:
         HF_TOKEN: ${{ secrets.HF_TOKEN }}
         # Build CUDA kernels only for the runner's RTX PRO 6000 (sm_120), not the image's ~6 archs.
         TORCH_CUDA_ARCH_LIST: "12.0"
+        # Where the gpu lane keeps the CUDA wheels it would otherwise compile (see noxfile.py).
+        WHEELHOUSE: /root/.cache/wheelhouse
     steps:
       - name: Install git
         # The vllm container ships without git; needed for a real checkout (correct
@@ -78,19 +80,61 @@ jobs:
         run: apt-get update && apt-get install -y git
       - uses: actions/checkout@v6
       - uses: nv-gha-runners/setup-proxy-cache@main
-      - uses: ./.github/actions/cache-extensions
+      - id: jit-cache
+        uses: ./.github/actions/cache-extensions
         with:
-          cache-key: rtxpro6000-${{ matrix.container_image }}
+          # One entry per lane: lanes sharing an image (gpu_megatron and the megatron_bridge
+          # example) would otherwise share one immutable, partially filled entry.
+          cache-key: rtxpro6000-${{ matrix.example }}-${{ matrix.container_image }}
       - name: Setup environment variables
         run: |
           echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu" >> $GITHUB_ENV
+      # The gpu lane compiles mamba_ssm, causal-conv1d and fast-hadamard-transform, 6-8 min. The
+      # nightly run on main builds them once per container image and ISO week (so new releases
+      # are picked up) and saves the wheels; every other run restores them. Without a cached
+      # wheelhouse the gpu session compiles them itself.
+      - name: Get the ISO week
+        id: week
+        if: matrix.example == 'gpu'
+        run: echo "iso=$(date -u +%G-W%V)" >> "$GITHUB_OUTPUT"
+      - name: Restore CUDA wheels
+        id: wheelhouse
+        if: matrix.example == 'gpu'
+        uses: actions/cache/restore@v6
+        with:
+          path: /root/.cache/wheelhouse
+          key: wheelhouse-${{ matrix.container_image }}-${{ steps.week.outputs.iso }}-${{ hashFiles('noxfile.py') }}
+      - name: Build CUDA wheels
+        if: matrix.example == 'gpu' && github.ref == 'refs/heads/main' && steps.wheelhouse.outputs.cache-hit != 'true'
+        run: python3 -m pip install nox && nox -s gpu_wheels
+      - name: Save CUDA wheels
+        if: matrix.example == 'gpu' && github.ref == 'refs/heads/main' && steps.wheelhouse.outputs.cache-hit != 'true'
+        uses: actions/cache/save@v6
+        with:
+          path: /root/.cache/wheelhouse
+          key: ${{ steps.wheelhouse.outputs.cache-primary-key }}
       - name: Run gpu tests
         env:
           COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml
           COVERAGE_FILE: ${{ github.workspace }}/.coverage
+          # sys.monitoring records the same lines as the default C tracer at a fraction of the
+          # overhead on import-heavy worker processes. Needs Python 3.12+.
+          COVERAGE_CORE: sysmon
+          # Tests marked `release` (real-weight downloads and the like) run on nightly and
+          # on-demand runs, not on pull requests.
+          PYTEST_ADDOPTS: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') && '--run-release' || '' }}
         run: |
           # Use `python3` (the vllm image has no `python` on PATH)
           python3 -m pip install nox && nox -s ${{ matrix.example }}
+      - name: Save JIT caches
+        # Also after a failed or cancelled run, so the next run starts warm
+        if: always()
+        uses: ./.github/actions/cache-extensions/save
+        with:
+          ext-key: ${{ steps.jit-cache.outputs.ext-key }}
+          ext-hit: ${{ steps.jit-cache.outputs.ext-hit }}
+          triton-key: ${{ steps.jit-cache.outputs.triton-key }}
+          triton-hit: ${{ steps.jit-cache.outputs.triton-hit }}
       - name: Upload GPU coverage to Codecov
         uses: codecov/codecov-action@v7
         with:

.github/workflows/regression_tests.yml

@@ -60,6 +60,9 @@ jobs:
         env:
           COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml
           COVERAGE_FILE: ${{ github.workspace }}/.coverage
+          # sys.monitoring records the same lines as the default C tracer at a fraction of the
+          # overhead on import-heavy worker processes. Needs Python 3.12+.
+          COVERAGE_CORE: sysmon
         run: python -m pip install nox && nox -s regression
       - name: Upload regression coverage to Codecov
         uses: codecov/codecov-action@v7

.github/workflows/unit_tests.yml

@@ -69,6 +69,12 @@ jobs:
           COVERAGE_FILE: ${{ github.workspace }}/.coverage
           # Everything but PRs also runs the slow tests marked ``manual`` (e.g. AutoNAS).
           PYTEST_EXTRA: ${{ github.event_name != 'pull_request' && '--run-manual' || '' }}
+          # sys.monitoring records the same lines as the default C tracer at a fraction of the overhead
+          COVERAGE_CORE: sysmon
+          # One single-threaded xdist worker per vCPU (`-n auto` counts physical cores); torch's
+          # default threads per worker oversubscribe the runner and run slower than serial.
+          PYTEST_ADDOPTS: -n 4 --dist loadfile
+          OMP_NUM_THREADS: "1"
         run: pip install nox uv && nox -s "unit-3.12(torch_214, tf_latest)" -- $PYTEST_EXTRA
       - name: Upload coverage reports to Codecov
         uses: codecov/codecov-action@v7
@@ -146,6 +152,9 @@ jobs:
         with:
           python-version: ${{ matrix.python_version }}
       - name: Run unit tests
+        env:
+          PYTEST_ADDOPTS: -n 4 --dist loadfile
+          OMP_NUM_THREADS: "1"
         run: pip install nox uv && nox -s "${{ matrix.nox_session }}"
   partial-install:
     if: needs.check-file-changes.outputs.any_changed == 'true'

docs/source/deployment/3_unified_hf.rst

@@ -144,8 +144,8 @@ prompts, asserting that each returns non-empty output.
 Two limits are worth stating plainly, because they bound what any ✅ below can mean:
 
 * **These are declared cases, not PR-gated coverage.** The suite is marked ``release`` and collects
-  only when pytest is given ``--run-release``, which no workflow in ``.github/workflows/`` currently
-  passes. A green check on a pull request does not mean these cases ran.
+  only when pytest is given ``--run-release``, which no example workflow in ``.github/workflows/``
+  currently passes. A green check on a pull request does not mean these cases ran.
 * **Each case is a load-and-generate smoke check on the text path.** It does not verify accuracy,
   image or audio inputs, diffusion output, or that speculative decoding actually engages.
 

examples/torch_onnx/torch_quant_to_onnx.py

@@ -27,7 +27,6 @@
 
 import timm
 import torch
-import torch.multiprocessing as mp
 import torch.nn.functional as F
 from datasets import load_dataset
 from download_example_onnx import export_to_onnx
@@ -56,9 +55,6 @@
 """
 
 
-mp.set_start_method("spawn", force=True)  # Needed for data loader with multiple workers
-
-
 _FP8_CONV_OVERRIDE: list = [
     {
         "parent_class": "nn.Conv2d",
@@ -311,16 +307,17 @@ def load_calibration_data(model, data_size, batch_size, device, with_labels=Fals
     calib_tensor = [transforms(img) for img in images]
     calib_tensor = [t.to(device) for t in calib_tensor]
 
+    # The samples already sit on ``device``; worker processes would only add spawn and import cost.
     if with_labels:
         labels = dataset["train"][:data_size]["label"]
         labels = torch.tensor(labels, device=device)
         calib_dataset = [{"image": img, "label": lbl} for img, lbl in zip(calib_tensor, labels)]
         return torch.utils.data.DataLoader(
-            calib_dataset, batch_size=batch_size, shuffle=True, num_workers=4
+            calib_dataset, batch_size=batch_size, shuffle=True, num_workers=0
         )
     else:
         return torch.utils.data.DataLoader(
-            calib_tensor, batch_size=batch_size, shuffle=True, num_workers=4
+            calib_tensor, batch_size=batch_size, shuffle=True, num_workers=0
         )
 
 
@@ -761,6 +758,10 @@ def build_trt_engine(onnx_path, optimization_level=4):
         f"--onnx={onnx_path}",
         "--stronglyTyped",
         f"--builderOptimizationLevel={optimization_level}",
+        # One inference iteration shows the engine runs; the default 3 s timed benchmark is not shown.
+        "--warmUp=0",
+        "--duration=0",
+        "--iterations=1",
     ]
     print(f"\nBuilding TensorRT engine: {' '.join(cmd)}")
     try:

noxfile.py

@@ -77,11 +77,12 @@ def unit(session, torch_ver, tf_ver):
     if run_manual:
         # A separate process: some manual tests (e.g. multi-threaded from_pretrained) leak global
         # state that breaks later tests. No coverage, so it cannot overwrite the main run's data.
+        # Serial too (CI's PYTEST_ADDOPTS asks for xdist): a worker would run them in some order.
         session.run(
             "python",
             # Caller args first: the last ``-m`` wins, and it must be ours.
             *("-m", "pytest", "tests/unit", *pytest_args, "--run-manual", "-m", "manual"),
-            env={**_CPU_ONLY_ENV, "COVERAGE_PROCESS_START": ""},
+            env={**_CPU_ONLY_ENV, "COVERAGE_PROCESS_START": "", "PYTEST_ADDOPTS": ""},
         )
 
 
@@ -113,6 +114,25 @@ def partial_unit(session, subset):
 # existing Python environment (e.g. /opt/venv in NeMo) instead of an isolated one.
 # Use `python -m pip/pytest` to ensure the container's active venv Python is used,
 # not a stale PATH entry (e.g. NeMo container has pip → /usr/local/bin/pip but python → /opt/venv/bin/python).
+
+# CUDA packages the gpu session compiles against the container's torch, which takes ~7 min together
+# (wheel file prefix -> requirement). CI builds them once with `nox -s gpu_wheels` into $WHEELHOUSE
+# (an Actions cache); the gpu session installs the wheels found there instead of compiling again.
+_WHEELHOUSE = os.environ.get("WHEELHOUSE")
+_CUDA_BUILDS = {
+    "fast_hadamard_transform": "git+https://github.com/Dao-AILab/fast-hadamard-transform.git",
+    # Latest *released* sdists
+    "mamba_ssm": "mamba_ssm",
+    "causal_conv1d": "causal-conv1d",
+}
+
+
+def _cuda_build(name):
+    """The wheel of ``name`` from the wheelhouse if it has one, else the requirement to build it from."""
+    wheels = sorted(glob.glob(f"{_WHEELHOUSE}/{name}-*.whl")) if _WHEELHOUSE else []
+    return wheels[-1] if wheels else _CUDA_BUILDS[name]
+
+
 # Container: nvcr.io/nvidia/pytorch:26.07-py3 or later with CUDA 13
 @nox.session(venv_backend="none")
 def gpu(session):
@@ -124,7 +144,7 @@ def gpu(session):
         "pip",
         "install",
         "--no-build-isolation",
-        "git+https://github.com/Dao-AILab/fast-hadamard-transform.git",
+        _cuda_build("fast_hadamard_transform"),
     )
     session.run("python", "-m", "pip", "install", "-e", ".[all,dev-test,dev-fla]")
     session.run(
@@ -133,13 +153,32 @@ def gpu(session):
         "pip",
         "install",
         "--no-build-isolation",
-        # Install the latest *released* sdists (built against the container torch)
-        "mamba_ssm",
-        "causal-conv1d",
+        _cuda_build("mamba_ssm"),
+        _cuda_build("causal_conv1d"),
     )
     session.run("python", "-m", "pytest", "tests/gpu", *_cov_args())
 
 
+# Container: same as gpu
+@nox.session(venv_backend="none")
+def gpu_wheels(session):
+    """Build the CUDA wheels of the gpu session into $WHEELHOUSE."""
+    if not _WHEELHOUSE:
+        session.error("Set WHEELHOUSE to the directory the wheels are built into")
+    for requirement in _CUDA_BUILDS.values():
+        session.run(
+            "python",
+            "-m",
+            "pip",
+            "wheel",
+            "--no-build-isolation",
+            "--no-deps",
+            "-w",
+            _WHEELHOUSE,
+            requirement,
+        )
+
+
 # Container: nvcr.io/nvidia/nemo:26.08 or later
 @nox.session(venv_backend="none")
 def gpu_megatron(session):

pyproject.toml

@@ -362,8 +362,9 @@ skips = [
 [tool.pytest.ini_options]
 # Default additional options
 # Show a short test summary info for all except passed tests with -ra flag
-# print execution time for 50 slowest tests and generate coverage reports
-addopts = "-v -ra --instafail --cov-report=term-missing --cov-report=html --cov-report=xml:coverage.xml --cov-config=pyproject.toml --durations=50 --strict-markers"
+# print execution time for 50 slowest tests and write the coverage.xml that CI uploads
+# (add --cov-report=html or --cov-report=term-missing locally; each adds 5-20 s to the end of a run)
+addopts = "-v -ra --instafail --cov-report=xml:coverage.xml --cov-config=pyproject.toml --durations=50 --strict-markers"
 pythonpath = ["tests/"]
 # Apply per-test timeouts (see tests/conftest.py) to the test call only, not fixture setup/teardown
 timeout_func_only = true