CI follow-up: cache source-built wheels and extensions, bump CI containers, trim slow GPU tests

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

@@ -1,22 +1,32 @@
 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')
-          }}
+        # 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'

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

@@ -0,0 +1,22 @@
+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:
+    # 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 }}

.github/workflows/_example_tests_runner.yml

@@ -59,8 +59,10 @@ 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 }}
+          # 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
@@ -113,6 +115,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: |

.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]"
@@ -169,6 +169,8 @@ jobs:
       contents: read
     secrets: inherit
     with:
+      # 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]"

.github/workflows/gpu_tests.yml

@@ -45,21 +45,16 @@ jobs:
         include:
           - example: gpu
             timeout: 60
-            container_image: nvcr.io/nvidia/pytorch:26.07-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
-          # Covers the RoutedExperts layout (0.24+)
           - example: gpu_vllm
-            timeout: 15
+            timeout: 20
             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:
@@ -79,18 +74,49 @@ 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 }}
+          # 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
+        if: matrix.example == 'gpu'
+        uses: actions/cache/restore@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
       - name: Run gpu tests
         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 }}
+      - 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: 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, 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
+          key: wheelhouse-${{ matrix.container_image }}-${{ hashFiles('noxfile.py') }}
       - name: Upload GPU coverage to Codecov
         uses: codecov/codecov-action@v7
         with:

.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

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:])
 
 

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

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
 

modelopt/torch/quantization/qtensor/nvfp4_tensor.py

@@ -46,6 +46,24 @@ 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)
+    # 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):
     """Implements the NVFP4 quantization on tensors for more efficient storage or computation.
 
@@ -294,11 +312,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 +377,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

modelopt/torch/quantization/utils/core_utils.py

@@ -828,20 +828,37 @@ 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.
+    """
+    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()
+        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
+        group.reshard = lambda: None
+    try:
         yield
+    finally:
+        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):

modelopt/torch/speculative/plugins/hf_eagle.py

@@ -147,7 +147,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: