build: update development PyTorch image to 26.09

.gitlab/stages/01.build.yml

@@ -62,12 +62,12 @@ test:pre_build_image:
       - IMAGE: CI_MCORE_DEV_IMAGE
         FILE: Dockerfile.ci.dev
         IMAGE_TYPE: dev
-        BASE_IMAGE: nvcr.io/nvidia/pytorch:26.08-py3
+        BASE_IMAGE: nvcr.io/nvidia/pytorch:26.09-py3
         PLATFORM: amd64
       - IMAGE: CI_MCORE_DEV_IMAGE
         FILE: Dockerfile.ci.dev
         IMAGE_TYPE: dev
-        BASE_IMAGE: nvcr.io/nvidia/pytorch:26.08-py3
+        BASE_IMAGE: nvcr.io/nvidia/pytorch:26.09-py3
         PLATFORM: arm64
       - IMAGE: UTILITY_IMAGE
         FILE: Dockerfile.linting

docker/.ngc_version.dev

@@ -1 +1 @@
-nvcr.io/nvidia/pytorch:26.08-py3
+nvcr.io/nvidia/pytorch:26.09-py3

docker/Dockerfile.ci.dev

@@ -17,7 +17,8 @@ ENV UV_LINK_MODE=copy
 
 RUN bash -ex <<"EOF"
     apt-get update
-    apt-get install -y --no-install-recommends gettext python3-venv psmisc uuid-runtime
+    # DeepGEMM/DeepJIT requires elfutils/libdwfl.h for its exception support.
+    apt-get install -y --no-install-recommends gettext python3-venv psmisc uuid-runtime libdw-dev
     apt-get clean
     python -m venv /opt/jet
     ARCH=$(uname -m)
@@ -88,6 +89,7 @@ RUN --mount=type=bind,source=docker/common/install_nccl.sh,target=/opt/install_n
 EOF
 
 COPY pyproject.toml uv.lock /workspace/
+COPY docker/patches/mamba-cxx20.patch /workspace/mamba-cxx20.patch
 COPY megatron/core/__init__.py /workspace/megatron/core/
 COPY megatron/core/package_info.py /workspace/megatron/core/
 ARG IMAGE_TYPE=dev
@@ -131,6 +133,7 @@ RUN --mount=type=cache,id=uv-${TARGETARCH}-${NVTE_CUDA_ARCHS},target=/root/.cach
         --no-install-package torch \
         --no-install-package torchvision \
         --no-install-package triton \
+        --no-install-package mamba-ssm \
         --no-install-package transformer-engine-cu12 \
         --no-install-package nvidia-cublas-cu12 \
         --no-install-package nvidia-cuda-cupti-cu12 \
@@ -161,6 +164,25 @@ RUN --mount=type=cache,id=uv-${TARGETARCH}-${NVTE_CUDA_ARCHS},target=/root/.cach
     cp -a "${NCCL_EP_JIT_HEADERS}/../nccl_ep.h" /opt/nccl-ep/include/
     test -f /opt/nccl-ep/include/nccl_ep/device/ht_ep.cuh
     test -f /opt/nccl-ep/include/nccl_ep.h
+
+    # Backport state-spaces/mamba#1000 without changing the pinned runtime sources.
+    # PyTorch's ATen headers now require C++20 for both host and CUDA compilation.
+    MAMBA_REV=$(sed -n 's/^mamba-ssm = .*rev = "\([^"]*\)".*/\1/p' /workspace/pyproject.toml)
+    test -n "${MAMBA_REV}"
+    MAMBA_BUILD_DIR=$(mktemp -d)
+    trap 'rm -rf "${MAMBA_BUILD_DIR}"' EXIT
+    git -C "${MAMBA_BUILD_DIR}" init
+    git -C "${MAMBA_BUILD_DIR}" fetch --depth 1 https://github.com/state-spaces/mamba.git "${MAMBA_REV}"
+    git -C "${MAMBA_BUILD_DIR}" checkout --detach "${MAMBA_REV}"
+    git -C "${MAMBA_BUILD_DIR}" apply /workspace/mamba-cxx20.patch
+    MAMBA_FORCE_BUILD=TRUE uv pip install --no-deps --no-build-isolation -v "${MAMBA_BUILD_DIR}"
+    # Keep the source revision and applied patch after deleting the build checkout.
+    mkdir -p /opt/mamba-build-info
+    git -C "${MAMBA_BUILD_DIR}" rev-parse HEAD > /opt/mamba-build-info/revision
+    cp /workspace/mamba-cxx20.patch /opt/mamba-build-info/
+    rm -rf "${MAMBA_BUILD_DIR}"
+    trap - EXIT
+
 EOF
 
 # Reuse the wheel without copying a second archive into the final image.

docker/common/install.sh

@@ -86,7 +86,8 @@ main() {
 
     # Install tools
     apt-get update
-    apt-get install -y wget curl git cmake
+    # DeepGEMM/DeepJIT requires elfutils/libdwfl.h for its exception support.
+    apt-get install -y wget curl git cmake libdw-dev
 
     # Install CUDA
     if [[ "$BASE_IMAGE" == "ubuntu" ]]; then

docker/patches/mamba-cxx20.patch

@@ -0,0 +1,33 @@
+Backport of https://github.com/state-spaces/mamba/commit/653923ce8fb0d47cdd9bcfd5904a0f1d58f91274
+Build the existing pinned Mamba sources with the standard required by PyTorch's ATen headers.
+
+diff --git a/setup.py b/setup.py
+--- a/setup.py
++++ b/setup.py
+@@ -210,10 +210,10 @@ def append_nvcc_threads(nvcc_extra_args):
+     if HIP_BUILD:
+
+         extra_compile_args = {
+-            "cxx": ["-O3", "-std=c++17"],
++            "cxx": ["-O3", "-std=c++20"],
+             "nvcc": [
+                 "-O3",
+-                "-std=c++17",
++                "-std=c++20",
+                 f"--offload-arch={os.getenv('HIP_ARCHITECTURES', 'native')}",
+                 "-U__CUDA_NO_HALF_OPERATORS__",
+                 "-U__CUDA_NO_HALF_CONVERSIONS__",
+@@ -223,11 +223,11 @@ def append_nvcc_threads(nvcc_extra_args):
+         }
+     else:
+         extra_compile_args = {
+-            "cxx": ["-O3", "-std=c++17"],
++            "cxx": ["-O3", "-std=c++20"],
+             "nvcc": append_nvcc_threads(
+                 [
+                     "-O3",
+-                    "-std=c++17",
++                    "-std=c++20",
+                     "-U__CUDA_NO_HALF_OPERATORS__",
+                     "-U__CUDA_NO_HALF_CONVERSIONS__",
+                     "-U__CUDA_NO_BFLOAT16_OPERATORS__",

megatron/core/fusions/fused_bias_swiglu.py

@@ -3,14 +3,44 @@
 
 # pylint: disable=missing-function-docstring, missing-class-docstring
 
+from functools import cache
 from typing import Optional
 
 import torch
 import torch.nn.functional as F
+from torch._subclasses.fake_tensor import unset_fake_temporarily
 
 from megatron.core.jit import jit_fuser
 from megatron.core.utils import nvtx_decorator
 
+
+@cache
+def _probe_clamp_boundaries():
+    """Read the installed scalar-clamp subgradient using only CPU constants."""
+    # PyTorch changed this within the 2.14 development series, so a version check
+    # cannot distinguish builds. FunctionsManual.cpp at b2c75dd062 uses strict
+    # inequalities; 4fdf77b940 uses inclusive ones. Probe both scalar-clamp forms
+    # once on CPU, including calls under no_grad or inference_mode.
+    # This zero-input probe needs real CPU constants, even during fake tracing.
+    with unset_fake_temporarily(), torch.inference_mode(False), torch.enable_grad():
+        gate = torch.tensor(1.0, dtype=torch.float32, device="cpu", requires_grad=True)
+        linear = torch.tensor([-1.0, 1.0], dtype=torch.float32, device="cpu", requires_grad=True)
+        gate_grad, linear_grad = torch.autograd.grad(
+            gate.clamp(max=1.0) + linear.clamp(min=-1.0, max=1.0).sum(), (gate, linear)
+        )
+        gradients = [gate_grad.item(), *linear_grad.tolist()]
+    if gradients not in ([0.0, 0.0, 0.0], [1.0, 1.0, 1.0]):
+        raise RuntimeError(f"Unsupported scalar-clamp boundary gradients: {gradients}")
+    return gradients[0] == 1.0
+
+
+@torch.compiler.assume_constant_result
+def _clamp_includes_boundaries():
+    # Delay autograd initialization until this clamped backward is actually used.
+    # Keep the cached probe outside tracing and specialize only its Python bool.
+    return _probe_clamp_boundaries()
+
+
 ###### BIAS SWIGLU FUSION/ NO AUTOGRAD ################
 
 
@@ -163,14 +193,20 @@ def clamped_swiglu_back(g, y, clamp_value):
     y_1, y_2 = torch.chunk(y.to(torch.float32), 2, -1)
     y_1c = y_1.clamp(min=None, max=clamp_value)
     y_2c = y_2.clamp(min=-clamp_value, max=clamp_value)
+    if _clamp_includes_boundaries():
+        gate_mask = y_1 <= clamp_value
+        linear_mask = (y_2 >= -clamp_value) & (y_2 <= clamp_value)
+    else:
+        gate_mask = y_1 < clamp_value
+        linear_mask = (y_2 > -clamp_value) & (y_2 < clamp_value)
     res = torch.cat(
         (
             g
             * torch.sigmoid(y_1c)
             * (1 + y_1c * (1 - torch.sigmoid(y_1c)))
             * y_2c
-            * (y_1 <= clamp_value).to(g.dtype),
-            g * F.silu(y_1c) * ((y_2 >= -clamp_value) & (y_2 <= clamp_value)).to(g.dtype),
+            * gate_mask.to(g.dtype),
+            g * F.silu(y_1c) * linear_mask.to(g.dtype),
         ),
         -1,
     )

megatron/core/ssm/ssm_inference.py

@@ -223,18 +223,15 @@ def ssm_dynamic_inference(
         if decode_req_count > 0:
             seq_len = 1 + context.num_speculative_tokens
             decode_token_count = decode_req_count * seq_len
-            if context.batch_invariant_mode:
-                # Batch-invariant execution may include token-only rows to preserve
-                # model-wide M alignment. Those rows do not represent requests and
-                # must not be passed to the recurrent decode kernels.
-                assert decode_token_count <= zxBCdt.shape[0], (
-                    "Batch-invariant SSM metadata describes more decode tokens "
-                    f"({decode_token_count}) than the input projection contains "
-                    f"({zxBCdt.shape[0]})."
-                )
-                zxBCdt_decode = zxBCdt[:decode_token_count]
-            else:
-                zxBCdt_decode = zxBCdt[:decode_token_count] if prefill_req_count > 0 else zxBCdt
+            # Token-only padding can come from kernel-level batch invariance even
+            # when the model/context flag is off. Only metadata-backed requests
+            # belong in the recurrent decode kernels.
+            assert decode_token_count <= zxBCdt.shape[0], (
+                "SSM metadata describes more decode tokens "
+                f"({decode_token_count}) than the input projection contains "
+                f"({zxBCdt.shape[0]})."
+            )
+            zxBCdt_decode = zxBCdt[:decode_token_count]
             # Reshape from [N*S, 1, d] to [N, S, d] for the decode kernels.
             zxBCdt_decode = zxBCdt_decode.squeeze(1).view(decode_req_count, seq_len, -1)
             y_decode = self.ssm_decode(
@@ -282,16 +279,15 @@ def ssm_dynamic_inference(
         else:
             raise RuntimeError("Dynamic inference called with 0 decode and 0 prefill requests")
 
-        if context.batch_invariant_mode:
-            # Restore the projection's token-only padding before the output projection.
-            # Its row count can be TP-local, unlike the context's global token count.
-            padding_token_count = zxBCdt.shape[0] - y.shape[0]
-            assert padding_token_count >= 0, (
-                "Batch-invariant SSM produced more token rows "
-                f"({y.shape[0]}) than the input projection contained ({zxBCdt.shape[0]})."
-            )
-            if padding_token_count > 0:
-                y = torch.cat((y, y.new_zeros(padding_token_count, *y.shape[1:])), dim=0)
+        # Restore the projection's token-only padding before the output projection.
+        # Its row count can be TP-local, unlike the context's global token count.
+        padding_token_count = zxBCdt.shape[0] - y.shape[0]
+        assert padding_token_count >= 0, (
+            "SSM produced more token rows "
+            f"({y.shape[0]}) than the input projection contained ({zxBCdt.shape[0]})."
+        )
+        if padding_token_count > 0:
+            y = torch.cat((y, y.new_zeros(padding_token_count, *y.shape[1:])), dim=0)
 
         # Zero padding positions to avoid corrupting quantization amax calculations.
         if is_using_quantization_scales(self.config):

megatron/core/tensor_parallel/random.py

@@ -448,16 +448,53 @@ def get_all_rng_states():
         return {}
 
 
+_CUDAGRAPH_NEEDS_GENERATOR_REGISTRATION = None
+
+
+def _probe_cudagraph_needs_generator_registration() -> bool:
+    """Capture an RNG op on an unregistered generator and report whether capture rejects it."""
+    generator = torch.Generator(device="cuda")
+    graph = torch.cuda.CUDAGraph()
+    try:
+        with torch.cuda.graph(graph, capture_error_mode="thread_local"):
+            torch.rand(1, device="cuda", generator=generator)
+    except RuntimeError:
+        # Without lazy registration, capture fails with "RNG op during graph capture but
+        # generator is not registered with the capturing graph". Treat any capture failure
+        # as needing registration: on PyTorch with lazy registration, registering anyway
+        # only costs a deprecation warning, while skipping it on older builds is fatal.
+        return True
+    finally:
+        del graph
+    return False
+
+
 def cudagraph_needs_generator_registration() -> bool:
     """Whether generators must be registered with a `torch.cuda.CUDAGraph` before capture.
 
-    PyTorch >= 2.14 (pytorch/pytorch#176753) lazily registers every generator whose Philox
+    PyTorch with pytorch/pytorch#176753 lazily registers every generator whose Philox
     state is consumed during capture, and `CUDAGraph.register_generator_state()` became a
     deprecated no-op that prints a warning on *every* call. Skip the explicit registration
     there: it does nothing, and with one call per layer, per graph and per generator it floods
     stderr (tens of thousands of lines per rank for dynamic inference with CUDA graphs).
+
+    The version string cannot identify that change: 2.14.0a0 nightlies built before it landed
+    still require explicit registration. Builds older than 2.14.0a0 always require it; for
+    newer builds, a small capture probe decides, and its result is cached.
     """
-    return not is_torch_min_version("2.14.0a0")
+    global _CUDAGRAPH_NEEDS_GENERATOR_REGISTRATION
+    if _CUDAGRAPH_NEEDS_GENERATOR_REGISTRATION is not None:
+        return _CUDAGRAPH_NEEDS_GENERATOR_REGISTRATION
+    if not is_torch_min_version("2.14.0a0"):
+        needs_registration = True
+    elif torch.cuda.is_current_stream_capturing():
+        # The probe cannot capture while another capture is active; register to be safe
+        # and probe on a later call.
+        return True
+    else:
+        needs_registration = _probe_cudagraph_needs_generator_registration()
+    _CUDAGRAPH_NEEDS_GENERATOR_REGISTRATION = needs_registration
+    return needs_registration
 
 
 def prime_cuda_rng_states_for_graph_capture() -> None:

megatron/training/checkpointing.py

@@ -3033,9 +3033,9 @@ def load_checkpoint(
         )
         # Same as the torch_dist branch: optimizer load templates select checkpoint-era keys by
         # version. fsdp_dtensor state has no dtype-keyed FQNs today; keep both paths identical.
-        optim_sd_kwargs['metadata']['checkpoint_version'] = (
-            state_dict.get('checkpoint_version') or 0
-        )
+        optim_sd_kwargs['metadata']['checkpoint_version'] = (state_dict or {}).get(
+            'checkpoint_version'
+        ) or 0
 
         # Megatron-FSDP materializes optimizer slots with a dummy zero-gradient
         # step while building a loading state dict. A normal full resume

pyproject.toml

@@ -255,7 +255,7 @@ requires-dist = ["torch", "packaging", "ninja"]
 flash_mla = [
     { git = "https://github.com/deepseek-ai/FlashMLA", rev = "nv_dev" },
 ]
-deep_gemm = { git = "https://github.com/deepseek-ai/DeepGEMM.git", rev = "714dd1a4a980f7937a74343d19a8eba4fe321480" }
+deep_gemm = { git = "https://github.com/deepseek-ai/DeepGEMM.git", rev = "ea2b7805b97956268bed1adf94591f428ded6cc1" }
 transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "6ea2a74a9e98c99e6d7b164a33775cc457520027" }
 nemo-run = { git = "https://github.com/NVIDIA-NeMo/Run.git", rev = "e3935393a290aed1822af52139b4b8ee270fed1f" }
 emerging_optimizers = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git", rev = "v0.3.0" }

tests/functional_tests/shell_test_utils/_run_training.sh

@@ -90,6 +90,9 @@ if [[ "$BEFORE_SCRIPT" != null ]]; then
     eval "$BEFORE_SCRIPT"
 fi
 
+# Keep functional validation on FA2 while using the unpatched NGC 26.09 stack.
+export NVTE_FLASH_ATTN_V4=0
+
 # Exit earlier to leave time for properly saving checkpoint
 if [[ "$IS_NEMO_TEST" == "true" ]]; then
     PARAMS=()
@@ -187,6 +190,9 @@ fi
 
 # Extract training params
 PARAMS=("${PARAMS[@]}" "${TRAINING_PARAMS_ARRAY[@]}")
+if [[ "$IS_NEMO_TEST" != "true" ]]; then
+    PARAMS+=("--flash-attention-version" "2")
+fi
 
 # Set PYTHONPATH
 export PYTHONPATH="$(pwd):${PYTHONPATH:-}"