From b7b7657227b9b3ece8af2f02a4c7f061ba7b07e8 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 04:07:35 -0700
Subject: [PATCH 01/34] Worker pool: report a failed rank before its teardown;
 move the helper test to unit

A failing rank blocked in megatron_worker_teardown's barrier never reported, so the 30 s failure grace
period could not start and the job waited for the pytest timeout instead. Also doc/comment fixes.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 tests/_test_utils/torch/distributed/utils.py  |  8 ++++-
 .../speculative/plugins/test_hf_dflash.py     |  2 +-
 .../utils/test_distributed_worker_pool.py     | 35 +++++++------------
 tests/gpu_megatron/conftest.py                | 12 ++++---
 tests/unit/torch/utils/test_get_live_pool.py  | 35 +++++++++++++++++++
 5 files changed, 62 insertions(+), 30 deletions(-)
 create mode 100644 tests/unit/torch/utils/test_get_live_pool.py

diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py
index ea954efbcf5..effc979154e 100644
--- a/tests/_test_utils/torch/distributed/utils.py
+++ b/tests/_test_utils/torch/distributed/utils.py
@@ -187,11 +187,16 @@ def _worker_loop(rank, world_size, backend, port, cmd_queue, result_queue, teard
             fn, args, kwargs = cmd
             status = "ok"
             tb = None
+            reported = False
             try:
                 fn(rank, world_size, *args, **kwargs)
             except Exception:
                 status = "error"
                 tb = traceback.format_exc()
+                # Report before teardown: it may block on a collective whose peer is still stuck in
+                # ``fn``, and the parent's failure grace period starts when it hears of a failure.
+                result_queue.put((status, rank, tb))
+                reported = True
             finally:
                 if teardown_fn is not None:
                     try:
@@ -201,7 +206,8 @@ def _worker_loop(rank, world_size, backend, port, cmd_queue, result_queue, teard
                         status = "error"
                         teardown_tb = traceback.format_exc()
                         tb = (tb + "\n" if tb else "") + f"[teardown] {teardown_tb}"
-            result_queue.put((status, rank, tb))
+            if not reported:
+                result_queue.put((status, rank, tb))
 
         dist.destroy_process_group()
 
diff --git a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py
index 69f0a117b63..078bf3f6b5a 100644
--- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py
@@ -287,7 +287,7 @@ def sharded_forward(*args, input_ids=None, **kwargs):
         model._base_model.forward = sharded_forward
 
     def test_generate_gathers_sharded_hidden_states(self):
-        # 8 layers so the default target_layer_ids ([1, 5]) straddle the split; 4 gives [1, 1].
+        # 8 layers so the default target_layer_ids land in different halves; 4 collapses them.
         model = get_tiny_llama(num_hidden_layers=8)
         mtsp.convert(model, [("dflash", get_dflash_config(block_size=4))])
         model = model.to("cuda:1").eval()
diff --git a/tests/gpu/torch/utils/test_distributed_worker_pool.py b/tests/gpu/torch/utils/test_distributed_worker_pool.py
index 7ae3a966278..166ce10e92e 100644
--- a/tests/gpu/torch/utils/test_distributed_worker_pool.py
+++ b/tests/gpu/torch/utils/test_distributed_worker_pool.py
@@ -17,21 +17,24 @@
 
 import queue
 import time
-from types import SimpleNamespace
 
 import pytest
 import torch.distributed as dist
-from _test_utils.torch.distributed.utils import DistributedWorkerPool, get_live_pool
+from _test_utils.torch.distributed.utils import DistributedWorkerPool
 
 
 def _hang(rank, size):
     time.sleep(600)
 
 
-def _rank_0_fails_and_rank_1_waits_for_it(rank, size):
+def _rank_0_fails_and_rank_1_hangs(rank, size):
     if rank == 0:
         raise ValueError("rank 0 failed")
-    dist.barrier()  # never completes: its partner left
+    time.sleep(600)
+
+
+def _teardown_barrier(rank, size):
+    dist.barrier()  # like megatron_worker_teardown: blocks while a peer is still in the job
 
 
 def test_a_hung_job_kills_the_pool_instead_of_stalling_later_runs(capfd):
@@ -52,30 +55,16 @@ def test_a_hung_job_kills_the_pool_instead_of_stalling_later_runs(capfd):
         pool.shutdown()
 
 
-def test_a_failed_rank_does_not_make_its_peers_wait_out_the_job_timeout():
-    pool = DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=None)
+def test_a_failed_rank_is_reported_even_when_its_teardown_blocks_on_a_hung_peer():
+    pool = DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=_teardown_barrier)
     try:
         pool.failure_grace = 2
+        pool.job_timeout = 30  # without the early report the parent would wait this long
         start = time.monotonic()
         with pytest.raises(RuntimeError, match="rank 0 failed"):
-            pool.run(_rank_0_fails_and_rank_1_waits_for_it)
+            pool.run(_rank_0_fails_and_rank_1_hangs)
 
-        assert time.monotonic() - start < 60  # not the 600 s job timeout
+        assert time.monotonic() - start < 25
         assert pool.broken
     finally:
         pool.shutdown()
-
-
-def test_a_killed_pool_is_replaced_and_a_live_one_is_reused():
-    pools = {}
-    made = []
-
-    def make():
-        made.append(SimpleNamespace(broken=False))
-        return made[-1]
-
-    first = get_live_pool(pools, 2, make)
-    assert get_live_pool(pools, 2, make) is first
-    first.broken = True
-    second = get_live_pool(pools, 2, make)
-    assert second is not first and len(made) == 2
diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py
index 5760b9fe22d..51261e30023 100644
--- a/tests/gpu_megatron/conftest.py
+++ b/tests/gpu_megatron/conftest.py
@@ -47,10 +47,10 @@ def _prebuild_quant_cuda_extensions():
     extensions during session setup and let tests fall back to on-demand JIT if needed.
 
     Doing it here in session setup (``pyproject`` sets ``timeout_func_only``) keeps the
-    build off the per-test clock and, unlike the ``_extensions/test_torch_extensions.py``
-    prebuild tests, runs regardless of test selection/ordering (e.g. ``-k`` filters) and
-    is not itself capped by a per-test timeout. Worker subprocesses then load the cached
-    .so from the shared ``TORCH_EXTENSIONS_DIR``.
+    build off the per-test clock, runs regardless of test selection/ordering (e.g. ``-k``
+    filters) and is not itself capped by a per-test timeout. Worker subprocesses then load
+    the cached .so from the shared ``TORCH_EXTENSIONS_DIR``. (The extension tests themselves
+    live in ``tests/gpu/_extensions``.)
     """
     ext.precompile()
 
@@ -99,7 +99,9 @@ def _pool_cache():
     instead of spawning duplicates — e.g. on a 2-GPU runner ``dist_workers`` and
     ``dist_workers_size_2`` are both size 2. The cache is module-scoped and torn down at
     module end, so workers are never reused across modules (avoids cross-test
-    state contamination).
+    state contamination). The ``dist_workers*`` fixtures are deliberately function-scoped so
+    ``get_live_pool`` can swap in a new pool after one was killed; they therefore cannot be
+    requested from a module- or session-scoped fixture.
     """
     pools: dict[int, DistributedWorkerPool] = {}
     yield pools
diff --git a/tests/unit/torch/utils/test_get_live_pool.py b/tests/unit/torch/utils/test_get_live_pool.py
new file mode 100644
index 00000000000..1431a9e2a52
--- /dev/null
+++ b/tests/unit/torch/utils/test_get_live_pool.py
@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""``get_live_pool`` hands back a live pool and replaces a killed one."""
+
+from types import SimpleNamespace
+
+from _test_utils.torch.distributed.utils import get_live_pool
+
+
+def test_a_killed_pool_is_replaced_and_a_live_one_is_reused():
+    pools = {}
+    made = []
+
+    def make():
+        made.append(SimpleNamespace(broken=False))
+        return made[-1]
+
+    first = get_live_pool(pools, 2, make)
+    assert get_live_pool(pools, 2, make) is first
+    first.broken = True
+    second = get_live_pool(pools, 2, make)
+    assert second is not first and len(made) == 2

From e5d149e7ac90821b7477e7497863478e96bbc857 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 04:07:41 -0700
Subject: [PATCH 02/34] CI: cache source-built wheels in the gpu job and bump
 the CI containers

- gpu job: pin and cache mamba_ssm / causal-conv1d / fast-hadamard-transform wheels per container image
  (~10 min of compilation per run); noxfile builds a wheel only if the wheelhouse lacks it
- Bump pytorch 26.05/26.07 -> 26.09, trtllm 1.3.0rc20 -> rc29, vllm v0.30.0 -> v0.31.0 (legacy v0.20.0 kept);
  nemo 26.08 is already the newest public. The tensorrt:26.05 images used by the ONNX/diffusers examples stay
  (torch-tensorrt needs libnvinfer.so.10)
- benchmark.py: tolerate TensorRT builds without NetworkDefinitionCreationFlag.EXPLICIT_BATCH (removed in 11)

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/example_tests.yml           |  6 +--
 .github/workflows/gpu_tests.yml               | 15 ++++--
 .github/workflows/regression_tests.yml        |  2 +-
 .../onnx/quantization/autotune/benchmark.py   |  5 +-
 noxfile.py                                    | 50 +++++++++++--------
 5 files changed, 48 insertions(+), 30 deletions(-)

diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml
index 8cc432a22ba..5a1b919f3c4 100644
--- a/.github/workflows/example_tests.yml
+++ b/.github/workflows/example_tests.yml
@@ -128,7 +128,7 @@ jobs:
       contents: read
     secrets: inherit
     with:
-      docker_image: "nvcr.io/nvidia/pytorch:26.07-py3"
+      docker_image: "nvcr.io/nvidia/pytorch:26.09-py3"
       example: ${{ matrix.example }}
       timeout_minutes: 30
       pip_install_extras: "[hf,dev-test]"
@@ -148,7 +148,7 @@ jobs:
       contents: read
     secrets: inherit
     with:
-      docker_image: "vllm/vllm-openai:v0.30.0"
+      docker_image: "vllm/vllm-openai:v0.31.0"
       container_options: "--shm-size=2gb --user 0"
       example: ${{ matrix.example }}
       timeout_minutes: 30
@@ -169,7 +169,7 @@ jobs:
       contents: read
     secrets: inherit
     with:
-      docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20"
+      docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc29"
       example: ${{ matrix.example }}
       pip_install_extras: "[hf,dev-test]"
       runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index 8e73f0f123a..d871ff263b2 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -45,19 +45,17 @@ jobs:
         include:
           - example: gpu
             timeout: 60
-            # Pinned to 26.05: benchmark.py uses trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH,
-            # which newer TensorRT (26.06) removed. Bump once the source is updated for TensorRT 10.
-            container_image: nvcr.io/nvidia/pytorch:26.05-py3
+            container_image: nvcr.io/nvidia/pytorch:26.09-py3
           - example: gpu_megatron
             timeout: 60
             container_image: nvcr.io/nvidia/nemo:26.08
           - example: gpu_trtllm
             timeout: 15
-            container_image: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20
+            container_image: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc29
           # Covers the RoutedExperts layout (0.24+)
           - example: gpu_vllm
             timeout: 15
-            container_image: docker.io/vllm/vllm-openai:v0.30.0
+            container_image: docker.io/vllm/vllm-openai:v0.31.0
           # Covers the legacy FusedMoE / SharedFusedMoE registration branches (vLLM < 0.24)
           - example: gpu_vllm
             timeout: 15
@@ -83,6 +81,13 @@ jobs:
       - uses: ./.github/actions/cache-extensions
         with:
           cache-key: rtxpro6000-${{ matrix.container_image }}
+      - name: Cache wheels built from source
+        # mamba_ssm and fast-hadamard-transform take ~10 min to compile (see noxfile.py `gpu`)
+        if: matrix.example == 'gpu'
+        uses: actions/cache@v6
+        with:
+          path: /root/.cache/wheelhouse
+          key: wheelhouse-${{ matrix.container_image }}-${{ hashFiles('noxfile.py') }}
       - name: Setup environment variables
         run: |
           echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu" >> $GITHUB_ENV
diff --git a/.github/workflows/regression_tests.yml b/.github/workflows/regression_tests.yml
index 001c4e6af4a..803572fb681 100644
--- a/.github/workflows/regression_tests.yml
+++ b/.github/workflows/regression_tests.yml
@@ -43,7 +43,7 @@ jobs:
     runs-on: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
     timeout-minutes: 15
     container:
-      image: nvcr.io/nvidia/pytorch:26.07-py3
+      image: nvcr.io/nvidia/pytorch:26.09-py3
       env:
         GIT_DEPTH: 1000 # For correct version
         PIP_CONSTRAINT: "" # Disable pip constraint for upgrading packages
diff --git a/modelopt/onnx/quantization/autotune/benchmark.py b/modelopt/onnx/quantization/autotune/benchmark.py
index ba5cf1142bf..12ca67e7f16 100644
--- a/modelopt/onnx/quantization/autotune/benchmark.py
+++ b/modelopt/onnx/quantization/autotune/benchmark.py
@@ -358,7 +358,10 @@ def __init__(
         trt.init_libnvinfer_plugins(self.trt_logger, "")
         self._plugin_registry = trt.get_plugin_registry()
 
-        self.network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
+        # Newer TensorRT removed EXPLICIT_BATCH: networks are always explicit-batch there.
+        self.network_flags = 0
+        if hasattr(trt.NetworkDefinitionCreationFlag, "EXPLICIT_BATCH"):
+            self.network_flags |= 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
         self.network_flags |= 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
         self._timing_cache = None
         self._load_timing_cache()
diff --git a/noxfile.py b/noxfile.py
index edd0544e1cd..b6b8d9d6234 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -76,7 +76,7 @@ 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.
+        # state that breaks later tests. No --cov here, so this run writes no coverage data.
         session.run(
             "python",
             # Caller args first: the last ``-m`` wins, and it must be ours.
@@ -113,32 +113,42 @@ 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 compiled against the container's torch, which takes minutes each. Pinned so the CI
+# cache of the built wheels (keyed on this file and the container image, see gpu_tests.yml) stays
+# valid: bump a pin deliberately, which also invalidates that cache.
+_SOURCE_BUILT_WHEELS = {
+    "fast_hadamard_transform": "git+https://github.com/Dao-AILab/fast-hadamard-transform.git@e7706faf8d1c3b9f241e36860640ad1dac644ede",
+    "mamba_ssm": "mamba_ssm==2.3.2.post1",
+    "causal_conv1d": "causal-conv1d==1.7.0",
+}
+
+
+def _install_built_from_source(session, *names):
+    """Install the named ``_SOURCE_BUILT_WHEELS``, building a wheel only if the wheelhouse lacks it.
+
+    CI caches ``MODELOPT_WHEELHOUSE`` per container image, so only the first run on an image builds.
+    """
+    wheelhouse = os.environ.get("MODELOPT_WHEELHOUSE", os.path.expanduser("~/.cache/wheelhouse"))
+    for name in names:
+        if not glob.glob(os.path.join(wheelhouse, f"{name}-*.whl")):
+            session.run(
+                *("python", "-m", "pip", "wheel", "--no-build-isolation", "--no-deps"),
+                *("--wheel-dir", wheelhouse, _SOURCE_BUILT_WHEELS[name]),
+            )
+    wheels = [glob.glob(os.path.join(wheelhouse, f"{name}-*.whl"))[0] for name in names]
+    session.run("python", "-m", "pip", "install", *wheels)
+
+
 # Container: nvcr.io/nvidia/pytorch:26.01-py3 or later
 @nox.session(venv_backend="none")
 def gpu(session):
+    _install_built_from_source(session, "fast_hadamard_transform")
+    session.run("python", "-m", "pip", "install", "-e", ".[all,dev-test,dev-fla]")
     # tests/gpu/_extensions/test_onnx_extensions.py fails for newer containers
     # until https://github.com/tbenthompson/cppimport/pull/98
-    session.run(
-        "python",
-        "-m",
-        "pip",
-        "install",
-        "--no-build-isolation",
-        "git+https://github.com/Dao-AILab/fast-hadamard-transform.git",
-    )
-    session.run("python", "-m", "pip", "install", "-e", ".[all,dev-test,dev-fla]")
     session.run("python", "-m", "pip", "uninstall", "-y", "cupy-cuda12x")
     session.run("python", "-m", "pip", "install", "cupy-cuda13x")
-    session.run(
-        "python",
-        "-m",
-        "pip",
-        "install",
-        "--no-build-isolation",
-        # Install the latest *released* sdists (built against the container torch)
-        "mamba_ssm",
-        "causal-conv1d",
-    )
+    _install_built_from_source(session, "mamba_ssm", "causal_conv1d")
     session.run("python", "-m", "pytest", "tests/gpu", *_cov_args())
 
 

From 4aedd5facb967002cfc79169bcbb1d948ceba68f Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 04:07:46 -0700
Subject: [PATCH 03/34] Remove the duplicate gpu_megatron copy of the CUDA
 extension tests

tests/gpu_megatron/_extensions/test_torch_extensions.py was identical to tests/gpu/_extensions/; the megatron conftest already prebuilds the extensions.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 tests/gpu_megatron/_extensions/test_torch_extensions.py | 1 -
 1 file changed, 1 deletion(-)
 delete mode 120000 tests/gpu_megatron/_extensions/test_torch_extensions.py

diff --git a/tests/gpu_megatron/_extensions/test_torch_extensions.py b/tests/gpu_megatron/_extensions/test_torch_extensions.py
deleted file mode 120000
index a07099dfd85..00000000000
--- a/tests/gpu_megatron/_extensions/test_torch_extensions.py
+++ /dev/null
@@ -1 +0,0 @@
-../../gpu/_extensions/test_torch_extensions.py
\ No newline at end of file

From 95b4913c1c5761c77be09dfbb9dcc6afcf2a2299 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 04:08:46 -0700
Subject: [PATCH 04/34] gpu job: save the wheel cache even when the tests fail
 or time out

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/gpu_tests.yml | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index d871ff263b2..ee4711e4b37 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -81,10 +81,11 @@ jobs:
       - uses: ./.github/actions/cache-extensions
         with:
           cache-key: rtxpro6000-${{ matrix.container_image }}
-      - name: Cache wheels built from source
+      - name: Restore wheels built from source
         # mamba_ssm and fast-hadamard-transform take ~10 min to compile (see noxfile.py `gpu`)
+        id: wheelhouse
         if: matrix.example == 'gpu'
-        uses: actions/cache@v6
+        uses: actions/cache/restore@v6
         with:
           path: /root/.cache/wheelhouse
           key: wheelhouse-${{ matrix.container_image }}-${{ hashFiles('noxfile.py') }}
@@ -98,6 +99,13 @@ jobs:
         run: |
           # Use `python3` (the vllm image has no `python` on PATH)
           python3 -m pip install nox && nox -s ${{ matrix.example }}
+      - name: Save wheels built from source
+        # Not tied to the tests passing: the wheels are valid even if a test fails or the job times out.
+        if: always() && matrix.example == 'gpu' && steps.wheelhouse.outputs.cache-hit != 'true'
+        uses: actions/cache/save@v6
+        with:
+          path: /root/.cache/wheelhouse
+          key: wheelhouse-${{ matrix.container_image }}-${{ hashFiles('noxfile.py') }}
       - name: Upload GPU coverage to Codecov
         uses: codecov/codecov-action@v7
         with:

From 7e8cfd19d785b45908ee0c554d593a64acb1b3f9 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 04:31:28 -0700
Subject: [PATCH 05/34] CI: save the CUDA extension cache even when the job
 fails or times out

The nightly gpu job never found its extension cache in a month of runs: actions/cache saves only when the job succeeds, and the nightly was failing or timing out every night, so the main-branch cache was never populated.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/actions/cache-extensions/action.yml | 12 ++++++++++--
 .github/actions/save-extensions/action.yml  | 16 ++++++++++++++++
 .github/workflows/_example_tests_runner.yml |  7 +++++++
 .github/workflows/gpu_tests.yml             |  7 +++++++
 4 files changed, 40 insertions(+), 2 deletions(-)
 create mode 100644 .github/actions/save-extensions/action.yml

diff --git a/.github/actions/cache-extensions/action.yml b/.github/actions/cache-extensions/action.yml
index 7ddf65de8ce..bc7584e813f 100644
--- a/.github/actions/cache-extensions/action.yml
+++ b/.github/actions/cache-extensions/action.yml
@@ -1,18 +1,26 @@
 name: Cache extensions
-description: Cache and reuse JIT-compiled extensions (e.g. CUDA extensions) across runs.
+description: Restore JIT-compiled extensions (e.g. CUDA extensions) cached by an earlier run. Pair with ./.github/actions/save-extensions.
 
 inputs:
   cache-key:
     description: Environment discriminator for the cache key (e.g. GPU arch + container image).
     required: true
 
+outputs:
+  cache-hit:
+    description: Whether the extensions were restored from the exact key.
+    value: ${{ steps.cache.outputs.cache-hit }}
+  cache-key:
+    description: The key to pass to ./.github/actions/save-extensions.
+    value: ${{ steps.cache.outputs.cache-primary-key }}
+
 runs:
   using: composite
   steps:
     - shell: bash
       run: echo "TORCH_EXTENSIONS_DIR=/root/.cache/torch_extensions" >> "$GITHUB_ENV"
     - id: cache
-      uses: actions/cache@v6
+      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')
diff --git a/.github/actions/save-extensions/action.yml b/.github/actions/save-extensions/action.yml
new file mode 100644
index 00000000000..bc639f28804
--- /dev/null
+++ b/.github/actions/save-extensions/action.yml
@@ -0,0 +1,16 @@
+name: Save extensions
+description: Save the JIT-compiled extensions restored by ./.github/actions/cache-extensions. Run it with `always()` so a failing or timed-out job still
+  populates the cache (actions/cache itself only saves when the job succeeds).
+
+inputs:
+  cache-key:
+    description: The `cache-key` output of ./.github/actions/cache-extensions.
+    required: true
+
+runs:
+  using: composite
+  steps:
+    - uses: actions/cache/save@v6
+      with:
+        path: /root/.cache/torch_extensions
+        key: ${{ inputs.cache-key }}
diff --git a/.github/workflows/_example_tests_runner.yml b/.github/workflows/_example_tests_runner.yml
index e3378865b65..d3237f1b147 100644
--- a/.github/workflows/_example_tests_runner.yml
+++ b/.github/workflows/_example_tests_runner.yml
@@ -59,6 +59,7 @@ jobs:
       - uses: actions/checkout@v6
       - uses: nv-gha-runners/setup-proxy-cache@main
       - uses: ./.github/actions/cache-extensions
+        id: ext-cache
         with:
           cache-key: rtxpro6000-${{ inputs.docker_image }}
       - name: Setup environment variables
@@ -113,6 +114,12 @@ jobs:
             pytest_args+=(-n 4)
           fi
           python3 -m pytest tests/examples/${{ inputs.example }} --cov "${pytest_args[@]}"
+      - name: Save extensions
+        # Not tied to the tests passing, see the action description
+        if: always() && steps.ext-cache.outputs.cache-hit != 'true'
+        uses: ./.github/actions/save-extensions
+        with:
+          cache-key: ${{ steps.ext-cache.outputs.cache-key }}
       - name: Flag allowed failure
         if: ${{ inputs.allow_failure && steps.run_tests.outcome == 'failure' }}
         run: |
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index ee4711e4b37..8e73dcd13bc 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -79,6 +79,7 @@ jobs:
       - uses: actions/checkout@v6
       - uses: nv-gha-runners/setup-proxy-cache@main
       - uses: ./.github/actions/cache-extensions
+        id: ext-cache
         with:
           cache-key: rtxpro6000-${{ matrix.container_image }}
       - name: Restore wheels built from source
@@ -99,6 +100,12 @@ jobs:
         run: |
           # Use `python3` (the vllm image has no `python` on PATH)
           python3 -m pip install nox && nox -s ${{ matrix.example }}
+      - name: Save extensions
+        # Not tied to the tests passing, see the action description
+        if: always() && steps.ext-cache.outputs.cache-hit != 'true'
+        uses: ./.github/actions/save-extensions
+        with:
+          cache-key: ${{ steps.ext-cache.outputs.cache-key }}
       - name: Save wheels built from source
         # Not tied to the tests passing: the wheels are valid even if a test fails or the job times out.
         if: always() && matrix.example == 'gpu' && steps.wheelhouse.outputs.cache-hit != 'true'

From 8c5fadce1cf193482c011a6b1d2ddc24a696deac Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 04:38:31 -0700
Subject: [PATCH 06/34] EAGLE: compile _eagle_forward without CUDA graphs

mode="max-autotune" enables inductor CUDA graphs, whose recording intermittently fails in training (cudaStreamCaptureStatusInvalidated, 'graph recording observed an input tensor deallocate'); it failed test_llama_eagle3 on 5 of the last 10 nightlies and on the pytorch:26.09 dispatch run. max-autotune-no-cudagraphs keeps the autotuning.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 modelopt/torch/speculative/plugins/hf_eagle.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/speculative/plugins/hf_eagle.py b/modelopt/torch/speculative/plugins/hf_eagle.py
index 94399912983..1789e55aff2 100644
--- a/modelopt/torch/speculative/plugins/hf_eagle.py
+++ b/modelopt/torch/speculative/plugins/hf_eagle.py
@@ -154,7 +154,9 @@ def _activate_torch_compile(self):
 
         compile_targets: list[tuple[str, dict[str, Any]]] = [
             ("_prepare_eagle_inputs", {}),
-            ("_eagle_forward", {"mode": "max-autotune"}),
+            # No CUDA graphs: "max-autotune" enables them, and recording them intermittently fails
+            # in training (cudaStreamCaptureStatusInvalidated / "graph recording observed ...").
+            ("_eagle_forward", {"mode": "max-autotune-no-cudagraphs"}),
             ("_eagle_loss", {"fullgraph": True}),
         ]
         for name, kwargs in compile_targets:

From 613e40e4c2f84c185d02865d17f9f645c61aa52a Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 05:09:41 -0700
Subject: [PATCH 07/34] gpu job: build mamba_ssm and causal-conv1d as C++20

Their setup.py hardcodes -std=c++17, which fails against the torch headers of pytorch:26.09 ('ATen requires C++20'). Patch the standard when building the wheel; verified to compile and import against torch 2.13 headers.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 noxfile.py | 35 +++++++++++++++++++++++++++++++----
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/noxfile.py b/noxfile.py
index b6b8d9d6234..1263bbfac97 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -26,6 +26,8 @@
 import glob
 import os
 import shutil
+import tarfile
+import tempfile
 
 import nox
 
@@ -123,6 +125,34 @@ def partial_unit(session, subset):
 }
 
 
+# Their setup.py files hardcode -std=c++17, but the torch headers of recent containers need C++20.
+_BUILD_AS_CXX20 = ("mamba_ssm", "causal_conv1d")
+
+
+def _build_wheel(session, name, wheelhouse):
+    pip_wheel = ("python", "-m", "pip", "wheel", "--no-build-isolation", "--no-deps")
+    pip_wheel += ("--wheel-dir", wheelhouse)
+    requirement = _SOURCE_BUILT_WHEELS[name]
+    if name not in _BUILD_AS_CXX20:
+        session.run(*pip_wheel, requirement)
+        return
+    with tempfile.TemporaryDirectory() as tmp:
+        session.run(
+            *("python", "-m", "pip", "download", "--no-deps", "--no-binary", ":all:"),
+            *("--no-build-isolation", "--dest", tmp, requirement),
+        )
+        sdist = glob.glob(os.path.join(tmp, "*.tar.gz"))[0]
+        with tarfile.open(sdist) as tar:
+            tar.extractall(tmp, filter="data")
+        source = sdist.removesuffix(".tar.gz")
+        setup_py = os.path.join(source, "setup.py")
+        with open(setup_py) as f:
+            patched = f.read().replace("-std=c++17", "-std=c++20")
+        with open(setup_py, "w") as f:
+            f.write(patched)
+        session.run(*pip_wheel, source)
+
+
 def _install_built_from_source(session, *names):
     """Install the named ``_SOURCE_BUILT_WHEELS``, building a wheel only if the wheelhouse lacks it.
 
@@ -131,10 +161,7 @@ def _install_built_from_source(session, *names):
     wheelhouse = os.environ.get("MODELOPT_WHEELHOUSE", os.path.expanduser("~/.cache/wheelhouse"))
     for name in names:
         if not glob.glob(os.path.join(wheelhouse, f"{name}-*.whl")):
-            session.run(
-                *("python", "-m", "pip", "wheel", "--no-build-isolation", "--no-deps"),
-                *("--wheel-dir", wheelhouse, _SOURCE_BUILT_WHEELS[name]),
-            )
+            _build_wheel(session, name, wheelhouse)
     wheels = [glob.glob(os.path.join(wheelhouse, f"{name}-*.whl"))[0] for name in names]
     session.run("python", "-m", "pip", "install", *wheels)
 

From ae6f163c4667e2f2b09ffc2d8332d042393cf214 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 05:09:43 -0700
Subject: [PATCH 08/34] NVFP4: convert cutlass-swizzled block scales in
 ModelOpt instead of importing them from TensorRT-LLM

TensorRT-LLM 1.3.0rc29 removed tensorrt_llm._torch.auto_deploy, so the dequantize/export paths (and gpu_trtllm collection) failed with ModuleNotFoundError. The conversion is a tensor permutation: implement the inverse of swizzle_nvfp4_scales locally, with a round-trip test.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 modelopt/torch/export/unified_export_hf.py    | 14 ++-----
 .../quantization/qtensor/nvfp4_tensor.py      | 37 ++++++++++---------
 .../quantization/backends/test_nvfp4_gemm.py  |  4 +-
 .../torch/quantization/test_nvfp4_tensor.py   | 22 +++++++++++
 4 files changed, 47 insertions(+), 30 deletions(-)

diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index f605ddaf636..51f9392028b 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -72,7 +72,10 @@
 from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer
 from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor
 from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper
-from modelopt.torch.quantization.qtensor.nvfp4_tensor import _cast_per_block_scale_to_fp8
+from modelopt.torch.quantization.qtensor.nvfp4_tensor import (
+    _cast_per_block_scale_to_fp8,
+    cutlass_fp4_scale_to_modelopt_fp4_scale,
+)
 from modelopt.torch.quantization.utils import (
     fsdp2_aware_weight_update,
     module_name_maps,
@@ -591,15 +594,6 @@ def _compressed_per_block_scale(
     scale = getattr(weight_quantizer, "_scale", None)
     if scale is None or not (scale.dtype == torch.uint8 and scale.ndim == 1):
         return scale
-    try:
-        from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import (
-            cutlass_fp4_scale_to_modelopt_fp4_scale,
-        )
-    except ImportError as e:
-        raise ImportError(
-            "This weight was compressed by TensorRT-LLM, so its NVFP4 block scale is "
-            "cutlass-swizzled, but tensorrt_llm cannot be imported to convert it for export."
-        ) from e
     return cutlass_fp4_scale_to_modelopt_fp4_scale(scale, weight.metadata["shape"][-2:])
 
 
diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py
index b3095e0ca98..80894db888a 100644
--- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py
+++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py
@@ -46,6 +46,23 @@ def _cast_per_block_scale_to_fp8(
     return per_block_scale.clamp(min=2**-9, max=E4M3_MAX).to(torch.float8_e4m3fn)
 
 
+def cutlass_fp4_scale_to_modelopt_fp4_scale(
+    cutlass_scale: torch.Tensor, weight_shape: tuple[int, int]
+) -> torch.Tensor:
+    """Undo the cuBLAS/cutlass tiled swizzle of NVFP4 block scales (TensorRT-LLM ``fp4_quantize``).
+
+    The swizzled scale holds 128x4 tiles, each stored as ``[row % 32, row // 32, col]``; this is the
+    inverse of the rearrangement in ``swizzle_nvfp4_scales``. Returns the ``[m, n // 16]`` E4M3
+    scales of an ``[m, n]`` weight with block size 16.
+    """
+    m, n = weight_shape
+    n //= 16
+    row_blocks, col_blocks = -(-m // 128), -(-n // 4)
+    tiles = cutlass_scale.reshape(row_blocks, col_blocks, 32, 4, 4)
+    padded = tiles.permute(0, 3, 2, 1, 4).reshape(row_blocks * 128, col_blocks * 4)
+    return padded[:m, :n].view(torch.float8_e4m3fn)
+
+
 class NVFP4QTensor(BaseQuantizedTensor):
     """Implements the NVFP4 quantization on tensors for more efficient storage or computation.
 
@@ -294,11 +311,6 @@ def quantize(
             try:
                 import tensorrt_llm  # noqa: F401
 
-                # Make sure this utils is available for dequantize
-                from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import (
-                    cutlass_fp4_scale_to_modelopt_fp4_scale,  # noqa: F401
-                )
-
                 packed_weight, weights_scaling_factor = torch.ops.trtllm.fp4_quantize(
                     input, 1.0 / weights_scaling_factor_2, block_size, False
                 )
@@ -364,18 +376,9 @@ def _unpack_tensor(input: torch.Tensor):
         # Get scales from kwargs
         if kwarg["scale"].dtype == torch.uint8 and kwarg["scale"].ndim == 1:
             # If quantization is done by trtllm, convert cutlass fp4 scale to modelopt fp4 scale
-            try:
-                from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import (
-                    cutlass_fp4_scale_to_modelopt_fp4_scale,
-                )
-
-                kwarg["scale"] = cutlass_fp4_scale_to_modelopt_fp4_scale(
-                    kwarg["scale"], self.metadata["shape"][-2:]
-                )
-            except ImportError as e:
-                raise ImportError(
-                    "This tensor is quantized by trtllm, but tensorrt_llm cannot be imported."
-                ) from e
+            kwarg["scale"] = cutlass_fp4_scale_to_modelopt_fp4_scale(
+                kwarg["scale"], self.metadata["shape"][-2:]
+            )
 
         if fast:
             from modelopt.torch.kernels.quantization.gemm.fp4_kernel import fp4_dequantize
diff --git a/tests/gpu_trtllm/torch/quantization/backends/test_nvfp4_gemm.py b/tests/gpu_trtllm/torch/quantization/backends/test_nvfp4_gemm.py
index afa0aa1f2de..38c22ec3533 100644
--- a/tests/gpu_trtllm/torch/quantization/backends/test_nvfp4_gemm.py
+++ b/tests/gpu_trtllm/torch/quantization/backends/test_nvfp4_gemm.py
@@ -15,13 +15,11 @@
 import pytest
 import torch
 from _test_utils.torch.quantization.models import OneLayerLinear
-from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import (
-    cutlass_fp4_scale_to_modelopt_fp4_scale,
-)
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.quantization.backends.utils import fp4_compatible
 from modelopt.torch.quantization.qtensor import NVFP4QTensor
+from modelopt.torch.quantization.qtensor.nvfp4_tensor import cutlass_fp4_scale_to_modelopt_fp4_scale
 
 
 @pytest.mark.skipif(not fp4_compatible(), reason="FP4 is not supported on this GPU")
diff --git a/tests/unit/torch/quantization/test_nvfp4_tensor.py b/tests/unit/torch/quantization/test_nvfp4_tensor.py
index a2a977f45b1..a987bcafad3 100644
--- a/tests/unit/torch/quantization/test_nvfp4_tensor.py
+++ b/tests/unit/torch/quantization/test_nvfp4_tensor.py
@@ -24,6 +24,7 @@
 from modelopt.torch.quantization.qtensor.nvfp4_tensor import (
     NVFP4QTensor,
     _cast_per_block_scale_to_fp8,
+    cutlass_fp4_scale_to_modelopt_fp4_scale,
 )
 
 _FP8_E4M3FN_MIN = 2**-9  # 0.001953125 — smallest positive FP8 E4M3FN subnormal
@@ -127,3 +128,24 @@ def test_static_path_no_nan_when_block_amax_zero(self):
         assert (per_block_scale_f32 <= _FP8_E4M3FN_MAX).all(), (
             f"Static per-block scale exceeds FP8 max 448: {per_block_scale_f32.tolist()}"
         )
+
+
+def _cublas_swizzle(scale):
+    """Flat cuBLAS tiled layout of a ``[rows, cols]`` block-scale matrix (the 128x4-tile swizzle)."""
+    rows, cols = scale.shape
+    row_blocks, col_blocks = -(-rows // 128), -(-cols // 4)
+    padded = torch.zeros(row_blocks * 128, col_blocks * 4, dtype=scale.dtype)
+    padded[:rows, :cols] = scale
+    blocks = padded.view(row_blocks, 128, col_blocks, 4).permute(0, 2, 1, 3)
+    return blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1)
+
+
+@pytest.mark.parametrize("weight_shape", [(128, 64), (3, 16), (130, 80), (256, 128)])
+def test_cutlass_fp4_scale_to_modelopt_fp4_scale_undoes_the_swizzle(weight_shape):
+    m, n = weight_shape
+    scale = torch.randint(0, 256, (m, n // 16), dtype=torch.uint8)
+
+    restored = cutlass_fp4_scale_to_modelopt_fp4_scale(_cublas_swizzle(scale), weight_shape)
+
+    assert restored.dtype == torch.float8_e4m3fn
+    assert torch.equal(restored.view(torch.uint8), scale)

From e0dafde6ecc1836ee9371cbb1ab0a36a84636cb8 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 05:09:48 -0700
Subject: [PATCH 09/34] Keep vLLM at v0.30.0: v0.31.0 breaks the vLLM plugin

v0.31.0 fails test_kimi_gemm_rs_ar_bypasses_enabled_fake_quantizers (gemm_rs_ar has no .apply), test_flashinfer_metadata_builder_patch_stashes_common_metadata (FlashInferMetadataBuilder.cache_config) and times out test_tiny_llama_quantize; needs plugin changes.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/example_tests.yml | 2 +-
 .github/workflows/gpu_tests.yml     | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml
index 5a1b919f3c4..66688c093a8 100644
--- a/.github/workflows/example_tests.yml
+++ b/.github/workflows/example_tests.yml
@@ -148,7 +148,7 @@ jobs:
       contents: read
     secrets: inherit
     with:
-      docker_image: "vllm/vllm-openai:v0.31.0"
+      docker_image: "vllm/vllm-openai:v0.30.0"
       container_options: "--shm-size=2gb --user 0"
       example: ${{ matrix.example }}
       timeout_minutes: 30
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index 8e73dcd13bc..db7f43c563d 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -55,7 +55,7 @@ jobs:
           # Covers the RoutedExperts layout (0.24+)
           - example: gpu_vllm
             timeout: 15
-            container_image: docker.io/vllm/vllm-openai:v0.31.0
+            container_image: docker.io/vllm/vllm-openai:v0.30.0
           # Covers the legacy FusedMoE / SharedFusedMoE registration branches (vLLM < 0.24)
           - example: gpu_vllm
             timeout: 15

From 27235c78dca7850265683a564416b379509c851f Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 05:49:25 -0700
Subject: [PATCH 10/34] Address review: wheelhouse path, complete-only cache
 save, wheel lookup, force-build, teardown error, test margins

- gpu job: set MODELOPT_WHEELHOUSE=/root/.cache/wheelhouse (HOME is /github/home in container jobs, so the build wrote to a path the cache steps never saw)
- save the wheelhouse only once a .complete marker shows every wheel is built
- deterministic wheel lookup with a clear error; MAMBA_FORCE_BUILD / CAUSAL_CONV1D_FORCE_BUILD so the patched source is compiled
- do not build a teardown traceback that is never reported; widen the pool test's timing margin; 2026 copyright

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/gpu_tests.yml               | 12 +++++--
 .../onnx/quantization/autotune/benchmark.py   |  6 ++--
 noxfile.py                                    | 35 ++++++++++++++++---
 tests/_test_utils/torch/distributed/utils.py  |  8 ++---
 .../utils/test_distributed_worker_pool.py     |  4 +--
 tests/unit/torch/utils/test_get_live_pool.py  |  2 +-
 6 files changed, 50 insertions(+), 17 deletions(-)

diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index db7f43c563d..6315c26a455 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -97,6 +97,8 @@ jobs:
         env:
           COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml
           COVERAGE_FILE: ${{ github.workspace }}/.coverage
+          # Explicit: HOME is /github/home in container jobs, not the /root the cache steps use
+          MODELOPT_WHEELHOUSE: /root/.cache/wheelhouse
         run: |
           # Use `python3` (the vllm image has no `python` on PATH)
           python3 -m pip install nox && nox -s ${{ matrix.example }}
@@ -106,9 +108,15 @@ jobs:
         uses: ./.github/actions/save-extensions
         with:
           cache-key: ${{ steps.ext-cache.outputs.cache-key }}
+      - name: Check the wheelhouse is complete
+        id: wheelhouse-complete
+        if: always() && matrix.example == 'gpu'
+        run: test -f /root/.cache/wheelhouse/.complete && echo "complete=true" >> "$GITHUB_OUTPUT" || true
       - name: Save wheels built from source
-        # Not tied to the tests passing: the wheels are valid even if a test fails or the job times out.
-        if: always() && matrix.example == 'gpu' && steps.wheelhouse.outputs.cache-hit != 'true'
+        # Not tied to the tests passing, but only once every wheel is built (see noxfile.py `gpu`).
+        if: >-
+          always() && matrix.example == 'gpu' && steps.wheelhouse.outputs.cache-hit != 'true'
+          && steps.wheelhouse-complete.outputs.complete == 'true'
         uses: actions/cache/save@v6
         with:
           path: /root/.cache/wheelhouse
diff --git a/modelopt/onnx/quantization/autotune/benchmark.py b/modelopt/onnx/quantization/autotune/benchmark.py
index 12ca67e7f16..0e8b4025c76 100644
--- a/modelopt/onnx/quantization/autotune/benchmark.py
+++ b/modelopt/onnx/quantization/autotune/benchmark.py
@@ -358,11 +358,11 @@ def __init__(
         trt.init_libnvinfer_plugins(self.trt_logger, "")
         self._plugin_registry = trt.get_plugin_registry()
 
-        # Newer TensorRT removed EXPLICIT_BATCH: networks are always explicit-batch there.
-        self.network_flags = 0
+        self.network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
+        # TensorRT 11 removed EXPLICIT_BATCH (networks are always explicit-batch); drop this guard
+        # once older TensorRT is no longer supported.
         if hasattr(trt.NetworkDefinitionCreationFlag, "EXPLICIT_BATCH"):
             self.network_flags |= 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
-        self.network_flags |= 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
         self._timing_cache = None
         self._load_timing_cache()
         self._shape_configs = {}
diff --git a/noxfile.py b/noxfile.py
index 1263bbfac97..0d44a470543 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -127,6 +127,15 @@ def partial_unit(session, subset):
 
 # Their setup.py files hardcode -std=c++17, but the torch headers of recent containers need C++20.
 _BUILD_AS_CXX20 = ("mamba_ssm", "causal_conv1d")
+# Without these, setup.py may download a matching prebuilt wheel instead of compiling the patched source.
+_FORCE_BUILD = {
+    "mamba_ssm": {"MAMBA_FORCE_BUILD": "TRUE"},
+    "causal_conv1d": {"CAUSAL_CONV1D_FORCE_BUILD": "TRUE"},
+}
+
+
+def _wheelhouse():
+    return os.environ.get("MODELOPT_WHEELHOUSE", os.path.expanduser("~/.cache/wheelhouse"))
 
 
 def _build_wheel(session, name, wheelhouse):
@@ -150,19 +159,31 @@ def _build_wheel(session, name, wheelhouse):
             patched = f.read().replace("-std=c++17", "-std=c++20")
         with open(setup_py, "w") as f:
             f.write(patched)
-        session.run(*pip_wheel, source)
+        session.run(*pip_wheel, source, env=_FORCE_BUILD[name])
+
+
+def _find_wheel(wheelhouse, name):
+    # ``-[0-9]`` so ``causal_conv1d`` does not also match a ``causal_conv1d_cuda-*`` wheel.
+    found = sorted(glob.glob(os.path.join(wheelhouse, f"{name}-[0-9]*.whl")))
+    return found[-1] if found else None
 
 
 def _install_built_from_source(session, *names):
     """Install the named ``_SOURCE_BUILT_WHEELS``, building a wheel only if the wheelhouse lacks it.
 
-    CI caches ``MODELOPT_WHEELHOUSE`` per container image, so only the first run on an image builds.
+    CI caches the wheelhouse (``MODELOPT_WHEELHOUSE``) per container image, so only the first run on
+    an image builds.
     """
-    wheelhouse = os.environ.get("MODELOPT_WHEELHOUSE", os.path.expanduser("~/.cache/wheelhouse"))
+    wheelhouse = _wheelhouse()
+    wheels = []
     for name in names:
-        if not glob.glob(os.path.join(wheelhouse, f"{name}-*.whl")):
+        wheel = _find_wheel(wheelhouse, name)
+        if wheel is None:
             _build_wheel(session, name, wheelhouse)
-    wheels = [glob.glob(os.path.join(wheelhouse, f"{name}-*.whl"))[0] for name in names]
+            wheel = _find_wheel(wheelhouse, name)
+        if wheel is None:
+            raise RuntimeError(f"No wheel for {name!r} in {wheelhouse} after building it")
+        wheels.append(wheel)
     session.run("python", "-m", "pip", "install", *wheels)
 
 
@@ -176,6 +197,10 @@ def gpu(session):
     session.run("python", "-m", "pip", "uninstall", "-y", "cupy-cuda12x")
     session.run("python", "-m", "pip", "install", "cupy-cuda13x")
     _install_built_from_source(session, "mamba_ssm", "causal_conv1d")
+    # Every wheel is built: CI saves the wheelhouse only then, so a run that stopped halfway cannot
+    # freeze a partial wheelhouse under the cache key.
+    with open(os.path.join(_wheelhouse(), ".complete"), "w"):
+        pass
     session.run("python", "-m", "pytest", "tests/gpu", *_cov_args())
 
 
diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py
index effc979154e..4d57929d3dd 100644
--- a/tests/_test_utils/torch/distributed/utils.py
+++ b/tests/_test_utils/torch/distributed/utils.py
@@ -202,10 +202,10 @@ def _worker_loop(rank, world_size, backend, port, cmd_queue, result_queue, teard
                     try:
                         teardown_fn(rank, world_size)
                     except Exception as e:
-                        print(f"Error tearing down worker: {e}")
-                        status = "error"
-                        teardown_tb = traceback.format_exc()
-                        tb = (tb + "\n" if tb else "") + f"[teardown] {teardown_tb}"
+                        print(f"Error tearing down worker: {e}\n{traceback.format_exc()}")
+                        if not reported:  # else the job failure is already sent; the teardown error is printed above
+                            status = "error"
+                            tb = f"[teardown] {traceback.format_exc()}"
             if not reported:
                 result_queue.put((status, rank, tb))
 
diff --git a/tests/gpu/torch/utils/test_distributed_worker_pool.py b/tests/gpu/torch/utils/test_distributed_worker_pool.py
index 166ce10e92e..0983f790b8e 100644
--- a/tests/gpu/torch/utils/test_distributed_worker_pool.py
+++ b/tests/gpu/torch/utils/test_distributed_worker_pool.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -59,7 +59,7 @@ def test_a_failed_rank_is_reported_even_when_its_teardown_blocks_on_a_hung_peer(
     pool = DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=_teardown_barrier)
     try:
         pool.failure_grace = 2
-        pool.job_timeout = 30  # without the early report the parent would wait this long
+        pool.job_timeout = 120  # without the early report the parent would wait this long
         start = time.monotonic()
         with pytest.raises(RuntimeError, match="rank 0 failed"):
             pool.run(_rank_0_fails_and_rank_1_hangs)
diff --git a/tests/unit/torch/utils/test_get_live_pool.py b/tests/unit/torch/utils/test_get_live_pool.py
index 1431a9e2a52..2166d5ee560 100644
--- a/tests/unit/torch/utils/test_get_live_pool.py
+++ b/tests/unit/torch/utils/test_get_live_pool.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

From db6896798ff9a02d2efbf1d8df6b575a71e1732b Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 06:05:32 -0700
Subject: [PATCH 11/34] trtllm examples: install fla-core for TensorRT-LLM rc29

rc29's quickstart_multimodal.py imports fla (flash-linear-attention) through its Kimi model code, which the image does not ship; test_qwen_vl failed with ModuleNotFoundError.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/example_tests.yml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml
index 66688c093a8..161ef573788 100644
--- a/.github/workflows/example_tests.yml
+++ b/.github/workflows/example_tests.yml
@@ -171,7 +171,8 @@ jobs:
     with:
       docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc29"
       example: ${{ matrix.example }}
-      pip_install_extras: "[hf,dev-test]"
+      # dev-fla: rc29's multimodal quickstart imports `fla`, which the image does not ship
+      pip_install_extras: "[hf,dev-test,dev-fla]"
       runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
       allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), format(',{0},', matrix.example)) }}
 

From 4ffb17e586e601fd6bf9fecc6809c4088669bb51 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 06:14:01 -0700
Subject: [PATCH 12/34] Pin TensorRT-LLM to rc23; build the conv3d
 implicit-GEMM extension without a hardcoded C++17

- rc23 is the newest TensorRT-LLM release without the Kimi KDA models (rc24+), whose fla import the image does not ship and which breaks quickstart_multimodal.py; drop the dev-fla workaround
- implicit_gemm_cuda.py passed -std=c++17 to the JIT build, which fails against torch headers that need C++20 (pytorch:26.09: 165 failures); torch's load() adds the standard it needs

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/example_tests.yml                        | 7 ++++---
 .github/workflows/gpu_tests.yml                            | 2 +-
 .../torch/kernels/quantization/conv/implicit_gemm_cuda.py  | 3 ++-
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml
index 161ef573788..6fbb1d13a93 100644
--- a/.github/workflows/example_tests.yml
+++ b/.github/workflows/example_tests.yml
@@ -169,10 +169,11 @@ jobs:
       contents: read
     secrets: inherit
     with:
-      docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc29"
+      # rc23 is the newest without the Kimi KDA models (rc24+), whose `fla` import the image does not
+      # ship and which breaks quickstart_multimodal.py; rc29 also removed tensorrt_llm._torch.auto_deploy.
+      docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc23"
       example: ${{ matrix.example }}
-      # dev-fla: rc29's multimodal quickstart imports `fla`, which the image does not ship
-      pip_install_extras: "[hf,dev-test,dev-fla]"
+      pip_install_extras: "[hf,dev-test]"
       runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
       allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), format(',{0},', matrix.example)) }}
 
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index 6315c26a455..f3c51d03568 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -51,7 +51,7 @@ jobs:
             container_image: nvcr.io/nvidia/nemo:26.08
           - example: gpu_trtllm
             timeout: 15
-            container_image: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc29
+            container_image: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc23
           # Covers the RoutedExperts layout (0.24+)
           - example: gpu_vllm
             timeout: 15
diff --git a/modelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.py b/modelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.py
index 713b5f82f9b..8c6d3dfacbf 100644
--- a/modelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.py
+++ b/modelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.py
@@ -59,7 +59,8 @@ def _get_cuda_module():
                 "--use_fast_math",
                 "-lineinfo",
                 "--ptxas-options=-v",
-                "-std=c++17",
+                # No -std flag: torch's load() adds the C++ standard its headers need (C++20 in
+                # recent releases), and a hardcoded -std=c++17 fails to build there.
             ],
         )
     return _cuda_module

From a04bd58df2d0dd821171ed17a4831fc9b1b3742c Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 06:22:27 -0700
Subject: [PATCH 13/34] persistent_materialization: skip FSDP2 unshard only for
 the materialized layer

_disable_fsdp_unshard_reshard patched FSDPParamGroup.unshard/reshard on the class, so while one layer
was materialized every other FSDP module also skipped unshard: a full forward (layerwise calibration
collecting the next layer's inputs) ran with sharded DTensor params, and on torch 2.14 failed with
"'FSDPParam' object has no attribute '_unsharded_param'" (test_layerwise_calibrate_fsdp2 on pytorch:26.09).
Patch the param groups inside the layer instead; CPU/gloo regression test.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../torch/quantization/utils/core_utils.py    | 31 +++++++++++--------
 tests/unit/torch/quantization/test_utils.py   | 22 +++++++++++++
 2 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py
index 63f349dd10a..4923b687f7f 100644
--- a/modelopt/torch/quantization/utils/core_utils.py
+++ b/modelopt/torch/quantization/utils/core_utils.py
@@ -828,20 +828,25 @@ def _init_mp_dtypes(self) -> None:
 
 @contextmanager
 def _disable_fsdp_unshard_reshard(layer):
-    """Disable FSDP2 unshard/reshard if *layer* is FSDP-wrapped."""
-    if isinstance(layer, FSDPModule):
-        _pg_cls = torch.distributed.fsdp._fully_shard._fsdp_param_group.FSDPParamGroup
-        orig_unshard = _pg_cls.unshard
-        orig_reshard = _pg_cls.reshard
-        _pg_cls.unshard = lambda self, async_op=False: None
-        _pg_cls.reshard = lambda self: None
-        try:
-            yield
-        finally:
-            _pg_cls.unshard = orig_unshard
-            _pg_cls.reshard = orig_reshard
-    else:
+    """Make the FSDP2 param groups inside *layer* skip unshard/reshard (already gathered).
+
+    Patched per group instance, not on ``FSDPParamGroup``: the rest of the model still runs forward
+    (e.g. while collecting the next layer's inputs) and must unshard as usual.
+    """
+    groups = [
+        group
+        for module in layer.modules()
+        if isinstance(module, FSDPModule)
+        and (group := fully_shard.state(module)._fsdp_param_group) is not None
+    ]
+    for group in groups:
+        group.unshard = lambda async_op=False: None
+        group.reshard = lambda: None
+    try:
         yield
+    finally:
+        for group in groups:
+            del group.unshard, group.reshard
 
 
 def build_param_index(model):
diff --git a/tests/unit/torch/quantization/test_utils.py b/tests/unit/torch/quantization/test_utils.py
index 933c553ce12..53d664917fd 100644
--- a/tests/unit/torch/quantization/test_utils.py
+++ b/tests/unit/torch/quantization/test_utils.py
@@ -15,12 +15,15 @@
 
 import pytest
 import torch
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.fsdp import fully_shard
 
 from modelopt.torch.quantization.utils import (
     convert_quantization_axis_to_reduce_axis,
     reduce_amax,
     reduce_block_amax,
 )
+from modelopt.torch.quantization.utils.core_utils import persistent_materialization
 from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
 
 
@@ -333,3 +336,22 @@ def _forward_loop(m):
     assert "L1" not in call_log_for_layer1
 
     assert torch.allclose(inp2[0][0][0], torch.tensor([[10.0 + 1.0 + 2.0]]))
+
+
+def _forward_through_other_layers_while_one_is_materialized(rank, size):
+    mesh = init_device_mesh("cpu", (size,))
+    torch.manual_seed(0)
+    model = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8))
+    for layer in model:
+        fully_shard(layer, mesh=mesh)
+    fully_shard(model, mesh=mesh)
+    x = torch.randn(2, 8)
+    expected = model(x)
+    # Layerwise calibration runs the whole model to collect the next layer's inputs while one layer
+    # is materialized; the other FSDP modules must still unshard.
+    with persistent_materialization(model[0]):
+        torch.testing.assert_close(model(x), expected)
+
+
+def test_persistent_materialization_only_skips_unshard_for_its_layer(cpu_dist_workers):
+    cpu_dist_workers.run(_forward_through_other_layers_while_one_is_materialized)

From 878ffc54f0499fd5d1c99edb8c36c3c3f3b4c5a2 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 06:22:33 -0700
Subject: [PATCH 14/34] noxfile: drop the stale cppimport note

test_onnx_extensions.py passes on pytorch:26.09.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 noxfile.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/noxfile.py b/noxfile.py
index 0d44a470543..0a5001f8c8a 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -192,8 +192,6 @@ def _install_built_from_source(session, *names):
 def gpu(session):
     _install_built_from_source(session, "fast_hadamard_transform")
     session.run("python", "-m", "pip", "install", "-e", ".[all,dev-test,dev-fla]")
-    # tests/gpu/_extensions/test_onnx_extensions.py fails for newer containers
-    # until https://github.com/tbenthompson/cppimport/pull/98
     session.run("python", "-m", "pip", "uninstall", "-y", "cupy-cuda12x")
     session.run("python", "-m", "pip", "install", "cupy-cuda13x")
     _install_built_from_source(session, "mamba_ssm", "causal_conv1d")

From 50615b34a9af06708d5b64e108ff05e6b67c21ff Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 06:54:27 -0700
Subject: [PATCH 15/34] Address review: retire a pool after any failed job;
 restore nested FSDP unshard/reshard overrides

- A failed rank reports before its teardown, so even when every rank replied a worker may still be blocked; kill the pool on any failure (get_live_pool replaces it), dumping stacks only for hangs
- _disable_fsdp_unshard_reshard nests (persistent_materialization(writeback=False) enters it twice) and fully_shard([a, b]) shares a group: dedupe groups and restore the prior instance attributes instead of deleting them (the writeback=False path raised AttributeError)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../torch/quantization/utils/core_utils.py    | 17 ++++++++++++----
 tests/_test_utils/torch/distributed/utils.py  | 20 +++++++++++--------
 .../utils/test_distributed_worker_pool.py     | 15 ++++++++++++++
 tests/unit/torch/quantization/test_utils.py   | 11 ++++++----
 4 files changed, 47 insertions(+), 16 deletions(-)

diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py
index 4923b687f7f..df7c0ffae23 100644
--- a/modelopt/torch/quantization/utils/core_utils.py
+++ b/modelopt/torch/quantization/utils/core_utils.py
@@ -833,11 +833,16 @@ def _disable_fsdp_unshard_reshard(layer):
     Patched per group instance, not on ``FSDPParamGroup``: the rest of the model still runs forward
     (e.g. while collecting the next layer's inputs) and must unshard as usual.
     """
-    groups = [
-        group
+    groups = {
+        id(group): group  # ``fully_shard([a, b])`` gives several modules one group
         for module in layer.modules()
         if isinstance(module, FSDPModule)
         and (group := fully_shard.state(module)._fsdp_param_group) is not None
+    }.values()
+    # Contexts nest (``persistent_materialization(writeback=False)`` enters this twice), so restore
+    # whatever instance attributes were there before instead of deleting them.
+    saved = [
+        (group, {k: group.__dict__.get(k) for k in ("unshard", "reshard")}) for group in groups
     ]
     for group in groups:
         group.unshard = lambda async_op=False: None
@@ -845,8 +850,12 @@ def _disable_fsdp_unshard_reshard(layer):
     try:
         yield
     finally:
-        for group in groups:
-            del group.unshard, group.reshard
+        for group, prior in saved:
+            for name, value in prior.items():
+                if value is None:
+                    del group.__dict__[name]
+                else:
+                    setattr(group, name, value)
 
 
 def build_param_index(model):
diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py
index 4d57929d3dd..02d5f6e8869 100644
--- a/tests/_test_utils/torch/distributed/utils.py
+++ b/tests/_test_utils/torch/distributed/utils.py
@@ -217,11 +217,11 @@ def run(self, fn, *args, **kwargs):
         ``fn`` is called as ``fn(rank, world_size, *args, **kwargs)`` and must be picklable
         (top-level function or ``functools.partial`` of one).
 
-        If the wait is interrupted (e.g. by ``pytest-timeout``), or a rank fails and its peers do
-        not finish within ``failure_grace`` (they are typically blocked in a collective waiting for
-        it), the workers are still inside the job and every later ``run`` would hang until its own
-        timeout. The workers are killed instead and the pool is marked ``broken``; callers that
-        want to continue need a new pool.
+        If the job fails or the wait is interrupted (e.g. by ``pytest-timeout``), the workers may
+        still be inside the job or its teardown, typically blocked in a collective, and every later
+        ``run`` would hang until its own timeout. The workers are killed instead and the pool is
+        marked ``broken``; callers that want to continue need a new pool (``get_live_pool``). After
+        a rank fails, its peers get ``failure_grace`` seconds to report before the kill.
         """
         if self.broken:
             raise RuntimeError("Worker pool was killed after an earlier job hung; create a new one")
@@ -246,12 +246,16 @@ def run(self, fn, *args, **kwargs):
             raise
 
         if errors:
+            if not self.broken:
+                # Every rank replied, but a failed rank reports before its teardown, which can still
+                # be blocked; never hand these workers to another job.
+                self._abort(dump_stacks=False)
             raise RuntimeError("Worker(s) failed:\n" + "\n".join(errors))
 
-    def _abort(self):
-        """Print every worker's stack (to stderr, for the CI log), then kill the workers."""
+    def _abort(self, dump_stacks=True):
+        """Kill the workers, first printing their stacks (to stderr, for the CI log) if asked."""
         self.broken = True
-        if hasattr(signal, "SIGUSR1"):
+        if dump_stacks and hasattr(signal, "SIGUSR1"):
             for p in self._processes:
                 with contextlib.suppress(OSError):
                     os.kill(p.pid, signal.SIGUSR1)
diff --git a/tests/gpu/torch/utils/test_distributed_worker_pool.py b/tests/gpu/torch/utils/test_distributed_worker_pool.py
index 0983f790b8e..36dce35968f 100644
--- a/tests/gpu/torch/utils/test_distributed_worker_pool.py
+++ b/tests/gpu/torch/utils/test_distributed_worker_pool.py
@@ -33,6 +33,10 @@ def _rank_0_fails_and_rank_1_hangs(rank, size):
     time.sleep(600)
 
 
+def _fail(rank, size):
+    raise ValueError("job failed")
+
+
 def _teardown_barrier(rank, size):
     dist.barrier()  # like megatron_worker_teardown: blocks while a peer is still in the job
 
@@ -68,3 +72,14 @@ def test_a_failed_rank_is_reported_even_when_its_teardown_blocks_on_a_hung_peer(
         assert pool.broken
     finally:
         pool.shutdown()
+
+
+def test_a_failed_job_retires_the_pool_even_when_every_rank_replied():
+    """A failed rank reports before its teardown, which may still be blocked; never reuse it."""
+    pool = DistributedWorkerPool(world_size=1, backend="gloo", teardown_fn=None)
+    try:
+        with pytest.raises(RuntimeError, match="job failed"):
+            pool.run(_fail)
+        assert pool.broken
+    finally:
+        pool.shutdown()
diff --git a/tests/unit/torch/quantization/test_utils.py b/tests/unit/torch/quantization/test_utils.py
index 53d664917fd..1591ed51632 100644
--- a/tests/unit/torch/quantization/test_utils.py
+++ b/tests/unit/torch/quantization/test_utils.py
@@ -338,7 +338,7 @@ def _forward_loop(m):
     assert torch.allclose(inp2[0][0][0], torch.tensor([[10.0 + 1.0 + 2.0]]))
 
 
-def _forward_through_other_layers_while_one_is_materialized(rank, size):
+def _forward_through_other_layers_while_one_is_materialized(rank, size, writeback):
     mesh = init_device_mesh("cpu", (size,))
     torch.manual_seed(0)
     model = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8))
@@ -349,9 +349,12 @@ def _forward_through_other_layers_while_one_is_materialized(rank, size):
     expected = model(x)
     # Layerwise calibration runs the whole model to collect the next layer's inputs while one layer
     # is materialized; the other FSDP modules must still unshard.
-    with persistent_materialization(model[0]):
+    with persistent_materialization(model[0], writeback=writeback):
         torch.testing.assert_close(model(x), expected)
+    torch.testing.assert_close(model(x), expected)  # FSDP unshard/reshard restored on exit
 
 
-def test_persistent_materialization_only_skips_unshard_for_its_layer(cpu_dist_workers):
-    cpu_dist_workers.run(_forward_through_other_layers_while_one_is_materialized)
+# writeback=False nests the unshard/reshard override (via _fsdp2_unshard_context).
+@pytest.mark.parametrize("writeback", [True, False])
+def test_persistent_materialization_only_skips_unshard_for_its_layer(cpu_dist_workers, writeback):
+    cpu_dist_workers.run(_forward_through_other_layers_while_one_is_materialized, writeback)

From 635c76148dbe6f486fba5326d68cf225cd04daa7 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 07:54:50 -0700
Subject: [PATCH 16/34] NVFP4: return a contiguous scale from
 cutlass_fp4_scale_to_modelopt_fp4_scale

The slice of the padded tile layout is a strided view whenever padding was needed, and fp4_dequantize indexes the scale pointer flatly.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 modelopt/torch/quantization/qtensor/nvfp4_tensor.py | 3 ++-
 tests/unit/torch/quantization/test_nvfp4_tensor.py  | 1 +
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py
index 80894db888a..8726dc5c719 100644
--- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py
+++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py
@@ -60,7 +60,8 @@ def cutlass_fp4_scale_to_modelopt_fp4_scale(
     row_blocks, col_blocks = -(-m // 128), -(-n // 4)
     tiles = cutlass_scale.reshape(row_blocks, col_blocks, 32, 4, 4)
     padded = tiles.permute(0, 3, 2, 1, 4).reshape(row_blocks * 128, col_blocks * 4)
-    return padded[:m, :n].view(torch.float8_e4m3fn)
+    # contiguous(): the slice is a strided view, and fp4_dequantize indexes the scale flatly.
+    return padded[:m, :n].contiguous().view(torch.float8_e4m3fn)
 
 
 class NVFP4QTensor(BaseQuantizedTensor):
diff --git a/tests/unit/torch/quantization/test_nvfp4_tensor.py b/tests/unit/torch/quantization/test_nvfp4_tensor.py
index a987bcafad3..a7b64515e3a 100644
--- a/tests/unit/torch/quantization/test_nvfp4_tensor.py
+++ b/tests/unit/torch/quantization/test_nvfp4_tensor.py
@@ -148,4 +148,5 @@ def test_cutlass_fp4_scale_to_modelopt_fp4_scale_undoes_the_swizzle(weight_shape
     restored = cutlass_fp4_scale_to_modelopt_fp4_scale(_cublas_swizzle(scale), weight_shape)
 
     assert restored.dtype == torch.float8_e4m3fn
+    assert restored.is_contiguous()  # consumers index it flatly
     assert torch.equal(restored.view(torch.uint8), scale)

From f181890ba1828088e9c2d992d59812ba8611a8de Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 07:54:56 -0700
Subject: [PATCH 17/34] Keep TensorRT-LLM at rc20

No later release runs the multimodal quickstart that test_qwen_vl uses: rc21-rc23 import MultimodalConfig, which tensorrt_llm.llmapi does not export yet, and rc24+ import fla, which the image lacks.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/example_tests.yml | 6 +++---
 .github/workflows/gpu_tests.yml     | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml
index 6fbb1d13a93..e1b50be0984 100644
--- a/.github/workflows/example_tests.yml
+++ b/.github/workflows/example_tests.yml
@@ -169,9 +169,9 @@ jobs:
       contents: read
     secrets: inherit
     with:
-      # rc23 is the newest without the Kimi KDA models (rc24+), whose `fla` import the image does not
-      # ship and which breaks quickstart_multimodal.py; rc29 also removed tensorrt_llm._torch.auto_deploy.
-      docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc23"
+      # Newer releases break quickstart_multimodal.py (test_qwen_vl): rc21-rc23 import MultimodalConfig,
+      # which tensorrt_llm.llmapi does not export yet, and rc24+ import `fla`, which the image lacks.
+      docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20"
       example: ${{ matrix.example }}
       pip_install_extras: "[hf,dev-test]"
       runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index f3c51d03568..c589fd375fb 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -51,7 +51,7 @@ jobs:
             container_image: nvcr.io/nvidia/nemo:26.08
           - example: gpu_trtllm
             timeout: 15
-            container_image: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc23
+            container_image: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20
           # Covers the RoutedExperts layout (0.24+)
           - example: gpu_vllm
             timeout: 15

From 9d3d714ab6a7eaf632082265844fc5f7111e4469 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 13:24:13 -0700
Subject: [PATCH 18/34] gpu_megatron: destroy each test's process groups in the
 worker teardown

Megatron's destroy_model_parallel() only drops its references to the NCCL groups, so communicators
piled up in the persistent workers: ~20 groups and ~0.4 GiB of device memory per test (1221 groups /
24 GiB after 61 tests of test_megatron.py), slowing later tests until they hit the 120 s timeout
(the gpu_megatron nightly timeouts) and causing NCCL CUDA errors on smaller GPUs. With the groups
destroyed, test_megatron.py stays at 1 group / ~1-3 GiB and passes 131/131 in 4:54 on 2 GPUs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 tests/gpu_megatron/conftest.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py
index 51261e30023..6b6ab484475 100644
--- a/tests/gpu_megatron/conftest.py
+++ b/tests/gpu_megatron/conftest.py
@@ -79,6 +79,12 @@ def megatron_worker_teardown(rank, world_size):
             apex_destroy()
         except Exception as e:
             print(f"Error destroying model parallel with Apex: {e}")
+    # destroy_model_parallel() only drops Megatron's references to its NCCL groups, so every test's
+    # communicators (and their GPU buffers) piled up in the persistent workers (~20 groups and
+    # ~0.4 GiB per test), slowing later tests until they timed out. Destroy all but the default.
+    for group in list(torch.distributed.distributed_c10d._world.pg_map):
+        if group is not torch.distributed.GroupMember.WORLD:
+            torch.distributed.destroy_process_group(group)
     torch.cuda.empty_cache()
 
 

From 7fc059581353879391c6d41cb54a48f84f8c31df Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 22:57:24 -0700
Subject: [PATCH 19/34] Use a fixed tile config for the NVFP4 FP8 sweep kernel;
 drop unneeded xfails

Autotuning compiled all three configs on first use (~70s per dtype on the
pytorch:26.09 CI runner), pushing the MSE sweep tests past the 120s timeout.
BPP=64/nw=8 was fastest on B300 and within noise elsewhere.

The preprocess_data xfails xpassed in every gpu_megatron run over the last week.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../quantization/gemm/nvfp4_fp8_sweep.py      | 21 +++++-------
 .../plugins/test_megatron_preprocess_data.py  | 33 -------------------
 2 files changed, 8 insertions(+), 46 deletions(-)

diff --git a/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py b/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py
index a56588b21d8..c9ee0aa24bd 100644
--- a/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py
+++ b/modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py
@@ -24,8 +24,6 @@
 the per-block scale is the identity, so the kernel can use
 ``scale = candidate * global_amax / 6.0`` without an explicit FP8 cast — making it
 runnable on any CUDA GPU with Triton (no ``tl.float8e4nv`` requirement).
-
-Tile shape (``BLOCKS_PER_PROGRAM``) and ``num_warps`` are autotuned per ``N_BLOCKS``.
 """
 
 import torch
@@ -43,18 +41,13 @@
 ]
 
 
-# Selected from a (BLOCKS_PER_PROGRAM, num_warps) sweep on B300:
-#   BPP=16,nw=2: 6.06 ms   BPP=32,nw=4: 6.06 ms   BPP=64,nw=8: 5.08 ms
-# The smaller-tile entries cover cases where N_BLOCKS is small enough that BPP=64
-# would underfill the SMs.
-_FP8_SWEEP_AUTOTUNE_CONFIGS = [
-    triton.Config({"BLOCKS_PER_PROGRAM": 16}, num_warps=2),
-    triton.Config({"BLOCKS_PER_PROGRAM": 32}, num_warps=4),
-    triton.Config({"BLOCKS_PER_PROGRAM": 64}, num_warps=8),
-]
+# Fastest in a (BLOCKS_PER_PROGRAM, num_warps) sweep on B300 (BPP=16/32: 6.06 ms, BPP=64: 5.08 ms)
+# and within noise of the others on small shapes. Hard-coded since autotuning compiles every
+# candidate (~10 s each) on first use.
+_FP8_SWEEP_BLOCKS_PER_PROGRAM = 64
+_FP8_SWEEP_NUM_WARPS = 8
 
 
-@triton.autotune(configs=_FP8_SWEEP_AUTOTUNE_CONFIGS, key=["N_BLOCKS"])
 @triton.jit
 def _fp8_scale_sweep_kernel(
     x_ptr,  # [N_BLOCKS * BLOCK_SIZE], any float dtype (loaded as fp32)
@@ -148,7 +141,7 @@ def nvfp4_fp8_scale_sweep(
     candidates = fp8_scale_candidates(x.device).to(dtype=torch.float32)
     global_amax_f32 = global_amax.detach().to(device=x.device, dtype=torch.float32).reshape(1)
 
-    grid = lambda meta: (triton.cdiv(n_blocks, meta["BLOCKS_PER_PROGRAM"]),)
+    grid = (triton.cdiv(n_blocks, _FP8_SWEEP_BLOCKS_PER_PROGRAM),)
     with torch.cuda.device(x.device):
         _fp8_scale_sweep_kernel[grid](
             x_flat,
@@ -158,6 +151,8 @@ def nvfp4_fp8_scale_sweep(
             n_blocks,
             BLOCK_SIZE=block_size,
             NUM_CANDIDATES=int(candidates.numel()),
+            BLOCKS_PER_PROGRAM=_FP8_SWEEP_BLOCKS_PER_PROGRAM,
+            num_warps=_FP8_SWEEP_NUM_WARPS,
         )
     return best_amax
 
diff --git a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py
index 5dd0815080c..b676e1e3f85 100644
--- a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py
+++ b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py
@@ -40,15 +40,6 @@ def test_worker_context_falls_back_to_spawn_without_forkserver(monkeypatch):
     _prestart_forkserver()
 
 
-# Depends on the external HF datasets-server, which intermittently returns 5xx
-# (e.g. 503) on the /splits lookup. Allow-fail the transient outage while keeping
-# real assertion failures visible (raises=RuntimeError only, strict=False so a
-# normal pass doesn't xpass-fail). download_hf_dataset_as_jsonl already retries.
-@pytest.mark.xfail(
-    raises=RuntimeError,
-    strict=False,
-    reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
-)
 def test_megatron_preprocess_data_with_jsonl_path(tmp_path):
     input_jsonl = tmp_path / "minipile.jsonl"
     input_jsonl.write_text(
@@ -109,12 +100,6 @@ def test_megatron_preprocess_data_jsonl_stops_at_max_tokens(tmp_path):
     assert limited_size < full_size
 
 
-# Allow-fail transient HF datasets-server 5xx (see the /splits note above).
-@pytest.mark.xfail(
-    raises=RuntimeError,
-    strict=False,
-    reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
-)
 def test_megatron_preprocess_data_hf_split_stops_at_max_tokens(tmp_path):
     common_args = {
         "hf_dataset": "nanotron/minipile_100_samples",
@@ -141,12 +126,6 @@ def test_megatron_preprocess_data_hf_split_stops_at_max_tokens(tmp_path):
     assert limited_size < full_size
 
 
-# Allow-fail transient HF datasets-server 5xx (see the /splits note above).
-@pytest.mark.xfail(
-    raises=RuntimeError,
-    strict=False,
-    reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
-)
 def test_megatron_preprocess_data_hf_split_resume_uses_cached_token_count(tmp_path):
     args = {
         "hf_dataset": "nanotron/minipile_100_samples",
@@ -168,12 +147,6 @@ def test_megatron_preprocess_data_hf_split_resume_uses_cached_token_count(tmp_pa
     assert set(tmp_path.iterdir()) == cached_files
 
 
-# Allow-fail transient HF datasets-server 5xx (see the /splits note above).
-@pytest.mark.xfail(
-    raises=RuntimeError,
-    strict=False,
-    reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
-)
 @pytest.mark.parametrize(
     ("hf_dataset", "hf_split", "json_keys"),
     [
@@ -318,12 +291,6 @@ def test_megatron_preprocess_data_tool_calls_arguments_normalized(tmp_path):
     )
 
 
-# Allow-fail transient HF datasets-server 5xx (see the /splits note above).
-@pytest.mark.xfail(
-    raises=RuntimeError,
-    strict=False,
-    reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
-)
 def test_megatron_preprocess_data_hf_streaming_warning(tmp_path):
     # hf_streaming without hf_max_samples_per_split should warn and fall back to non-streaming
     with pytest.warns(UserWarning, match="hf_streaming"):

From 528c39ea1ca7105220a231d6bfbae07ed16331f9 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 16 Oct 2026 23:04:03 -0700
Subject: [PATCH 20/34] Test the FSDP2 unshard/reshard override with a shared
 param group and nesting

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 tests/unit/torch/quantization/test_utils.py | 25 ++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/tests/unit/torch/quantization/test_utils.py b/tests/unit/torch/quantization/test_utils.py
index 1591ed51632..4366a062ed2 100644
--- a/tests/unit/torch/quantization/test_utils.py
+++ b/tests/unit/torch/quantization/test_utils.py
@@ -23,7 +23,10 @@
     reduce_amax,
     reduce_block_amax,
 )
-from modelopt.torch.quantization.utils.core_utils import persistent_materialization
+from modelopt.torch.quantization.utils.core_utils import (
+    _disable_fsdp_unshard_reshard,
+    persistent_materialization,
+)
 from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
 
 
@@ -358,3 +361,23 @@ def _forward_through_other_layers_while_one_is_materialized(rank, size, writebac
 @pytest.mark.parametrize("writeback", [True, False])
 def test_persistent_materialization_only_skips_unshard_for_its_layer(cpu_dist_workers, writeback):
     cpu_dist_workers.run(_forward_through_other_layers_while_one_is_materialized, writeback)
+
+
+def _shared_group_override_restores_prior_attributes(rank, size):
+    mesh = init_device_mesh("cpu", (size,))
+    block = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8))
+    fully_shard(list(block), mesh=mesh)  # both linears share one FSDPParamGroup
+    group = fully_shard.state(block[0])._fsdp_param_group
+    prior_reshard = group.reshard  # pre-existing instance override must survive
+    group.reshard = prior_reshard
+    with _disable_fsdp_unshard_reshard(block):
+        with _disable_fsdp_unshard_reshard(block):
+            pass
+        assert group.unshard() is None  # outer no-op survives the inner exit
+        assert "unshard" in group.__dict__
+    assert "unshard" not in group.__dict__
+    assert group.__dict__["reshard"] is prior_reshard
+
+
+def test_disable_fsdp_unshard_reshard_shared_group_and_nesting(cpu_dist_workers):
+    cpu_dist_workers.run(_shared_group_override_restores_prior_attributes)

From 1203fcd76ea3626538645efc58e502c87678167a Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 17 Oct 2026 01:15:40 -0700
Subject: [PATCH 21/34] Key the torch-extension cache on the C++/CUDA sources
 only

Editing a Triton kernel invalidated the cache, so the gpu_vllm job rebuilt the
extensions and hit its 15-min timeout.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/actions/cache-extensions/action.yml | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/.github/actions/cache-extensions/action.yml b/.github/actions/cache-extensions/action.yml
index bc7584e813f..a25dcb5fafc 100644
--- a/.github/actions/cache-extensions/action.yml
+++ b/.github/actions/cache-extensions/action.yml
@@ -23,8 +23,10 @@ runs:
       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')
-          }}
+        # Only the JIT-built sources and their build scripts, so Triton-only edits keep the cache.
+        key: torch-ext-${{ inputs.cache-key }}-${{ hashFiles('modelopt/torch/kernels/quantization/**/*.cpp', 'modelopt/torch/kernels/quantization/**/*.cu',
+          'modelopt/torch/kernels/quantization/**/*.cuh', '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'

From 0ef079a96f97f543515a7fadbcc3d44d7598771c Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 17 Oct 2026 01:16:59 -0700
Subject: [PATCH 22/34] Drop the vLLM v0.20.0 gpu_vllm job

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/gpu_tests.yml | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index c589fd375fb..e2822ab67c2 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -52,14 +52,9 @@ jobs:
           - example: gpu_trtllm
             timeout: 15
             container_image: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20
-          # Covers the RoutedExperts layout (0.24+)
           - example: gpu_vllm
             timeout: 15
             container_image: docker.io/vllm/vllm-openai:v0.30.0
-          # Covers the legacy FusedMoE / SharedFusedMoE registration branches (vLLM < 0.24)
-          - example: gpu_vllm
-            timeout: 15
-            container_image: docker.io/vllm/vllm-openai:v0.20.0
     runs-on: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
     timeout-minutes: ${{ matrix.timeout }}
     container:

From f222d46ee99ac336711a861181524cf224043dcd Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 17 Oct 2026 01:55:43 -0700
Subject: [PATCH 23/34] Only skip FSDP unshard/reshard when the materialized
 layer is FSDP-wrapped

persistent_materialization gathers nothing for a layer that is not itself an
FSDPModule, so its FSDP submodules must keep unsharding (the previous behavior).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 modelopt/torch/quantization/utils/core_utils.py | 3 +++
 tests/unit/torch/quantization/test_utils.py     | 4 ++++
 2 files changed, 7 insertions(+)

diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py
index df7c0ffae23..c0f3884f41d 100644
--- a/modelopt/torch/quantization/utils/core_utils.py
+++ b/modelopt/torch/quantization/utils/core_utils.py
@@ -833,6 +833,9 @@ def _disable_fsdp_unshard_reshard(layer):
     Patched per group instance, not on ``FSDPParamGroup``: the rest of the model still runs forward
     (e.g. while collecting the next layer's inputs) and must unshard as usual.
     """
+    if not isinstance(layer, FSDPModule):  # nothing was gathered for its FSDP submodules
+        yield
+        return
     groups = {
         id(group): group  # ``fully_shard([a, b])`` gives several modules one group
         for module in layer.modules()
diff --git a/tests/unit/torch/quantization/test_utils.py b/tests/unit/torch/quantization/test_utils.py
index 4366a062ed2..0caab634562 100644
--- a/tests/unit/torch/quantization/test_utils.py
+++ b/tests/unit/torch/quantization/test_utils.py
@@ -367,6 +367,7 @@ def _shared_group_override_restores_prior_attributes(rank, size):
     mesh = init_device_mesh("cpu", (size,))
     block = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8))
     fully_shard(list(block), mesh=mesh)  # both linears share one FSDPParamGroup
+    fully_shard(block, mesh=mesh)
     group = fully_shard.state(block[0])._fsdp_param_group
     prior_reshard = group.reshard  # pre-existing instance override must survive
     group.reshard = prior_reshard
@@ -377,6 +378,9 @@ def _shared_group_override_restores_prior_attributes(rank, size):
         assert "unshard" in group.__dict__
     assert "unshard" not in group.__dict__
     assert group.__dict__["reshard"] is prior_reshard
+    # A layer that is not itself FSDP-wrapped was not gathered, so its FSDP submodules must unshard.
+    with _disable_fsdp_unshard_reshard(torch.nn.Sequential(*block)):
+        assert "unshard" not in group.__dict__
 
 
 def test_disable_fsdp_unshard_reshard_shared_group_and_nesting(cpu_dist_workers):

From f086740f168f2a3fa5bc25ac75aa86dfa9756d99 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 17 Oct 2026 03:24:43 -0700
Subject: [PATCH 24/34] Raise gpu_vllm timeout to 20 min

The v0.30.0 job passed all tests with a warm extension cache but hit the
15 min limit: the image pull alone takes ~4.5 min.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/gpu_tests.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index e2822ab67c2..b7bca3bf5e3 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -53,7 +53,7 @@ jobs:
             timeout: 15
             container_image: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20
           - example: gpu_vllm
-            timeout: 15
+            timeout: 20
             container_image: docker.io/vllm/vllm-openai:v0.30.0
     runs-on: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
     timeout-minutes: ${{ matrix.timeout }}

From 996f48eceed24bc7570484bcf4d6e0511953adad Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 17 Oct 2026 03:40:00 -0700
Subject: [PATCH 25/34] Run only 3 puzzletron models in CI and raise
 test_mamba_search_space timeout

Puzzletron is not actively maintained, so CI keeps Nemotron-3-Nano, gpt-oss
and Qwen3-VL and marks the other six models manual (~4 min of the gpu job).
test_mamba_search_space usually takes ~70s but hit 123s on slow runners.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 tests/gpu/torch/puzzletron/test_puzzletron.py  | 18 ++++++++++++------
 .../test_megatron_hybrid_dynamic_modules.py    |  2 ++
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/tests/gpu/torch/puzzletron/test_puzzletron.py b/tests/gpu/torch/puzzletron/test_puzzletron.py
index 4953db2c737..222054ea187 100644
--- a/tests/gpu/torch/puzzletron/test_puzzletron.py
+++ b/tests/gpu/torch/puzzletron/test_puzzletron.py
@@ -32,19 +32,25 @@
 #
 
 SEED = 1234
+_MANUAL = pytest.mark.manual(reason="~40s each; run with --run-manual")
 
 
 @pytest.mark.parametrize(
     ("hf_model_name", "converter", "hybrid_override_pattern", "has_moe_layers"),
     [
-        ("meta-llama/Llama-3.1-8B-Instruct", "llama", None, False),
-        ("meta-llama/Llama-3.2-3B-Instruct", "llama", None, False),
-        ("mistralai/Mistral-Small-24B-Instruct-2501", "mistral_small", None, False),
+        # Puzzletron is not actively maintained, so CI runs only Nemotron-3-Nano, gpt-oss and Qwen3-VL.
+        pytest.param("meta-llama/Llama-3.1-8B-Instruct", "llama", None, False, marks=_MANUAL),
+        pytest.param("meta-llama/Llama-3.2-3B-Instruct", "llama", None, False, marks=_MANUAL),
+        pytest.param(
+            "mistralai/Mistral-Small-24B-Instruct-2501", "mistral_small", None, False, marks=_MANUAL
+        ),
         ("nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16", "nemotron_h", "*E", True),
-        ("nvidia/NVIDIA-Nemotron-Nano-12B-v2", "nemotron_h_v2", "*-", False),
+        pytest.param(
+            "nvidia/NVIDIA-Nemotron-Nano-12B-v2", "nemotron_h_v2", "*-", False, marks=_MANUAL
+        ),
         ("openai/gpt-oss-20b", "gpt_oss", None, True),
-        ("Qwen/Qwen2.5-7B-Instruct", "qwen2", None, False),
-        ("Qwen/Qwen3-8B", "qwen3", None, False),
+        pytest.param("Qwen/Qwen2.5-7B-Instruct", "qwen2", None, False, marks=_MANUAL),
+        pytest.param("Qwen/Qwen3-8B", "qwen3", None, False, marks=_MANUAL),
         ("Qwen/Qwen3-VL-30B-A3B-Instruct", "qwen3_vl", None, True),
     ],
 )
diff --git a/tests/gpu_megatron/torch/nas/plugins/test_megatron_hybrid_dynamic_modules.py b/tests/gpu_megatron/torch/nas/plugins/test_megatron_hybrid_dynamic_modules.py
index 8c98f7127fa..55404619de9 100644
--- a/tests/gpu_megatron/torch/nas/plugins/test_megatron_hybrid_dynamic_modules.py
+++ b/tests/gpu_megatron/torch/nas/plugins/test_megatron_hybrid_dynamic_modules.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 
 
+import pytest
 import torch
 from _test_utils.import_helper import skip_if_no_mamba
 
@@ -133,6 +134,7 @@ def _test_mamba_search_space(rank, size):
     assert not any(named_dynamic_modules(model))
 
 
+@pytest.mark.timeout(150)  # usually ~70s, but up to 123s on slow runners
 def test_mamba_search_space(dist_workers):
     dist_workers.run(_test_mamba_search_space)
 

From f1311ea5f9e3d78ffd74a0d2d734a7953f6c6b57 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 17 Oct 2026 03:52:08 -0700
Subject: [PATCH 26/34] Key the extension cache per job and skip saving an
 empty one

Example and GPU jobs on the same image shared one key. An example job that
builds no extensions finished first and saved an empty cache, so the gpu and
gpu_megatron jobs restored it as an exact hit, recompiled every run, and
could never save their own.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/actions/save-extensions/action.yml  | 8 +++++++-
 .github/workflows/_example_tests_runner.yml | 3 ++-
 .github/workflows/gpu_tests.yml             | 3 ++-
 3 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/.github/actions/save-extensions/action.yml b/.github/actions/save-extensions/action.yml
index bc639f28804..09f03d45794 100644
--- a/.github/actions/save-extensions/action.yml
+++ b/.github/actions/save-extensions/action.yml
@@ -10,7 +10,13 @@ inputs:
 runs:
   using: composite
   steps:
-    - uses: actions/cache/save@v6
+    # A job cancelled before building anything must not claim the key with an empty cache:
+    # later runs would restore it as an exact hit and never save again.
+    - id: built
+      shell: bash
+      run: find /root/.cache/torch_extensions -name '*.so' 2>/dev/null | grep -q . && echo "any=true" >> "$GITHUB_OUTPUT" || true
+    - if: steps.built.outputs.any == 'true'
+      uses: actions/cache/save@v6
       with:
         path: /root/.cache/torch_extensions
         key: ${{ inputs.cache-key }}
diff --git a/.github/workflows/_example_tests_runner.yml b/.github/workflows/_example_tests_runner.yml
index d3237f1b147..63147221b81 100644
--- a/.github/workflows/_example_tests_runner.yml
+++ b/.github/workflows/_example_tests_runner.yml
@@ -61,7 +61,8 @@ jobs:
       - uses: ./.github/actions/cache-extensions
         id: ext-cache
         with:
-          cache-key: rtxpro6000-${{ inputs.docker_image }}
+          # Per example: jobs sharing an image build different extensions, and the first to save wins the key.
+          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
diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index b7bca3bf5e3..9b17b21b4f4 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -76,7 +76,8 @@ jobs:
       - uses: ./.github/actions/cache-extensions
         id: ext-cache
         with:
-          cache-key: rtxpro6000-${{ matrix.container_image }}
+          # Per job: jobs sharing an image build different extensions, and the first to save wins the key.
+          cache-key: rtxpro6000-${{ matrix.example }}-${{ matrix.container_image }}
       - name: Restore wheels built from source
         # mamba_ssm and fast-hadamard-transform take ~10 min to compile (see noxfile.py `gpu`)
         id: wheelhouse

From 8f58d407f12c89ec6ba839022b9682c7ce6c6adb Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Mon, 19 Oct 2026 22:10:11 -0700
Subject: [PATCH 27/34] Fix the new 2-GPU nightly failures: vLLM /dev/shm and
 the packed distill test's head_dim

- gpu_vllm: the TP2 tests from #2842 need 160 MiB of /dev/shm; the container had Docker's 64 MiB
- test_distill_llm_sft_hf_chat[packed]: the tiny model's head_dim=2 left only TE's unfused attention,
  which the sm_120 runner rejected for packed input at TP2; use head_dim=8

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/gpu_tests.yml                | 1 +
 tests/examples/megatron_bridge/test_distill.py | 6 +++++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index b196747a874..95a806c8203 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -59,6 +59,7 @@ jobs:
     timeout-minutes: ${{ matrix.timeout }}
     container:
       image: ${{ matrix.container_image }}
+      options: --shm-size=2gb # vLLM TP>1 tests on the 2-GPU runner need more than the default 64 MB
       env:
         GIT_DEPTH: 1000 # For correct version
         PIP_CONSTRAINT: "" # Disable pip constraint for upgrading packages
diff --git a/tests/examples/megatron_bridge/test_distill.py b/tests/examples/megatron_bridge/test_distill.py
index 6977061d282..3bdcc460c27 100644
--- a/tests/examples/megatron_bridge/test_distill.py
+++ b/tests/examples/megatron_bridge/test_distill.py
@@ -130,7 +130,11 @@ def test_distill_llm_sft_hf_chat(tmp_path, num_gpus, packed):
     ``packed`` runs an MoE student with conversations packed per micro-batch and the aux loss off.
     """
     create_dir = create_tiny_qwen3_moe_dir if packed else create_tiny_qwen3_dir
-    teacher_hf_path = create_dir(tmp_path, with_tokenizer=True)
+    # The default head_dim=2 leaves only TE's unfused attention, which the 2-GPU sm_120 runner
+    # rejects for packed (THD) input; head_dim=8 also qualifies the flash and cuDNN backends.
+    teacher_hf_path = create_dir(
+        tmp_path, with_tokenizer=True, **({"head_dim": 8} if packed else {})
+    )
     train_iters = 2
     gbs = 4
     # More records than train_iters * gbs so the sampler does not run dry; every 4th is multi-turn.

From f25ff984cbb40c6137f246f3e0a8789cfef80130 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Tue, 20 Oct 2026 00:41:39 -0700
Subject: [PATCH 28/34] Raise the gpu_vllm timeout to 30 min for the 2-GPU
 nightly

With /dev/shm fixed, the TP2 tests run on the 2-GPU runner (~3.5 min) and the dispatch run hit the
20 min limit at 51% of tests after a 4.5 min image pull.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/gpu_tests.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index 95a806c8203..906aa0081e8 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -53,7 +53,7 @@ jobs:
             timeout: 15
             container_image: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20
           - example: gpu_vllm
-            timeout: 20
+            timeout: 30 # ~13 min of tests on 1 GPU (PRs), ~16 min with the TP2 tests on 2 GPUs (nightly), plus image pull
             container_image: docker.io/vllm/vllm-openai:v0.30.0
     runs-on: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
     timeout-minutes: ${{ matrix.timeout }}

From 165cdd692a924d6ad7e99917d699717bdb2a29b0 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Tue, 20 Oct 2026 00:47:34 -0700
Subject: [PATCH 29/34] Raise the gpu_megatron timeout to 75 min for the 2-GPU
 nightly

The nightly runs every test on 2 GPUs and hit the 60 min limit at 89-99% of tests (Oct 20 nightly
and this branch's dispatch run); PR runs take ~42 min.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/gpu_tests.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index 906aa0081e8..1b3c5d50ecd 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -47,7 +47,7 @@ jobs:
             timeout: 60
             container_image: nvcr.io/nvidia/pytorch:26.09-py3
           - example: gpu_megatron
-            timeout: 60
+            timeout: 75 # ~42 min on 1 GPU (PRs); the 2-GPU nightly runs every test and took 46 to 60+ min
             container_image: nvcr.io/nvidia/nemo:26.08
           - example: gpu_trtllm
             timeout: 15

From 378b5b630abc9e6900850ee900db63debe114b0c Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Tue, 20 Oct 2026 00:48:43 -0700
Subject: [PATCH 30/34] Drop the comments on the gpu_vllm and gpu_megatron
 timeouts

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .github/workflows/gpu_tests.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml
index 1b3c5d50ecd..b73f85affeb 100644
--- a/.github/workflows/gpu_tests.yml
+++ b/.github/workflows/gpu_tests.yml
@@ -47,13 +47,13 @@ jobs:
             timeout: 60
             container_image: nvcr.io/nvidia/pytorch:26.09-py3
           - example: gpu_megatron
-            timeout: 75 # ~42 min on 1 GPU (PRs); the 2-GPU nightly runs every test and took 46 to 60+ min
+            timeout: 75
             container_image: nvcr.io/nvidia/nemo:26.08
           - example: gpu_trtllm
             timeout: 15
             container_image: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc20
           - example: gpu_vllm
-            timeout: 30 # ~13 min of tests on 1 GPU (PRs), ~16 min with the TP2 tests on 2 GPUs (nightly), plus image pull
+            timeout: 30
             container_image: docker.io/vllm/vllm-openai:v0.30.0
     runs-on: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
     timeout-minutes: ${{ matrix.timeout }}

From 49ab19f1e81467e3f3c8298c0065c99fa64cd2fa Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Tue, 20 Oct 2026 01:32:57 -0700
Subject: [PATCH 31/34] Fix unit CI: pin leptonai<0.47 for the launcher; halve
 the DFlash2 overfit test

- leptonai 0.47 (released 2026-10-20) removed leptonai.api.v1, which nemo-run 0.11.1 imports, so every
  launcher test failed to collect
- test_overfits_a_single_batch took 51-61 s on most unit runners today (2-7 s on others; 16 s locally
  on one thread) and timed out at 60 s; both accuracies pass 0.9 by step ~21, so run 30 steps

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 tests/unit/torch/speculative/plugins/test_hf_dflash2.py | 2 +-
 tools/launcher/pyproject.toml                           | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py
index 931a4a3ee46..97def429a1e 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash2.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash2.py
@@ -361,7 +361,7 @@ def test_overfits_a_single_batch(self):
 
         input_ids, attention_mask, labels = _make_batch(model.dflash_config.vocab_size)
         optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=5e-3)
-        for _ in range(60):
+        for _ in range(30):  # both accuracies pass 0.9 by step ~21
             out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
             optimizer.zero_grad()
             out.loss.backward()
diff --git a/tools/launcher/pyproject.toml b/tools/launcher/pyproject.toml
index 8ba825776d4..02adf5991fe 100644
--- a/tools/launcher/pyproject.toml
+++ b/tools/launcher/pyproject.toml
@@ -4,6 +4,7 @@ version = "0.1.0"
 description = "ModelOpt Launcher — submit quantization, training, and evaluation jobs to Slurm clusters"
 requires-python = ">=3.10"
 dependencies = [
+    "leptonai<0.47",  # 0.47 removed leptonai.api.v1, which nemo-run (<=0.11.1) imports
     "nemo-run>=0.8.0",
     "pyyaml",
 ]

From 3ae2c85ab7b87c0bbd1a0f650fc9e552a9c24ec4 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Tue, 20 Oct 2026 02:42:23 -0700
Subject: [PATCH 32/34] Skip two more Megatron quantization tests hit by the
 known Blackwell illegal-memory-access flake

test_fp8_real_quantize (twice, 2-GPU) and test_layerwise_calibrate_vlm_text_forward_loop[True] (PR CI,
1 GPU) failed with the intermittent CUDA illegal memory access that skip_flaky_on_blackwell already
covers for this file; test_megatron.py alone passed on the same 2-GPU runners in all 4 tries.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 tests/gpu_megatron/torch/quantization/plugins/test_megatron.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
index 9e2a7be8a0b..c53a87276e6 100644
--- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
+++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
@@ -691,6 +691,7 @@ def _test_fp8_real_quantize_helper(rank, size):
     assert real_quant_mem < cur_mem
 
 
+@skip_flaky_on_blackwell
 def test_fp8_real_quantize(dist_workers):
     dist_workers.run(_test_fp8_real_quantize_helper)
 
@@ -1725,6 +1726,7 @@ def forward_loop(_model):
     assert all(m.input_quantizer.amax is not None for m in calibrated)
 
 
+@skip_flaky_on_blackwell
 @pytest.mark.parametrize("drive_wrapper", [False, True])
 def test_layerwise_calibrate_vlm_text_forward_loop(dist_workers_size_1, drive_wrapper):
     dist_workers_size_1.run(partial(_test_layerwise_calibrate_vlm_text_forward_loop, drive_wrapper))

From 6bc011949c14ecb42968fbfe6aa1ab0532cadfd5 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Tue, 20 Oct 2026 08:40:56 -0700
Subject: [PATCH 33/34] gpu_megatron: stop destroying NCCL groups in the worker
 teardown; drop the new Blackwell skips

On the 1-GPU PR runners the branch hit CUDA illegal memory accesses in random test_megatron.py tests
(3 of 7 full-suite runs), never with main's conftest (0 of 2 bisect runs, and none on main).
Destroying each test's process groups in the teardown is the only functional difference, so revert
it; the 2-GPU full suite takes the same time without it (43:36 vs 44:21). The skips added for the
same crashes in 3ae2c85ab7 go too.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 tests/gpu_megatron/conftest.py                              | 6 ------
 .../torch/quantization/plugins/test_megatron.py             | 2 --
 2 files changed, 8 deletions(-)

diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py
index 6b6ab484475..51261e30023 100644
--- a/tests/gpu_megatron/conftest.py
+++ b/tests/gpu_megatron/conftest.py
@@ -79,12 +79,6 @@ def megatron_worker_teardown(rank, world_size):
             apex_destroy()
         except Exception as e:
             print(f"Error destroying model parallel with Apex: {e}")
-    # destroy_model_parallel() only drops Megatron's references to its NCCL groups, so every test's
-    # communicators (and their GPU buffers) piled up in the persistent workers (~20 groups and
-    # ~0.4 GiB per test), slowing later tests until they timed out. Destroy all but the default.
-    for group in list(torch.distributed.distributed_c10d._world.pg_map):
-        if group is not torch.distributed.GroupMember.WORLD:
-            torch.distributed.destroy_process_group(group)
     torch.cuda.empty_cache()
 
 
diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
index c53a87276e6..9e2a7be8a0b 100644
--- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
+++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
@@ -691,7 +691,6 @@ def _test_fp8_real_quantize_helper(rank, size):
     assert real_quant_mem < cur_mem
 
 
-@skip_flaky_on_blackwell
 def test_fp8_real_quantize(dist_workers):
     dist_workers.run(_test_fp8_real_quantize_helper)
 
@@ -1726,7 +1725,6 @@ def forward_loop(_model):
     assert all(m.input_quantizer.amax is not None for m in calibrated)
 
 
-@skip_flaky_on_blackwell
 @pytest.mark.parametrize("drive_wrapper", [False, True])
 def test_layerwise_calibrate_vlm_text_forward_loop(dist_workers_size_1, drive_wrapper):
     dist_workers_size_1.run(partial(_test_layerwise_calibrate_vlm_text_forward_loop, drive_wrapper))

From cb1cac93b0ebed943dd7edb38b00f02afeeedc75 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Tue, 20 Oct 2026 09:25:17 -0700
Subject: [PATCH 34/34] hf_eagle: note that CUDA-graph capture is currently off
 where rotary_emb init is hoisted

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 modelopt/torch/speculative/plugins/hf_eagle.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/hf_eagle.py b/modelopt/torch/speculative/plugins/hf_eagle.py
index 251f5776428..910d5907870 100644
--- a/modelopt/torch/speculative/plugins/hf_eagle.py
+++ b/modelopt/torch/speculative/plugins/hf_eagle.py
@@ -754,7 +754,8 @@ def forward(
                 base_outputs,
             )
 
-        # Init rotary_emb outside _eagle_forward so inv_freq is not captured by CUDAGraph.
+        # Init rotary_emb outside _eagle_forward so inv_freq is not captured by CUDAGraph. Capture is
+        # off for now (see _activate_torch_compile); this keeps re-enabling it safe.
         self.eagle_module._maybe_init_rope(device=eagle_input_hiddens.device)
 
         # ====Run eagle forward with extra training-time-test steps====
@@ -902,7 +903,8 @@ def pseudo_speculative_generate(
         else:
             eagle_input_hidden_states = base_model_hidden_states
 
-        # Init rotary_emb outside _eagle_forward so inv_freq is not captured by CUDAGraph.
+        # Init rotary_emb outside _eagle_forward so inv_freq is not captured by CUDAGraph. Capture is
+        # off for now (see _activate_torch_compile); this keeps re-enabling it safe.
         self.eagle_module._maybe_init_rope(device=eagle_input_hidden_states.device)
 
         draft_tokens = []