Reflect #4440 onto GPT models, not just hybrid

megatron/core/models/gpt/gpt_model.py

@@ -29,6 +29,7 @@
 from megatron.core.quantization.utils import get_quant_config_or_none
 from megatron.core.tensor_parallel import gather_from_sequence_parallel_region
 from megatron.core.transformer.enums import InferenceCudaGraphScope, ModelType
+from megatron.core.transformer.module import GraphableMegatronModule
 from megatron.core.transformer.moe.paged_stash import paged_stash_init_chunk_handler
 from megatron.core.transformer.multi_token_prediction import (
     MultiTokenPredictionBlock,
@@ -48,7 +49,7 @@
 logger = logging.getLogger(__name__)
 
 
-class GPTModel(LanguageModule):
+class GPTModel(LanguageModule, GraphableMegatronModule):
     """GPT Transformer language model.
 
     Args:
@@ -237,6 +238,8 @@ def __init__(
             pg_collection=self.pg_collection,
             vp_stage=vp_stage,
         )
+        if hasattr(self, 'cudagraph_manager') and hasattr(self.decoder, 'cudagraph_manager'):
+            del self.decoder.cudagraph_manager
 
         if self.mtp_process:
             self.mtp = MultiTokenPredictionBlock(
@@ -522,6 +525,42 @@ def preprocess_for_paged_stash(self):
             vp_size=self.config.virtual_pipeline_model_parallel_size, vp_stage=self.vp_stage
         )
 
+    def _should_call_local_cudagraph(self, *args, **kwargs):
+        """
+        Check if we should call the local cudagraph path.
+        """
+        if (
+            InferenceMode.is_active()
+            and hasattr(self, 'cudagraph_manager')
+            and (
+                kwargs.get('inference_context') is not None
+                or kwargs.get('inference_params') is not None
+            )
+            and self.config.inference_cuda_graph_scope == InferenceCudaGraphScope.block
+        ):
+            if kwargs['inference_context'].is_static_batching():
+                using_cuda_graph = kwargs['inference_context'].is_decode_only()
+            else:
+                using_cuda_graph = kwargs['inference_context'].using_cuda_graph_this_step()
+
+            if using_cuda_graph:
+                return True
+        return False
+
+    def __call__(self, *args, **kwargs):
+        if self._should_call_local_cudagraph(*args, **kwargs):
+            return super().__call__(*args, **kwargs)[0]
+        return super().__call__(*args, **kwargs)
+
+    def create_mcore_cudagraph_manager(self, config):
+        """
+        Create the cudagraph manager for the full iteration inference scope
+        """
+        if config.inference_cuda_graph_scope == InferenceCudaGraphScope.block:
+            from megatron.core.transformer.cuda_graphs import CudaGraphManager
+
+            self.cudagraph_manager = CudaGraphManager(config)
+
     def forward(
         self,
         input_ids: Tensor,

tests/unit_tests/inference/engines/test_dynamic_engine.py

@@ -213,6 +213,7 @@ def _assert_inference_cuda_graphs_disabled(env) -> None:
         assert env.engine.cuda_graph_impl == "full_iteration"
         assert env.engine.inference_cuda_graph_scope == InferenceCudaGraphScope.none
         assert env.engine.capture_stats is None
+        assert not hasattr(model, 'cudagraph_manager')
         assert not hasattr(model.decoder, 'cudagraph_manager')
         for layer in model.decoder.layers:
             assert not hasattr(layer, 'cudagraph_manager')
@@ -845,11 +846,10 @@ def test_simple(self, model_provider, num_cuda_graphs, inference_cuda_graph_scop
             assert env.engine.context.cuda_graph_batch_dimensions_list
             model = env.engine.controller.inference_wrapped_model.model
             if inference_cuda_graph_scope == InferenceCudaGraphScope.block:
-                # hybrid models attach cudagraph_manager to the model; others attach to the decoder
-                if model_provider == "hybrid":
-                    assert model.cudagraph_manager.cudagraph_runners
-                else:
-                    assert model.decoder.cudagraph_manager.cudagraph_runners
+                # GPT and hybrid models both own the block-scope graph at model level;
+                # GPT removes the decoder's fallback manager at construction.
+                assert model.cudagraph_manager.cudagraph_runners
+                assert not hasattr(model.decoder, 'cudagraph_manager')
             else:
                 # check if cudagraph runners are created at the layer level
                 for layer in model.decoder.layers:
@@ -959,7 +959,8 @@ def test_deprecated_full_iteration_inference_scope_matches_new_flag_runtime_beha
         model = env.engine.controller.inference_wrapped_model.model
         assert model.config.inference_cuda_graph_scope == InferenceCudaGraphScope.block
         assert model.config.cuda_graph_modules == []
-        assert model.decoder.cudagraph_manager.cudagraph_runners
+        assert model.cudagraph_manager.cudagraph_runners
+        assert not hasattr(model.decoder, 'cudagraph_manager')
         for layer in model.decoder.layers:
             assert not hasattr(layer, 'cudagraph_manager')