chore: nightly sync main into dev (30_07_2026)

[svcnvidia-nemo-ci]

Summary

Nightly sync of main into dev.

• 147 commits from main merged into dev.
• Merge commit 7cd77c636; parents origin/dev (95e4baf) + origin/main (ff6b92a).
• Python-only line stats: Python lines: +42496 / -5586 across 334 files (git diff --numstat origin/dev...HEAD -- '*.py'). The large + is dominated by new/renamed files main brought in (gated_delta_net restructure, experimental lite), not conflict edits.
• 62 files had merge conflicts, all resolved surgically (no global -X theirs).

Files where main's version was taken

Skill-sanctioned "Files to Override from Main" (reference dev-only args/APIs main removed/renamed) — taken via git checkout origin/main, then API-mismatch-audited:

• megatron/training/training.py
• megatron/training/utils/common_utils.py (dev split utils.py; is_gtp_remat_active re-export preserved in utils/__init__.py)
• megatron/core/optimizer/layer_wise_optimizer.py
• megatron/training/arguments.py, argument_utils.py, checkpointing.py, models/dist_utils.py (resolved toward main where they reference main-side arg/API shapes)

Dynamic-inference stack — taken from main because main strictly supersedes dev:

• megatron/core/inference/engines/dynamic_engine.py

• megatron/core/inference/contexts/dynamic_context.py

• tests/unit_tests/inference/test_mtp_cuda_graph_inference.py

  Main's mamba prefix caching (Mamba prefix caching fixes #5502, Ensure Mamba prefix cache snapshots are recorded for multi-chunk prompts #5952) builds on dev's Fix Mamba prefix match for chunked prefill #4758, and main's async scheduling (Extend dynamic inference asynchronous scheduling support #5939) extends dev's Implement async scheduling for dynamic inference #5453 (removed _validate_async_sched_support_for_request; SERIAL→LEGACY/ASYNC). Main's _can_schedule_* / CG-admission helpers were already auto-merged and call each other, so a piecemeal resolution would have produced Frankenstein code (orphaned dev-only mamba_caching_enabled, calls into the deleted validator). The merged engine calls main-only can_prepare_requests() and expects main's prepare_requests/resolve_requests signatures, so dynamic_context.py must also be main's.

  Preserved dev-only line: [dev] partial cuda graph support for thd format training. #4359's max_seqlen=query_emb.size(0) in dynamic_context.apply_rotary_emb_query (required — merged rope_utils.py is dev's, which asserts max_seqlen is not None).

gated_delta_net restructure (dev moved GatedDeltaNet into common.py; main added standalone gdn.py):

• common.py — kept dev's version (authoritative GatedDeltaNet, get_parameter_local_cp_headwise, _build_thd_cp_a2a_perm, _build_head_perm_for_split_sections).
• gdn.py — converted main's standalone module into a re-export shim so the ...gdn import path keeps working.
• __init__.py — re-exports updated to _headwise name; no stale get_parameter_local_cp importers remain.

Files deleted in dev but restored

• tests/unit_tests/inference/engines/test_cg_admission_gating.py — restored from main. Added by Inference: Cudagraph-aware admission gating in prefill scheduler #4870 (ancestor of both branches), present at merge-base + main, but absent on dev due to an accidental drop in a prior sync merge (a5e420afa), not a deliberate dev deletion. Main modified it via Extend dynamic inference asynchronous scheduling support #5939; it exercises exactly the CG-admission methods kept in the merged engine.

Kept deleted (dev intentionally removed): .github/workflows/multi-approval-bot.yml (deleted on dev, modified on main; dev's deletion intentional).

Dependency triple & CODEOWNERS

• pyproject.toml, uv.lock, docker/Dockerfile.ci.dev, .github/CODEOWNERS — all verified byte-identical to origin/dev in the merge result.
• mamba-ssm git source NOT added. Main declares a mamba-ssm git rev; dev pins mamba-ssm~=2.2 (PyPI). The triple has comprehensively diverged with dev ahead (TransformerEngine f031cf87 > main e7c550c5; emerging_optimizers v0.3.0 > v0.2.0; dev-only nvidia-cudnn-frontend, nvidia-resiliency-ext) — the skill's "keep dev's triple" case, not the narrow "main added one source" reconciliation. Merged mamba_ssm imports (selective_state_update, RMSNorm as RMSNormGated, ssd_combined, import mamba_ssm) are standard public 2.x API satisfied by dev's PyPI pin and guarded at runtime by megatron/core/utils.py version checks — no symbol requires main's rev. Avoids an unnecessary uv.lock regen.

Disposition of pre-push advisory findings

The advisory audit reported 3363 raw hits. These are exact-string, whitespace-sensitive diffs over two branches each ~150 commits past their merge-base, so the vast majority are cosmetic (black reflow, line-wrapping, trailing commas) or lines main legitimately evolved. Scoping to lines dev-added after merge-base, absent on main, absent in the merge (the only genuine-accident pattern) with whitespace normalization reduces this to 13 files, all resolved:

• 0 genuine merge accidents.
• 2 intentional main removals (documented):
  • megatron/post_training/model_builder.py — KD-state block mto.ModeloptStateManager(model).state_dict().pop() (present at base) deleted by main e97ac8360 (Various ModelOpt fixes: QAD test for CICD, use model builder config instead of model provider, allow loading teacher ckpt independently of student #4520). No dangling reference remains.
  • megatron/rl/rl_utils.py — packing_context.has_inference_logprobs = ... (present at base) removed by main c922805a4 (Add support for non-Gym multi-turn environments #5312), replaced with a locally-computed has_inference_logprobs. The packing_context attribute is no longer referenced.
• Remaining flags are false positives: correct unions (hybrid_layer_allocation.py VALID_LAYERS/MLA_ATTENTION hold both dev's CSA/HCA/WINDOW and main's MLA; test_optimizer.py import block has all symbols), reflowed calls (router.py z_loss record), superseded APIs (dev's per-param copy_tensor_to_quantized_param → main's batched copy_tensors_to_quantized_params, Perf: skip per-param copy_ dispatch in the MXFP8 param copy-back #6094; singular still defined for tests), and reflowed docstrings.

CODEOWNERS and the dep triple appear in the remerge-diff (resolution differed from naive auto-merge) but are verified identical to origin/dev in the final tree.

Commit signing

Signed off (-s). Cryptographic signing (-S) is disabled in this environment (USE_COMMIT_SIGNING=false, empty SSH_SIGNING_KEY, no key material), so /ok to test references the SHA explicitly.

Merge verification

• 0 conflict markers tree-wide; all 333 changed .py files ast.parse clean.
• All 12 gated_delta_net re-exported symbols resolve to staged common.py.
• black 24.10.0 + isort applied to merge-drifted files.

Conflict resolutions (remerge-diff)

Full git show --remerge-diff HEAD is 12,307 lines (dominated by golden-value JSON + uv.lock). Excerpt below covers the highest-value code resolutions (inference engine/context, gated_delta_net, rl_utils, model_builder, param_and_grad_buffer), each capped.

remerge-diff excerpt (key code files)

diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py
remerge CONFLICT (content): Merge conflict in megatron/core/distributed/param_and_grad_buffer.py
index 60459eed9..8e490f8c5 100644
--- a/megatron/core/distributed/param_and_grad_buffer.py
+++ b/megatron/core/distributed/param_and_grad_buffer.py
@@ -29,13 +29,9 @@ from ..fp4_utils import (
     modify_nvfp4_rowwise_storage,
 )
 from ..fp8_utils import (
-<<<<<<< 95e4bafeb ([dev] [fix] fix optimizer_cpu_offload with mark_keep_in_fp32 (#6124))
     _stage_param_to_bf16,
     copy_back_gathered_bf16_into_fp8_param,
-    copy_tensor_to_quantized_param,
-=======
     copy_tensors_to_quantized_params,
->>>>>>> ff6b92a7c (Fix --freeze-all-layers for new model builder path + unit tests (#5926))
     is_float8tensor,
     is_grouped_mxfp8tensor,
     is_grouped_tensor,
@@ -377,24 +373,15 @@ class _ParamAndGradBucketGroup:
         """Run post-processing after param all-gather completes."""
         if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag:
             for bucket in self.buckets:
-<<<<<<< 95e4bafeb ([dev] [fix] fix optimizer_cpu_offload with mark_keep_in_fp32 (#6124))
                 # Non-DistOpt LayerWise bucket has no param buffer (param_data is None); its weights
                 # were already requantized in finish_param_sync's non-DistOpt branch. The
                 # param-buffer copy/zero below only applies to buckets that own a param buffer.
                 if bucket.param_data is None:
                     continue
                 has_non_quantized_weight = False
-=======
-                if bucket.param_data is None:
-                    # LayerWise variable-size gather path: params are already updated via
-                    # unflatten + copy_ in finish_param_sync, and there is no param_data
... (truncated; 85 more lines in full remerge-diff) ...

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
remerge CONFLICT (content): Merge conflict in megatron/core/inference/contexts/dynamic_context.py
index 44662d9cf..b938b4f2c 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -36,7 +36,9 @@ from megatron.core.models.hybrid.hybrid_layer_allocation import (
 )
 from megatron.core.package_info import __version__ as mcore_version
 from megatron.core.transformer import MLATransformerConfig, TransformerConfig
+from megatron.core.transformer.enums import InferenceCudaGraphScope
 from megatron.core.transformer.moe.token_dispatcher_inference import (
+    InferenceAllGatherDispatcherBase,
     NCCLAllGatherDispatcher,
     NVLSAllGatherVDispatcher,
 )
@@ -49,7 +51,7 @@ from .attention_context.mha_metadata import GraphedMHAMetadata, NonGraphedMHAMet
 from .base_context import BaseInferenceContext
 from .gpu_view import ContextGPUView
 from .kv_block_allocator import KVBlockAllocator
-from .mamba_slot_allocator import MambaSlotAllocator
+from .mamba_slot_allocator import MAX_INTERMEDIATE_OFFSETS_PER_REQUEST, MambaSlotAllocator
 from .routing_metadata import RoutingMetadata
 
 try:
@@ -565,6 +567,8 @@ class DynamicInferenceContext(BaseInferenceContext):
 
         # Initialize context state.
         self.params_dtype = model_config.params_dtype
+        self.hidden_size = model_config.hidden_size
+        self.inference_cuda_graph_scope = model_config.inference_cuda_graph_scope
         self.max_sequence_length = inference_config.max_sequence_length
 
         # Block ids. With speculative decoding, blocks are pre-allocated when the
@@ -697,21 +701,6 @@ class DynamicInferenceContext(BaseInferenceContext):
... (truncated; 228 more lines in full remerge-diff) ...

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
remerge CONFLICT (content): Merge conflict in megatron/core/inference/engines/dynamic_engine.py
index 99f6030ab..c1668946a 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -1626,29 +1626,8 @@ class DynamicInferenceEngine(AbstractEngine):
         """
         return {"waits": self._prefix_coordination_waits}
 
-<<<<<<< 95e4bafeb ([dev] [fix] fix optimizer_cpu_offload with mark_keep_in_fp32 (#6124))
-    def _find_mamba_match_count(self, req: DynamicInferenceRequest) -> int:
-        """Find farthest block with cached Mamba state by iterating from the end.
-
-        Not all blocks have Mamba state cached in mamba_hash_to_block_id,
-        only divergence and last-aligned blocks do. Iterating from the end
-        finds the farthest block with cached state, which is the only one
-        needed for restore since Mamba state is cumulative.
-        """
-        if not req.precomputed_block_hashes:
-            return 0
-        mamba_map = self.context.mamba_slot_allocator.hash_to_block_id
-        for i in range(len(req.precomputed_block_hashes) - 1, -1, -1):
-            if req.precomputed_block_hashes[i] in mamba_map:
-                return i + 1
-        return 0
-
-    def schedule_waiting_requests(self):
-        """Tries to schedule any requests in the waiting pool."""
-=======
     def schedule_waiting_requests(self) -> None:
         """Try to schedule requests from the waiting pool."""
->>>>>>> ff6b92a7c (Fix --freeze-all-layers for new model builder path + unit tests (#5926))
         # Keep track of which requests get scheduled.
         waiting_before = set(self.waiting_request_ids)
... (truncated; 149 more lines in full remerge-diff) ...

diff --git a/megatron/core/ssm/gated_delta_net/common.py b/megatron/core/ssm/gated_delta_net/common.py
remerge CONFLICT (content): Merge conflict in megatron/core/ssm/gated_delta_net/common.py
index 07fd804bc..7e28691c1 100644
--- a/megatron/core/ssm/gated_delta_net/common.py
+++ b/megatron/core/ssm/gated_delta_net/common.py
@@ -1,29 +1,25 @@
-# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 # Copyright (c) 2025, Songlin Yang, Jan Kautz, Ali Hatamizadeh.
 
 # Some of this code was adopted from https://github.com/huggingface/transformers
 # This source code is licensed under the Apache license found in the
 # LICENSE file in the root directory of this source tree.
 
-# pylint: disable=unused-import
-
 import logging
 from dataclasses import dataclass
 from functools import lru_cache
-from typing import Callable, Optional, Protocol, Union
+from typing import Optional, Union
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch import Tensor
 
-<<<<<<< 95e4bafeb ([dev] [fix] fix optimizer_cpu_offload with mark_keep_in_fp32 (#6124)):megatron/core/ssm/gated_delta_net.py
 from megatron.core import tensor_parallel
 from megatron.core.context_parallel_layout import (
     contiguous_to_zigzag_chunks,
     zigzag_to_contiguous_chunks,
 )
-=======
... (truncated; 431 more lines in full remerge-diff) ...

diff --git a/megatron/core/ssm/gated_delta_net/gdn.py b/megatron/core/ssm/gated_delta_net/gdn.py
index 65d9dc7df..a94b17e0f 100644
--- a/megatron/core/ssm/gated_delta_net/gdn.py
+++ b/megatron/core/ssm/gated_delta_net/gdn.py
@@ -5,274 +5,15 @@
 # This source code is licensed under the Apache license found in the
 # LICENSE file in the root directory of this source tree.
 
-from functools import partial
-from typing import Optional
+# pylint: disable=unused-import
 
-import torch
-import torch.nn.functional as F
+"""GatedDeltaNet layer.
 
-from megatron.core import tensor_parallel
-from megatron.core.inference.contexts import BaseInferenceContext
-from megatron.core.packed_seq_params import PackedSeqParams
-from megatron.core.ssm.gated_delta_net.common import (
-    _build_head_perm_for_split_sections,
-    _build_thd_cp_a2a_perm,
-    _GDNBase,
-    causal_conv1d,
-    chunk_gated_delta_rule,
-    get_parameter_local_cp,
-    tensor_a2a_cp2hp,
-    torch_chunk_gated_delta_rule,
-)
-from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push
+The full ``GatedDeltaNet`` implementation lives in
+:mod:`megatron.core.ssm.gated_delta_net.common`; this module re-exports it so the
+``megatron.core.ssm.gated_delta_net.gdn`` import path keeps working.
+"""
... (truncated; 253 more lines in full remerge-diff) ...

diff --git a/megatron/post_training/model_builder.py b/megatron/post_training/model_builder.py
remerge CONFLICT (content): Merge conflict in megatron/post_training/model_builder.py
index 8887a6563..3b7ed3b96 100644
--- a/megatron/post_training/model_builder.py
+++ b/megatron/post_training/model_builder.py
@@ -515,13 +515,6 @@ def modelopt_gpt_hybrid_builder(
         # Additional tweaks needed for MCore.
         # (accounts for sharded state, pipeline parallel, and potentially skipping LM loss)
         mtd_mcore.adjust_distillation_model_for_mcore(model, distill_cfg)
-<<<<<<< 95e4bafeb ([dev] [fix] fix optimizer_cpu_offload with mark_keep_in_fp32 (#6124))
-        # Also remove KD mode state to prevent issues with re-conversion after restore.
-        mto.ModeloptStateManager(
-            model
-        ).state_dict().pop()  # TODO(aanoosheh): remove once fixed in ModelOpt
-=======
->>>>>>> ff6b92a7c (Fix --freeze-all-layers for new model builder path + unit tests (#5926))
 
     return model
 

diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py
remerge CONFLICT (content): Merge conflict in megatron/rl/rl_utils.py
index 0881a239d..8311d4640 100644
--- a/megatron/rl/rl_utils.py
+++ b/megatron/rl/rl_utils.py
@@ -278,19 +278,11 @@ def verify_model_weights_swap(
 
 @dataclass(slots=True)
 class RolloutStats:
-<<<<<<< 95e4bafeb ([dev] [fix] fix optimizer_cpu_offload with mark_keep_in_fp32 (#6124))
     rewards: list[list[float]]  # inner list is for a group
     env_ids: list[str]  # same length as len(rewards)
-    turn_lens: list[list[int]]  # token lengths of turns, grouped.
-    traj_lens: list[list[int]]  # all turns comprise one trajectory.
+    turn_lens: list[list[int]]  # tokens newly added by each turn, grouped.
+    traj_lens: list[list[int]]  # the final sequence is the full trajectory.
     num_turns: None | list[list[int]]  # num_turns per traj
-=======
-    rewards: list[list[float]] # inner list is for a group
-    env_ids: list[str] # same length as len(rewards)
-    turn_lens: list[list[int]] # tokens newly added by each turn, grouped.
-    traj_lens: list[list[int]] # the final sequence is the full trajectory.
-    num_turns: None | list[list[int]] # num_turns per traj
->>>>>>> ff6b92a7c (Fix --freeze-all-layers for new model builder path + unit tests (#5926))
     advantages: None | list[list[float]]
     min_piold_to_inf_prob: None | float
     max_piold_to_inf_prob: None | float
@@ -966,9 +958,7 @@ def compute_group_stats(
             group_kv_epoch.append([epoch for turn in rollout.kv_cache_epoch for _, epoch in turn])
             # completed_epochs is per-turn, so it cannot be masked per-rollout downstream
             if rollout.trajectory:
-                group_completed_epochs.extend(
-                    turn[-1][1] for turn in rollout.policy_epoch
-                )
... (truncated; 319 more lines in full remerge-diff) ...

[svcnvidia-nemo-ci]

/ok to test a62c10c

[svcnvidia-nemo-ci]

/ok to test 201c517

[svcnvidia-nemo-ci]

/ok to test 1241aa1