Enable CUDA graphs for MTP inference

examples/inference/gpt/gpt_dynamic_inference.py

@@ -286,10 +286,6 @@ def main():
     )
     initialize_megatron()
 
-    # Start Nsight profiler.
-    if os.environ.get("NSIGHT_PREFIX"):
-        torch.cuda.cudart().cudaProfilerStart()
-
     level_str = os.getenv("LOG_LEVEL", "INFO").upper()
     level = getattr(logging, level_str, logging.INFO)
     logging.basicConfig(level=level, force=True)
@@ -350,28 +346,53 @@ def main():
     print(setup_prefix)
     print("~~~")
 
+    # Warmup: run one untimed iteration so CUDA caches, JIT kernels, and
+    # allocator pools are ready before the measured runs.
+    if args.inference_repeat_n > 1:
+        print("Running warmup iteration ...")
+        engine.reset()
+        run_inference(requests, engine)
+        torch.cuda.synchronize()
+        engine.reset()
+
+    # Start CUDA profiler after warmup so nsys traces only the measured runs.
+    if os.environ.get("NSIGHT_PREFIX"):
+        torch.cuda.cudart().cudaProfilerStart()
+
     # Run and time test, optionally `args.inference_repeat_n` times.
     throughputs = []
+    cuda_start_event = torch.cuda.Event(enable_timing=True)
+    cuda_end_event = torch.cuda.Event(enable_timing=True)
     for _ in range(args.inference_repeat_n):
 
         # Reset engine.
         engine.reset()
 
         torch.cuda.reset_peak_memory_stats()
 
-        # Trial.
+        # Synchronize before starting the timer to avoid measuring stale GPU work.
+        torch.cuda.synchronize()
+
+        # Trial — use both wall-clock and CUDA events for accurate GPU timing.
         t = get_curr_time()
+        cuda_start_event.record()
         result = run_inference(requests, engine)
+        cuda_end_event.record()
         step_times = result["step_times"]
         add_times = result["add_times"]
         output_times = result["output_times"]
         total_output_tokens = result["total_output_tokens"]
         torch.cuda.synchronize()
         total_time = get_curr_time() - t
+        cuda_elapsed_ms = cuda_start_event.elapsed_time(cuda_end_event)
         stats = torch.cuda.memory_stats()
         throughput = total_output_tokens / total_time
         throughputs.append(throughput)
 
+    # Stop CUDA profiler after measured runs.
+    if os.environ.get("NSIGHT_PREFIX"):
+        torch.cuda.cudart().cudaProfilerStop()
+
     # Validate all requests finished.
     for request in requests:
         assert request.state == "finished", f"request.state == '{request.state}' != 'finished'."
@@ -505,19 +526,17 @@ def escape_str(s):
         #     f"count [ p {p_count}, d {d_count} ]."
         # )
         capture_str = f"{engine.capture_stats['time']:.2f} sec" if engine.capture_stats else "--"
+        cuda_throughput = total_output_tokens / (cuda_elapsed_ms / 1000.0)
         print(
             f"{setup_prefix} … " f"throughput: {throughput:.3f} tok/s … ",
             f"total time: {total_time:.3f}s … "
+            f"cuda time: {cuda_elapsed_ms:.1f}ms ({cuda_throughput:.3f} tok/s) … "
             f"mem {peak_alloc_gb:.1f}/{peak_resvd_gb:.1f} GB … "
             f"steps: {engine.context.step_count:d} … "
             f"capture {capture_str}",
         )
         print("~~~")
 
-    # Stop Nsight profiler.
-    if os.environ.get("NSIGHT_PREFIX"):
-        torch.cuda.cudart().cudaProfilerStop()
-
 
 if __name__ == "__main__":
     main()

examples/inference/gpt/gpt_dynamic_inference_12b.sh

@@ -5,7 +5,6 @@
 set -u
 
 # Libraries.
-pip install simpy
 pip install sentencepiece
 pip install tiktoken
 

examples/inference/gpt/gpt_dynamic_inference_357m.sh

@@ -5,7 +5,6 @@
 set -u
 
 # Libraries.
-pip install simpy
 pip install sentencepiece
 pip install tiktoken
 

examples/inference/gpt/utils.py

@@ -106,18 +106,13 @@ def get_time_offsets(
 
     random.seed(seed)
 
-    import simpy  # Guard against this import in test case
-
-    # Generate random time offsets.
-    def arrival(r):
-        while True:
-            yield env.timeout(random.expovariate(r))
-            time_offsets.append(env.now)
-
+    # Generate Poisson arrival times by accumulating exponential inter-arrival intervals.
     time_offsets = []
-    env = simpy.Environment()
-    env.process(arrival(incoming_requests_per_sec))
-    env.run(incoming_requests_duration)
+    current_time = 0.0
+    while current_time < incoming_requests_duration:
+        current_time += random.expovariate(incoming_requests_per_sec)
+        if current_time < incoming_requests_duration:
+            time_offsets.append(current_time)
 
     # Ensure at least a single request.
     if len(time_offsets) == 0:

megatron/core/inference/batch_dimensions_utils.py

@@ -85,6 +85,10 @@ def is_valid(
         Returns:
             True if the config is valid, False otherwise
         """
+        # A dimension with no tokens serves no requests.
+        if self.token_count <= 0:
+            return False
+
         # Check if total requests exceed maximum
         if self.prefill_req_count + self.decode_req_count > max_requests:
             return False

megatron/core/inference/engines/dynamic_engine.py

@@ -50,7 +50,7 @@
     unset_inference_cuda_graphed_iteration_for_ep_inference,
 )
 from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.core.transformer.cuda_graphs import delete_cuda_graphs
+from megatron.core.transformer.cuda_graphs import delete_cuda_graphs, graph_capture
 from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
 from megatron.core.utils import (
@@ -64,6 +64,7 @@
     nvtx_range_pop,
     nvtx_range_push,
     trace_async_exceptions,
+    unwrap_model,
 )
 
 from .async_zmq_communicator import AsyncZMQCommunicator
@@ -365,6 +366,21 @@ def create_cuda_graphs(self, reset_context: bool = True):
             unwrapped_model = controller.inference_wrapped_model.model
             set_inference_cuda_graphed_iteration_for_ep_inference(unwrapped_model)
 
+        # MTP warmup preparation: capture MTP CUDA graphs alongside the
+        # decoder graphs within the same loop rather than in a separate pass.
+        unwrapped = unwrap_model(controller.inference_wrapped_model.model)
+        mtp_warmup_enabled = (
+            controller.num_mtp_heads > 0
+            and (controller.num_speculative_tokens or 0) > 0
+            and hasattr(unwrapped, 'mtp')
+        )
+        if mtp_warmup_enabled:
+            tp_size = get_pg_size(controller.inference_wrapped_model.tp_group)
+            sp_enabled = model_config.sequence_parallel and tp_size > 1
+            mtp_pass_depth = not unwrapped.mtp.mtp_use_repeated_layer
+            mtp_warmup_depths = range(controller._num_mtp_depths) if mtp_pass_depth else [None]
+            mtp_seen_batch_sizes = set()
+
         tbar = enumerate(context.cuda_graph_batch_dimensions_list)
         if HAVE_TQDM:
             tbar = tqdm(tbar, total=len(context.cuda_graph_batch_dimensions_list))
@@ -390,12 +406,39 @@ def create_cuda_graphs(self, reset_context: bool = True):
             # Forward pass -> logits.
             controller._dynamic_step_forward_logits(input_ids, position_ids)
 
+            # MTP CUDA graph warmup for this batch dimension.
+            if mtp_warmup_enabled:
+                n = cuda_graph_batch_dimension.req_count
+                if sp_enabled:
+                    n += (tp_size - n % tp_size) % tp_size
+                if n > 0 and n not in mtp_seen_batch_sizes:
+                    mtp_seen_batch_sizes.add(n)
+                    device = torch.cuda.current_device()
+                    batch_dim = n // tp_size if sp_enabled else n
+                    # Use zeros (not empty) — garbage token IDs cause OOB embedding lookups during graph capture/replay.
+                    for depth in mtp_warmup_depths:
+                        with graph_capture():
+                            unwrapped.compute_mtp_single_step(
+                                hidden_states=torch.zeros(
+                                    (batch_dim, 1, model_config.hidden_size),
+                                    device=device,
+                                    dtype=model_config.params_dtype,
+                                ),
+                                next_token_ids=torch.zeros((1, n), device=device, dtype=torch.long),
+                                position_ids=torch.zeros((1, n), device=device, dtype=torch.int64),
+                                depth=depth,
+                            )
+
             context.reset()
 
         # Disable inference dispatcher after graph capture
         if is_inference_optimized_ep:
             unset_inference_cuda_graphed_iteration_for_ep_inference(unwrapped_model)
 
+        if mtp_warmup_enabled and mtp_seen_batch_sizes:
+            controller._has_mtp_cuda_graphs = True
+            logging.info("> MTP CUDA graph warmup: %d batch size(s)", len(mtp_seen_batch_sizes))
+
         # Memory usage.
         time_end = time.time()
         mem_stats_end = torch.cuda.memory_stats()