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

@@ -20,7 +20,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)
@@ -66,6 +67,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
@@ -109,6 +111,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 \
@@ -139,6 +142,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
 
 ENV NCCL_EXTENSIONS_VERSION=${NCCL_EXTENSIONS_VERSION}
@@ -212,6 +234,21 @@ RUN bash -ex <<"EOF"
     rm -rf DeepEP
 EOF
 
+# NGC 26.09 bundles FA4 beta24, which calls a helper removed by newer Quack.
+# Backport Dao-AILab/flash-attention#2787 onto the bundled package.
+COPY docker/patches/flash-attn-quack.patch /opt/flash-attn-build-info/quack.patch
+RUN bash -ex <<"EOF"
+    if [ "$IMAGE_TYPE" = "lts" ]; then
+        exit 0
+    fi
+    FA4_SITE_PACKAGES=$(python -c 'from importlib.metadata import distribution; d = distribution("flash-attn-4"); assert d.version == "4.0.0b24", f"Re-evaluate the FA4 Quack backport for version {d.version}"; print(d.locate_file(""))')
+    patch --batch --forward --fuzz=0 -p1 -d "${FA4_SITE_PACKAGES}" < /opt/flash-attn-build-info/quack.patch
+    python -c 'import cutlass.cute as cute; from flash_attn.cute import utils; assert callable(cute.arch.sub_packed_f32x2)'
+    printf 'flash-attn-4==4.0.0b24\nbackport=145b1010051dbfd4bdc41a0ae55d495b08d7a458\n' > /opt/flash-attn-build-info/revision
+    sha256sum "${FA4_SITE_PACKAGES}/flash_attn/cute/utils.py" \
+        "${FA4_SITE_PACKAGES}/flash_attn/cute/flash_bwd_sm100.py" > /opt/flash-attn-build-info/patched-files.sha256
+EOF
+
 # Install the local project after expensive dependency builds so README-only
 # changes do not invalidate those layers.
 COPY README.md /workspace/

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/flash-attn-quack.patch

@@ -0,0 +1,47 @@
+diff --git a/flash_attn/cute/flash_bwd_sm100.py b/flash_attn/cute/flash_bwd_sm100.py
+index bded7c05303..5cf6fb506c7 100644
+--- a/flash_attn/cute/flash_bwd_sm100.py
++++ b/flash_attn/cute/flash_bwd_sm100.py
+@@ -14,7 +14,6 @@
+ import cutlass.utils.blackwell_helpers as sm100_utils_basic
+ from cutlass.pipeline import PipelineAsync
+
+-import quack.activation
+ from quack import layout_utils
+ from flash_attn.cute import utils
+ from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
+@@ -3268,7 +3267,7 @@ def compute_loop(
+                                 utils.shuffle_sync(tSrdPsum, offset=2 * v + 1),
+                             )
+                         tdPrdP_cur[2 * v], tdPrdP_cur[2 * v + 1] = (
+-                            quack.activation.sub_packed_f32x2(
++                            cute.arch.sub_packed_f32x2(
+                                 (tdPrdP_cur[2 * v], tdPrdP_cur[2 * v + 1]), dPsum_pair
+                             )
+                         )
+diff --git a/flash_attn/cute/utils.py b/flash_attn/cute/utils.py
+index 9feafab5f05..5391fa7cfed 100644
+--- a/flash_attn/cute/utils.py
++++ b/flash_attn/cute/utils.py
+@@ -17,8 +17,6 @@
+ from cutlass.cute.runtime import from_dlpack
+
+
+-import quack.activation
+-
+ _MIXER_ATTRS = ("__vec_size__",)
+
+
+@@ -780,10 +778,10 @@ def ex2_emulation_2(
+     xy_rounded = cute.arch.add_packed_f32x2(xy_clamped, (fp32_round_int, fp32_round_int), rnd="rm")
+     # The integer floor of x & y are now in the last 8 bits of xy_rounded
+     # We want the next 2 ops to round to nearest even. The rounding mode is important.
+-    xy_rounded_back = quack.activation.sub_packed_f32x2(
++    xy_rounded_back = cute.arch.sub_packed_f32x2(
+         xy_rounded, (fp32_round_int, fp32_round_int)
+     )
+-    xy_frac = quack.activation.sub_packed_f32x2(xy_clamped, xy_rounded_back)
++    xy_frac = cute.arch.sub_packed_f32x2(xy_clamped, xy_rounded_back)
+     xy_frac_ex2 = evaluate_polynomial_2(*xy_frac, POLY_EX2[poly_degree], loc=loc, ip=ip)
+     x_out = combine_int_frac_ex2(xy_rounded[0], xy_frac_ex2[0], loc=loc, ip=ip)
+     y_out = combine_int_frac_ex2(xy_rounded[1], xy_frac_ex2[1], loc=loc, ip=ip)

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

@@ -27,11 +27,7 @@
     get_tensor_model_parallel_rank,
 )
 from megatron.core.tensor_observation import suspend_tensor_observations
-from megatron.core.utils import (
-    is_te_min_version,
-    is_torch_min_version,
-    safely_set_viewless_tensor_data,
-)
+from megatron.core.utils import is_te_min_version, safely_set_viewless_tensor_data
 
 # ---------------------------------------------------------------------------
 # C++ extension: zero-copy storage sharing for CheckpointWithoutOutput
@@ -451,13 +447,15 @@ def get_all_rng_states():
 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
-    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).
+    NGC 26.08's exact PyTorch build implements lazy registration from pytorch/pytorch#176753;
+    explicit registration is a deprecated no-op that floods stderr. Other 2.14 prereleases,
+    including NGC 26.09, require registration again, so a minimum version is insufficient.
+    Skip only the verified no-op build and conservatively register on every other build.
     """
-    return not is_torch_min_version("2.14.0a0")
+    if str(torch.__version__) != "2.14.0a0+4fdf77b940.nv26.08":
+        return True
+    git_version = getattr(torch.version, "git_version", None)
+    return git_version not in (None, "Unknown") and not git_version.startswith("4fdf77b940")
 
 
 def prime_cuda_rng_states_for_graph_capture() -> None:

megatron/core/transformer/attention.py

@@ -1093,6 +1093,19 @@ def flash_decode_and_prefill(
         need_lse = softmax_offset is not None
 
         use_fa4, use_fa3 = self._resolve_flash_version()
+        fa4_num_splits = 1 if self.batch_invariant_mode else 0
+        if use_fa4 and q.is_cuda and torch.cuda.get_device_capability(q.device)[0] == 9:
+            # FA4's automatic heuristic can select SplitKV, which Hopper does not support.
+            fa4_num_splits = 1
+            if (
+                self.flash_attention_version is None
+                and k.shape[1] % 256 == 0
+                and (q.shape[-1] % 16 != 0 or v.shape[-1] % 16 != 0)
+            ):
+                # Hopper's FA4 non-TMA paged path requires 16-aligned head dimensions.
+                # These page sizes are supported by FA2; retain FA4 for smaller pages,
+                # which can use its TMA path and are not necessarily supported by FA2.
+                use_fa4, use_fa3 = False, HAVE_FA3
 
         # Flash attn kernel.
         if not is_decode_only:
@@ -1116,7 +1129,8 @@ def flash_decode_and_prefill(
                     softmax_scale=softmax_scale,
                     causal=True,
                     window_size=window_size,
-                    num_splits=0 if not self.batch_invariant_mode else 1,
+                    num_splits=fa4_num_splits,
+                    return_lse=need_lse,
                     **softcap_kwargs,
                 )
             elif use_fa3:
@@ -1259,7 +1273,8 @@ def flash_decode_and_prefill(
                         softmax_scale=softmax_scale,
                         causal=True,
                         window_size=window_size,
-                        num_splits=0 if not self.batch_invariant_mode else 1,
+                        num_splits=fa4_num_splits,
+                        return_lse=need_lse,
                         **softcap_kwargs,
                     )
                     if need_lse:

megatron/training/checkpointing.py

@@ -3007,9 +3007,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