Enable inference CUDA graphs with Shortcut MoE and wide residuals

megatron/core/ssm/ssm_inference.py

@@ -57,6 +57,9 @@
     tensor_get_slice_after,
     tensor_merge,
 )
+from megatron.core.transformer.custom_layers.batch_invariant_kernels import (
+    is_batch_invariant_mode_enabled,
+)
 from megatron.core.utils import is_using_quantization_scales
 
 
@@ -287,6 +290,12 @@ def ssm_dynamic_inference(
             f"({y.shape[0]}) than the input projection contained ({zxBCdt.shape[0]})."
         )
         if padding_token_count > 0:
+            # Graph bucket alignment reads the process-global kernel switch;
+            # the context's SSM batch-invariant path can also require token padding.
+            assert context.batch_invariant_mode or is_batch_invariant_mode_enabled(), (
+                "Token-only SSM padding requires batch-invariant SSM execution "
+                "or the global batch-invariant kernel mode."
+            )
             y = torch.cat((y, y.new_zeros(padding_token_count, *y.shape[1:])), dim=0)
 
         # Zero padding positions to avoid corrupting quantization amax calculations.

megatron/core/transformer/transformer_config.py

@@ -1665,9 +1665,20 @@ def __post_init__(self):
         if self.wide_residual is not None:
             if self.enable_mhc_connections:
                 raise ValueError("wide_residual and enable_mhc_connections are mutually exclusive.")
-            if self.cuda_graph_impl != "none" or self.enable_cuda_graph or self.external_cuda_graph:
+            if (
+                self.cuda_graph_impl != "none" or self.enable_cuda_graph or self.external_cuda_graph
+            ) and not (
+                self.cuda_graph_impl == "local"
+                and not self.enable_cuda_graph
+                and not self.external_cuda_graph
+                and normalize_inference_cuda_graph_scope(
+                    self.inference_cuda_graph_scope, self.cuda_graph_impl
+                )
+                == InferenceCudaGraphScope.block
+            ):
                 raise NotImplementedError(
-                    "wide_residual does not yet support CUDA graphs; use cuda_graph_impl='none'."
+                    "wide_residual supports CUDA graphs only with cuda_graph_impl='local' "
+                    "and inference_cuda_graph_scope='block'."
                 )
             if self.pipeline_model_parallel_size > 1:
                 raise NotImplementedError(
@@ -3401,8 +3412,16 @@ def _scope_to_str(s):
         ), 'cuda_graph_modules must be empty when cuda_graph_impl="full_iteration".'
 
         assert not (
-            self.moe_shortcut_connection and self.cuda_graph_impl != "none"
-        ), "CUDA graphs are not supported with moe_shortcut_connection."
+            self.moe_shortcut_connection
+            and self.cuda_graph_impl != "none"
+            and not (
+                self.cuda_graph_impl == "local"
+                and self.inference_cuda_graph_scope == InferenceCudaGraphScope.block
+            )
+        ), (
+            "moe_shortcut_connection supports CUDA graphs only with "
+            "cuda_graph_impl='local' and inference_cuda_graph_scope='block'."
+        )
 
         if self.cuda_graph_impl != "none":
 

megatron/training/arguments.py

@@ -1965,8 +1965,16 @@ def validate_args(args, defaults={}):
     assert not (
         args.cuda_graph_impl == "full_iteration" and args.cuda_graph_modules
     ), '--cuda-graph-modules must be empty when --cuda-graph-impl=full_iteration.'
-    assert not (args.moe_shortcut_connection and args.cuda_graph_impl != "none"), (
-        "CUDA graphs are not supported with --moe-shortcut-connection."
+    assert not (
+        args.moe_shortcut_connection
+        and args.cuda_graph_impl != "none"
+        and not (
+            args.cuda_graph_impl == "local"
+            and args.inference_cuda_graph_scope == InferenceCudaGraphScope.block
+        )
+    ), (
+        "--moe-shortcut-connection supports CUDA graphs only with "
+        "--cuda-graph-impl=local and --inference-cuda-graph-scope=block."
     )
 
     if args.multi_latent_attention:

tests/unit_tests/inference/engines/test_shortcut_moe_cuda_graph_e2e.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Exercise Shortcut MoE and wide residuals through the inference graph manager."""
+
+import torch
+
+from megatron.core import parallel_state
+from megatron.core.activations import squared_relu
+from megatron.core.inference.symmetric_memory import SymmetricMemoryManager
+from megatron.core.models.hybrid.hybrid_layer_specs import (
+    wide_residual_gated_delta_product_inference_stack_spec,
+)
+from megatron.core.models.hybrid.hybrid_model import HybridModel
+from megatron.core.models.hybrid.shortcut_block import ShortcutMoEBlock
+from megatron.core.transformer.attention import Attention
+from megatron.core.transformer.enums import AttnBackend
+from megatron.core.transformer.moe.token_dispatcher_inference import NVLSAllGatherVDispatcher
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.transformer.wide_residual_config import WideResidualConfig
+from tests.unit_tests.inference.engines.test_gdp_cuda_graph_e2e import (
+    FLASH_ATTENTION_VERSION,
+    MAX_SEQ_LEN,
+    VOCAB_SIZE,
+)
+from tests.unit_tests.inference.engines.test_gdp_cuda_graph_e2e import (
+    TestGDPCudaGraphE2E as _GDPGraphTest,
+)
+from tests.unit_tests.test_utilities import Utils
+
+
+class TestShortcutMoECudaGraphE2E(_GDPGraphTest):
+    """Greedy tokens must match eager inference for a GDP/Shortcut-MoE pair."""
+
+    @classmethod
+    def setup_class(cls):
+        Utils.initialize_distributed()
+        Utils.initialize_model_parallel(
+            expert_model_parallel_size=torch.distributed.get_world_size()
+        )
+
+    @classmethod
+    def teardown_class(cls):
+        try:
+            NVLSAllGatherVDispatcher._delete_buffers()
+            SymmetricMemoryManager.destroy()
+        finally:
+            super().teardown_class()
+
+    def _create_model(self):
+        torch.manual_seed(1234)
+        config = TransformerConfig(
+            params_dtype=torch.bfloat16,
+            num_layers=3,
+            hidden_size=256,
+            num_attention_heads=16,
+            ffn_hidden_size=512,
+            activation_func=squared_relu,
+            mamba_num_heads=8,
+            mamba_head_dim=32,
+            mamba_num_groups=8,
+            mamba_state_dim=64,
+            gdp_num_householder=2,
+            num_moe_experts=8,
+            moe_router_topk=2,
+            moe_router_pre_softmax=True,
+            moe_router_dtype="fp32",
+            moe_token_dispatcher_type="alltoall",
+            moe_shortcut_connection=True,
+            moe_shared_expert_intermediate_size=512,
+            moe_shortcut_post_norm=True,
+            wide_residual=WideResidualConfig(num_streams=3),
+            use_cpu_initialization=True,
+            transformer_impl="inference_optimized",
+            normalization="RMSNorm",
+            cuda_graph_impl="local",
+            inference_cuda_graph_scope="block",
+            inference_rng_tracker=True,
+            tensor_model_parallel_size=1,
+            expert_model_parallel_size=parallel_state.get_expert_model_parallel_world_size(),
+            pipeline_model_parallel_size=1,
+            pipeline_dtype=torch.bfloat16,
+            add_bias_linear=False,
+            is_hybrid_model=True,
+            attention_dropout=0.0,
+            attention_backend=AttnBackend.flash,
+            flash_attention_version=FLASH_ATTENTION_VERSION,
+        )
+        model = HybridModel(
+            config=config,
+            hybrid_stack_spec=wide_residual_gated_delta_product_inference_stack_spec,
+            vocab_size=VOCAB_SIZE,
+            max_sequence_length=MAX_SEQ_LEN,
+            parallel_output=True,
+            hybrid_layer_pattern="ME*",
+            pre_process=parallel_state.is_pipeline_first_stage(),
+            post_process=parallel_state.is_pipeline_last_stage(),
+        ).cuda()
+        assert any(isinstance(layer, ShortcutMoEBlock) for layer in model.decoder.layers)
+        for param in model.parameters():
+            param.data = param.data.to(config.params_dtype)
+        model.eval()
+        for module in model.modules():
+            if isinstance(module, Attention):
+                module.batch_invariant_mode = True
+        return model

tests/unit_tests/inference/test_mimo_text_checkpoint_detection.py

@@ -0,0 +1,29 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""MIMO checkpoints select the text backbone and its checkpoint prefix."""
+
+from argparse import Namespace
+
+import pytest
+
+from examples.inference.advanced import vlm_dynamic_inference
+
+
+@pytest.mark.parametrize("provider", ["nemotron-moe-vlm", "nemotron-moe-mistral-vit"])
+def test_mimo_text_detection_selects_text_backbone(monkeypatch, provider):
+    """Model detection selects the text backbone without constructing a VLM."""
+    args = Namespace(model_provider="gpt")
+    saved = Namespace(model_provider=provider, mimo_llm_tp=1)
+    monkeypatch.setattr(vlm_dynamic_inference, "load_args_from_checkpoint", lambda _: (args, saved))
+    # Detection also reads tensor metadata on newer main; keep this unit test independent of I/O.
+    monkeypatch.setattr(
+        vlm_dynamic_inference,
+        "_checkpoint_tensor_keys",
+        lambda _: [
+            "language_model.module.module.decoder.layers.0.self_attention.linear_qkv.weight"
+        ],
+        raising=False,
+    )
+    assert not vlm_dynamic_inference._detect_vlm_from_checkpoint(args)
+    assert args.model_provider == "hybrid"
+    assert args.checkpoint_model_prefix == "language_model.module.module."

tests/unit_tests/ssm/test_gdp_dynamic_inference.py

@@ -64,6 +64,7 @@
 from megatron.core.inference.utils import InferenceMode
 from megatron.core.models.hybrid.hybrid_layer_specs import gated_delta_product_inference_stack_spec
 from megatron.core.models.hybrid.hybrid_model import HybridModel
+from megatron.core.ssm import ssm_inference
 from megatron.core.ssm.gated_delta_product import GatedDeltaProductMixer
 from megatron.core.ssm.ops.gdp.chunk import chunk_gated_delta_product_varlen
 from megatron.core.ssm.ops.gdp.fused_recurrent import fused_recurrent_gated_delta_rule_update
@@ -217,19 +218,25 @@ def out_proj(self, y):
 
 
 @pytest.mark.parametrize(
-    ("batch_invariant_mode", "num_requests", "tokens_per_request", "padded_token_count"),
-    [
-        (False, 40, 1, 40),
-        (True, 40, 1, 40),
-        (True, 40, 1, 64),
-        (False, 20, 3, 60),
-        (True, 20, 3, 64),
-    ],
+    ("batch_invariant_mode", "global_batch_invariant_mode"),
+    [(False, False), (False, True), (True, False), (True, True)],
+)
+@pytest.mark.parametrize(
+    ("num_requests", "tokens_per_request", "padded_token_count"),
+    [(40, 1, 40), (40, 1, 64), (20, 3, 60), (20, 3, 64)],
 )
-def test_decode_ssm_preserves_batch_invariant_token_padding(
-    batch_invariant_mode, num_requests, tokens_per_request, padded_token_count
+def test_decode_ssm_preserves_token_padding(
+    monkeypatch,
+    batch_invariant_mode,
+    global_batch_invariant_mode,
+    num_requests,
+    tokens_per_request,
+    padded_token_count,
 ):
-    """Only batch-invariant token-only rows bypass SSM decode."""
+    """Token-only padding requires at least one of the independent switches."""
+    monkeypatch.setattr(
+        ssm_inference, "is_batch_invariant_mode_enabled", lambda: global_batch_invariant_mode
+    )
     metadata_token_count = num_requests * tokens_per_request
     projected = torch.arange(padded_token_count * 4, dtype=torch.float32).reshape(
         padded_token_count, 1, 4
@@ -250,6 +257,13 @@ def test_decode_ssm_preserves_batch_invariant_token_padding(
         padding_slice=slice(metadata_token_count, padded_token_count),
     )
 
+    if padded_token_count > metadata_token_count and not (
+        batch_invariant_mode or global_batch_invariant_mode
+    ):
+        with pytest.raises(AssertionError, match="Token-only SSM padding requires"):
+            mixer.ssm_dynamic_inference(torch.empty(0), context)
+        return
+
     output, bias = mixer.ssm_dynamic_inference(torch.empty(0), context)
 
     decode_input, received_indices = mixer.decode_inputs[0]

tests/unit_tests/transformer/test_residual_recompute.py

@@ -499,7 +499,7 @@ def test_rejects_invalid_block_size(self, block_size):
             )
 
     def test_inherits_wide_residual_cuda_graph_guard(self):
-        with pytest.raises(NotImplementedError, match="wide_residual does not yet support"):
+        with pytest.raises(NotImplementedError, match="wide_residual supports CUDA graphs only"):
             TransformerConfig(
                 num_layers=2,
                 hidden_size=8,

tests/unit_tests/transformer/test_transformer_config.py

@@ -190,8 +190,8 @@ def test_attn_logit_softcapping_rejects_invalid_values(softcap):
         pytest.param(
             {"num_moe_experts": 2, "moe_shortcut_connection": True, "cuda_graph_impl": "local"},
             AssertionError,
-            "CUDA graphs are not supported",
-            id="cuda-graphs",
+            "supports CUDA graphs only",
+            id="cuda-graphs-layer-scope",
         ),
     ],
 )
@@ -207,6 +207,22 @@ def test_shortcut_rejects_incompatible_configurations(overrides, error, message)
         )
 
 
+def test_shortcut_accepts_local_block_cuda_graphs():
+    config = TransformerConfig(
+        num_layers=2,
+        hidden_size=128,
+        num_attention_heads=4,
+        num_moe_experts=2,
+        moe_router_topk=1,
+        moe_router_pre_softmax=True,
+        moe_shortcut_connection=True,
+        cuda_graph_impl="local",
+        inference_cuda_graph_scope="block",
+    )
+    assert config.cuda_graph_impl == "local"
+    assert config.inference_cuda_graph_scope.name == "block"
+
+
 def _make_shortcut_memory_config(**overrides):
     kwargs = dict(
         num_layers=2,

tests/unit_tests/transformer/test_wide_residual_connection.py

@@ -718,9 +718,13 @@ def test_static_subclass_is_rejected_for_mtp_auxiliary_stacks(self):
         ids=("local", "transformer_engine", "full_iteration", "legacy_local", "legacy_te"),
     )
     def test_wide_residual_rejects_cuda_graphs(self, config_overrides):
-        with pytest.raises(NotImplementedError, match="does not yet support CUDA graphs"):
+        with pytest.raises(NotImplementedError, match="supports CUDA graphs only"):
             _wide_config(**config_overrides)
 
+    def test_wide_residual_accepts_local_block_cuda_graphs(self):
+        config = _wide_config(cuda_graph_impl="local", inference_cuda_graph_scope="block")
+        assert config.inference_cuda_graph_scope.name == "block"
+
 
 class TestWideResidualLayerIntegration:
     @staticmethod