From 85efc1f3ead8c0900f442e81c6a0cbaf28c56fb5 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Wed, 14 Jan 2026 15:06:17 -0800
Subject: [PATCH 01/30] WIP engine and context cleanup

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 .../inference/gpt/gpt_dynamic_inference.py    | 274 +++++-------------
 .../gpt_dynamic_inference_with_coordinator.py | 111 +++----
 examples/inference/gpt/utils.py               | 154 ++++++----
 .../inference/contexts/dynamic_context.py     | 132 ++++-----
 .../core/inference/contexts/static_context.py |   6 +-
 .../core/inference/engines/dynamic_engine.py  |  25 +-
 .../core/inference/engines/static_engine.py   |  13 +-
 .../abstract_model_inference_wrapper.py       | 221 +++-----------
 .../gpt/gpt_inference_wrapper.py              |   8 +-
 .../inference_wrapper_config.py               |  66 -----
 .../t5/t5_inference_wrapper.py                |   7 +-
 .../text_generation_controller.py             |  65 +++--
 megatron/core/transformer/attention.py        |   5 +-
 .../core/transformer/transformer_config.py    |   9 +
 megatron/rl/inference/megatron.py             |  20 +-
 megatron/training/arguments.py                |   5 +-
 tools/run_dynamic_text_generation_server.py   |  25 +-
 tools/run_inference_performance_test.py       |  63 +---
 18 files changed, 409 insertions(+), 800 deletions(-)
 delete mode 100644 megatron/core/inference/model_inference_wrappers/inference_wrapper_config.py

diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py
index 7eeee21562c..6bd4a3df037 100644
--- a/examples/inference/gpt/gpt_dynamic_inference.py
+++ b/examples/inference/gpt/gpt_dynamic_inference.py
@@ -8,12 +8,12 @@
 import pickle
 import sys
 import warnings
-import torch
 from argparse import ArgumentParser
 from collections import defaultdict
-from functools import partial
+from typing import Dict, List, Optional, Tuple
+
+import torch
 from tqdm import tqdm
-from typing import Dict, List, Tuple, Optional
 
 sys.path.append(
     os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
@@ -27,14 +27,15 @@
     build_requests,
     get_curr_time,
     get_global_peak_memory_stats_bytes,
+    get_model,
+)
+from megatron.core.inference.contexts.attention_context.mamba_metadata import (
+    MambaInferenceStateConfig,
 )
 from megatron.core.inference.contexts.dynamic_context import (
     ContextOverflowError,
     DynamicInferenceContext,
 )
-from megatron.core.inference.contexts.attention_context.mamba_metadata import (
-    MambaInferenceStateConfig,
-)
 from megatron.core.inference.engines import DynamicInferenceEngine, EngineSuspendedError
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
     GPTInferenceWrapper,
@@ -50,19 +51,18 @@
 sys.path.append(
     os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
 )
-from megatron.training import get_args, get_model as _get_model, get_tokenizer, initialize_megatron
-from megatron.training.checkpointing import load_checkpoint
-from model_provider import model_provider
-from gpt_builders import gpt_builder
-from mamba_builders import mamba_builder
+import logging
 
 from megatron.core.utils import configure_nvtx_profiling
-import logging
+from megatron.training import get_args
+from megatron.training import get_model as _get_model
+from megatron.training import get_tokenizer, initialize_megatron
 
 torch.serialization.add_safe_globals([io.BytesIO])
 torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunState])
 torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunDiagnostic])
 
+
 def add_dynamic_inference_args(parser: ArgumentParser) -> ArgumentParser:
     """Dynamic inference arguments."""
 
@@ -75,68 +75,37 @@ def add_dynamic_inference_args(parser: ArgumentParser) -> ArgumentParser:
         help="Load checkpoint with `strict=False`.",
     )
     group.add_argument(
-        "--termination-id", type=int, default=None,
+        "--termination-id",
+        type=int,
+        default=None,
         help="Termination ID that overrides `tokenizer.eod`.",
     )
     group.add_argument(
-        "--suspend-resume-interval", type=int, default=None,
+        "--suspend-resume-interval",
+        type=int,
+        default=None,
         help="Suspend and resume the dynamic engine every "
         "`suspend_resume_interval` steps. This is used to tet the suspend/resume "
         "system.",
     )
     group.add_argument(
-        "--inference-repeat-n", type=int, default=1,
-        help="Repeat inference iterations N times for benchmarking."
+        "--inference-repeat-n",
+        type=int,
+        default=1,
+        help="Repeat inference iterations N times for benchmarking.",
     )
     group.add_argument(
         "--throughput-check-only",
         action='store_true',
         default=False,
-        help="If true, only run throughput check without verifying outputs."
+        help="If true, only run throughput check without verifying outputs.",
     )
 
     return parser
 
 
-def get_model() -> MegatronModule:
-    """Initialize model and load checkpoint."""
-
-    args = get_args()
-
-    if args.model_provider == "gpt":
-        model_builder = gpt_builder
-    elif args.model_provider == "mamba":
-        model_builder = mamba_builder
-    else:
-        raise ValueError(f"Invalid model provider {args.model_provider}")
-
-    # Build model.
-    model = _get_model(
-        partial(model_provider, model_builder),
-        wrap_with_ddp=False
-    )
-
-    # Load checkpoint.
-    assert args.load is not None
-    args.exit_on_missing_checkpoint = True
-    load_checkpoint(
-        ddp_model=model,
-        optimizer=None,
-        opt_param_scheduler=None,
-        strict=not args.inference_ckpt_non_strict,
-    )
-
-    # No virtual PP.
-    assert len(model) == 1, "Above condition should have caught this"
-    model = model[0]
-
-    # Eval mode.
-    model.eval()
-
-    return model
-
-
 def get_inference_context(
+    model,
     requests: List[Request],
     sampling_params: Optional[SamplingParams] = None,
     calculate_max_sequence_length_from_requests: bool = True,
@@ -146,87 +115,16 @@ def get_inference_context(
 
     args = get_args()
 
+    overrides = None
+
     # Max sequence length.
     if calculate_max_sequence_length_from_requests:
-        max_gen_length = sampling_params.num_tokens_to_generate    
+        max_gen_length = sampling_params.num_tokens_to_generate
         max_context_length = max(len(r.prompt_tokens) for r in requests)
         max_sequence_length = max_context_length + max_gen_length
-    else:
-        max_sequence_length = args.inference_max_seq_length
-
-    metrics_writer = None
-    if args.inference_logging_step_interval > 0 and args.inference_wandb_logging:
-        metrics_writer = get_wandb_writer()
-
-    # Inference context.
-    context = DynamicInferenceContext(
-        params_dtype=args.params_dtype,
-        num_layers=args.num_layers // args.pipeline_model_parallel_size,
-        kv_channels=args.kv_channels,
-        num_attention_heads=(
-            args.num_query_groups if args.group_query_attention else args.num_attention_heads
-        ),
-        max_sequence_length=max_sequence_length,
-        num_cuda_graphs=(
-            args.inference_dynamic_batching_num_cuda_graphs
-            if args.cuda_graph_impl == "local"
-            else None
-        ),
-        block_size_tokens=args.inference_dynamic_batching_block_size,
-        buffer_size_gb=args.inference_dynamic_batching_buffer_size_gb,
-        max_requests=args.inference_dynamic_batching_max_requests,
-        max_tokens=args.inference_dynamic_batching_max_tokens,
-        tensor_model_parallel_size=args.tensor_model_parallel_size,
-        pipeline_model_parallel_size=args.pipeline_model_parallel_size,
-        materialize_only_last_token_logits=not args.return_log_probs,
-        mamba_inference_state_config=mamba_inference_state_config,
-        cache_mla_latent=args.multi_latent_attention and args.cache_mla_latents,
-        kv_lora_rank=args.kv_lora_rank if args.multi_latent_attention else None,
-        qk_pos_emb_head_dim=args.qk_pos_emb_head_dim,
-        use_cuda_graphs_for_non_decode_steps=not args.decode_only_cuda_graphs,
-        use_flashinfer_fused_rope=args.use_flashinfer_fused_rope,
-        unified_memory_level=args.inference_dynamic_batching_unified_memory_level,
-        cuda_graph_max_tokens=args.inference_dynamic_batching_cuda_graph_max_tokens,
-        cuda_graph_mixed_prefill_count=args.inference_dynamic_batching_cuda_graph_mixed_prefill_count,
-        metrics_writer=metrics_writer,
-    )
-
-    return context
+        overrides = {"max_sequence_length": max_sequence_length}
 
-
-def get_inference_controller(
-    model: MegatronModule, context: DynamicInferenceContext
-) -> TextGenerationController:
-    """Buid text generation controller, which manages the model inference context.
-
-    Args:
-        model (MegatronModule): Megatron GPT model.
-        context (DynamicInferenceContext): Context for managing KV cache blocks.
-
-    Return:
-        (TextGenerationController) Inference text generation controller.
-    """
-
-    args = get_args()
-    if args.legacy_tokenizer:
-        tokenizer = get_tokenizer()
-    else:
-        tokenizer = build_tokenizer(args)
-
-    # Wrap model in inference wrapper.
-    model = GPTInferenceWrapper(model, args, context)
-
-    # Note: the following is taken from AbstractModelInferenceWrapper.prep_model_for_inference().
-    from megatron.core import parallel_state
-
-    model.model_is_pipeline_parallel = not (
-        parallel_state.is_pipeline_first_stage() and parallel_state.is_pipeline_last_stage()
-    )
-
-    # Text generation controller.
-    controller = TextGenerationController(model, tokenizer)
-
-    return controller
+    return DynamicInferenceContext.from_model_and_args(model, args, overrides)
 
 
 def run_inference(
@@ -281,11 +179,7 @@ def _add_request():
         """
         nonlocal num_requests_added
         _request = requests[num_requests_added]
-        engine.add_request(
-            num_requests_added,
-            _request.prompt_text,
-            _request.sampling_params,
-        )
+        engine.add_request(num_requests_added, _request.prompt_text, _request.sampling_params)
         _request.time_start = get_curr_time()
         _request.state = "started"
         num_requests_added += 1
@@ -302,10 +196,9 @@ def _add_request():
                 _add_request()
         else:
             # Add deterministic number of requests (generally used for debugging).
-            for i in range(min(
-                args.incoming_requests_per_step,
-                num_requests_total - num_requests_added,
-            )):
+            for i in range(
+                min(args.incoming_requests_per_step, num_requests_total - num_requests_added)
+            ):
                 _add_request()
         add_times.append(get_curr_time() - add_start)
 
@@ -315,11 +208,11 @@ def _add_request():
             result = engine.step_modern()
         except EngineSuspendedError as e:
             result = e
-            pass # ignore error in order to call 'engine.resume()' below.
+            pass  # ignore error in order to call 'engine.resume()' below.
         attempted_step_count += 1
 
         # After step, we lost track of last iteration's is_decode_only, so we need to get it from the engine
-        is_decode_only = engine.is_decode_only 
+        is_decode_only = engine.is_decode_only
 
         # Test suspending and resuming engine.
         if args.suspend_resume_interval is not None:
@@ -332,9 +225,9 @@ def _add_request():
             # Resume, 0+ attempted steps later.
             if (
                 attempted_step_count > 0
-                and
-                (attempted_step_count - args.suspend_resume_interval // 2)
-                    % args.suspend_resume_interval == 0
+                and (attempted_step_count - args.suspend_resume_interval // 2)
+                % args.suspend_resume_interval
+                == 0
             ):
                 print("**** step %d/%d ... resume." % (engine.step_count, attempted_step_count))
                 engine.resume()
@@ -346,7 +239,9 @@ def _add_request():
         # Record cuda_graph_request_count.
         cuda_graph_request_count = result["cuda_graph_request_count"]
         if args.cuda_graph_impl == "local" and cuda_graph_request_count is not None:
-            cuda_graph_request_count_map[cuda_graph_request_count] = cuda_graph_request_count_map.get(cuda_graph_request_count, 0) + 1
+            cuda_graph_request_count_map[cuda_graph_request_count] = (
+                cuda_graph_request_count_map.get(cuda_graph_request_count, 0) + 1
+            )
 
         # Update requests.
         active_request_ids = result["active_request_ids"]
@@ -404,11 +299,11 @@ def _add_request():
         engine.resume()
 
     return {
-        "step_times" : step_times,
-        "add_times" : add_times,
-        "output_times" : output_times,
-        "total_output_tokens" : total_output_tokens,
-        "cuda_graph_request_count_map" : cuda_graph_request_count_map,
+        "step_times": step_times,
+        "add_times": add_times,
+        "output_times": output_times,
+        "total_output_tokens": total_output_tokens,
+        "cuda_graph_request_count_map": cuda_graph_request_count_map,
     }
 
 
@@ -424,9 +319,9 @@ def main():
     # 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) 
+
+    level_str = os.getenv("LOG_LEVEL", "INFO").upper()
+    level = getattr(logging, level_str, logging.INFO)
     logging.basicConfig(level=level, force=True)
 
     configure_nvtx_profiling(True)
@@ -452,20 +347,14 @@ def main():
         termination_id=args.termination_id if args.termination_id is not None else tokenizer.eod,
         top_n_logprobs=args.top_n_logprobs,
         stop_words=args.stop_words,
-    ) 
+    )
 
     model = get_model()
 
-    mamba_inference_state_config = get_mamba_inference_state_config_from_model(model)
-
     # Requests, context, controller.
     requests = build_requests(args, tokenizer, sampling_params)
-    context = get_inference_context(
-        requests,
-        sampling_params,
-        mamba_inference_state_config=mamba_inference_state_config,
-    )
-    controller = get_inference_controller(model, context)
+    context = get_inference_context(model, requests=requests, sampling_params=sampling_params)
+    controller = TextGenerationController.from_model_and_args(model, args, context)
 
     # Validate all context_length's <= max_tokens.
     if args.disable_chunked_prefill:
@@ -473,21 +362,14 @@ def main():
         for request_idx, request in enumerate(requests):
             if len(request.prompt_tokens) > context.max_tokens:
                 invalid_prompt_length_map[request_idx] = len(request.prompt_tokens)
-        assert not invalid_prompt_length_map, (
-            "request idxs with prompts longer than context.max_tokens: "
-            ", ".join(f"{k}({v})" for k, v in invalid_prompt_length_map.items())
+        assert (
+            not invalid_prompt_length_map
+        ), "request idxs with prompts longer than context.max_tokens: " ", ".join(
+            f"{k}({v})" for k, v in invalid_prompt_length_map.items()
         )
 
     # Inference engine.
-    engine = DynamicInferenceEngine(
-        controller,
-        context,
-        enable_cuda_graph=args.cuda_graph_impl == "local",
-        random_seed=args.seed,
-        track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events,
-        enable_chunked_prefill=not args.disable_chunked_prefill,
-        inference_logging_step_interval=args.inference_logging_step_interval,
-    )
+    engine = DynamicInferenceEngine.from_model_and_args(model, args, controller, context)
 
     setup_prefix = build_dynamic_engine_setup_prefix(args, model, context, requests)
     print("~~~")
@@ -518,14 +400,13 @@ def main():
 
     # Validate all requests finished.
     for request in requests:
-        assert request.state == "finished", (
-            f"request.state == '{request.state}' != 'finished'."
-        )
+        assert request.state == "finished", f"request.state == '{request.state}' != 'finished'."
 
     peak_mem_stats = get_global_peak_memory_stats_bytes()
 
     # Print unique prompts + outputs.
     if torch.distributed.get_rank() == 0:
+
         def escape_str(s):
             return s.replace("\n", "\\n")
 
@@ -543,7 +424,9 @@ def escape_str(s):
             # ---- Prompt summary line ----
             prompt_len = len(requests[request_idxs[0]].prompt_tokens)
             escaped_prompt_text = escape_str(prompt_text)
-            print(f"{unique_idx+1}/{len(unique_prompt_map)} [n {len(request_idxs)}, l {prompt_len}] {escaped_prompt_text}")
+            print(
+                f"{unique_idx+1}/{len(unique_prompt_map)} [n {len(request_idxs)}, l {prompt_len}] {escaped_prompt_text}"
+            )
 
             # ---- Group all outputs for this prompt ----
             output_map = defaultdict(list)
@@ -557,16 +440,16 @@ def escape_str(s):
                     # Use hash of prompt + generated text in case engine was
                     # suspended and resumed, which misaligns boundary between
                     # prompt and generated tokens.
-                    o_hash = hashlib.sha256(
-                        (prompt_text + output_text).encode()
-                    ).hexdigest()[:6]
+                    o_hash = hashlib.sha256((prompt_text + output_text).encode()).hexdigest()[:6]
                     o_len = len(requests[output_request_idxs[0]].output_tokens)
                     escaped_output_text = escape_str(output_text)
                 else:
                     o_hash = "--"
                     o_len = 0
                     escaped_output_text = "--"
-                print(f"  >>>> [n {len(output_request_idxs)}, {o_len} tokens, hash {o_hash}] {escaped_output_text}")
+                print(
+                    f"  >>>> [n {len(output_request_idxs)}, {o_len} tokens, hash {o_hash}] {escaped_output_text}"
+                )
                 text_hashes.append(o_hash)
 
         # Write results to JSON. Primarily used for functional testing.
@@ -582,14 +465,16 @@ def escape_str(s):
                         "generated_text": req.output_text,
                         "generated_tokens": req.output_tokens,
                         "latency": req.time_end - req.time_start,
-                        "cuda_graph_request_count_map" : result["cuda_graph_request_count_map"],
-                        "step_count" : engine.step_count,
-                        "top_n_logprobs" : getattr(req, 'generated_top_n_logprobs', None),
-                        "prompt_top_n_logprobs" : getattr(req, 'prompt_top_n_logprobs', None),
+                        "cuda_graph_request_count_map": result["cuda_graph_request_count_map"],
+                        "step_count": engine.step_count,
+                        "top_n_logprobs": getattr(req, 'generated_top_n_logprobs', None),
+                        "prompt_top_n_logprobs": getattr(req, 'prompt_top_n_logprobs', None),
                     }
                     if req.sampling_params.return_log_probs:
                         result_dict["prompt_logprobs"] = getattr(req, 'prompt_log_probs', None)
-                        result_dict["generated_logprobs"] = getattr(req, 'generated_log_probs', None)
+                        result_dict["generated_logprobs"] = getattr(
+                            req, 'generated_log_probs', None
+                        )
                         result_dict["logprobs"] = getattr(req, 'logprobs', None)
                     json_results[req.request_id] = result_dict
 
@@ -621,7 +506,7 @@ def escape_str(s):
         d_count = len(d_times)
 
         p_mean = p_total / p_count
-        d_mean = d_total / d_count if d_count != 0 else 0.
+        d_mean = d_total / d_count if d_count != 0 else 0.0
 
         # Commented out for now as the step/add/output times are not calculated correctly.
         # print(
@@ -633,18 +518,13 @@ def escape_str(s):
         #     f"mean [ p {p_mean:.3f}s, d {d_mean:.3f}s ], "
         #     f"count [ p {p_count}, d {d_count} ]."
         # )
-        capture_str = (
-            f"{engine.capture_stats['time']:.2f} sec"
-            if engine.capture_stats else
-            "--"
-        )
+        capture_str = f"{engine.capture_stats['time']:.2f} sec" if engine.capture_stats else "--"
         print(
-            f"{setup_prefix} … "
-            f"throughput: {throughput:.3f} tok/s … ",
+            f"{setup_prefix} … " f"throughput: {throughput:.3f} tok/s … ",
             f"total time: {total_time:.3f}s … "
             f"mem {peak_alloc_gb:.1f}/{peak_resvd_gb:.1f} GB … "
             f"steps: {engine.step_count:d} … "
-            f"capture {capture_str}"
+            f"capture {capture_str}",
         )
         print("~~~")
 
diff --git a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
index f354b122a7e..97117b9dbe1 100644
--- a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
+++ b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
@@ -2,36 +2,31 @@
 
 import asyncio
 import json
+import logging
 import os
 import time
+import warnings
+from collections import defaultdict
+from typing import List
+
 import torch
 import torch.distributed as dist
-from collections import defaultdict
 from tqdm import tqdm
-from typing import List
-import warnings
-import logging
 
-from examples.inference.gpt.gpt_dynamic_inference import (
-    add_dynamic_inference_args,
-    get_inference_context,
-    get_inference_controller,
-    get_model,
-)
 from examples.inference.gpt.utils import (
-    Request, 
-    build_dynamic_engine_setup_prefix, 
+    Request,
+    add_dynamic_inference_args,
+    add_common_inference_args,
+    build_dynamic_engine_setup_prefix,
     build_requests,
-    add_common_inference_args
+    get_model
 )
-
 from megatron.core import parallel_state
 from megatron.core.inference.engines import DynamicInferenceEngine
 from megatron.core.inference.inference_client import InferenceClient
 from megatron.core.inference.inference_request import DynamicInferenceRequestRecord
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.utils import get_mamba_inference_state_config_from_model
-
 from megatron.training import get_args, get_tokenizer, initialize_megatron
 from megatron.training.arguments import parse_args
 
@@ -39,6 +34,7 @@
 
 logging.basicConfig(level=logging.INFO, force=True)
 
+
 async def main(
     engine: DynamicInferenceEngine,
     requests: List[Request],
@@ -51,14 +47,13 @@ async def main(
             "Sampling parameters are specified per request.",
             DeprecationWarning,
         )
-        
+
     # once you call engine.start_listening_to_data_parallel_coordinator,
     # the engine will start accepting requests from the data parallel coordinator.
     # and processing them in an asyncio coroutine.
-    
+
     await engine.start_listening_to_data_parallel_coordinator(
-        inference_coordinator_port=port,
-        launch_inference_coordinator=True,
+        inference_coordinator_port=port, launch_inference_coordinator=True
     )
 
     args = get_args()
@@ -68,14 +63,11 @@ async def main(
         # Since the client doesn't directly call engine.async_step here, we test
         # the suspend-resume system ~4 times.
         suspend_resume_interval = max(1, len(requests) // 4)
-        suspend_idxs = set(range(
-            suspend_resume_interval,
-            len(requests) + 1,
-            suspend_resume_interval,
-        ))
+        suspend_idxs = set(
+            range(suspend_resume_interval, len(requests) + 1, suspend_resume_interval)
+        )
         resume_idxs = set(
-            min(len(requests), i + suspend_resume_interval // 2)
-            for i in suspend_idxs
+            min(len(requests), i + suspend_resume_interval // 2) for i in suspend_idxs
         )
     else:
         suspend_idxs = set()
@@ -97,7 +89,10 @@ async def main(
             current_time = time.time_ns() / 10**9
             if args.incoming_requests_per_step is None:
                 # Only add requests that have arrived at the current time.
-                while num_requests_added < num_requests_total and requests[num_requests_added].time_arrival <= current_time:
+                while (
+                    num_requests_added < num_requests_total
+                    and requests[num_requests_added].time_arrival <= current_time
+                ):
                     request = requests[num_requests_added]
                     # These add-request calls will queue up the request on a zmq socket and return
                     # instantaneously. They will return an asyncio future which can be awaited for
@@ -113,10 +108,9 @@ async def main(
 
             else:
                 # Add deterministic number of requests (generally used for debugging).
-                for i in range(min(
-                    args.incoming_requests_per_step,
-                    num_requests_total - num_requests_added
-                )):
+                for i in range(
+                    min(args.incoming_requests_per_step, num_requests_total - num_requests_added)
+                ):
                     # Change sampling parameters to force different generation lengths.
                     request = requests[num_requests_added]
                     n = request.sampling_params.num_tokens_to_generate
@@ -134,7 +128,7 @@ async def main(
                 break
             # Relinquish control since there are no more requests to add at the moment. This allows the engine to run.
             await asyncio.sleep(0)
-        
+
         # While we wait for the requests to complete, the engine runs in the background.
         results: List[DynamicInferenceRequestRecord] = await asyncio.gather(*futures)
 
@@ -169,16 +163,19 @@ async def main(
                 req = record.merge()
                 unique_prompt_map[req.prompt].append(req)
             for idx, (prompt_text, reqs) in enumerate(unique_prompt_map.items()):
-                print(f"%d/%d. prompt '%s' ... [%d] output '%s'." % (
-                    idx,
-                    len(unique_prompt_map),
-                    prompt_text.replace("\n", "\\n"),
-                    len(reqs),
-                    reqs[0].generated_text.replace("\n", "\\n"),
-                ))
+                print(
+                    f"%d/%d. prompt '%s' ... [%d] output '%s'."
+                    % (
+                        idx,
+                        len(unique_prompt_map),
+                        prompt_text.replace("\n", "\\n"),
+                        len(reqs),
+                        reqs[0].generated_text.replace("\n", "\\n"),
+                    )
+                )
 
         # kill the engines and suspend the client
-        # Right now, we can only call stop when all requests are done. 
+        # Right now, we can only call stop when all requests are done.
         # Todo: Make this explicit in the Client class....
         await client.stop_engines()
         client.stop()
@@ -189,7 +186,7 @@ async def main(
 
 
 if __name__ == "__main__":
-    # enable inference mode in the very beginning as some fp-8 optimizations
+    # enable inference mode in the very beginning as some fp8 optimizations
     # check for it.
     with torch.inference_mode():
         initialize_megatron(
@@ -212,31 +209,13 @@ async def main(
             ),
         )
 
-        # Requests, context, conroller.
         model = get_model()
-        mamba_inference_state_config = get_mamba_inference_state_config_from_model(model)
+
         requests = (
             build_requests(args, tokenizer, sampling_params) if dist.get_rank() == 0 else None
         )
 
-        context = get_inference_context(
-            None,
-            None,
-            calculate_max_sequence_length_from_requests=False,
-            mamba_inference_state_config=mamba_inference_state_config,
-        )
-
-        controller = get_inference_controller(model, context)
-
-        # Inference engine.
-        engine = DynamicInferenceEngine(
-            controller,
-            context,
-            enable_cuda_graph=args.cuda_graph_impl == "local",
-            random_seed=args.seed,
-            enable_chunked_prefill=not args.disable_chunked_prefill,
-            inference_logging_step_interval=args.inference_logging_step_interval,
-        )
+        engine = DynamicInferenceEngine.from_model_and_args(model, args)
 
         if dist.get_rank() == 0:
             setup_prefix = build_dynamic_engine_setup_prefix(args, model, context, requests)
@@ -248,14 +227,8 @@ async def main(
         if os.environ.get("NSIGHT_PREFIX"):
             torch.cuda.cudart().cudaProfilerStart()
 
-        asyncio.run(
-            main(
-                engine,
-                requests,
-                args.inference_coordinator_port,
-            )
-        )
+        asyncio.run(main(engine, requests, args.inference_coordinator_port))
 
         # Stop Nsight profiler.
         if os.environ.get("NSIGHT_PREFIX"):
-            torch.cuda.cudart().cudaProfilerStop()
\ No newline at end of file
+            torch.cuda.cudart().cudaProfilerStop()
diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py
index fa269606708..52fd10791a8 100644
--- a/examples/inference/gpt/utils.py
+++ b/examples/inference/gpt/utils.py
@@ -1,21 +1,63 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import copy
-import json
 import itertools
+import json
 import random
 import time
-import torch
 from argparse import ArgumentParser, Namespace
-from tqdm import tqdm
+from functools import partial
 from typing import Any, List, Optional
 
-from megatron.core.inference.inference_request import DynamicInferenceRequest
+import torch
+from tqdm import tqdm
+
+from gpt_builders import gpt_builder
+from mamba_builders import mamba_builder
 from megatron.core.inference.contexts import DynamicInferenceContext
 from megatron.core.inference.contexts.dynamic_context import get_mem_size_str
+from megatron.core.inference.inference_request import DynamicInferenceRequest
+from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.transformer.module import MegatronModule
+from megatron.training import get_args
+from megatron.training import get_model as _get_model
+from megatron.training.checkpointing import load_checkpoint
+from model_provider import model_provider
 
-from megatron.core.inference.sampling_params import SamplingParams
+
+def get_model() -> MegatronModule:
+    """Initialize model and load checkpoint."""
+
+    args = get_args()
+
+    if args.model_provider == "gpt":
+        model_builder = gpt_builder
+    elif args.model_provider == "mamba":
+        model_builder = mamba_builder
+    else:
+        raise ValueError(f"Invalid model provider {args.model_provider}")
+
+    # Build model.
+    model = _get_model(partial(model_provider, model_builder), wrap_with_ddp=False)
+
+    # Load checkpoint.
+    assert args.load is not None
+    args.exit_on_missing_checkpoint = True
+    load_checkpoint(
+        ddp_model=model,
+        optimizer=None,
+        opt_param_scheduler=None,
+        strict=not args.inference_ckpt_non_strict,
+    )
+
+    # No virtual PP.
+    assert len(model) == 1, "Above condition should have caught this"
+    model = model[0]
+
+    # Eval mode.
+    model.eval()
+
+    return model
 
 
 def add_common_inference_args(parser: ArgumentParser) -> ArgumentParser:
@@ -68,7 +110,8 @@ def add_common_inference_args(parser: ArgumentParser) -> ArgumentParser:
     )
     group.add_argument(
         "--incoming-requests-per-step",
-        type=int, default=None,
+        type=int,
+        default=None,
         help="Add a deterministic number of requests per step. This arg is "
         "prioritized over `--incoming-requests-per-sec` below (which is non-"
         "deterministic). Note that the number of requests added per step is "
@@ -91,16 +134,10 @@ def add_common_inference_args(parser: ArgumentParser) -> ArgumentParser:
         "total number of requests. Set to -1 to add all requests together.",
     )
     group.add_argument(
-        "--model-provider",
-        choices=["mamba", "gpt"],
-        default="gpt",
-        help="Model provider",
+        "--model-provider", choices=["mamba", "gpt"], default="gpt", help="Model provider"
     )
     group.add_argument(
-        "--skip-prompt-log-probs",
-        action='store_true',
-        default=False,
-        help='Skip prompt log probs.',
+        "--skip-prompt-log-probs", action='store_true', default=False, help='Skip prompt log probs.'
     )
     group.add_argument(
         "--stop-words",
@@ -112,10 +149,7 @@ def add_common_inference_args(parser: ArgumentParser) -> ArgumentParser:
         'separated by space. Example: --stop-words "\\n\\n" "END" "###"',
     )
     group.add_argument(
-        "--output-path",
-        type=str,
-        default=None,
-        help="Path to save generations as JSON",
+        "--output-path", type=str, default=None, help="Path to save generations as JSON"
     )
     group.add_argument(
         "--output-every-n-results",
@@ -148,8 +182,7 @@ def add_common_inference_args(parser: ArgumentParser) -> ArgumentParser:
         "--no-record-throughput",
         action='store_false',
         dest="record_throughput",
-        help="Disable throughput recording in --output-file"
-        
+        help="Disable throughput recording in --output-file",
     )
 
     return parser
@@ -162,9 +195,10 @@ def get_default_sampling_params(termination_id: int = None):
         top_p=0.0,
         return_log_probs=False,
         num_tokens_to_generate=30,
-        termination_id = termination_id,
+        termination_id=termination_id,
     )
 
+
 def get_curr_time() -> float:
     """Get synchronized time across ranks."""
     curr_time = torch.cuda.LongTensor([time.time_ns()])
@@ -188,7 +222,13 @@ class Request:
         tokenizer (Any): Tokenizer for tokenizing the prompt.
     """
 
-    def __init__(self, prompt_text: str, time_offset: float, tokenizer: Any, sampling_params: SamplingParams = None):
+    def __init__(
+        self,
+        prompt_text: str,
+        time_offset: float,
+        tokenizer: Any,
+        sampling_params: SamplingParams = None,
+    ):
         self.prompt_text = prompt_text
         self.prompt_tokens = tokenizer.tokenize(prompt_text)
         self.output_text = None
@@ -198,7 +238,11 @@ def __init__(self, prompt_text: str, time_offset: float, tokenizer: Any, samplin
         self.time_start = None
         self.time_end = None
         self.state = "not-started"
-        self.sampling_params: SamplingParams = sampling_params if sampling_params is not None else get_default_sampling_params(tokenizer.eod)
+        self.sampling_params: SamplingParams = (
+            sampling_params
+            if sampling_params is not None
+            else get_default_sampling_params(tokenizer.eod)
+        )
         self.sampling_params = copy.deepcopy(self.sampling_params)
 
     def __str__(self) -> str:
@@ -225,10 +269,10 @@ def get_time_offsets(
 
     # if num_requests is not None:
     incoming_requests_duration = num_requests / incoming_requests_per_sec
-    incoming_requests_duration *= 2 # extra margin, to accomodate time sampling
+    incoming_requests_duration *= 2  # extra margin, to accomodate time sampling
 
     random.seed(seed)
-    
+
     import simpy  # Guard against this import in test case
 
     # Generate random time offsets.
@@ -241,14 +285,14 @@ def arrival(r):
     env = simpy.Environment()
     env.process(arrival(incoming_requests_per_sec))
     env.run(incoming_requests_duration)
-    
+
     # Ensure at least a single request.
     if len(time_offsets) == 0:
         time_offsets = [0.0]
 
     # Ensure first time is 0.
     time_offsets = [to - time_offsets[0] for to in time_offsets]
-    
+
     # Truncate to num_requests.
     assert len(time_offsets) >= num_requests
     time_offsets = time_offsets[:num_requests]
@@ -257,7 +301,7 @@ def arrival(r):
 
 
 def get_cli_requests(
-        args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None
+    args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None
 ) -> list[Request]:
 
     # Get time offsets.
@@ -269,7 +313,7 @@ def get_cli_requests(
     )
 
     # Init requests.
-    requests = [Request(p, t, tokenizer, sampling_params) for p,t in zip(args.prompts, t_offsets)]
+    requests = [Request(p, t, tokenizer, sampling_params) for p, t in zip(args.prompts, t_offsets)]
     return requests
 
 
@@ -289,18 +333,14 @@ def get_synthetic_requests(
     # Build prompts with expected lengths.
     assert (
         len(args.num_tokens_to_prompt) == 2
-        and
-        args.num_tokens_to_prompt[1] >= args.num_tokens_to_prompt[0]
+        and args.num_tokens_to_prompt[1] >= args.num_tokens_to_prompt[0]
     )
     max_prompt_length = args.num_tokens_to_prompt[1]
     max_prompt_text = "hi " * max_prompt_length
     max_prompt_tokens = tokenizer.tokenize(max_prompt_text)
-    prompt_lengths = [
-        random.randint(*args.num_tokens_to_prompt)
-        for _ in time_offsets
-    ]
-    prompt_tokens_list = [ max_prompt_tokens[:l] for l in prompt_lengths ]
-    prompt_texts = [ tokenizer.detokenize(tt) for tt in prompt_tokens_list ]
+    prompt_lengths = [random.randint(*args.num_tokens_to_prompt) for _ in time_offsets]
+    prompt_tokens_list = [max_prompt_tokens[:l] for l in prompt_lengths]
+    prompt_texts = [tokenizer.detokenize(tt) for tt in prompt_tokens_list]
 
     # Init requests.
     assert len(prompt_texts) == len(time_offsets)
@@ -340,16 +380,15 @@ def get_requests_from_file(
 
     # Get time offsets.
     time_offsets: list[float] = get_time_offsets(
-        args.seed,
-        args.incoming_requests_per_step,
-        args.incoming_requests_per_sec,
-        len(prompts),
+        args.seed, args.incoming_requests_per_step, args.incoming_requests_per_sec, len(prompts)
     )
 
     # Init requests.
     requests = [
         Request(p, t, tokenizer, sp)
-        for p, t, sp in tqdm(zip(prompts, time_offsets, sampling_params_list), "init requests", total=len(prompts))
+        for p, t, sp in tqdm(
+            zip(prompts, time_offsets, sampling_params_list), "init requests", total=len(prompts)
+        )
     ]
 
     return requests
@@ -411,19 +450,21 @@ def build_dynamic_engine_setup_prefix(
 
     # Prompt description
     prompt_src_str = (
-        "cli" if args.prompts else
-        "file" if args.prompt_file else
-        f"synth({', '.join(map(str, args.num_tokens_to_prompt))})"
+        "cli"
+        if args.prompts
+        else (
+            "file"
+            if args.prompt_file
+            else f"synth({', '.join(map(str, args.num_tokens_to_prompt))})"
+        )
     )
     request_str = (
-        f"requests: {prompt_src_str}, "
-        f"n {len(requests):d}, g {args.num_tokens_to_generate:d}, "
+        f"requests: {prompt_src_str}, " f"n {len(requests):d}, g {args.num_tokens_to_generate:d}, "
     )
     request_str += (
-        f"dur {args.incoming_requests_duration:.1e} "
-        f"r/sec {args.incoming_requests_per_sec:.1e}"
-        if args.incoming_requests_per_step is None else
-        f"r/step {args.incoming_requests_per_step}"
+        f"dur {args.incoming_requests_duration:.1e} " f"r/sec {args.incoming_requests_per_sec:.1e}"
+        if args.incoming_requests_per_step is None
+        else f"r/step {args.incoming_requests_per_step}"
     )
 
     # Buffer limits config
@@ -433,14 +474,7 @@ def build_dynamic_engine_setup_prefix(
         f"[r {context.max_active_requests}, t {context.max_tokens}]"
     )
 
-    parts = [
-        get_model_size_str(model),
-        "dynamic",
-        cg_str,
-        uvm_str,
-        request_str,
-        buffer_limits_str,
-    ]
+    parts = [get_model_size_str(model), "dynamic", cg_str, uvm_str, request_str, buffer_limits_str]
 
     return " | ".join(parts)
 
@@ -456,4 +490,4 @@ def get_global_peak_memory_stats_bytes() -> dict:
         t = torch.tensor([peak_alloc], device="cuda", dtype=torch.int64)
         torch.distributed.all_reduce(t, op=torch.distributed.ReduceOp.MAX)
         peak_alloc = int(t[0].item())
-    return {"mem-max-allocated-bytes": peak_alloc}
\ No newline at end of file
+    return {"mem-max-allocated-bytes": peak_alloc}
diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 10fb335addd..4b7e47b9d22 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -4,7 +4,7 @@
 import math
 import warnings
 from contextlib import nullcontext
-from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple
 
 import torch
 import torch.nn.functional as F
@@ -17,9 +17,6 @@
     InferenceBatchDimensions,
 )
 from megatron.core.inference.inference_request import DynamicInferenceRequest
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.unified_memory import (
     UnifiedMemoryUnsupportedError,
@@ -30,9 +27,14 @@
 from megatron.core.package_info import __version__ as mcore_version
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.ssm.mamba_hybrid_layer_allocation import get_layer_maps_from_layer_type_list
-from megatron.core.transformer import TransformerConfig
+from megatron.core.transformer import MLATransformerConfig, TransformerConfig
 from megatron.core.utils import divide as core_divide
-from megatron.core.utils import get_attr_wrapped_model, get_pg_size, internal_api
+from megatron.core.utils import (
+    get_attr_wrapped_model,
+    get_mamba_inference_state_config_from_model,
+    get_pg_size,
+    internal_api,
+)
 
 from .attention_context.mamba_metadata import MambaInferenceStateConfig, MambaMetadata
 from .attention_context.mha_metadata import GraphedMHAMetadata, NonGraphedMHAMetadata
@@ -205,10 +207,7 @@ class DynamicInferenceContext(BaseInferenceContext):
     any unassigned blocks equate to unused space.
 
     Args:
-        params_dtype (torch.dtype): Dtype used for KV cache.
-        num_layers (int): Number of layers on this pipeline parallel rank.
-        kv_channels (int): Hidden dimension per attention head.
-        num_attention_heads (int): Number of attention heads.
+        model_config (TransformerConfig): Model config.
         max_sequence_length (int): Max possible sequence length (prompt + output)
             that will occur.
         buffer_size_gb (float): Buffer size reserved on the GPU for the KV cache.
@@ -216,6 +215,8 @@ class DynamicInferenceContext(BaseInferenceContext):
             utilized, resulting in a total buffer size of `2 * buffer_size_gb`.
             Regardless of total buffer size, the KV cache is conceptually divided
             into 50% active requests and 50% paused requests.
+        mamba_inference_state_config (Optional[MambaInferenceStateConfig]): The Mamba
+            inference state config if the model is a hybrid model.
         max_requests (int): Max number of active requests to use for
             decode-only forward passes. This value is primarily limited by the
             combination of `buffer_size_gb` and `max_sequence_length`.
@@ -223,7 +224,6 @@ class DynamicInferenceContext(BaseInferenceContext):
             primarily limited by prefill activation memory usage. (Defaults to
             16384).
         block_size_tokens (int): Size of KV cache block size.
-        tensor_model_parallel_size (Optional[int]): Tensor model parallel size.
         num_cuda_graphs (Optional[int]): Maximum number of cuda graphs to capture,
             where the cuda graph batch sizes range from 1 to `max_active_requests`
             (as computed below). Due to rounding, the actual number of cuda graphs
@@ -231,8 +231,6 @@ class DynamicInferenceContext(BaseInferenceContext):
         materialize_only_last_token_logits (Optional[bool]): Whether to only
             materialize logits for the last token. This should be set to False
             if returning log probs.
-        mamba_inference_state_config (Optional[MambaInferenceStateConfig]): The Mamba
-            inference state config if the model is a hybrid model.
         use_cuda_graphs_for_non_decode_steps (bool): If True, use cuda graphs for non-decode
             engine steps.
         unified_memory_level (Optional[int]): Set unified memory usage within the
@@ -254,24 +252,16 @@ class DynamicInferenceContext(BaseInferenceContext):
     def __init__(
         self,
         *,
-        params_dtype: torch.dtype,
-        num_layers: int,
-        kv_channels: int,
-        num_attention_heads: int,
+        model_config: TransformerConfig,
         max_sequence_length: int,
         buffer_size_gb: float,
-        max_requests: int = None,
+        mamba_inference_state_config: Optional[MambaInferenceStateConfig] = None,
+        pg_collection: Optional[ProcessGroupCollection] = None,
+        max_requests: Optional[int] = None,
         max_tokens: int = DEFAULT_MAX_TOKENS,
         block_size_tokens: int = 256,
-        tensor_model_parallel_size: Optional[int] = None,
-        pipeline_model_parallel_size: Optional[int] = None,
-        pg_collection: Optional[ProcessGroupCollection] = None,
-        cache_mla_latent: bool = False,
-        kv_lora_rank: Optional[int] = None,
-        qk_pos_emb_head_dim: Optional[int] = None,
         num_cuda_graphs: Optional[int] = None,
         materialize_only_last_token_logits: Optional[bool] = True,
-        mamba_inference_state_config: Optional[MambaInferenceStateConfig] = None,
         use_cuda_graphs_for_non_decode_steps: bool = True,
         use_flashinfer_fused_rope: bool = False,
         unified_memory_level: Optional[int] = 0,
@@ -282,7 +272,9 @@ def __init__(
     ):
         super().__init__(materialize_only_last_token_logits=materialize_only_last_token_logits)
 
-        self.cache_mla_latent = cache_mla_latent
+        self.cache_mla_latent = (
+            isinstance(model_config, MLATransformerConfig) and model_config.cache_mla_latents
+        )
         if self.cache_mla_latent:
             assert (
                 block_size_tokens == 64
@@ -300,26 +292,29 @@ def __init__(
         self.metrics_writer = metrics_writer
 
         # Per partition num heads and hidden size.
-        projection_size = kv_channels * num_attention_heads
-        if tensor_model_parallel_size is None:
+        num_attention_heads = model_config.num_query_groups or model_config.num_attention_heads
+        projection_size = model_config.kv_channels * num_attention_heads
+        if model_config.tensor_model_parallel_size is None:
+            assert pg_collection is not None
             tp_size = (
                 get_pg_size(pg_collection.tp)
                 if pg_collection is not None
                 else parallel_state.get_tensor_model_parallel_world_size()
             )
         else:
-            tp_size = tensor_model_parallel_size
+            tp_size = model_config.tensor_model_parallel_size
         self.hidden_size_per_attention_head = core_divide(projection_size, num_attention_heads)
         self.num_attention_heads_per_partition = core_divide(num_attention_heads, tp_size)
 
-        if pipeline_model_parallel_size is None:
+        if model_config.pipeline_model_parallel_size is None:
+            assert pg_collection is not None
             pp_size = (
                 get_pg_size(pg_collection.pp)
                 if pg_collection is not None
                 else parallel_state.get_pipeline_model_parallel_world_size()
             )
         else:
-            pp_size = pipeline_model_parallel_size
+            pp_size = model_config.pipeline_model_parallel_size
 
         # Cache the PP group we should use for PP collectives inside the context.
         # If the model provides a pg_collection with a pp group, prefer it.
@@ -358,7 +353,7 @@ def __init__(
             self.layer_map = attention_layer_map | mamba_layer_map
         else:
             # The layer map is the identity function for pure Transformer models.
-            self.num_attention_layers = num_layers
+            self.num_attention_layers = model_config.num_layers // pp_size
             self.num_mamba_layers = 0
             (self.mamba_conv_states_shape, self.mamba_ssm_states_shape) = (None, None)
             self.layer_map = {i: i for i in range(self.num_attention_layers)}
@@ -369,7 +364,7 @@ def __init__(
             )
 
         # Block size tokens, bytes.
-        dtype_size_bytes = params_dtype.itemsize
+        dtype_size_bytes = model_config.params_dtype.itemsize
         self.block_size_tokens = block_size_tokens
         if self.cache_mla_latent:
             #   one vector  c_t  (rank)  +  optional RoPE phase slice
@@ -466,7 +461,7 @@ def __init__(
         self.request_metadata_types = request_metadata_types
 
         # Initialize context state.
-        self.params_dtype = params_dtype
+        self.params_dtype = model_config.params_dtype
         self.max_sequence_length = max_sequence_length
 
         # Request and token counts.
@@ -735,25 +730,16 @@ def round_up_tokens(cls, value, tp_size=None):
         return token_rounder * int(math.ceil(int(value) / token_rounder))
 
     @classmethod
-    def from_config(
-        cls,
-        inference_config: InferenceWrapperConfig,
-        model,
-        max_batch_size: int,
-        buffer_size_gb: float = 40,
-        num_cuda_graphs: int = None,
-        mamba_inference_state_config: Optional[MambaInferenceStateConfig] = None,
-        unified_memory_level: int = 0,
-    ):
+    def from_model_and_args(cls, model, args, overrides: Optional[Dict[str, Any]] = None):
         """
-        Instantiate a `DynamicInferenceContext` from a `TransformerConfig` and an `InferenceWrapperConfig`.
+        Instantiate a `DynamicInferenceContext` from a model and command-line args.
         """
-        # TODO: Add other necessary configs from inference_config
+        config = model.config
 
         # Max sequence length.
         position_embedding_type = get_attr_wrapped_model(model, "position_embedding_type")
         model_max_seq_len = get_attr_wrapped_model(model, "max_sequence_length")
-        inf_max_seq_len = inference_config.inference_max_seq_length
+        inf_max_seq_len = args.inference_max_seq_length
 
         if position_embedding_type == "learned_absolute":
             # When using absolute position embeddings, it is critical that the
@@ -767,28 +753,38 @@ def from_config(
                 max_sequence_length = model_max_seq_len
             assert max_batch_size <= model_max_seq_len
         else:
-            max_sequence_length = (
-                inference_config.inference_max_seq_length or model_config.max_sequence_length
+            max_sequence_length = inf_max_seq_len
+        if args.inference_dynamic_batching_max_requests is not None:
+            max_sequence_length = max(
+                max_sequence_length, args.inference_dynamic_batching_max_requests
             )
-        max_sequence_length = max(max_sequence_length, max_batch_size)
-
-        # Context.
-        model_config = model.config
-        return cls(
-            params_dtype=inference_config.params_dtype,
-            num_layers=model_config.num_layers // model_config.pipeline_model_parallel_size,
-            kv_channels=model_config.kv_channels,
-            num_attention_heads=model_config.num_query_groups,
-            tensor_model_parallel_size=model_config.tensor_model_parallel_size,
-            pipeline_model_parallel_size=model_config.pipeline_model_parallel_size,
-            max_sequence_length=max_sequence_length,
-            buffer_size_gb=buffer_size_gb,
-            materialize_only_last_token_logits=False,
-            num_cuda_graphs=num_cuda_graphs,
-            use_flashinfer_fused_rope=None,
-            mamba_inference_state_config=mamba_inference_state_config,
-            unified_memory_level=unified_memory_level,
-        )
+
+        mamba_inference_state_config = get_mamba_inference_state_config_from_model(model)
+
+        kwargs = {
+            "model_config": config,
+            "max_sequence_length": max_sequence_length,
+            "mamba_inference_state_config": mamba_inference_state_config,
+            "num_cuda_graphs": (
+                args.inference_dynamic_batching_num_cuda_graphs
+                if args.cuda_graph_impl == "local"
+                else None
+            ),
+            "block_size_tokens": args.inference_dynamic_batching_block_size,
+            "buffer_size_gb": args.inference_dynamic_batching_buffer_size_gb,
+            "max_requests": args.inference_dynamic_batching_max_requests,
+            "max_tokens": args.inference_dynamic_batching_max_tokens,
+            "materialize_only_last_token_logits": not args.return_log_probs,
+            "use_flashinfer_fused_rope": args.use_flashinfer_fused_rope,
+            "unified_memory_level": args.inference_dynamic_batching_unified_memory_level,
+            "cuda_graph_max_tokens": args.inference_dynamic_batching_cuda_graph_max_tokens,
+            "cuda_graph_mixed_prefill_count": args.inference_dynamic_batching_cuda_graph_mixed_prefill_count,
+        }
+
+        if overrides is not None:
+            kwargs.update(overrides)
+
+        return cls(**kwargs)
 
     @classmethod
     def round_up_requests(cls, value, tp_size=None):
diff --git a/megatron/core/inference/contexts/static_context.py b/megatron/core/inference/contexts/static_context.py
index 8c83d2f09b3..ba41b0c2401 100644
--- a/megatron/core/inference/contexts/static_context.py
+++ b/megatron/core/inference/contexts/static_context.py
@@ -1,8 +1,6 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
+from megatron.core.transformer.transformer_config import TransformerConfig
 
 from .base_context import BaseInferenceContext
 
@@ -28,7 +26,7 @@ def __init__(
         self.decode_mode = False
 
     @classmethod
-    def from_config(cls, config: InferenceWrapperConfig) -> "StaticInferenceContext":
+    def from_config(cls, config: TransformerConfig) -> "StaticInferenceContext":
         """Initialize context from a config."""
         max_batch_size = config.inference_max_requests
         max_sequence_length = config.inference_max_seq_length
diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index c7698b8a4bb..d1885bf0d6c 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -226,6 +226,29 @@ def __init__(
         # Create cuda graphs.
         self.create_cuda_graphs()
 
+    @classmethod
+    def from_model_and_args(
+        cls,
+        model,
+        args,
+        context: Optional[DynamicInferenceContext] = None,
+        controller: Optional[TextGenerationController] = None,
+    ):
+        if context is None:
+            context = DynamicInferenceContext.from_model_and_args(model, args)
+        if controller is None:
+            controller = TextGenerationController.from_model_and_args(model, args)
+
+        return cls(
+            context,
+            controller,
+            enable_cuda_graph=args.cuda_graph_impl == "local",
+            random_seed=args.seed,
+            track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events,
+            enable_chunked_prefill=not args.disable_chunked_prefill,
+            inference_logging_step_interval=args.inference_logging_step_interval,
+        )
+
     def reset(self) -> None:
         """Reset by removing all requests and reset all state."""
 
@@ -281,8 +304,6 @@ def create_cuda_graphs(self, reset_context: bool = True):
         context = self.context
         controller = self.controller
 
-        config = controller.inference_wrapped_model.inference_wrapper_config
-
         time_start = time.time()
         mem_stats_start = torch.cuda.memory_stats()
 
diff --git a/megatron/core/inference/engines/static_engine.py b/megatron/core/inference/engines/static_engine.py
index d4c61965d2b..661962b37a2 100644
--- a/megatron/core/inference/engines/static_engine.py
+++ b/megatron/core/inference/engines/static_engine.py
@@ -42,8 +42,6 @@ class StaticInferenceEngine(AbstractEngine):
             controller that will be used to define how to preprocess prompts, generate
             outputs and detokenizer the output tokens.
         max_batch_size (int, optional): The maximum number of requests to process at once.
-            Will be set from the InferenceWrapperConfig in `text_generation_controller` by
-            default.
         random_seed (int, optional): Use a random seed if you want deterministic
             results. Defaults to None.
     """
@@ -69,13 +67,12 @@ def __init__(
                 DeprecationWarning,
             )
 
-        inference_wrapper_config = (
-            text_generation_controller.inference_wrapped_model.inference_wrapper_config
-        )
         self.controller = text_generation_controller
+        self.inference_wrapped_model = self.controller.inference_wrapped_model
+        self.config = self.inference_wrapped_model.config
         self.random_seed = random_seed or 1234
 
-        inference_max_batch_size = inference_wrapper_config.inference_max_requests
+        inference_max_batch_size = self.config.inference_max_requests
         if max_batch_size is None:
             max_batch_size = inference_max_batch_size
         elif max_batch_size > inference_max_batch_size:
@@ -91,10 +88,10 @@ def __init__(
         self.scheduler = Scheduler(max_batch_size=max_batch_size)
 
         # Store original context in case we need to fall back to legacy static engine
-        original_context = text_generation_controller.inference_wrapped_model.inference_context
+        original_context = self.inference_wrapped_model.inference_context
 
         mamba_inference_state_config = get_mamba_inference_state_config_from_model(
-            text_generation_controller.inference_wrapped_model.model
+            self.inference_wrapped_model.model
         )
 
         try:
diff --git a/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py b/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py
index 6a17de685bf..ae8e8742fad 100644
--- a/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py
+++ b/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py
@@ -1,8 +1,6 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 
 import abc
-import math
-import warnings
 from typing import Any, Dict, Iterable, Optional, Union
 
 import torch
@@ -15,27 +13,22 @@
     send_to_next_pipeline_rank,
 )
 from megatron.core.inference.contexts import BaseInferenceContext
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.models.gpt.gpt_model import GPTModel
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.utils import get_attr_wrapped_model, get_model_config
 
 
-# pylint: disable=line-too-long
 class AbstractModelInferenceWrapper(abc.ABC):
     """Abstract inference wrapper
 
     Extend this to create a version for your model.
 
-    The wrapper prepares the model for inference, provides the required input data and runs the forward pass.
+    The wrapper prepares the model for inference, provides the required input data and
+    runs the forward pass.
 
     Args:
         model (Union[GPTModel, LegacyGPTModel]): The actual GPT model (MCore
             or MLM).
-        inference_wrapper_config (InferenceWrapperConfig): Has info like
-            hidden size, vocab size etc.
         inference_context (BaseInferenceContext): Context for managing KV
             cache and other inference params.
         pg_collection (ProcessGroupCollection): Process groups for model communication.
@@ -44,30 +37,18 @@ class AbstractModelInferenceWrapper(abc.ABC):
     def __init__(
         self,
         model: Union['LegacyGPTModel', GPTModel],  # type: ignore[name-defined]
-        inference_wrapper_config: InferenceWrapperConfig,
-        inference_context: Optional[BaseInferenceContext] = None,
+        inference_context: BaseInferenceContext,
         pg_collection: Optional[ProcessGroupCollection] = None,
     ):
         assert not isinstance(
             model, Iterable
         ), 'interleaving schedule is not supported for inference'
         self.model = model
-        self.inference_wrapper_config = inference_wrapper_config
+        self.config = get_model_config(self.model)
         self.pipeline_communication_dtype = (
-            torch.float
-            if self.inference_wrapper_config.fp32_residual_connection
-            else self.inference_wrapper_config.params_dtype
+            torch.float if self.config.fp32_residual_connection else self.config.params_dtype
         )
-        model_config = get_model_config(self.model)
-        self.sequence_parallel = model_config.sequence_parallel
-
-        if inference_context is None:
-            warnings.warn(
-                "`inference_context` must be passed in as an argument starting in `megatron-core` 0.13."
-            )
-            from megatron.core.inference.contexts import StaticInferenceContext
-
-            inference_context = StaticInferenceContext.from_config(inference_wrapper_config)
+        self.sequence_parallel = self.config.sequence_parallel
 
         self.inference_context = inference_context
 
@@ -78,40 +59,18 @@ def __init__(
         self.pp_group = pg_collection.pp
         self.tp_size = torch.distributed.get_world_size(self.tp_group)
 
-        if self.inference_wrapper_config.fp8 is not None:
+        if self.config.fp8 is not None:
             self.model = prepare_model_for_fp8_inference(self.model)
 
-    @property
-    def inference_params(self):
-        """Getter for deprecated `inference_params`."""
-        warnings.warn(
-            "`inference_params` renamed to `inference_context`, and will be removed in `megatron-core` 0.13."
-        )
-        return self.inference_context
-
-    @inference_params.setter
-    def inference_params(self, value):
-        """Setter for deprecated `inference_params`."""
-        warnings.warn(
-            "`inference_params` renamed to `inference_context`, and will be removed in `megatron-core` 0.13."
-        )
-        self.inference_context = value
+        # TODO(ksanthanam): Add support for fp4
 
-    def prep_model_for_inference(self, prompts_tokens: Optional[torch.Tensor] = None):
+    def prep_model_for_inference(self):
         """A utility function for preparing model for inference
 
         The function gets called once before the auto regressive inference loop.
         It puts the model in eval mode.
 
-        Args:
-            prompts_tokens (torch.Tensor, optional): Deprecated, will be removed in `megatron-core` 0.13
         """
-        if prompts_tokens is not None:
-            warnings.warn(
-                "Passing `prompts_tokens` is deprecated and this argument will be ignored."
-                "This parameter will be removed in `megatron-core` 0.13."
-            )
-
         self.model.eval()
 
         # For TP only model both is_pp_first_stage and _is_pp_last_stage returns True
@@ -121,23 +80,13 @@ def prep_model_for_inference(self, prompts_tokens: Optional[torch.Tensor] = None
 
         self.inference_context.reset()
 
-    @abc.abstractmethod
-    def prep_inference_input(self, prompt_tokens) -> Dict[str, Any]:
-        """Prepares the inference input data.
-
-        Args:
-            prompts_tokens (torch.Tensor): A tensor of shape [batch_size, max_seq_len]
-
-        Returns:
-            A dict with all the inference input needed for the batch.
-        """
-        raise NotImplementedError()
-
     @abc.abstractmethod
     def get_batch_for_context_window(self, *args, **kwargs) -> Dict[str, Any]:
         """Returns the input data for inference
 
-        This function gets called iteratively in the inference loop . It can be used to extract relevant input from the prompt tokens, attention mask etc. required for each step in inference.
+        This function gets called iteratively in the inference loop.
+        It can be used to extract relevant input from the prompt tokens, attention mask etc.
+        required for each step in inference.
 
         """
         raise NotImplementedError()
@@ -183,15 +132,16 @@ def _get_batch_size_and_seq_len(
         self, tokens: torch.Tensor, recv_buffer_seq_len: Optional[int] = None
     ):
         """
-        Returns the batch size and sequence length based on the tokens tensor and recv_buffer_seq_len.
+        Returns the batch size and sequence length based on the tokens tensor and
+        recv_buffer_seq_len.
 
         Args:
             tokens (torch.Tensor): The input tensor of shape (batch_size, seq_len).
             recv_buffer_seq_len (int, optional): An optional recv buffer sequence length.
 
         Returns:
-            tuple: A tuple (batch_size, seq_len), where batch_size is the first dimension of tokens
-                   and seq_len is either the second dimension or recv_buffer_seq_len.
+            tuple: A tuple (batch_size, seq_len), where batch_size is the first dimension of
+                tokens and seq_len is either the second dimension or recv_buffer_seq_len.
         """
         batch_size = tokens.shape[0]
         seq_len = recv_buffer_seq_len if recv_buffer_seq_len is not None else tokens.shape[1]
@@ -204,7 +154,7 @@ def _allocate_recv_buffer(self, batch_size, seq_len):
             # sequence parallelism. Static batching does not support sequence parallelism
             # except for the MoE layers which is handled separately.
             seq_len = seq_len // self.tp_size
-        recv_size = (seq_len, batch_size, self.inference_wrapper_config.hidden_size)
+        recv_size = (seq_len, batch_size, self.config.hidden_size)
         return torch.empty(
             recv_size, dtype=self.pipeline_communication_dtype, device=torch.cuda.current_device()
         )
@@ -214,10 +164,12 @@ def forward_pass_without_pipeline_parallel(
     ) -> torch.Tensor:
         """Utility to carry out simple forward pass for TP or no model parallel models
 
-        Runs a very simple forward pass for model. Used  in the case of models without any parallelism or only tensor parallelism.
+        Runs a very simple forward pass for model. Used  in the case of models without any
+        parallelism or only tensor parallelism.
 
         Args:
-            inference_input (Dict[str, Any]): A dict containg the inputs for the gpt model [tokens, position ids, attention mask]
+            inference_input (Dict[str, Any]): A dict containg the inputs for the gpt model
+                [tokens, position ids, attention mask]
 
         Returns:
             torch.Tensor: The output logits of shape [batch_size, seq_len, padded_vocab_size]
@@ -228,16 +180,18 @@ def forward_pass_without_pipeline_parallel(
 
         return logits
 
-    def forward_pass_with_pipeline_parallel_small_input_batch(
+    def forward_pass_with_pipeline_parallel(
         self, inference_input: Dict[str, Any], recv_buffer_seq_len: Optional[int] = None
     ) -> torch.Tensor:
-        """Utility to carry out forward pass for PP models with very small inputs
+        """Utility to carry out forward pass for PP models
 
-        If a model is pipeline parallel, yet, the input global batch is very small, we compute a foward pass on the entire global batch, rather than splitting it up into micro batches and doing something more complex as in the forward_pass_with_pipeline_parallel_large_input_batch method
+        TODO: Add support for asynchronous microbatches
 
         Args:
-            inference_input (Dict[str, Any]): A dict containing the inputs for the gpt model [tokens, position ids, attention mask]
-            recv_buffer_seq_len (int): An optional sequence length for the pipeline parallel recv buffer.
+            inference_input (Dict[str, Any]): A dict containing the inputs for the gpt model
+                [tokens, position ids, attention mask]
+            recv_buffer_seq_len (int): An optional sequence length for the pipeline parallel
+                recv buffer.
 
         Returns:
             torch.Tensor: The output logits of shape [batch_size, seq_len, padded_vocab_size]
@@ -268,98 +222,8 @@ def forward_pass_with_pipeline_parallel_small_input_batch(
             logits = output_tensor
 
             # Explicitly cast logits to expected dtype
-            logits = logits.to(self.inference_wrapper_config.params_dtype)
-
-        return logits
-
-    def forward_pass_with_pipeline_parallel_large_input_batch(
-        self, inference_input: Dict[str, Any], recv_buffer_seq_len=None
-    ) -> torch.Tensor:
-        """Utility to carry out forward pass PP models.
-
-        Runs the forward pass for models which are pipeline parallel.
-        This is more complex than forward_pass_with_pipeline_parallel_small_input_batch because
-        this splits the global batch into small micro batches and runs them through the model.
-
-        Args:
-            inference_input (Dict[str, Any]): A dict containg the inputs for the gpt model [tokens, position ids, attention mask]
-            recv_buffer_seq_len (int): An optional sequence length for the pipeline parallel recv buffer.
-
-        Returns:
-            torch.Tensor: The output logits of shape [batch_size, seq_len, padded_vocab_size]
-        """
-        tokens = inference_input["tokens"]
-        position_ids = inference_input["position_ids"]
-        attention_mask = inference_input["attention_mask"]
-        materialize_only_last_token_logits = (
-            self.inference_context.materialize_only_last_token_logits
-        )
-
-        micro_batch_size = max(
-            1,
-            self.inference_wrapper_config.inference_batch_times_seqlen_threshold // tokens.size(1),
-        )
-        batch_size, seq_len = self._get_batch_size_and_seq_len(tokens, recv_buffer_seq_len)
-        # Round up to account for the last partial micro batch if present
-        num_micro_batches = math.ceil(batch_size / micro_batch_size)
-
-        logits = None
-        # Preallocate memory for output logits.
-        if is_pipeline_last_stage(self.pp_group):
-            logits_seq_len = 1 if materialize_only_last_token_logits else seq_len
-            logits = torch.empty(
-                (batch_size, logits_seq_len, self.inference_wrapper_config.padded_vocab_size),
-                dtype=self.pipeline_communication_dtype,
-                device=torch.cuda.current_device(),
-            )
-
-        recv_buffer = None
-        if not is_pipeline_first_stage(self.pp_group):
-            recv_buffer = self._allocate_recv_buffer(micro_batch_size, seq_len)
-        for micro_batch_index in range(num_micro_batches):
-            start = micro_batch_index * micro_batch_size
-            end = min(start + micro_batch_size, batch_size)
-            tokens2use = tokens[start:end, ...]
-            position_ids2use = position_ids[start:end, ...]
-            current_micro_batch_size = end - start
-
-            # Need to change recv buffer shape for the last partial microbatch (if exists)
-            if current_micro_batch_size != micro_batch_size:
-                recv_buffer = self._allocate_recv_buffer(current_micro_batch_size, seq_len)
-
-            if not is_pipeline_first_stage(self.pp_group):
-                recv_from_prev_pipeline_rank_(recv_buffer, self.pp_group)
-
-            self.model.set_input_tensor(recv_buffer)
-
-            output_tensor = self._forward(
-                {
-                    "tokens": tokens2use,
-                    "position_ids": position_ids2use,
-                    "attention_mask": attention_mask,
-                    "inference_context": self.inference_context,
-                }
-            )
-
-            if not is_pipeline_last_stage(self.pp_group):
-                send_to_next_pipeline_rank(output_tensor, self.pp_group)
-
-            self.inference_context.batch_size_offset += current_micro_batch_size
-
-            if is_pipeline_last_stage(self.pp_group):
-                assert logits is not None
-                logits[start:end, ...] = output_tensor
-
-        # Explicitly cast logits to expected dtype
-        if is_pipeline_last_stage(self.pp_group):
-            assert logits is not None
-            logits = logits.to(self.inference_wrapper_config.params_dtype)
-
-        # Once done with all micro batches, we reset batch size offset and seq len offset
-        self.inference_context.increment_sequence_len_offset(seq_len)
-        self.inference_context.reset_batch_size_offset()
+            logits = logits.to(self.config.params_dtype)
 
-        # NOTE: Only returns the logits on the last pipeline stage
         return logits
 
     @torch.inference_mode()
@@ -368,14 +232,18 @@ def run_one_forward_step(
     ) -> torch.Tensor:
         """The forward pass of the model for inference
 
-        Appropriate utility is called for the forward pass depending on the type of model parallelism used
+        Appropriate utility is called for the forward pass depending on the type of model
+        parallelism used
 
         Args:
-            inference_input (Dict[str, Any]): A dict containing the inputs for the gpt model [tokens, position ids, attention mask]
-            recv_buffer_seq_len (int): An optional sequence length for the pipeline parallel recv buffer.
+            inference_input (Dict[str, Any]): A dict containing the inputs for the gpt model
+                [tokens, position ids, attention mask]
+            recv_buffer_seq_len (int): An optional sequence length for the pipeline parallel
+                recv buffer.
 
         Returns:
-            torch.Tensor: The output logits of shape [batch_size, seq_len, padded_vocab_size]. The logits are returned only in the last pipeline stage for PP models.
+            torch.Tensor: The output logits of shape [batch_size, seq_len, padded_vocab_size].
+            The logits are returned only in the last pipeline stage for PP models.
         """
         # Check if we are in a PP model
         if not (is_pipeline_first_stage(self.pp_group) and is_pipeline_last_stage(self.pp_group)):
@@ -383,19 +251,6 @@ def run_one_forward_step(
             current_batch_size, seq_len = self._get_batch_size_and_seq_len(
                 tokens, recv_buffer_seq_len
             )
-            # If input batch is large, we need to split into micro batches and run the forward pass
-            if (
-                current_batch_size * seq_len
-                > self.inference_wrapper_config.inference_batch_times_seqlen_threshold
-                and self.inference_wrapper_config.inference_batch_times_seqlen_threshold != -1
-            ):
-                return self.forward_pass_with_pipeline_parallel_large_input_batch(
-                    inference_input, recv_buffer_seq_len
-                )
-            else:
-                # If input batch is very small we can do a simple forward pass on the entire global batch
-                return self.forward_pass_with_pipeline_parallel_small_input_batch(
-                    inference_input, recv_buffer_seq_len
-                )
+            return self.forward_pass_with_pipeline_parallel(inference_input, recv_buffer_seq_len)
         else:
             return self.forward_pass_without_pipeline_parallel(inference_input)
diff --git a/megatron/core/inference/model_inference_wrappers/gpt/gpt_inference_wrapper.py b/megatron/core/inference/model_inference_wrappers/gpt/gpt_inference_wrapper.py
index ba89fbc2f6c..031eecfd27f 100644
--- a/megatron/core/inference/model_inference_wrappers/gpt/gpt_inference_wrapper.py
+++ b/megatron/core/inference/model_inference_wrappers/gpt/gpt_inference_wrapper.py
@@ -7,9 +7,6 @@
 from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
     AbstractModelInferenceWrapper,
 )
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.inference.utils import get_attention_mask
 from megatron.core.models.gpt import GPTModel
 from megatron.core.process_groups_config import ProcessGroupCollection
@@ -25,8 +22,6 @@ class GPTInferenceWrapper(AbstractModelInferenceWrapper):
 
     Args:
         model (GPTModel): The GPT model (MCore or legacy)
-        inference_wrapper_config (InferenceWrapperConfig): Has info like hidden size, vocab
-            size, etc.
         inference_context (BaseInferenceContext): Manages KV cache, and tracks
             sequence/token/batch offsets.
         pg_collection (ProcessGroupCollection): Process groups for model communication.
@@ -36,11 +31,10 @@ class GPTInferenceWrapper(AbstractModelInferenceWrapper):
     def __init__(
         self,
         model: GPTModel,
-        inference_wrapper_config: InferenceWrapperConfig,
         inference_context: Optional[BaseInferenceContext] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
     ):
-        super().__init__(model, inference_wrapper_config, inference_context, pg_collection)
+        super().__init__(model, inference_context, pg_collection)
 
     def prep_inference_input(self, prompts_tokens: torch.Tensor) -> Dict[str, Any]:
         """Prepares the inference input data.
diff --git a/megatron/core/inference/model_inference_wrappers/inference_wrapper_config.py b/megatron/core/inference/model_inference_wrappers/inference_wrapper_config.py
deleted file mode 100644
index 5d89085add2..00000000000
--- a/megatron/core/inference/model_inference_wrappers/inference_wrapper_config.py
+++ /dev/null
@@ -1,66 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
-from dataclasses import dataclass
-from typing import Optional
-
-import torch
-
-
-@dataclass
-class InferenceWrapperConfig:
-    """Config for the model inference wrapper
-
-    NOTE : All the arguments here are obtained from arguments.py file
-    """
-
-    hidden_size: int
-    """Receive happens between the layers during PP with size [seq_len, batch_size, hidden_size]"""
-
-    params_dtype: torch.dtype
-    """Can be torch.float or torch.half if --fp16 is used, or torch.bfloat16 if --bf16 is used"""
-
-    inference_batch_times_seqlen_threshold: int
-    """if (batch-size * sequence-length) is smaller than this threshold then we will not pipeline 
-    the batch."""
-
-    padded_vocab_size: int
-    """The final padded vocab size (Padded to make it divisible by 
-    --make-vocab-size-divisible-by value)"""
-
-    inference_max_requests: int = 8
-    """ Maximum number of requests for inference (prefill & decode). Necessary for CUDA graphs. """
-
-    inference_max_seq_length: int = 2560
-    """ Maximum sequence length for inference (prefill & decode). Necessary for CUDA graphs. """
-
-    fp32_residual_connection: bool = False
-    """Move residual connections to fp32. Obtained from arguments.py"""
-
-    nccl_all_reduce_for_prefill: bool = False
-    """When using symmetric all reduce kernels we keep the default all reduces for nccl. 
-    This can be more effecient for large prefill sizes"""
-
-    fp8: Optional[str] = None
-    """If set, enables the use of FP8 precision through Transformer Engine. There are 2 predefined
-    choices (1) 'e4m3' uniformly uses e4m3 for all FP8 tensors, (2) 'hybrid' uses e4m3 for all FP8
-    activation and weight tensors and e5m2 for all FP8 output activation gradient tensors."""
-
-    moe_pad_experts_for_cuda_graph_inference: bool = False
-    """Some MoE routers have a D2H sync that will break cuda graphs.  If this flag is set the router
-    will switch to dropping and padding during decode time which does not have a D2H sync. The 
-    capacity factor is set to the max that an expert could see during inference so no tokens are
-    actually dropped. """
-
-    def add_attributes(self, attribute_value_pair: dict):
-        """Utility to add more attributes to inference params
-
-        Use this method to pass in a custom dictionary to add more configs to the instance created.
-        Use as follows:
-        c = InferenceWrapperConfig
-        c.add_attributes({'precision':'fp32'})
-
-        Args:
-            attribute_value_pair (dict): A dictionary containing attributes as the key names and
-            corresponding values.
-        """
-        for key, value in attribute_value_pair.items():
-            setattr(self, key, value)
diff --git a/megatron/core/inference/model_inference_wrappers/t5/t5_inference_wrapper.py b/megatron/core/inference/model_inference_wrappers/t5/t5_inference_wrapper.py
index 2ae1e2ade6f..c773ab507a3 100644
--- a/megatron/core/inference/model_inference_wrappers/t5/t5_inference_wrapper.py
+++ b/megatron/core/inference/model_inference_wrappers/t5/t5_inference_wrapper.py
@@ -11,9 +11,6 @@
 from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
     AbstractModelInferenceWrapper,
 )
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.models.T5 import T5Model
 from megatron.core.utils import get_attr_wrapped_model
 
@@ -27,7 +24,6 @@ class T5InferenceWrapper(AbstractModelInferenceWrapper):
 
     Args:
         model (T5Model): The T5 model (MCore or legacy)
-        inference_wrapper_config (InferenceWrapperConfig): The command line arguments that were passed
         inference_context (BaseInferenceContext): Manages KV cache, and tracks
             sequence/token/batch offsets.
         use_local (bool): Whether  the T5 model's transformer impl
@@ -37,11 +33,10 @@ class T5InferenceWrapper(AbstractModelInferenceWrapper):
     def __init__(
         self,
         model: T5Model,
-        inference_wrapper_config: InferenceWrapperConfig,
         inference_context: Optional[BaseInferenceContext] = None,
         use_local: bool = False,
     ):
-        super().__init__(model, inference_wrapper_config, inference_context)
+        super().__init__(model, inference_context)
         self.use_local = use_local
 
     def prep_inference_input(
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index f757d4b539d..74bad4f571d 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -13,23 +13,34 @@
 from torch import Tensor
 from torch.distributed import ProcessGroup
 
+from megatron.core import parallel_state
 from megatron.core.inference.async_stream import AsyncStream
 from megatron.core.inference.communication_utils import (
     broadcast_from_last_pipeline_stage,
     is_pipeline_first_stage,
     is_pipeline_last_stage,
 )
+from megatron.core.inference.contexts.base_context import BaseInferenceContext
 from megatron.core.inference.contexts.dynamic_context import MaxSequenceLengthOverflowError
 from megatron.core.inference.inference_request import InferenceRequest, Status
 from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
     AbstractModelInferenceWrapper,
 )
+from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
+    GPTInferenceWrapper,
+)
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.utils import get_attention_mask, set_decode_expert_padding
+from megatron.core.tokenizers.text.utils.build_tokenizer import build_tokenizer
 from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.transformer.moe.moe_layer import BaseMoELayer
 from megatron.core.transformer.utils import set_model_to_sequence_parallel
-from megatron.core.utils import get_asyncio_loop, get_model_config, unwrap_model
+from megatron.core.utils import (
+    get_asyncio_loop,
+    get_attr_wrapped_model,
+    get_model_config,
+    unwrap_model,
+)
 
 try:
     import transformer_engine as te  # pylint: disable=unused-import
@@ -62,6 +73,7 @@ def __init__(
         pp_group: ProcessGroup = None,
     ):
         self.inference_wrapped_model = inference_wrapped_model
+        self.model_config = self.inference_wrapped_model.model.config
         self.tokenizer = tokenizer
 
         self.pp_group = pp_group
@@ -71,13 +83,28 @@ def __init__(
             is_pipeline_first_stage(self.pp_group) and is_pipeline_last_stage(self.pp_group)
         )
 
-        model_config = get_model_config(self.inference_wrapped_model.model)
         self.sampling_rng = torch.Generator(device=torch.cuda.current_device())
-        self.sampling_rng.manual_seed(model_config.inference_sampling_seed)
+        self.sampling_rng.manual_seed(self.model_config.inference_sampling_seed)
 
         if self.inference_wrapped_model.inference_context.is_dynamic_batching():
             self._init_dynamic_sampling_tensors()
 
+    @classmethod
+    def from_model_and_args(
+        cls,
+        model,
+        args,
+        context: BaseInferenceContext,
+        model_inference_wrapper_cls: type[AbstractModelInferenceWrapper] = GPTInferenceWrapper,
+    ):
+        tokenizer = build_tokenizer(args)
+        # TODO(ksanthanam): Condition this on model type?
+        model = model_inference_wrapper_cls(model, context)
+        model.model_is_pipeline_parallel = not (
+            parallel_state.is_pipeline_first_stage() and parallel_state.is_pipeline_last_stage()
+        )
+        return cls(model, tokenizer)
+
     def set_stop_word_finished_ids_callback(self, callback):
         """Set a callback to get request IDs that should be marked as finished due to stop words.
 
@@ -98,9 +125,9 @@ def _init_dynamic_sampling_tensors(self):
         self._get_stop_word_finished_ids_callback = None
 
         device = torch.cuda.current_device()
-        logits_dtype = self.inference_wrapped_model.inference_wrapper_config.params_dtype
+        logits_dtype = self.inference_wrapped_model.config.params_dtype
         # Use padded vocab size because tokenizer vocab size might pad to nearest power of 2.
-        vocab_size = self.inference_wrapped_model.inference_wrapper_config.padded_vocab_size
+        vocab_size = get_attr_wrapped_model(self.inference_wrapped_model.model, "vocab_size")
 
         self._sampling_backend = "torch"
         self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
@@ -505,7 +532,6 @@ def _dynamic_step_context_init(
             position_ids (Tensor): The active position IDs.
         """
         context = self.inference_wrapped_model.inference_context
-        inference_wrapper_config = self.inference_wrapped_model.inference_wrapper_config
         active_request_slice = slice(context.paused_request_count, context.total_request_count)
 
         # Remove Float16Module wrapper if it exists
@@ -517,11 +543,11 @@ def _dynamic_step_context_init(
 
         # If using symmetric kernels and we are using using nccl
         # for prefill turn off symmetric kernels
-        symmetric_ar_type = model_config.symmetric_ar_type
-        nccl_all_reduce_for_prefill = inference_wrapper_config.nccl_all_reduce_for_prefill
+        symmetric_ar_type = self.model_config.symmetric_ar_type
+        nccl_all_reduce_for_prefill = self.model_config.nccl_all_reduce_for_prefill
         # Turning on/off MoE padding for cuda-graphs
         moe_pad_experts_for_cuda_graph_inference = (
-            inference_wrapper_config.moe_pad_experts_for_cuda_graph_inference
+            self.model_config.moe_pad_experts_for_cuda_graph_inference
         )
         if moe_pad_experts_for_cuda_graph_inference:
             if context.using_cuda_graph_this_step():
@@ -569,8 +595,6 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
             input_ids (Tensor): The input token IDs.
             position_ids (Tensor): The position IDs.
         """
-        inference_wrapper_config = self.inference_wrapped_model.inference_wrapper_config
-
         context = self.inference_wrapped_model.inference_context
         active_request_count = context.total_request_count - context.paused_request_count
 
@@ -585,7 +609,7 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
                 if context.materialize_only_last_token_logits
                 else input_ids.shape[1]
             )
-            vocab_size = inference_wrapper_config.padded_vocab_size
+            vocab_size = get_attr_wrapped_model(self.inference_wrapped_model.model, "vocab_size")
             logits_shape = [1, logits_seq_len, vocab_size]
 
             if is_pipeline_last_stage(self.pp_group):
@@ -593,7 +617,7 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
 
             logits = broadcast_from_last_pipeline_stage(
                 logits_shape,
-                dtype=inference_wrapper_config.params_dtype,
+                dtype=self.model_config.params_dtype,
                 tensor=logits,
                 pp_group=self.pp_group,
             )
@@ -1024,9 +1048,8 @@ def generate_all_output_tokens_static_batch(
         # Pad batch tokens if necessary
         batch_size = len(active_requests)
         max_sequence_length = max_prompt_length_in_batch + sampling_params.num_tokens_to_generate
-        inference_wrapper_config = self.inference_wrapped_model.inference_wrapper_config
-        inference_max_batch_size = inference_wrapper_config.inference_max_requests
-        inference_max_sequence_length = inference_wrapper_config.inference_max_seq_length
+        inference_max_batch_size = self.model_config.inference_max_requests
+        inference_max_sequence_length = self.model_config.inference_max_seq_length
         padded_batch_size = inference_max_batch_size if enable_cuda_graph else batch_size
         if padded_batch_size > inference_max_batch_size:
             raise ValueError(
@@ -1068,7 +1091,7 @@ def generate_all_output_tokens_static_batch(
 
         # Use padded vocab size because tokenizer vocab size might not include padding
         # to nearest power of 2
-        vocab_size = inference_wrapper_config.padded_vocab_size
+        vocab_size = get_attr_wrapped_model(self.inference_wrapped_model.model, "vocab_size")
 
         # Check whether early termination is enabled
         no_early_termination = getattr(sampling_params, "no_early_termination", False)
@@ -1130,14 +1153,14 @@ def generate_all_output_tokens_static_batch(
 
             # If using symmetric kernels and we are using using nccl
             # for prefill turn off symmetric kernels
-            symmetric_ar_type = model_config.symmetric_ar_type
-            nccl_all_reduce_for_prefill = inference_wrapper_config.nccl_all_reduce_for_prefill
+            symmetric_ar_type = self.model_config.symmetric_ar_type
+            nccl_all_reduce_for_prefill = self.model_config.nccl_all_reduce_for_prefill
             if symmetric_ar_type is not None and nccl_all_reduce_for_prefill:
                 unwrapped_model.set_symmetric_ar(None)
 
             # Turning off MoE padding for prefill
             moe_pad_experts_for_cuda_graph_inference = (
-                inference_wrapper_config.moe_pad_experts_for_cuda_graph_inference
+                self.model_config.moe_pad_experts_for_cuda_graph_inference
             )
             if moe_pad_experts_for_cuda_graph_inference:
                 set_decode_expert_padding(unwrapped_model, False)
@@ -1219,7 +1242,7 @@ def generate_all_output_tokens_static_batch(
                     # and then broadcast the sampled tokens rather than broadcasting the raw logits.
                     logits = broadcast_from_last_pipeline_stage(
                         [batch_size, logits_seq_len, vocab_size],
-                        dtype=inference_wrapper_config.params_dtype,
+                        dtype=self.model_config.params_dtype,
                         tensor=logits,
                         pp_group=self.pp_group,
                     )
diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py
index 8265ee83ff5..2a5b2eab0f1 100644
--- a/megatron/core/transformer/attention.py
+++ b/megatron/core/transformer/attention.py
@@ -608,7 +608,7 @@ def flash_decode_and_prefill(
                     k_new=None,
                     v_new=None,
                     qv=None,
-                    out=None,
+                    out_=None,
                     cu_seqlens_q=cu_seqlens_q,
                     cu_seqlens_k=None,
                     cu_seqlens_k_new=None,
@@ -627,7 +627,8 @@ def flash_decode_and_prefill(
                     v_descale=None,
                     softmax_scale=softmax_scale,
                     causal=True,
-                    window_size=(-1, -1),
+                    window_size_left=-1,
+                    window_size_right=-1,
                     attention_chunk=0,
                     softcap=0.0,
                     rotary_interleaved=True,
diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index b07e3ed4d9d..391e72974f8 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -611,6 +611,12 @@ class TransformerConfig(ModelParallelConfig):
     the expert capacity length, effective only after the moe_expert_capacity_factor is set. The
     default setting is False."""
 
+    moe_pad_experts_for_cuda_graph_inference: bool = False
+    """moe_pad_experts_for_cuda_graph_inference (bool): If True, the router will switch to dropping
+    and padding during decode time which does not have a D2H sync. The capacity factor is set to the
+    max that an expert could see during inference so no tokens are actually dropped. The default
+    setting is False."""
+
     moe_token_drop_policy: str = "probs"
     """The policy to drop tokens. Can be either "probs" or "position". If "probs", the tokens with
     the lowest probabilities will be dropped. If "position", tokens at the end of each batch will
@@ -741,6 +747,9 @@ class TransformerConfig(ModelParallelConfig):
     symmetric_ar_type: Optional[str] = None
     """Type of symmetric all reduce to use"""
 
+    nccl_all_reduce_for_prefill: bool = False
+    """If True, use NCCL all-reduce kernels when symmetric all-reduce is enabled."""
+
     use_inference_optimized_layers: bool = False
     """If True, use inference optimized transformer layers during inference."""
 
diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py
index 126da94e3a6..a4d9ae52448 100644
--- a/megatron/rl/inference/megatron.py
+++ b/megatron/rl/inference/megatron.py
@@ -16,9 +16,6 @@
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
     GPTInferenceWrapper,
 )
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.text_generation_controllers.simple_text_generation_controller import (
     SimpleTextGenerationController,
@@ -66,20 +63,7 @@ def get_static_inference_engine(args: Namespace, model: MegatronModule) -> Abstr
     """
     tokenizer = get_tokenizer()
 
-    inference_wrapper_config = InferenceWrapperConfig(
-        hidden_size=args.hidden_size,
-        inference_batch_times_seqlen_threshold=args.inference_batch_times_seqlen_threshold,
-        fp32_residual_connection=args.fp32_residual_connection,
-        params_dtype=args.params_dtype,
-        padded_vocab_size=args.padded_vocab_size,
-        inference_max_seq_length=args.inference_max_seq_length,
-        inference_max_requests=(
-            args.inference_max_batch_size if args.inference_max_batch_size is not None else 1
-        ),
-        nccl_all_reduce_for_prefill=args.nccl_all_reduce_for_prefill,
-    )
-
-    inference_wrapped_model = GPTInferenceWrapper(model, inference_wrapper_config)
+    inference_wrapped_model = GPTInferenceWrapper(model)
     pg_collection = get_attr_wrapped_model(model, "pg_collection")
     pp_group = pg_collection.pp
     text_generation_controller = SimpleTextGenerationController(
@@ -90,7 +74,7 @@ def get_static_inference_engine(args: Namespace, model: MegatronModule) -> Abstr
     return MCoreEngine(
         text_generation_controller=text_generation_controller,
         max_batch_size=(
-            args.inference_max_batch_size if args.inference_max_batch_size is not None else 1
+            args.inference_max_reqeusts if args.inference_max_requests is not None else 1
         ),
     )
 
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 262c4ce79b2..f89ec449a64 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1556,13 +1556,10 @@ def _add_inference_args(parser):
                        dest='use_legacy_static_engine')
     group.add_argument('--inference-max-requests', type=int, default=8,
                        help='Maximum number of requests for inference.',
-                       dest='inference_max_batch_size')
+                       dest='inference_max_requests')
     group.add_argument('--inference-max-seq-length', type=int, default=2560,
                        help='Maximum sequence length expected for inference (prefill + decode).',
                        dest='inference_max_seq_length')
-    group.add_argument('--inference-max-batch-size', type=int, default=None,
-                       help='Maximum batch size for inference.',
-                       dest='inference_max_batch_size')
     group.add_argument('--inference-dynamic-batching',
                        action='store_true', default=False,
                        help='Enable dynamic batching mode.')
diff --git a/tools/run_dynamic_text_generation_server.py b/tools/run_dynamic_text_generation_server.py
index c1d582c53dd..67c0c6c30f3 100644
--- a/tools/run_dynamic_text_generation_server.py
+++ b/tools/run_dynamic_text_generation_server.py
@@ -8,8 +8,6 @@
 
 from examples.inference.gpt.gpt_dynamic_inference import (
     add_dynamic_inference_args,
-    get_inference_context,
-    get_inference_controller,
     get_model,
 )
 from megatron.core.inference.engines import DynamicInferenceEngine
@@ -80,33 +78,12 @@ async def run_text_generation_server(
         args = get_args()
         model = get_model()
 
-        if args.legacy_tokenizer:
-            tokenizer = get_tokenizer()
-        else:
-            tokenizer = build_tokenizer(args)
-
-        mamba_inference_state_config = get_mamba_inference_state_config_from_model(model)
-
         # Enable return_log_probs to allow prompt logprobs computation for echo=True requests
         # This sets materialize_only_last_token_logits=False in the inference context,
         # which is required for lm-eval compatibility (loglikelihood evaluation tasks)
         args.return_log_probs = True
 
-        context = get_inference_context(
-            None,
-            None,
-            calculate_max_sequence_length_from_requests=False,
-            mamba_inference_state_config=mamba_inference_state_config,
-        )
-
-        controller = get_inference_controller(model, context)
-
-        engine = DynamicInferenceEngine(
-            controller,
-            context,
-            enable_cuda_graph=args.cuda_graph_impl == "local",
-            random_seed=args.seed,
-        )
+        engine = DynamicInferenceEngine.from_model_and_args(model, args)
 
         asyncio.run(
             run_text_generation_server(
diff --git a/tools/run_inference_performance_test.py b/tools/run_inference_performance_test.py
index 32d61444530..a429a91724d 100644
--- a/tools/run_inference_performance_test.py
+++ b/tools/run_inference_performance_test.py
@@ -76,22 +76,8 @@ def get_inference_engine(args: argparse.Namespace, model: MegatronModule) -> Abs
     """
     tokenizer = get_tokenizer()
 
-    inference_wrapper_config = InferenceWrapperConfig(
-        hidden_size=args.hidden_size,
-        inference_batch_times_seqlen_threshold=args.inference_batch_times_seqlen_threshold,
-        fp32_residual_connection=args.fp32_residual_connection,
-        params_dtype=args.params_dtype,
-        padded_vocab_size=args.padded_vocab_size,
-        inference_max_requests=args.inference_max_batch_size,
-        inference_max_seq_length=args.inference_max_seq_length,
-        nccl_all_reduce_for_prefill=args.nccl_all_reduce_for_prefill,
-        moe_pad_experts_for_cuda_graph_inference=args.moe_pad_experts_for_cuda_graph_inference,
-    )
-
-    mamba_inference_state_config = get_mamba_inference_state_config_from_model(model)
-
     if args.engine_type == "static":
-        inference_wrapped_model = GPTInferenceWrapper(model, inference_wrapper_config)
+        inference_wrapped_model = GPTInferenceWrapper(model)
         inference_wrapped_model.model_is_pipeline_parallel = not (
             mpu.is_pipeline_first_stage() and mpu.is_pipeline_last_stage()
         )
@@ -100,52 +86,7 @@ def get_inference_engine(args: argparse.Namespace, model: MegatronModule) -> Abs
         )
         return StaticInferenceEngine(text_generation_controller=text_generation_controller)
     elif args.engine_type == "dynamic":
-        context = DynamicInferenceContext(
-            params_dtype=args.params_dtype,
-            num_layers=args.num_layers,
-            kv_channels=args.kv_channels,
-            num_attention_heads=(
-                args.num_query_groups if args.group_query_attention else args.num_attention_heads
-            ),
-            max_sequence_length=args.inference_max_seq_length,
-            num_cuda_graphs=(
-                args.inference_dynamic_batching_num_cuda_graphs
-                if args.cuda_graph_impl == "local"
-                else None
-            ),
-            buffer_size_gb=args.inference_dynamic_batching_buffer_size_gb,
-            buffer_guaranteed_fraction=args.inference_dynamic_batching_buffer_guaranteed_fraction,
-            buffer_overflow_factor=args.inference_dynamic_batching_buffer_overflow_factor,
-            max_requests_override=args.inference_dynamic_batching_max_requests_override,
-            max_tokens_override=args.inference_dynamic_batching_max_tokens_override,
-            block_size_tokens=args.inference_dynamic_batching_block_size,
-            tensor_model_parallel_size=args.tensor_model_parallel_size,
-            pipeline_model_parallel_size=args.pipeline_model_parallel_size,
-            materialize_only_last_token_logits=not args.return_log_probs,
-            mamba_inference_state_config=mamba_inference_state_config,
-            cache_mla_latent=args.multi_latent_attention and args.cache_mla_latents,
-            kv_lora_rank=args.kv_lora_rank if args.multi_latent_attention else None,
-            qk_pos_emb_head_dim=args.qk_pos_emb_head_dim,
-            use_cuda_graphs_for_non_decode_steps=not args.decode_only_cuda_graphs,
-            use_flashinfer_fused_rope=args.use_flashinfer_fused_rope,
-            unified_memory_level=args.inference_dynamic_batching_unified_memory_level,
-        )
-        inference_wrapped_model = GPTInferenceWrapper(
-            model, inference_wrapper_config, inference_context=context
-        )
-        inference_wrapped_model.model_is_pipeline_parallel = not (
-            mpu.is_pipeline_first_stage() and mpu.is_pipeline_last_stage()
-        )
-        text_generation_controller = TextGenerationController(
-            inference_wrapped_model=inference_wrapped_model, tokenizer=tokenizer
-        )
-        return DynamicInferenceEngine(
-            text_generation_controller,
-            context,
-            termination_id=-1,
-            enable_cuda_graph=args.cuda_graph_impl == "local",
-            random_seed=args.seed,
-        )
+        return DynamicInferenceEngine.from_model_and_args(model, args)
 
 
 async def generate(

From ca8b2780ccd96b0db03ccc26a00e638c2de34058 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Thu, 15 Jan 2026 02:25:24 -0800
Subject: [PATCH 02/30] Tests pass

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 .../inference/gpt/gpt_static_inference.py     |  47 ++-----
 .../inference/contexts/dynamic_context.py     |  31 ++---
 .../core/inference/contexts/static_context.py |   9 --
 .../core/inference/engines/dynamic_engine.py  |  20 ++-
 .../core/inference/engines/static_engine.py   |  40 +++---
 .../simple_text_generation_controller.py      |   5 -
 .../text_generation_controller.py             |  46 +++++--
 megatron/core/ssm/mamba_layer.py              |   5 +-
 megatron/rl/inference/megatron.py             | 118 +++---------------
 .../contexts/test_dynamic_context.py          |  41 ++++--
 .../inference/engines/test_dynamic_engine.py  |  23 +---
 .../inference/engines/test_static_engine.py   |  18 +--
 .../gpt/test_gpt_inference_wrapper.py         |  57 +--------
 .../t5/test_t5_inference_wrapper.py           |  17 +--
 .../test_model_inference_wrapper_config.py    |  21 ----
 .../inference/test_wandb_logging.py           |  18 +--
 ...oder_decoder_text_generation_controller.py |  17 +--
 ....py => test_text_generation_controller.py} |  28 +----
 .../test_vlm_text_generation_controller.py    |  17 +--
 .../models/test_gpt_model_batch_invariant.py  |  42 ++-----
 tests/unit_tests/models/test_mamba_model.py   |   5 +-
 tools/run_inference_performance_test.py       |  18 ++-
 tools/run_text_generation_server.py           |  27 +---
 23 files changed, 199 insertions(+), 471 deletions(-)
 delete mode 100644 megatron/core/inference/text_generation_controllers/simple_text_generation_controller.py
 delete mode 100644 tests/unit_tests/inference/model_inference_wrappers/test_model_inference_wrapper_config.py
 rename tests/unit_tests/inference/text_generation_controllers/{test_simple_text_generation_controller.py => test_text_generation_controller.py} (97%)

diff --git a/examples/inference/gpt/gpt_static_inference.py b/examples/inference/gpt/gpt_static_inference.py
index 03a60927ab2..906748d5d2e 100644
--- a/examples/inference/gpt/gpt_static_inference.py
+++ b/examples/inference/gpt/gpt_static_inference.py
@@ -1,9 +1,6 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 
 import os
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from model_provider import model_provider
 from gpt_builders import gpt_builder
 from mamba_builders import mamba_builder
@@ -23,9 +20,6 @@
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
     GPTInferenceWrapper,
 )
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.text_generation_controllers.text_generation_controller import (
     TextGenerationController,
@@ -49,6 +43,7 @@
 from megatron.training.checkpointing import load_checkpoint
 from megatron.training.initialize import initialize_megatron
 
+
 def add_static_inference_args(parser):
     """Static inference arguments."""
 
@@ -83,30 +78,16 @@ def get_inference_engine(args: Namespace, model: MegatronModule) -> StaticInfere
         tokenizer = get_tokenizer()
     else:
         tokenizer = build_tokenizer(args)
-    inference_wrapper_config = InferenceWrapperConfig(
-        hidden_size=args.hidden_size,
-        inference_batch_times_seqlen_threshold=args.inference_batch_times_seqlen_threshold,
-        fp32_residual_connection=args.fp32_residual_connection,
-        params_dtype=args.params_dtype,
-        padded_vocab_size=args.padded_vocab_size,
-        inference_max_requests=args.inference_max_batch_size,
-        inference_max_seq_length=args.inference_max_seq_length,
-        nccl_all_reduce_for_prefill=args.nccl_all_reduce_for_prefill,
-        fp8=args.fp8,
-        moe_pad_experts_for_cuda_graph_inference = args.moe_pad_experts_for_cuda_graph_inference
-    )
-
-    inference_context = StaticInferenceContext.from_config(inference_wrapper_config)
-
-    inference_wrapped_model = GPTInferenceWrapper(
-        model, inference_wrapper_config, inference_context
+    inference_context = StaticInferenceContext(
+        args.inference_max_requests, args.inference_max_seq_length
     )
+    inference_wrapped_model = GPTInferenceWrapper(model, inference_context)
     text_generation_controller = TextGenerationController(
         inference_wrapped_model=inference_wrapped_model, tokenizer=tokenizer
     )
     engine_kwargs = {
-        "text_generation_controller" : text_generation_controller,
-        "legacy" : args.use_legacy_static_engine,
+        "text_generation_controller": text_generation_controller,
+        "legacy": args.use_legacy_static_engine,
     }
     if not args.use_legacy_static_engine:
         engine_kwargs["buffer_size_gb"] = args.inference_dynamic_batching_buffer_size_gb
@@ -165,12 +146,6 @@ def main():
 
     args = get_args()
 
-    if args.max_batch_size is not None:
-        warnings.warn(
-            f"`--max-batch-size` has been deprecated in favor of `--inference-max-requests`."
-        )
-        args.inference_max_batch_size = max(args.max_batch_size, args.inference_max_batch_size)
-
     # Set up model and load checkpoint
     if args.model_provider == "gpt":
         model_builder = gpt_builder
@@ -246,14 +221,15 @@ def main():
         from collections import defaultdict
 
         unique_prompt_map = defaultdict(list)
-        for result_idx, result in enumerate(results):
+        for result_idx, record in enumerate(results):
+            result = record.requests[0]
             unique_prompt_map[result.prompt].append(result_idx)
 
         # Print unique prompts + outputs.
         for unique_idx, (prompt_text, result_idxs) in enumerate(unique_prompt_map.items()):
             result_idx = result_idxs[0]
-            result = results[result_idx]
-            generated_text = result.generated_text.replace("\n", "\\n")
+            record = results[result_idx]
+            generated_text = record.requests[0].generated_text.replace("\n", "\\n")
             print(
                 f"{unique_idx}/{len(unique_prompt_map)} [{len(result_idxs)}]. {prompt_text} "
                 f"... {generated_text}"
@@ -276,7 +252,7 @@ def main():
                 )
             ),
             len(requests),
-            args.inference_max_batch_size,
+            args.inference_max_requests,
             stats["allocated_bytes.all.peak"] / (1024**3),
             stats["reserved_bytes.all.peak"] / (1024**3),
             latency,
@@ -293,6 +269,5 @@ def main():
         torch.distributed.destroy_process_group()
 
 
-
 if __name__ == "__main__":
     main()
diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 4b7e47b9d22..0d2a40a51e4 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -294,25 +294,15 @@ def __init__(
         # Per partition num heads and hidden size.
         num_attention_heads = model_config.num_query_groups or model_config.num_attention_heads
         projection_size = model_config.kv_channels * num_attention_heads
-        if model_config.tensor_model_parallel_size is None:
-            assert pg_collection is not None
-            tp_size = (
-                get_pg_size(pg_collection.tp)
-                if pg_collection is not None
-                else parallel_state.get_tensor_model_parallel_world_size()
-            )
+        if pg_collection is not None:
+            tp_size = get_pg_size(pg_collection.tp)
         else:
             tp_size = model_config.tensor_model_parallel_size
         self.hidden_size_per_attention_head = core_divide(projection_size, num_attention_heads)
         self.num_attention_heads_per_partition = core_divide(num_attention_heads, tp_size)
 
-        if model_config.pipeline_model_parallel_size is None:
-            assert pg_collection is not None
-            pp_size = (
-                get_pg_size(pg_collection.pp)
-                if pg_collection is not None
-                else parallel_state.get_pipeline_model_parallel_world_size()
-            )
+        if pg_collection is not None:
+            pp_size = get_pg_size(pg_collection.pp)
         else:
             pp_size = model_config.pipeline_model_parallel_size
 
@@ -732,7 +722,16 @@ def round_up_tokens(cls, value, tp_size=None):
     @classmethod
     def from_model_and_args(cls, model, args, overrides: Optional[Dict[str, Any]] = None):
         """
-        Instantiate a `DynamicInferenceContext` from a model and command-line args.
+        Instantiate a `DynamicInferenceContext` from the model and args.
+
+        Args:
+            model: The Megatron model instance.
+            args: The arguments object.
+            overrides (Optional[Dict[str, Any]]): A dictionary of values to override
+                the default arguments derived from `model` and `args`.
+
+        Returns:
+            DynamicInferenceContext: The initialized inference context.
         """
         config = model.config
 
@@ -760,11 +759,13 @@ def from_model_and_args(cls, model, args, overrides: Optional[Dict[str, Any]] =
             )
 
         mamba_inference_state_config = get_mamba_inference_state_config_from_model(model)
+        pg_collection = get_attr_wrapped_model(model, "pg_collection")
 
         kwargs = {
             "model_config": config,
             "max_sequence_length": max_sequence_length,
             "mamba_inference_state_config": mamba_inference_state_config,
+            "pg_collection": pg_collection,
             "num_cuda_graphs": (
                 args.inference_dynamic_batching_num_cuda_graphs
                 if args.cuda_graph_impl == "local"
diff --git a/megatron/core/inference/contexts/static_context.py b/megatron/core/inference/contexts/static_context.py
index ba41b0c2401..98ba8b5185d 100644
--- a/megatron/core/inference/contexts/static_context.py
+++ b/megatron/core/inference/contexts/static_context.py
@@ -1,7 +1,5 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 
-from megatron.core.transformer.transformer_config import TransformerConfig
-
 from .base_context import BaseInferenceContext
 
 
@@ -25,13 +23,6 @@ def __init__(
         self.key_value_memory_dict = {}
         self.decode_mode = False
 
-    @classmethod
-    def from_config(cls, config: TransformerConfig) -> "StaticInferenceContext":
-        """Initialize context from a config."""
-        max_batch_size = config.inference_max_requests
-        max_sequence_length = config.inference_max_seq_length
-        return cls(max_batch_size, max_sequence_length)
-
     def swap_key_value_dict(self, batch_idx):
         "swap between batches"
         if len(self.key_value_memory_dict) == 0:
diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index d1885bf0d6c..e56d7fd528d 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -231,17 +231,31 @@ def from_model_and_args(
         cls,
         model,
         args,
-        context: Optional[DynamicInferenceContext] = None,
         controller: Optional[TextGenerationController] = None,
+        context: Optional[DynamicInferenceContext] = None,
     ):
+        """
+        Initializes a `DynamicInferenceEngine` from the model and args.
+
+        Args:
+            model: The Megatron model instance.
+            args: The arguments object.
+            controller (Optional[TextGenerationController]): An optional existing
+                controller. If None, one is created from the model and args.
+            context (Optional[DynamicInferenceContext]): An optional existing
+                context. If None, one is created from the model and args.
+
+        Returns:
+            DynamicInferenceEngine: The initialized inference engine.
+        """
         if context is None:
             context = DynamicInferenceContext.from_model_and_args(model, args)
         if controller is None:
-            controller = TextGenerationController.from_model_and_args(model, args)
+            controller = TextGenerationController.from_model_and_args(model, args, context)
 
         return cls(
-            context,
             controller,
+            context,
             enable_cuda_graph=args.cuda_graph_impl == "local",
             random_seed=args.seed,
             track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events,
diff --git a/megatron/core/inference/engines/static_engine.py b/megatron/core/inference/engines/static_engine.py
index 661962b37a2..196fc1054da 100644
--- a/megatron/core/inference/engines/static_engine.py
+++ b/megatron/core/inference/engines/static_engine.py
@@ -8,7 +8,7 @@
 import torch
 
 from megatron.core.inference.async_stream import AsyncStream
-from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
+from megatron.core.inference.contexts import DynamicInferenceContext, StaticInferenceContext
 from megatron.core.inference.engines.abstract_engine import AbstractEngine
 from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine
 from megatron.core.inference.inference_request import InferenceRequest
@@ -72,49 +72,53 @@ def __init__(
         self.config = self.inference_wrapped_model.config
         self.random_seed = random_seed or 1234
 
-        inference_max_batch_size = self.config.inference_max_requests
+        # Store original context in case we need to fall back to legacy static engine
+        original_context = self.inference_wrapped_model.inference_context
+        assert original_context is not None
+        assert isinstance(original_context, StaticInferenceContext)
+
         if max_batch_size is None:
-            max_batch_size = inference_max_batch_size
-        elif max_batch_size > inference_max_batch_size:
+            max_batch_size = original_context.max_batch_size
+        elif max_batch_size > original_context.max_batch_size:
             warnings.warn(
                 f"Engine `max_batch_size` ({max_batch_size}) > "
-                f"`inference_max_requests` in `inference_wrapper_config` "
-                f"({inference_max_batch_size}); setting `max_batch_size` to "
-                f"{inference_max_batch_size}",
+                f"`context.max_batch_size` in `inference_wrapped_model.inference_context` "
+                f"({original_context.max_batch_size}); setting `max_batch_size` to "
+                f"{original_context.max_batch_size}",
                 UserWarning,
             )
-            max_batch_size = inference_max_batch_size
+            max_batch_size = original_context.max_batch_size
 
         self.scheduler = Scheduler(max_batch_size=max_batch_size)
 
-        # Store original context in case we need to fall back to legacy static engine
-        original_context = self.inference_wrapped_model.inference_context
-
         mamba_inference_state_config = get_mamba_inference_state_config_from_model(
             self.inference_wrapped_model.model
         )
 
         try:
             if not legacy:
-                dynamic_context = DynamicInferenceContext.from_config(
-                    inference_config=inference_wrapper_config,
-                    model=text_generation_controller.inference_wrapped_model.model,
-                    max_batch_size=max_batch_size,
+                dynamic_context = DynamicInferenceContext(
+                    model_config=self.config,
+                    max_sequence_length=original_context.max_sequence_length,
                     buffer_size_gb=buffer_size_gb,
-                    num_cuda_graphs=1,
                     mamba_inference_state_config=mamba_inference_state_config,
+                    max_requests=max_batch_size,
+                    num_cuda_graphs=1,
+                    block_size_tokens=256,
+                    unified_memory_level=0,
                 )
+
                 self.controller.inference_wrapped_model.inference_context = dynamic_context
                 self.controller.inference_wrapped_model.prep_model_for_inference()
                 self.controller._init_dynamic_sampling_tensors()
 
                 self.dynamic_engine = DynamicInferenceEngine(
                     controller=self.controller,
-                    random_seed=self.random_seed,
                     context=dynamic_context,
-                    enable_cuda_graph=True,
+                    random_seed=self.random_seed,
                 )
         except Exception as e:
+            torch.distributed.breakpoint(0)
             # Get exception details for better debugging
             exception_msg = str(e) if str(e) else f"{type(e).__name__}: {repr(e)}"
             warnings.warn(
diff --git a/megatron/core/inference/text_generation_controllers/simple_text_generation_controller.py b/megatron/core/inference/text_generation_controllers/simple_text_generation_controller.py
deleted file mode 100644
index 340cadb48a9..00000000000
--- a/megatron/core/inference/text_generation_controllers/simple_text_generation_controller.py
+++ /dev/null
@@ -1,5 +0,0 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
-
-from megatron.core.inference.text_generation_controllers.text_generation_controller import (  # noqa: F401 # pylint: disable=unused-import
-    TextGenerationController as SimpleTextGenerationController,
-)
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index 74bad4f571d..a1fa8ed898a 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -22,6 +22,7 @@
 )
 from megatron.core.inference.contexts.base_context import BaseInferenceContext
 from megatron.core.inference.contexts.dynamic_context import MaxSequenceLengthOverflowError
+from megatron.core.inference.contexts.static_context import StaticInferenceContext
 from megatron.core.inference.inference_request import InferenceRequest, Status
 from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
     AbstractModelInferenceWrapper,
@@ -83,6 +84,17 @@ def __init__(
             is_pipeline_first_stage(self.pp_group) and is_pipeline_last_stage(self.pp_group)
         )
 
+        # Use padded vocab size because tokenizer vocab size might pad to nearest power of 2.
+        try:
+            self.vocab_size = get_attr_wrapped_model(
+                self.inference_wrapped_model.model, "vocab_size"
+            )
+        except RuntimeError as e:
+            # Handle LlaVa models
+            self.vocab_size = get_attr_wrapped_model(
+                self.inference_wrapped_model.model, "language_model"
+            ).vocab_size
+
         self.sampling_rng = torch.Generator(device=torch.cuda.current_device())
         self.sampling_rng.manual_seed(self.model_config.inference_sampling_seed)
 
@@ -97,6 +109,19 @@ def from_model_and_args(
         context: BaseInferenceContext,
         model_inference_wrapper_cls: type[AbstractModelInferenceWrapper] = GPTInferenceWrapper,
     ):
+        """
+        Initializes a `TextGenerationController` from the model and args.
+
+        Args:
+            model: The Megatron model instance.
+            args: The arguments object.
+            context (BaseInferenceContext): The inference context.
+            model_inference_wrapper_cls (type[AbstractModelInferenceWrapper]): The class
+                used to wrap the model for inference. Defaults to GPTInferenceWrapper.
+
+        Returns:
+            TextGenerationController: The initialized text generation controller.
+        """
         tokenizer = build_tokenizer(args)
         # TODO(ksanthanam): Condition this on model type?
         model = model_inference_wrapper_cls(model, context)
@@ -126,8 +151,6 @@ def _init_dynamic_sampling_tensors(self):
 
         device = torch.cuda.current_device()
         logits_dtype = self.inference_wrapped_model.config.params_dtype
-        # Use padded vocab size because tokenizer vocab size might pad to nearest power of 2.
-        vocab_size = get_attr_wrapped_model(self.inference_wrapped_model.model, "vocab_size")
 
         self._sampling_backend = "torch"
         self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
@@ -609,8 +632,7 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
                 if context.materialize_only_last_token_logits
                 else input_ids.shape[1]
             )
-            vocab_size = get_attr_wrapped_model(self.inference_wrapped_model.model, "vocab_size")
-            logits_shape = [1, logits_seq_len, vocab_size]
+            logits_shape = [1, logits_seq_len, self.vocab_size]
 
             if is_pipeline_last_stage(self.pp_group):
                 assert logits is not None and torch.Size(logits_shape) == logits.shape
@@ -1048,8 +1070,10 @@ def generate_all_output_tokens_static_batch(
         # Pad batch tokens if necessary
         batch_size = len(active_requests)
         max_sequence_length = max_prompt_length_in_batch + sampling_params.num_tokens_to_generate
-        inference_max_batch_size = self.model_config.inference_max_requests
-        inference_max_sequence_length = self.model_config.inference_max_seq_length
+        context = self.inference_wrapped_model.inference_context
+        assert isinstance(context, StaticInferenceContext)
+        inference_max_batch_size = context.max_batch_size
+        inference_max_sequence_length = context.max_sequence_length
         padded_batch_size = inference_max_batch_size if enable_cuda_graph else batch_size
         if padded_batch_size > inference_max_batch_size:
             raise ValueError(
@@ -1089,10 +1113,6 @@ def generate_all_output_tokens_static_batch(
             batch_size, device=torch.cuda.current_device()
         ).cuda()
 
-        # Use padded vocab size because tokenizer vocab size might not include padding
-        # to nearest power of 2
-        vocab_size = get_attr_wrapped_model(self.inference_wrapped_model.model, "vocab_size")
-
         # Check whether early termination is enabled
         no_early_termination = getattr(sampling_params, "no_early_termination", False)
         termination_id = -1 if no_early_termination else self.tokenizer.eod
@@ -1235,13 +1255,13 @@ def generate_all_output_tokens_static_batch(
                 if self.model_is_pipeline_parallel:
                     context_length = context_end_position - context_start_position
                     logits_seq_len = 1 if materialize_only_last_token_logits else context_length
-                    logits_shape = [batch_size, logits_seq_len, vocab_size]
+                    logits_shape = [batch_size, logits_seq_len, self.vocab_size]
                     if is_pipeline_last_stage(self.pp_group):
                         assert logits is not None and torch.Size(logits_shape) == logits.shape
                     # TODO(ksanthanam): Evaluate whether it makes more sense to sample on 1 rank
                     # and then broadcast the sampled tokens rather than broadcasting the raw logits.
                     logits = broadcast_from_last_pipeline_stage(
-                        [batch_size, logits_seq_len, vocab_size],
+                        [batch_size, logits_seq_len, self.vocab_size],
                         dtype=self.model_config.params_dtype,
                         tensor=logits,
                         pp_group=self.pp_group,
@@ -1271,7 +1291,7 @@ def generate_all_output_tokens_static_batch(
                 sampled_logits = self.sample_from_logits(
                     last_token_logits,
                     sampling_params,
-                    vocab_size,
+                    self.vocab_size,
                     generation_started=generation_started,
                     top_n_logprobs_dict=top_n_logprobs_dict,
                     logits=logits_for_top_n_prompt_logprobs,
diff --git a/megatron/core/ssm/mamba_layer.py b/megatron/core/ssm/mamba_layer.py
index e04fc1d226d..ac44b21967b 100644
--- a/megatron/core/ssm/mamba_layer.py
+++ b/megatron/core/ssm/mamba_layer.py
@@ -181,6 +181,9 @@ def _should_call_local_cudagraph(self, *args, **kwargs):
             and kwargs.get('attention_mask') is None
             and kwargs.get('inference_context') is not None
         ):
-            using_cuda_graph = kwargs['inference_context'].using_cuda_graph_this_step()
+            context = kwargs['inference_context']
+            using_cuda_graph = (context.is_static_batching() and context.is_decode_only()) or (
+                not context.is_static_batching() and context.using_cuda_graph_this_step()
+            )
             return using_cuda_graph
         return False
diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py
index a4d9ae52448..6e6053d1ca6 100644
--- a/megatron/rl/inference/megatron.py
+++ b/megatron/rl/inference/megatron.py
@@ -7,7 +7,6 @@
 import torch.distributed as dist
 from pydantic import PrivateAttr
 
-from megatron.core import parallel_state
 from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
 from megatron.core.inference.engines.abstract_engine import AbstractEngine
 from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine
@@ -17,19 +16,12 @@
     GPTInferenceWrapper,
 )
 from megatron.core.inference.sampling_params import SamplingParams
-from megatron.core.inference.text_generation_controllers.simple_text_generation_controller import (
-    SimpleTextGenerationController,
+from megatron.core.inference.text_generation_controllers.text_generation_controller import (
+    TextGenerationController,
 )
 from megatron.core.models.gpt.gpt_model import GPTModel
-from megatron.core.pipeline_parallel.utils import is_pp_first_stage, is_pp_last_stage
-from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols
 from megatron.core.transformer.module import MegatronModule
-from megatron.core.utils import (
-    get_attr_wrapped_model,
-    get_mamba_inference_state_config_from_model,
-    get_pg_size,
-    log_single_rank,
-)
+from megatron.core.utils import get_attr_wrapped_model, log_single_rank
 from megatron.training import get_wandb_writer
 from megatron.training.global_vars import get_args, get_tokenizer
 
@@ -66,10 +58,8 @@ def get_static_inference_engine(args: Namespace, model: MegatronModule) -> Abstr
     inference_wrapped_model = GPTInferenceWrapper(model)
     pg_collection = get_attr_wrapped_model(model, "pg_collection")
     pp_group = pg_collection.pp
-    text_generation_controller = SimpleTextGenerationController(
-        inference_wrapped_model=inference_wrapped_model,
-        tokenizer=tokenizer,
-        pp_group=pp_group,
+    text_generation_controller = TextGenerationController(
+        inference_wrapped_model=inference_wrapped_model, tokenizer=tokenizer, pp_group=pp_group
     )
     return MCoreEngine(
         text_generation_controller=text_generation_controller,
@@ -79,19 +69,14 @@ def get_static_inference_engine(args: Namespace, model: MegatronModule) -> Abstr
     )
 
 
-## This code is copied from tools/run_text_generation_server.py
 def get_dynamic_inference_engine(
     args: Namespace,
     model: MegatronModule,
     inference_logging_step_interval: int = 0,
-    metrics_writer = None
-) -> AbstractEngine:
-    """Get the relevant backend for running inference.
-
-    This function will automatically choose the TRTLLMBackend when possible,
-    and default to Mcore backend if the user does not specify any backends.
-    TRTLLMBackend is not implmented yet.
-
+    metrics_writer=None,
+) -> DynamicInferenceEngine:
+    """
+    Returns an inference engine.
     Args:
         args (Namespace): The user arguments parsed from command line
         model (MegatronModule): The megatron model.
@@ -99,81 +84,14 @@ def get_dynamic_inference_engine(
         metrics_writer: Metrics writer (wandb module) for logging.
 
     Returns:
-        AbstractBackend: The chosen backend
+        DynamicInferenceEngine: The inference engine
     """
-    tokenizer = get_tokenizer()
-
-    enable_cuda_graph = args.cuda_graph_impl == "local"
-
-    mamba_inference_state_config = get_mamba_inference_state_config_from_model(model)
-
-    # DynamicInferenceContext must use the inference model's TP / PP size, not the
-    # training TP / PP size from global args. The inference model may have a custom
-    # ProcessGroupCollection with a different TP / PP size.
-    pg_collection = get_attr_wrapped_model(model, "pg_collection")
-    tp_group = getattr(pg_collection, 'tp', None) if pg_collection is not None else None
-    if tp_group is not None:
-        inference_tp_size = get_pg_size(tp_group)
-    else:
-        inference_tp_size = args.tensor_model_parallel_size
-    pp_group = getattr(pg_collection, 'pp', None) if pg_collection is not None else None
-    if pp_group is not None:
-        inference_pp_size = get_pg_size(pp_group)
-    else:
-        inference_pp_size = args.pipeline_model_parallel_size
-
-    # Inference context.
-    inference_context = DynamicInferenceContext(
-        params_dtype=args.params_dtype,
-        num_layers=args.num_layers // inference_pp_size,
-        kv_channels=args.kv_channels,
-        num_attention_heads=(
-            args.num_query_groups if args.group_query_attention else args.num_attention_heads
-        ),
-        max_sequence_length=args.inference_max_seq_length,
-        num_cuda_graphs=(
-            args.inference_dynamic_batching_num_cuda_graphs if enable_cuda_graph else None
-        ),
-        block_size_tokens=args.inference_dynamic_batching_block_size,
-        buffer_size_gb=args.inference_dynamic_batching_buffer_size_gb,
-        max_requests=args.inference_dynamic_batching_max_requests,
-        max_tokens=args.inference_dynamic_batching_max_tokens,
-        pg_collection=pg_collection,  # TP/PP sizes are derived from the model's pg_collection.
-        materialize_only_last_token_logits=True,
-        mamba_inference_state_config=mamba_inference_state_config,
-        cache_mla_latent=args.multi_latent_attention and args.cache_mla_latents,
-        kv_lora_rank=args.kv_lora_rank if args.multi_latent_attention else None,
-        qk_pos_emb_head_dim=args.qk_pos_emb_head_dim,
-        use_cuda_graphs_for_non_decode_steps=not args.decode_only_cuda_graphs,
-        use_flashinfer_fused_rope=None,
-        unified_memory_level=args.inference_dynamic_batching_unified_memory_level,
-        cuda_graph_max_tokens=args.inference_dynamic_batching_cuda_graph_max_tokens,
-        cuda_graph_mixed_prefill_count=args.inference_dynamic_batching_cuda_graph_mixed_prefill_count,
-        metrics_writer=metrics_writer,
-    )
-
-    inference_wrapped_model = GPTInferenceWrapper(model, args, inference_context, pg_collection=pg_collection)
-
-    inference_wrapped_model.model_is_pipeline_parallel = not (
-        is_pp_first_stage(pg_collection.pp) and is_pp_last_stage(pg_collection.pp)
-    )
-
-    pp_group = getattr(pg_collection, "pp", None)
-    text_generation_controller = SimpleTextGenerationController(
-        inference_wrapped_model=inference_wrapped_model,
-        tokenizer=tokenizer,
-        pp_group=pp_group,
-    )
-
-    return DynamicInferenceEngine(
-        controller=text_generation_controller,
-        context=inference_context,
-        random_seed=args.seed,
-        track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events,
-        enable_chunked_prefill=not args.disable_chunked_prefill,
-        inference_logging_step_interval=inference_logging_step_interval,
-        pg_collection=pg_collection,
+    context = DynamicInferenceContext.from_model_and_args(
+        model, args, overrides={"metrics_writer": metrics_writer}
     )
+    controller = TextGenerationController.from_model_and_args(model, args, context)
+    engine = DynamicInferenceEngine.from_model_and_args(model, args, controller, context)
+    return engine
 
 
 class MegatronLocal(InferenceServer, ReturnsTokens, ReturnsRaw):
@@ -263,10 +181,12 @@ async def launch(cls, model: GPTModel, **kwargs):
             args, model, inference_logging_step_interval, metrics_writer
         )
         await inference_engine.start_listening_to_data_parallel_coordinator(
-            inference_coordinator_port=41521, launch_inference_coordinator=True,
+            inference_coordinator_port=41521, launch_inference_coordinator=True
         )
         if dist.get_rank() == 0:
-            # TODO: We have to do this only on the rank 0 process, should be fixed in the future when we have support for multiple inference clients. !2278
+            # TODO: We have to do this only on the rank 0 process,
+            # should be fixed in the future when we have support for
+            # multiple inference clients. !2278
             client = InferenceClient(inference_coordinator_port=41521)
             await client.start()
         else:
diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py
index 3b3ee09607d..996a419c2a8 100644
--- a/tests/unit_tests/inference/contexts/test_dynamic_context.py
+++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py
@@ -18,6 +18,7 @@
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+from megatron.core.transformer.transformer_config import TransformerConfig
 from tests.unit_tests.test_utilities import Utils
 
 
@@ -68,10 +69,12 @@ def _get_dynamic_context(
             mamba_inference_state_config = None
 
         dynamic_context = DynamicInferenceContext(
-            params_dtype=params_dtype,
-            num_layers=num_layers // self.pp_size,
-            kv_channels=kv_channels,
-            num_attention_heads=num_attention_heads,
+            model_config=TransformerConfig(
+                params_dtype=params_dtype,
+                num_layers=num_layers,
+                kv_channels=kv_channels,
+                num_attention_heads=num_attention_heads,
+            ),
             max_sequence_length=max_sequence_length,
             num_cuda_graphs=None,
             use_cuda_graphs_for_non_decode_steps=True,
@@ -1211,22 +1214,36 @@ def test_pipeline_parallel_uneven_layers(self):
 
         rank = parallel_state.get_pipeline_model_parallel_rank()
 
+        mamba_conv_states_shape = (544, 4)
+        mamba_ssm_states_shape = (8, 64, 16)
+
         if rank == 0:
-            local_num_layers = 12
+            mamba_inference_state_config = MambaInferenceStateConfig(
+                [Symbols.MAMBA] + [Symbols.ATTENTION] * 4,
+                mamba_conv_states_shape,
+                mamba_ssm_states_shape,
+            )
         else:
-            local_num_layers = 4
+            mamba_inference_state_config = MambaInferenceStateConfig(
+                [Symbols.MAMBA] * 4 + [Symbols.ATTENTION],
+                mamba_conv_states_shape,
+                mamba_ssm_states_shape,
+            )
 
         context = DynamicInferenceContext(
-            params_dtype=torch.float32,
-            num_layers=local_num_layers,
-            kv_channels=64,
-            num_attention_heads=8,
+            model_config=TransformerConfig(
+                params_dtype=torch.float32,
+                num_layers=10,
+                kv_channels=64,
+                num_attention_heads=8,
+                pipeline_model_parallel_size=pp_size,
+                tensor_model_parallel_size=1,
+                pipeline_dtype=torch.float32,
+            ),
             max_sequence_length=128,
             buffer_size_gb=0.1,
             block_size_tokens=16,
             max_tokens=1024,
-            pipeline_model_parallel_size=pp_size,
-            tensor_model_parallel_size=1,
             unified_memory_level=0,
         )
 
diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py
index d7c74b24036..7d163e7c94b 100644
--- a/tests/unit_tests/inference/engines/test_dynamic_engine.py
+++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py
@@ -28,9 +28,6 @@
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
     GPTInferenceWrapper,
 )
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.text_generation_controllers.text_generation_controller import (
     TextGenerationController,
@@ -217,11 +214,7 @@ def _build_inference_context(
 
         # Inference context.
         context = DynamicInferenceContext(
-            params_dtype=transformer_config.params_dtype,
-            num_layers=transformer_config.num_layers
-            // transformer_config.pipeline_model_parallel_size,
-            kv_channels=transformer_config.kv_channels,
-            num_attention_heads=transformer_config.num_query_groups,
+            model_config=transformer_config,
             max_sequence_length=test_config.max_sequence_length,
             num_cuda_graphs=test_config.num_cuda_graphs,
             use_cuda_graphs_for_non_decode_steps=True,
@@ -229,8 +222,6 @@ def _build_inference_context(
             block_size_tokens=test_config.context_block_size_tokens,
             max_requests=test_config.context_max_requests,
             max_tokens=test_config.context_max_tokens,
-            tensor_model_parallel_size=transformer_config.tensor_model_parallel_size,
-            pipeline_model_parallel_size=transformer_config.pipeline_model_parallel_size,
             mamba_inference_state_config=mamba_inference_state_config,
             materialize_only_last_token_logits=test_config.materialize_only_last_token_logits,
             use_flashinfer_fused_rope=None,  # default to using flash-infer if available
@@ -377,16 +368,6 @@ def _build_test_env(cls, test_config):
 
         mamba_inference_state_config = get_mamba_inference_state_config_from_model(model)
 
-        # Inference config.
-        inference_config = InferenceWrapperConfig(
-            hidden_size=transformer_config.hidden_size,
-            inference_batch_times_seqlen_threshold=400,
-            fp32_residual_connection=False,
-            params_dtype=transformer_config.params_dtype,
-            fp8=transformer_config.fp8,
-            padded_vocab_size=test_config.vocab_size,
-        )
-
         # Inference context.
         inference_context = cls._build_inference_context(
             test_config=test_config,
@@ -396,7 +377,7 @@ def _build_test_env(cls, test_config):
         )
 
         # Inference model wrapper.
-        inference_wrapped_model = GPTInferenceWrapper(model, inference_config, inference_context)
+        inference_wrapped_model = GPTInferenceWrapper(model, inference_context)
 
         # Note: the following is taken from AbstractModelInferenceWrapper.prep_model_for_inference().
         inference_wrapped_model.model_is_pipeline_parallel = not (
diff --git a/tests/unit_tests/inference/engines/test_static_engine.py b/tests/unit_tests/inference/engines/test_static_engine.py
index 03b3712e39a..c7d7f223d6e 100644
--- a/tests/unit_tests/inference/engines/test_static_engine.py
+++ b/tests/unit_tests/inference/engines/test_static_engine.py
@@ -20,9 +20,6 @@
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
     GPTInferenceWrapper,
 )
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.text_generation_controllers.text_generation_controller import (
     TextGenerationController,
@@ -85,20 +82,11 @@ def setup_engine(
         ).cuda()
         gpt_model.to(inference_config_params_dtype)
 
-        inference_wrapper_config = InferenceWrapperConfig(
-            hidden_size=self.hidden_size,
-            inference_batch_times_seqlen_threshold=400,
-            inference_max_requests=self.batch_size,
-            fp32_residual_connection=False,
-            params_dtype=inference_config_params_dtype,
-            padded_vocab_size=self.vocab_size,
+        inference_context = StaticInferenceContext(
+            max_batch_size=self.batch_size, max_sequence_length=self.sequence_length
         )
 
-        inference_context = StaticInferenceContext.from_config(inference_wrapper_config)
-
-        inference_wrapped_model = GPTInferenceWrapper(
-            gpt_model, inference_wrapper_config, inference_context
-        )
+        inference_wrapped_model = GPTInferenceWrapper(gpt_model, inference_context)
         self.mock_tokenizer = mock.Mock()
         # Set required tokenizer attributes before engine creation
         self.mock_tokenizer.vocab_size = self.vocab_size
diff --git a/tests/unit_tests/inference/model_inference_wrappers/gpt/test_gpt_inference_wrapper.py b/tests/unit_tests/inference/model_inference_wrappers/gpt/test_gpt_inference_wrapper.py
index 07afebe1067..086fc45118c 100644
--- a/tests/unit_tests/inference/model_inference_wrappers/gpt/test_gpt_inference_wrapper.py
+++ b/tests/unit_tests/inference/model_inference_wrappers/gpt/test_gpt_inference_wrapper.py
@@ -10,9 +10,6 @@
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
     GPTInferenceWrapper,
 )
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.models.gpt.gpt_layer_specs import (
     get_gpt_layer_local_spec,
     get_gpt_layer_with_transformer_engine_spec,
@@ -53,27 +50,15 @@ def setup_model(self, tensor_parallel_size, pipeline_parallel_size):
             post_process=parallel_state.is_pipeline_last_stage(),
         ).cuda()
 
-        inference_wrapper_config = InferenceWrapperConfig(
-            hidden_size=hidden_size,
-            inference_batch_times_seqlen_threshold=20,
-            inference_max_requests=self.batch_size,
-            fp32_residual_connection=False,
-            params_dtype=torch.float,
-            padded_vocab_size=self.vocab_size,
-        )
+        inference_context = StaticInferenceContext(self.batch_size, self.sequence_length)
 
-        inference_context = StaticInferenceContext.from_config(inference_wrapper_config)
-
-        self.inference_wrapped_model = GPTInferenceWrapper(
-            gpt_model, inference_wrapper_config, inference_context
-        )
+        self.inference_wrapped_model = GPTInferenceWrapper(gpt_model, inference_context)
 
     def teardown_method(self, method):
         Utils.destroy_model_parallel()
 
-    # This will call the inference_wrapped_model.forward_pass_with_pipeline_parallel_small_input_batch()
     @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
-    def test_inference_pipeline_parallel_small_size(self, materialize_only_last_token_logits):
+    def test_inference_pipeline_parallel(self, materialize_only_last_token_logits):
         self.setup_model(tensor_parallel_size=2, pipeline_parallel_size=2)
 
         batch_prompt_tokens = (
@@ -107,42 +92,6 @@ def test_inference_pipeline_parallel_small_size(self, materialize_only_last_toke
                 self.vocab_size,
             ), f"Shape mismatch . Expected {(self.batch_size, logits_seq_len, self.vocab_size)}, but got {logits.shape}"
 
-    # This will call the inference_wrapped_model.forward_pass_with_pipeline_parallel_large_input_batch()
-    @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
-    def test_inference_pipeline_parallel_large_size(self, materialize_only_last_token_logits):
-        self.setup_model(tensor_parallel_size=2, pipeline_parallel_size=2)
-
-        batch_prompt_tokens = (
-            torch.randint(low=0, high=self.vocab_size, size=(self.batch_size, self.sequence_length))
-            .int()
-            .cuda()
-        )
-        self.inference_wrapped_model.prep_model_for_inference()
-        self.inference_wrapped_model.inference_context.materialize_only_last_token_logits = (
-            materialize_only_last_token_logits
-        )
-
-        inference_input = self.inference_wrapped_model.prep_inference_input(
-            prompts_tokens=batch_prompt_tokens
-        )
-
-        inference_input_for_context_window = (
-            self.inference_wrapped_model.get_batch_for_context_window(inference_input, 0, 10)
-        )
-
-        logits_seq_len = 1 if materialize_only_last_token_logits else 10
-
-        logits = self.inference_wrapped_model.run_one_forward_step(
-            inference_input_for_context_window
-        )
-
-        if parallel_state.is_pipeline_last_stage():
-            assert logits.shape == (
-                self.batch_size,
-                logits_seq_len,
-                self.vocab_size,
-            ), f"Shape mismatch . Expected {(self.batch_size, logits_seq_len, self.vocab_size)}, but got {logits.shape}"
-
     @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False])
     def test_inference_only_tensor_parallel(self, materialize_only_last_token_logits):
         self.setup_model(tensor_parallel_size=4, pipeline_parallel_size=1)
diff --git a/tests/unit_tests/inference/model_inference_wrappers/t5/test_t5_inference_wrapper.py b/tests/unit_tests/inference/model_inference_wrappers/t5/test_t5_inference_wrapper.py
index 36d5187b5eb..1c167f1a98e 100644
--- a/tests/unit_tests/inference/model_inference_wrappers/t5/test_t5_inference_wrapper.py
+++ b/tests/unit_tests/inference/model_inference_wrappers/t5/test_t5_inference_wrapper.py
@@ -7,9 +7,6 @@
 
 from megatron.core import parallel_state
 from megatron.core.inference.contexts import StaticInferenceContext
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.inference.model_inference_wrappers.t5.t5_inference_wrapper import (
     T5InferenceWrapper,
 )
@@ -77,19 +74,9 @@ def setup_model(self, tensor_parallel_size, pipeline_parallel_size):
             add_decoder=True,
         ).cuda()
 
-        inference_wrapper_config = InferenceWrapperConfig(
-            hidden_size=hidden_size,
-            inference_batch_times_seqlen_threshold=-1,
-            fp32_residual_connection=False,
-            params_dtype=torch.float,
-            padded_vocab_size=self.vocab_size,
-        )
+        inference_context = StaticInferenceContext(max_batch_size=8, max_sequence_length=2560)
 
-        inference_context = StaticInferenceContext.from_config(inference_wrapper_config)
-
-        self.inference_wrapped_model = T5InferenceWrapper(
-            t5_model, inference_wrapper_config, inference_context
-        )
+        self.inference_wrapped_model = T5InferenceWrapper(t5_model, inference_context)
 
     def teardown_method(self, method):
         Utils.destroy_model_parallel()
diff --git a/tests/unit_tests/inference/model_inference_wrappers/test_model_inference_wrapper_config.py b/tests/unit_tests/inference/model_inference_wrappers/test_model_inference_wrapper_config.py
deleted file mode 100644
index 794634760d0..00000000000
--- a/tests/unit_tests/inference/model_inference_wrappers/test_model_inference_wrapper_config.py
+++ /dev/null
@@ -1,21 +0,0 @@
-import torch
-
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
-
-
-class TestModelInferenceWrapperConfig:
-
-    def test_inference_config(self):
-        inference_config = InferenceWrapperConfig(
-            hidden_size=10,
-            inference_batch_times_seqlen_threshold=10,
-            padded_vocab_size=10,
-            params_dtype=torch.float,
-            fp32_residual_connection=False,
-        )
-        inference_config.add_attributes({"abc": 45})
-        assert (
-            inference_config.abc == 45
-        ), f"min tokens not set correctly. it is {inference_config.min_tokens}"
diff --git a/tests/unit_tests/inference/test_wandb_logging.py b/tests/unit_tests/inference/test_wandb_logging.py
index 1d5d054b80e..4c8088b92e7 100644
--- a/tests/unit_tests/inference/test_wandb_logging.py
+++ b/tests/unit_tests/inference/test_wandb_logging.py
@@ -15,6 +15,7 @@
     TextGenerationController,
 )
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+from megatron.core.transformer.transformer_config import TransformerConfig
 from tests.unit_tests.test_utilities import Utils
 
 
@@ -54,10 +55,12 @@ def _get_dynamic_context(
     ):
         """Helper to create a DynamicInferenceContext."""
         return DynamicInferenceContext(
-            params_dtype=params_dtype,
-            num_layers=num_layers,
-            kv_channels=kv_channels,
-            num_attention_heads=num_attention_heads,
+            model_config=TransformerConfig(
+                params_dtype=params_dtype,
+                num_layers=num_layers,
+                kv_channels=kv_channels,
+                num_attention_heads=num_attention_heads,
+            ),
             max_sequence_length=max_sequence_length,
             num_cuda_graphs=None,
             buffer_size_gb=buffer_size_gb,
@@ -229,10 +232,9 @@ def test_paused_requests_in_stats(self):
         """Test that paused requests are correctly reflected in stats."""
         set_rounder(1)
         dynamic_context = DynamicInferenceContext(
-            params_dtype=torch.float32,
-            num_layers=2,
-            kv_channels=64,
-            num_attention_heads=8,
+            model_config=TransformerConfig(
+                params_dtype=torch.float32, num_layers=2, kv_channels=64, num_attention_heads=8
+            ),
             max_sequence_length=128,
             num_cuda_graphs=None,
             buffer_size_gb=0.01,  # Small buffer to force pausing
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_encoder_decoder_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_encoder_decoder_text_generation_controller.py
index 93a208710fc..9bf2183c0fd 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_encoder_decoder_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_encoder_decoder_text_generation_controller.py
@@ -12,9 +12,6 @@
 
 from megatron.core.inference.contexts import StaticInferenceContext
 from megatron.core.inference.inference_request import InferenceRequest, Status
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.inference.model_inference_wrappers.t5.t5_inference_wrapper import (
     T5InferenceWrapper,
 )
@@ -85,19 +82,9 @@ def setup_method(self, method):
             add_decoder=True,
         ).cuda()
 
-        inference_wrapper_config = InferenceWrapperConfig(
-            hidden_size=hidden_size,
-            inference_batch_times_seqlen_threshold=-1,
-            fp32_residual_connection=False,
-            params_dtype=torch.float,
-            padded_vocab_size=self.vocab_size,
-        )
+        inference_context = StaticInferenceContext(max_batch_size=8, max_sequence_length=2560)
 
-        inference_context = StaticInferenceContext.from_config(inference_wrapper_config)
-
-        inference_wrapped_model = T5InferenceWrapper(
-            t5_model, inference_wrapper_config, inference_context
-        )
+        inference_wrapped_model = T5InferenceWrapper(t5_model, inference_context)
 
         self.mock_tokenizer = mock.Mock()
 
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_simple_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
similarity index 97%
rename from tests/unit_tests/inference/text_generation_controllers/test_simple_text_generation_controller.py
rename to tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index 0885401e7a0..61ea73f823b 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_simple_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -24,9 +24,6 @@
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
     GPTInferenceWrapper,
 )
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.text_generation_controllers.text_generation_controller import (
     TextGenerationController,
@@ -100,26 +97,13 @@ def setup_model(
         if dtype == torch.bfloat16:
             gpt_model = Float16Module(gpt_model.config, gpt_model)
 
-        inference_wrapper_config = InferenceWrapperConfig(
-            hidden_size=self.hidden_size,
-            inference_batch_times_seqlen_threshold=-1,
-            inference_max_seq_length=2048,
-            inference_max_requests=16 if fp8 else self.batch_size,
-            fp32_residual_connection=False,
-            params_dtype=dtype,
-            padded_vocab_size=self.vocab_size,
-        )
-
         if static:
-            inference_context = StaticInferenceContext.from_config(inference_wrapper_config)
+            inference_context = StaticInferenceContext(
+                max_batch_size=16 if fp8 else self.batch_size, max_sequence_length=2048
+            )
         else:
             inference_context = DynamicInferenceContext(
-                params_dtype=dtype,
-                num_layers=transformer_config.num_layers // pipeline_model_parallel_size,
-                kv_channels=transformer_config.kv_channels,
-                num_attention_heads=transformer_config.num_attention_heads,
-                tensor_model_parallel_size=transformer_config.tensor_model_parallel_size,
-                pipeline_model_parallel_size=transformer_config.pipeline_model_parallel_size,
+                model_config=transformer_config,
                 max_sequence_length=2048,
                 buffer_size_gb=0.2,
                 materialize_only_last_token_logits=False,
@@ -128,9 +112,7 @@ def setup_model(
                 unified_memory_level=0,  # unit tests currently broken with UVM
             )
 
-        inference_wrapped_model = GPTInferenceWrapper(
-            gpt_model, inference_wrapper_config, inference_context
-        )
+        inference_wrapped_model = GPTInferenceWrapper(gpt_model, inference_context)
 
         inference_wrapped_model.model_is_pipeline_parallel = not (
             parallel_state.is_pipeline_first_stage() and parallel_state.is_pipeline_last_stage()
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_vlm_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_vlm_text_generation_controller.py
index 31bf415ba56..326e9b35873 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_vlm_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_vlm_text_generation_controller.py
@@ -13,9 +13,6 @@
 
 from megatron.core.inference.contexts import StaticInferenceContext
 from megatron.core.inference.inference_request import InferenceRequest, Status, VLMInferenceRequest
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.inference.model_inference_wrappers.multimodal.vlm_inference_wrapper import (
     VLMInferenceWrapper,
 )
@@ -92,19 +89,9 @@ def setup_method(self, method):
         self.image_token_index = self.model.image_token_index
         self.model = Float16Module(self.model.config, self.model)
 
-        inference_wrapper_config = InferenceWrapperConfig(
-            hidden_size=self.language_hidden_size,
-            inference_batch_times_seqlen_threshold=-1,
-            fp32_residual_connection=False,
-            params_dtype=torch.float,
-            padded_vocab_size=self.language_vocab_size,
-        )
+        inference_context = StaticInferenceContext(max_batch_size=8, max_sequence_length=2560)
 
-        inference_context = StaticInferenceContext.from_config(inference_wrapper_config)
-
-        inference_wrapped_model = VLMInferenceWrapper(
-            self.model, inference_wrapper_config, inference_context
-        )
+        inference_wrapped_model = VLMInferenceWrapper(self.model, inference_context)
 
         self.mock_tokenizer = mock.Mock()
 
diff --git a/tests/unit_tests/models/test_gpt_model_batch_invariant.py b/tests/unit_tests/models/test_gpt_model_batch_invariant.py
index ead9125e5ec..3468ef55b4c 100644
--- a/tests/unit_tests/models/test_gpt_model_batch_invariant.py
+++ b/tests/unit_tests/models/test_gpt_model_batch_invariant.py
@@ -10,12 +10,9 @@
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
     GPTInferenceWrapper,
 )
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.inference.sampling_params import SamplingParams
-from megatron.core.inference.text_generation_controllers.simple_text_generation_controller import (
-    SimpleTextGenerationController,
+from megatron.core.inference.text_generation_controllers.text_generation_controller import (
+    TextGenerationController,
 )
 from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
 from megatron.core.models.gpt.gpt_model import GPTModel
@@ -91,6 +88,8 @@ def _build_flash_attn_bik_model(seq_len: int, vocab_size: int, hidden_size: int
         normalization="RMSNorm",
         params_dtype=torch.bfloat16,
         attention_backend=AttnBackend.flash,
+        fp32_residual_connection=False,
+        nccl_all_reduce_for_prefill=False,
     )
     cfg.fp16 = False
     cfg.bf16 = True
@@ -184,10 +183,7 @@ def test_dynamic_engine_matches_batched_forward_rl(self):
         inference_model = Float16Module(base_model.config, base_model).cuda().eval()
 
         ctx = DynamicInferenceContext(
-            params_dtype=torch.bfloat16,
-            num_layers=base_model.config.num_layers,
-            kv_channels=base_model.config.kv_channels,
-            num_attention_heads=base_model.config.num_attention_heads,
+            model_config=base_model.config,
             max_sequence_length=seq_len,
             buffer_size_gb=0.125,
             block_size_tokens=16,
@@ -197,19 +193,9 @@ def test_dynamic_engine_matches_batched_forward_rl(self):
             unified_memory_level=0,
         )
 
-        wrapper_cfg = InferenceWrapperConfig(
-            hidden_size=base_model.config.hidden_size,
-            inference_batch_times_seqlen_threshold=-1,
-            fp32_residual_connection=False,
-            params_dtype=torch.bfloat16,
-            padded_vocab_size=vocab_size,
-            inference_max_seq_length=seq_len,
-            inference_max_requests=8,
-            nccl_all_reduce_for_prefill=False,
-        )
-        wrapper = GPTInferenceWrapper(inference_model, wrapper_cfg, ctx)
+        wrapper = GPTInferenceWrapper(inference_model, ctx)
         tokenizer = DummyTokenizer(vocab_size=vocab_size, bos=None, eod=vocab_size - 1, pad=0)
-        controller = SimpleTextGenerationController(wrapper, tokenizer)
+        controller = TextGenerationController(wrapper, tokenizer)
         engine = DynamicInferenceEngine(
             controller=controller, context=ctx, enable_cuda_graph=False, random_seed=123
         )
@@ -286,19 +272,9 @@ def _run_engine_with_order(order):
                 unified_memory_level=0,
             )
 
-            wrapper_cfg = InferenceWrapperConfig(
-                hidden_size=base_model.config.hidden_size,
-                inference_batch_times_seqlen_threshold=-1,
-                fp32_residual_connection=False,
-                params_dtype=torch.bfloat16,
-                padded_vocab_size=vocab_size,
-                inference_max_seq_length=seq_len,
-                inference_max_requests=8,
-                nccl_all_reduce_for_prefill=False,
-            )
-            wrapper = GPTInferenceWrapper(inference_model, wrapper_cfg, ctx)
+            wrapper = GPTInferenceWrapper(inference_model, ctx)
             tokenizer = DummyTokenizer(vocab_size=vocab_size, bos=None, eod=vocab_size - 1, pad=0)
-            controller = SimpleTextGenerationController(wrapper, tokenizer)
+            controller = TextGenerationController(wrapper, tokenizer)
             engine = DynamicInferenceEngine(
                 controller=controller, context=ctx, enable_cuda_graph=False, random_seed=123
             )
diff --git a/tests/unit_tests/models/test_mamba_model.py b/tests/unit_tests/models/test_mamba_model.py
index 6c0e4a563cd..ab3f9789f79 100644
--- a/tests/unit_tests/models/test_mamba_model.py
+++ b/tests/unit_tests/models/test_mamba_model.py
@@ -284,10 +284,7 @@ def test_dynamic_inference_padding_with_fp8(self):
         )
 
         inference_context = DynamicInferenceContext(
-            params_dtype=config.params_dtype,
-            num_layers=config.num_layers,
-            kv_channels=config.hidden_size // config.num_attention_heads,
-            num_attention_heads=config.num_attention_heads,
+            model_config=self.model.config,
             max_sequence_length=self.model.module.max_sequence_length,
             buffer_size_gb=1.0,
             block_size_tokens=256,
diff --git a/tools/run_inference_performance_test.py b/tools/run_inference_performance_test.py
index a429a91724d..9f3e38514a3 100644
--- a/tools/run_inference_performance_test.py
+++ b/tools/run_inference_performance_test.py
@@ -10,16 +10,13 @@
 
 from gpt_builders import gpt_builder
 from mamba_builders import mamba_builder
-from megatron.core.inference.contexts import DynamicInferenceContext
+from megatron.core.inference.contexts import DynamicInferenceContext, StaticInferenceContext
 from megatron.core.inference.engines import DynamicInferenceEngine, StaticInferenceEngine
 from megatron.core.inference.engines.abstract_engine import AbstractEngine
 from megatron.core.inference.inference_request import InferenceRequest
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
     GPTInferenceWrapper,
 )
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.text_generation_controllers.text_generation_controller import (
     TextGenerationController,
@@ -77,7 +74,10 @@ def get_inference_engine(args: argparse.Namespace, model: MegatronModule) -> Abs
     tokenizer = get_tokenizer()
 
     if args.engine_type == "static":
-        inference_wrapped_model = GPTInferenceWrapper(model)
+        context = StaticInferenceContext(
+            args.inference_max_requests, args.inference_max_sequence_length
+        )
+        inference_wrapped_model = GPTInferenceWrapper(model, context)
         inference_wrapped_model.model_is_pipeline_parallel = not (
             mpu.is_pipeline_first_stage() and mpu.is_pipeline_last_stage()
         )
@@ -173,9 +173,7 @@ def generate_dynamic(
         request_id = REQUEST_ID
         REQUEST_ID += 1
         prompt_tokens = request.prompt_tokens
-        inference_engine.add_request(
-            request_id, prompt_tokens, request.inference_parameters,
-        )
+        inference_engine.add_request(request_id, prompt_tokens, request.inference_parameters)
 
     start_time = time.perf_counter()
     all_finished_requests = []
@@ -292,9 +290,7 @@ def main():
                 prompts=args.prompts, inference_requests=requests, sampling_params=sampling_params
             )
         elif args.engine_type == "dynamic":
-            results: List[InferenceRequest] = generate_dynamic(
-                args, requests, inference_engine,
-            )
+            results: List[InferenceRequest] = generate_dynamic(args, requests, inference_engine)
     end_time = time.perf_counter()
     latency = end_time - start_time
 
diff --git a/tools/run_text_generation_server.py b/tools/run_text_generation_server.py
index 350173dc16f..89c1cfa5b86 100644
--- a/tools/run_text_generation_server.py
+++ b/tools/run_text_generation_server.py
@@ -22,9 +22,6 @@
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
     GPTInferenceWrapper,
 )
-from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
-    InferenceWrapperConfig,
-)
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.text_generation_controllers.text_generation_controller import (
     TextGenerationController,
@@ -63,27 +60,15 @@ def get_inference_engine(args: Namespace, model: MegatronModule) -> AbstractEngi
 
     tokenizer = get_tokenizer()
 
-    inference_wrapper_config = InferenceWrapperConfig(
-        hidden_size=args.hidden_size,
-        inference_batch_times_seqlen_threshold=args.inference_batch_times_seqlen_threshold,
-        fp32_residual_connection=args.fp32_residual_connection,
-        params_dtype=args.params_dtype,
-        padded_vocab_size=args.padded_vocab_size,
-        inference_max_seq_length=args.inference_max_seq_length,
-        inference_max_requests=args.inference_max_batch_size,
-        nccl_all_reduce_for_prefill=args.nccl_all_reduce_for_prefill,
-        moe_pad_experts_for_cuda_graph_inference = args.moe_pad_experts_for_cuda_graph_inference
-    )
-    inference_context = StaticInferenceContext.from_config(inference_wrapper_config)
+    inference_context = StaticInferenceContext(args.inference_max_requests, args.inference_max_sequence_length)
     inference_wrapped_model = GPTInferenceWrapper(
-        model, inference_wrapper_config, inference_context
+        model, inference_context
     )
     text_generation_controller = TextGenerationController(
         inference_wrapped_model=inference_wrapped_model, tokenizer=tokenizer
     )
     return StaticInferenceEngine(
         text_generation_controller=text_generation_controller,
-        max_batch_size=args.inference_max_batch_size,
     )
 
 
@@ -166,14 +151,6 @@ def main(model_type: str = "gpt"):
     model = model[0]
     model.eval()
 
-    if args.max_batch_size is not None:
-        assert args.inference_max_batch_size is not None
-        args.inference_max_batch_size = max(args.inference_max_batch_size, args.max_batch_size)
-        warnings.warn(
-            "`--max-batch-size` has been deprecated in favor of `--inference-max-requests`, "
-            f"setting maximum batch size to {args.inference_max_batch_size}"
-        )
-
     inference_engine = get_inference_engine(args, model)
 
     if args.cuda_graph_impl == "local":

From 01a82378dec60d715f6b3b050a6b9216ee9ee9f5 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Thu, 15 Jan 2026 02:38:05 -0800
Subject: [PATCH 03/30] Undo attention.py change

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 megatron/core/transformer/attention.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py
index 2a5b2eab0f1..8265ee83ff5 100644
--- a/megatron/core/transformer/attention.py
+++ b/megatron/core/transformer/attention.py
@@ -608,7 +608,7 @@ def flash_decode_and_prefill(
                     k_new=None,
                     v_new=None,
                     qv=None,
-                    out_=None,
+                    out=None,
                     cu_seqlens_q=cu_seqlens_q,
                     cu_seqlens_k=None,
                     cu_seqlens_k_new=None,
@@ -627,8 +627,7 @@ def flash_decode_and_prefill(
                     v_descale=None,
                     softmax_scale=softmax_scale,
                     causal=True,
-                    window_size_left=-1,
-                    window_size_right=-1,
+                    window_size=(-1, -1),
                     attention_chunk=0,
                     softcap=0.0,
                     rotary_interleaved=True,

From 6da70fb1324815660fd5082632b95b1298f99360 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Thu, 15 Jan 2026 08:33:54 -0800
Subject: [PATCH 04/30] Fix unit tests

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 tests/unit_tests/models/test_gpt_model.py              | 10 ++++++----
 .../models/test_gpt_model_batch_invariant.py           |  5 +----
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/tests/unit_tests/models/test_gpt_model.py b/tests/unit_tests/models/test_gpt_model.py
index cf3bd40ee4b..8c8b1be638f 100644
--- a/tests/unit_tests/models/test_gpt_model.py
+++ b/tests/unit_tests/models/test_gpt_model.py
@@ -392,10 +392,12 @@ def test_dynamic_inference_padding_with_fp8(self):
         config = self.gpt_model.config
 
         inference_context = DynamicInferenceContext(
-            params_dtype=config.params_dtype,
-            num_layers=config.num_layers,
-            kv_channels=config.hidden_size // config.num_attention_heads,
-            num_attention_heads=config.num_attention_heads,
+            model_config=TransformerConfig(
+                params_dtype=config.params_dtype,
+                num_layers=config.num_layers,
+                kv_channels=config.hidden_size // config.num_attention_heads,
+                num_attention_heads=config.num_attention_heads,
+            ),
             max_sequence_length=self.gpt_model.module.max_sequence_length,
             buffer_size_gb=1.0,
             block_size_tokens=256,
diff --git a/tests/unit_tests/models/test_gpt_model_batch_invariant.py b/tests/unit_tests/models/test_gpt_model_batch_invariant.py
index 3468ef55b4c..9284e3a6acc 100644
--- a/tests/unit_tests/models/test_gpt_model_batch_invariant.py
+++ b/tests/unit_tests/models/test_gpt_model_batch_invariant.py
@@ -259,10 +259,7 @@ def test_dynamic_engine_is_batch_invariant(self):
 
         def _run_engine_with_order(order):
             ctx = DynamicInferenceContext(
-                params_dtype=torch.bfloat16,
-                num_layers=base_model.config.num_layers,
-                kv_channels=base_model.config.kv_channels,
-                num_attention_heads=base_model.config.num_attention_heads,
+                model_config=based_model.config,
                 max_sequence_length=seq_len,
                 buffer_size_gb=0.125,
                 block_size_tokens=16,

From 2f67856b3c0a21b4708e64aa05d7bc4eaa033407 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Tue, 20 Jan 2026 15:54:59 -0800
Subject: [PATCH 05/30] Fix typo and clean up run_inference_performance_test.py

---
 megatron/rl/inference/megatron.py       |  2 +-
 tools/run_inference_performance_test.py | 78 +++++--------------------
 2 files changed, 14 insertions(+), 66 deletions(-)

diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py
index 6e6053d1ca6..8d94c7f8e3b 100644
--- a/megatron/rl/inference/megatron.py
+++ b/megatron/rl/inference/megatron.py
@@ -64,7 +64,7 @@ def get_static_inference_engine(args: Namespace, model: MegatronModule) -> Abstr
     return MCoreEngine(
         text_generation_controller=text_generation_controller,
         max_batch_size=(
-            args.inference_max_reqeusts if args.inference_max_requests is not None else 1
+            args.inference_max_requests if args.inference_max_requests is not None else 1
         ),
     )
 
diff --git a/tools/run_inference_performance_test.py b/tools/run_inference_performance_test.py
index 9f3e38514a3..ff08cd13dbd 100644
--- a/tools/run_inference_performance_test.py
+++ b/tools/run_inference_performance_test.py
@@ -57,7 +57,6 @@ def add_inference_benchmarking_args(parser):
     group.add_argument(
         "--benchmark-profile", action="store_true", default=False, help="If set, profile"
     )
-    group.add_argument('--stream', action="store_true", default=False, help="If set, stream tokens")
     return parser
 
 
@@ -88,53 +87,6 @@ def get_inference_engine(args: argparse.Namespace, model: MegatronModule) -> Abs
     elif args.engine_type == "dynamic":
         return DynamicInferenceEngine.from_model_and_args(model, args)
 
-
-async def generate(
-    inference_engine: Union[StaticInferenceEngine, DynamicInferenceEngine],
-    sampling_params: SamplingParams,
-    prompts: List[str],
-    inference_requests: List[InferenceRequest] = None,
-) -> List[InferenceRequest]:
-    async def collect_stream(prompt, request_id, stream_generator):
-        async for output in stream_generator:
-            pass
-
-    if inference_requests is None:
-        assert prompts is not None
-        inference_requests = [None for _ in range(len(prompts))]
-    elif prompts is None:
-        assert inference_requests is not None
-        tokenizer = get_tokenizer()
-        prompts = [tokenizer.detokenize(request.prompt_tokens) for request in inference_requests]
-
-    request_ids: List[int] = [
-        inference_engine.add_request(
-            prompt=prompt,
-            inference_request=inference_request,
-            inference_parameters=sampling_params,
-            streaming=True,
-        )
-        for prompt, inference_request in zip(prompts, inference_requests)
-    ]
-    stream_generators = [
-        inference_engine.get_stream_generator(request_id) for request_id in request_ids
-    ]
-
-    tasks = [
-        asyncio.create_task(collect_stream(prompt, request_id, stream_generator))
-        for (prompt, request_id, stream_generator) in zip(prompts, request_ids, stream_generators)
-    ]
-
-    await inference_engine.run_engine_async()
-    await asyncio.gather(*tasks)
-
-    results: List[InferenceRequest] = [
-        inference_engine.scheduler.completed_request_pool[request_id] for request_id in request_ids
-    ]
-
-    return results
-
-
 def get_random_prompt_tokens(tokenizer, num_input_tokens) -> List[int]:
     # Get the set of special token IDs to exclude
     special_token_ids = set()
@@ -178,7 +130,7 @@ def generate_dynamic(
     start_time = time.perf_counter()
     all_finished_requests = []
     while inference_engine.has_unfinished_requests():
-        result = inference_engine.step(verbose=False)
+        result = inference_engine.step()
         finished_requests = result["finished_requests"]
         for request in finished_requests:
             req_id = request.request_id
@@ -243,7 +195,7 @@ def main():
     requests = []
     if args.num_input_tokens is not None:
         assert args.prompts is None
-        batch_size = args.inference_max_batch_size
+        batch_size = args.inference_max_requests
         for i in range(batch_size):
             prompt_tokens = get_random_prompt_tokens(tokenizer, args.num_input_tokens)
             requests.append(
@@ -266,7 +218,7 @@ def main():
                 )
             )
 
-    if args.cuda_graph_impl == "local":
+    if args.cuda_graph_impl == "local" and args.engine_type == "static":
         print(f"Running warmup for CUDA graphs...")
         warmup_sampling_params = SamplingParams(num_tokens_to_generate=10)
         warmup_sampling_params.add_attributes({"no_early_termination": True})
@@ -276,21 +228,16 @@ def main():
         torch.cuda.cudart().cudaProfilerStart()
 
     start_time = time.perf_counter()
-    if args.stream:
-        if args.engine_type == "dynamic":
-            raise NotImplementedError("Streaming not supported with DynamicInferenceEngine")
-        results: List[InferenceRequest] = asyncio.run(
-            generate(
-                inference_engine, sampling_params, prompts=args.prompts, inference_requests=requests
-            )
+    if args.engine_type == "static":
+        results: List[InferenceRequest] = inference_engine.generate(
+            prompts=args.prompts, inference_requests=requests, sampling_params=sampling_params
         )
     else:
-        if args.engine_type == "static":
-            results: List[InferenceRequest] = inference_engine.generate(
-                prompts=args.prompts, inference_requests=requests, sampling_params=sampling_params
-            )
-        elif args.engine_type == "dynamic":
-            results: List[InferenceRequest] = generate_dynamic(args, requests, inference_engine)
+        prompts = [request.prompt_tokens for request in requests]
+        results: List[InferenceRequest] = inference_engine.generate(
+            prompts=prompts, sampling_params=sampling_params
+        )
+
     end_time = time.perf_counter()
     latency = end_time - start_time
 
@@ -300,7 +247,8 @@ def main():
         torch.cuda.cudart().cudaProfilerStop()
 
     if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
-        for idx, result in enumerate(results):
+        for idx, record in enumerate(results):
+            result = record.requests[0]
             print(f' \n------------- RESULT FOR PROMPT {idx} --------------- ')
             generated_log_probs = result.generated_log_probs
             result_dict = {

From b801c9c13e16f1ce91ee22796cc0915e5311767c Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Wed, 21 Jan 2026 12:59:43 -0800
Subject: [PATCH 06/30] Address reviewer feedback

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 .../inference/gpt/gpt_dynamic_inference.py    | 15 +-----
 .../inference/contexts/dynamic_context.py     | 17 +------
 .../core/inference/engines/dynamic_engine.py  | 47 +++++++++++++----
 .../text_generation_controller.py             | 19 +++----
 megatron/rl/inference/megatron.py             | 51 +------------------
 5 files changed, 49 insertions(+), 100 deletions(-)

diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py
index 245e7be5aa7..28f6a35a227 100644
--- a/examples/inference/gpt/gpt_dynamic_inference.py
+++ b/examples/inference/gpt/gpt_dynamic_inference.py
@@ -424,13 +424,7 @@ def escape_str(s):
             # ---- Prompt summary line ----
             prompt_len = len(requests[request_idxs[0]].prompt_tokens)
             escaped_prompt_text = escape_str(prompt_text)
-<<<<<<< HEAD
-            print(
-                f"{unique_idx+1}/{len(unique_prompt_map)} [n {len(request_idxs)}, l {prompt_len}] {escaped_prompt_text}"
-            )
-=======
             print(f"\n{unique_idx+1}/{len(unique_prompt_map)} [n {len(request_idxs)}, l {prompt_len}] {escaped_prompt_text}")
->>>>>>> upstream/main
 
             # ---- Group all outputs for this prompt ----
             output_map = defaultdict(list)
@@ -457,13 +451,8 @@ def escape_str(s):
                     o_hash = "--"
                     o_len = 0
                     escaped_output_text = "--"
-<<<<<<< HEAD
-                print(
-                    f"  >>>> [n {len(output_request_idxs)}, {o_len} tokens, hash {o_hash}] {escaped_output_text}"
-                )
-=======
-                print(f"  >>>> [n {len(output_request_idxs)}, {o_len} tokens, hash {o_hash}{', <evicted>' if evicted else ''}] {escaped_output_text}")
->>>>>>> upstream/main
+                print(f"  >>>> [n {len(output_request_idxs)}, {o_len} tokens, hash {o_hash}"
+                f"{', <evicted>' if evicted else ''}] {escaped_output_text}")
                 text_hashes.append(o_hash)
 
         # Write results to JSON. Primarily used for functional testing.
diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 03154a311eb..0502a0bf829 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -4,7 +4,7 @@
 import math
 import warnings
 from contextlib import nullcontext
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple
+from typing import Any, Dict, List, Optional, Sequence, Tuple
 
 import torch
 import torch.nn.functional as F
@@ -60,17 +60,6 @@
 except ImportError:
     HAVE_FLASHINFER = False
 
-try:
-    import wandb  # pylint: disable=unused-import
-
-    HAVE_WANDB = True
-except ImportError:
-    HAVE_WANDB = False
-    wandb = None
-
-if TYPE_CHECKING:
-    import wandb as WandbModule
-
 
 class ContextOverflowError(Exception):
     """Base exception for when a new request does not fit.
@@ -244,7 +233,6 @@ class DynamicInferenceContext(BaseInferenceContext):
             levels will be included to control other tensors within the context.
         use_flashinfer_fused_rope (bool): If True, use flashinfer's fused rope implementation.
             If None, defaults to using flash-infer if available.
-        metrics_writer (Optional['WandbModule']): Wandb module for writing metrics.
         request_metadata_types (Optional[List[Tuple[str, torch.dtype, bool]]]): A list of the
             per-request metadata types to track. Each entry is a tuple consisting of the string
             label, the target dtype, and whether to store the data on GPU.
@@ -273,7 +261,6 @@ def __init__(
         unified_memory_level: Optional[int] = 0,
         cuda_graph_max_tokens: Optional[int] = None,
         cuda_graph_mixed_prefill_count: Optional[int] = 16,
-        metrics_writer: Optional['WandbModule'] = None,
         request_metadata_types: Optional[List[Tuple[str, torch.dtype, bool]]] = None,
         persist_cuda_graphs: Optional[bool] = False,
     ):
@@ -296,8 +283,6 @@ def __init__(
                 DeprecationWarning,
             )
 
-        self.metrics_writer = metrics_writer
-
         # Per partition num heads and hidden size.
         num_attention_heads = model_config.num_query_groups or model_config.num_attention_heads
         projection_size = model_config.kv_channels * num_attention_heads
diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index f8f14bf27a8..e56f73999ea 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -48,6 +48,7 @@
     get_pg_size,
     get_pg_src_rank,
     internal_api,
+    log_single_rank,
     trace_async_exceptions,
 )
 
@@ -146,6 +147,7 @@ def __init__(
         *,
         track_paused_request_events: bool = False,
         enable_chunked_prefill: bool = True,
+        metrics_writer: Optional['WandbModule'] = None,
         inference_logging_step_interval: int = 0,
         pg_collection: Optional[ProcessGroupCollection] = None,
     ):
@@ -182,6 +184,7 @@ def __init__(
         self.random_seed = random_seed
         self.track_paused_request_events = track_paused_request_events
         self.enable_chunked_prefill = enable_chunked_prefill
+        self.metrics_writer = metrics_writer
         self.inference_logging_step_interval = inference_logging_step_interval
         self.unified_memory_level = context.unified_memory_level
         self.persist_cuda_graphs = context.persist_cuda_graphs
@@ -200,12 +203,12 @@ def __init__(
         )
 
         # Configure wandb to use separate step counter for inference metrics (only once)
-        if self.inference_logging_step_interval > 0 and self.context.metrics_writer is not None:
+        if self.inference_logging_step_interval > 0 and self.metrics_writer is not None:
             logging.info(
                 f"\033[1;93m[INFERENCE]\033[0m "
                 f"\033[1;95mLogging inference metrics to wandb (rank {self.rank})\033[0m"
             )
-            if HAVE_WANDB and self.context.metrics_writer.__name__ == "wandb":
+            if HAVE_WANDB and self.metrics_writer.__name__ == "wandb":
                 # Make all inference/* metrics use inference_step as their x-axis
                 # This allows inference and training to have independent step counters
                 context.metrics_writer.define_metric(
@@ -254,14 +257,40 @@ def from_model_and_args(
         if controller is None:
             controller = TextGenerationController.from_model_and_args(model, args, context)
 
+        # The model may have a custom ProcessGroupCollection with a different TP / PP size.
+        pg_collection = get_attr_wrapped_model(model, "pg_collection")
+
+        # Get inference logging configuration from args
+        log_inference_wandb = args.inference_wandb_logging
+        inference_logging_step_interval = args.inference_logging_step_interval
+
+        # Get metrics writer if logging is enabled and on the logging rank
+        # Use the same rank convention as training (last rank logs)
+        metrics_writer = None
+        if (
+            inference_logging_step_interval > 0
+            and log_inference_wandb
+            and args.rank == (args.world_size - 1)
+        ):
+            metrics_writer = get_wandb_writer()
+            if metrics_writer is None:
+                log_single_rank(
+                    logger,
+                    logging.WARNING,
+                    "WARNING: --rl-inference-logging-step-interval is set but no metrics writer "
+                    "wandb module is available. Inference logging will be disabled.",
+                )
+
         return cls(
             controller,
             context,
-            enable_cuda_graph=args.cuda_graph_impl == "local",
-            random_seed=args.seed,
+            enable_cuda_graph=model.config.cuda_graph_impl == "local",
+            random_seed=amodel.config.seed,
             track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events,
             enable_chunked_prefill=not args.disable_chunked_prefill,
+            metrics_writer=metrics_writer,
             inference_logging_step_interval=args.inference_logging_step_interval,
+            pg_collection=pg_collection,
         )
 
     def reset(self) -> None:
@@ -1187,7 +1216,7 @@ async def async_forward(self) -> Tuple[Dict, Dict, float, int]:
             self.inference_logging_step_interval > 0
             and self.step_count > 0
             and self.step_count % self.inference_logging_step_interval == 0
-            and self.context.metrics_writer is not None
+            and self.metrics_writer is not None
         ):
             kvcache_util_stats = self.context.get_kvcache_utilization_stats()
         else:
@@ -1312,12 +1341,10 @@ async def async_bookkeep(
                 else:
                     metrics[f'inference/{key}'] = value
 
-            if HAVE_WANDB and self.context.metrics_writer.__name__ == "wandb":
-                self.context.metrics_writer.log(metrics, commit=True)
+            if HAVE_WANDB and self.metrics_writer.__name__ == "wandb":
+                self.metrics_writer.log(metrics, commit=True)
             else:
-                raise ValueError(
-                    f"Unsupported metrics writer type: {type(self.context.metrics_writer)}"
-                )
+                raise ValueError(f"Unsupported metrics writer type: {type(self.metrics_writer)}")
 
         # Print context state.
         if (
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index 0212aa01274..e9eefa573c0 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -13,7 +13,6 @@
 from torch import Tensor
 from torch.distributed import ProcessGroup
 
-from megatron.core import parallel_state
 from megatron.core.inference.async_stream import AsyncStream
 from megatron.core.inference.communication_utils import (
     broadcast_from_last_pipeline_stage,
@@ -32,6 +31,7 @@
 )
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.utils import get_attention_mask, set_decode_expert_padding
+from megatron.core.models.multimodal.llava_model import LLaVAModel
 from megatron.core.tokenizers.text.utils.build_tokenizer import build_tokenizer
 from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.transformer.moe.moe_layer import BaseMoELayer
@@ -85,15 +85,15 @@ def __init__(
         )
 
         # Use padded vocab size because tokenizer vocab size might pad to nearest power of 2.
-        try:
-            self.vocab_size = get_attr_wrapped_model(
-                self.inference_wrapped_model.model, "vocab_size"
-            )
-        except RuntimeError as e:
-            # Handle LlaVa models
+        if isinstance(self.inference_wrapped_model.model, LLaVAModel):
+            # TODO(ksanthanam): Consider deprecating this check if LLaVAModel is no longer used
             self.vocab_size = get_attr_wrapped_model(
                 self.inference_wrapped_model.model, "language_model"
             ).vocab_size
+        else:
+            self.vocab_size = get_attr_wrapped_model(
+                self.inference_wrapped_model.model, "vocab_size"
+            )
 
         self.sampling_rng = torch.Generator(device=torch.cuda.current_device())
         self.sampling_rng.manual_seed(self.model_config.inference_sampling_seed)
@@ -123,11 +123,8 @@ def from_model_and_args(
             TextGenerationController: The initialized text generation controller.
         """
         tokenizer = build_tokenizer(args)
-        # TODO(ksanthanam): Condition this on model type?
         model = model_inference_wrapper_cls(model, context)
-        model.model_is_pipeline_parallel = not (
-            parallel_state.is_pipeline_first_stage() and parallel_state.is_pipeline_last_stage()
-        )
+        model.model_is_pipeline_parallel = model.config.pipeline_parallel_size > 1
         return cls(model, tokenizer)
 
     def set_stop_word_finished_ids_callback(self, callback):
diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py
index 8d94c7f8e3b..6f10f504768 100644
--- a/megatron/rl/inference/megatron.py
+++ b/megatron/rl/inference/megatron.py
@@ -68,32 +68,6 @@ def get_static_inference_engine(args: Namespace, model: MegatronModule) -> Abstr
         ),
     )
 
-
-def get_dynamic_inference_engine(
-    args: Namespace,
-    model: MegatronModule,
-    inference_logging_step_interval: int = 0,
-    metrics_writer=None,
-) -> DynamicInferenceEngine:
-    """
-    Returns an inference engine.
-    Args:
-        args (Namespace): The user arguments parsed from command line
-        model (MegatronModule): The megatron model.
-        inference_logging_step_interval (int): Step interval for logging inference metrics.
-        metrics_writer: Metrics writer (wandb module) for logging.
-
-    Returns:
-        DynamicInferenceEngine: The inference engine
-    """
-    context = DynamicInferenceContext.from_model_and_args(
-        model, args, overrides={"metrics_writer": metrics_writer}
-    )
-    controller = TextGenerationController.from_model_and_args(model, args, context)
-    engine = DynamicInferenceEngine.from_model_and_args(model, args, controller, context)
-    return engine
-
-
 class MegatronLocal(InferenceServer, ReturnsTokens, ReturnsRaw):
     """Interface to use MCoreEngine directly as an inference engine."""
 
@@ -156,30 +130,7 @@ async def launch(cls, model: GPTModel, **kwargs):
                 "WARNING: Tokenizer has no BOS token so prompt will not have BOS token",
             )
 
-        # Get inference logging configuration from args
-        log_inference_wandb = args.inference_wandb_logging
-        inference_logging_step_interval = args.inference_logging_step_interval
-
-        # Get metrics writer if logging is enabled and on the logging rank
-        # Use the same rank convention as training (last rank logs)
-        metrics_writer = None
-        if (
-            inference_logging_step_interval > 0
-            and log_inference_wandb
-            and args.rank == (args.world_size - 1)
-        ):
-            metrics_writer = get_wandb_writer()
-            if metrics_writer is None:
-                log_single_rank(
-                    logger,
-                    logging.WARNING,
-                    "WARNING: --rl-inference-logging-step-interval is set but no metrics writer "
-                    "wandb module is available. Inference logging will be disabled.",
-                )
-
-        inference_engine: DynamicInferenceEngine = get_dynamic_inference_engine(
-            args, model, inference_logging_step_interval, metrics_writer
-        )
+        inference_engine: DynamicInferenceEngine = DynamicInferenceEngine.from_model_and_args(model, args)
         await inference_engine.start_listening_to_data_parallel_coordinator(
             inference_coordinator_port=41521, launch_inference_coordinator=True
         )

From cff3c1efeb095a7621f851d90e2351ad1a30670e Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Wed, 21 Jan 2026 13:20:31 -0800
Subject: [PATCH 07/30] Fix copyright

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 .../model_inference_wrappers/t5/test_t5_inference_wrapper.py    | 2 ++
 .../test_encoder_decoder_text_generation_controller.py          | 2 ++
 .../test_vlm_text_generation_controller.py                      | 2 ++
 3 files changed, 6 insertions(+)

diff --git a/tests/unit_tests/inference/model_inference_wrappers/t5/test_t5_inference_wrapper.py b/tests/unit_tests/inference/model_inference_wrappers/t5/test_t5_inference_wrapper.py
index 1c167f1a98e..eb06f6ed78b 100644
--- a/tests/unit_tests/inference/model_inference_wrappers/t5/test_t5_inference_wrapper.py
+++ b/tests/unit_tests/inference/model_inference_wrappers/t5/test_t5_inference_wrapper.py
@@ -1,3 +1,5 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
 from argparse import Namespace
 from copy import deepcopy
 from unittest import mock
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_encoder_decoder_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_encoder_decoder_text_generation_controller.py
index 9bf2183c0fd..5bd39ec1324 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_encoder_decoder_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_encoder_decoder_text_generation_controller.py
@@ -1,3 +1,5 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
 import random
 import string
 import time
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_vlm_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_vlm_text_generation_controller.py
index 326e9b35873..50db5cc0afc 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_vlm_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_vlm_text_generation_controller.py
@@ -1,3 +1,5 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
 import copy
 import os
 import random

From b6d9c930ca16122bcf3ca25904347650e4ab3c9b Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Wed, 21 Jan 2026 22:14:39 -0800
Subject: [PATCH 08/30] Add back deprecated args with exception

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 .../inference/contexts/dynamic_context.py     | 26 +++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 0502a0bf829..472d2b11d41 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -263,9 +263,35 @@ def __init__(
         cuda_graph_mixed_prefill_count: Optional[int] = 16,
         request_metadata_types: Optional[List[Tuple[str, torch.dtype, bool]]] = None,
         persist_cuda_graphs: Optional[bool] = False,
+        params_dtype: Optional[Any] = None,  # Deprecated
+        num_layers: Optional[Any] = None,  # Deprecated
+        kv_channels: Optional[Any] = None,  # Deprecated
+        num_attention_heads: Optional[Any] = None,  # Deprecated
+        tensor_model_parallel_size: Optional[Any] = None,  # Deprecated
+        pipeline_model_parallel_size: Optional[Any] = None,  # Deprecated
+        cache_mla_latent: Optional[Any] = None,  # Deprecated
+        kv_lora_rank: Optional[Any] = None,  # Deprecate
+        qk_pos_emb_head_dim: Optional[Any] = None,  # Deprecated
     ):
         super().__init__(materialize_only_last_token_logits=materialize_only_last_token_logits)
 
+        deprecated_params = [
+            params_dtype,
+            num_layers,
+            kv_channels,
+            num_attention_heads,
+            tensor_model_parallel_size,
+            pipeline_model_parallel_size,
+            cache_mla_latent,
+            kv_lora_rank,
+            qk_pos_emb_head_dim,
+        ]
+        if any(param is not None for param in deprecated_params):
+            raise TypeError(
+                "Passing `TransformerConfig` arguments directly is deprecated. "
+                "Please pass `model_config` instead."
+            )
+
         self.cache_mla_latent = (
             isinstance(model_config, MLATransformerConfig) and model_config.cache_mla_latents
         )

From 98d5acc31ea6ae71272dc9752269097ef939ab96 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Wed, 21 Jan 2026 22:21:26 -0800
Subject: [PATCH 09/30] Add deprecation exception for GPTInferenceWrapper

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 .../model_inference_wrappers/gpt/gpt_inference_wrapper.py      | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/megatron/core/inference/model_inference_wrappers/gpt/gpt_inference_wrapper.py b/megatron/core/inference/model_inference_wrappers/gpt/gpt_inference_wrapper.py
index 031eecfd27f..2a1f10daa1e 100644
--- a/megatron/core/inference/model_inference_wrappers/gpt/gpt_inference_wrapper.py
+++ b/megatron/core/inference/model_inference_wrappers/gpt/gpt_inference_wrapper.py
@@ -33,7 +33,10 @@ def __init__(
         model: GPTModel,
         inference_context: Optional[BaseInferenceContext] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
+        inference_wrapper_config: Optional[Any] = None,  # Deprecated
     ):
+        if inference_wrapper_config is not None:
+            raise TypeError("Passing `inference_wrapper_config` is deprecated.")
         super().__init__(model, inference_context, pg_collection)
 
     def prep_inference_input(self, prompts_tokens: torch.Tensor) -> Dict[str, Any]:

From ec2fad20dff0c71a95348260e4fd3f35c5d6b2ce Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Thu, 22 Jan 2026 09:31:18 -0800
Subject: [PATCH 10/30] Fix wandb test

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 tests/unit_tests/inference/test_wandb_logging.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/tests/unit_tests/inference/test_wandb_logging.py b/tests/unit_tests/inference/test_wandb_logging.py
index f39729f5fbd..7d111ac1d93 100644
--- a/tests/unit_tests/inference/test_wandb_logging.py
+++ b/tests/unit_tests/inference/test_wandb_logging.py
@@ -51,7 +51,6 @@ def _get_dynamic_context(
         max_sequence_length=512,
         buffer_size_gb=0.03,
         block_size_tokens=128,
-        metrics_writer=None,
     ):
         """Helper to create a DynamicInferenceContext."""
         return DynamicInferenceContext(
@@ -65,7 +64,6 @@ def _get_dynamic_context(
             num_cuda_graphs=None,
             buffer_size_gb=buffer_size_gb,
             block_size_tokens=block_size_tokens,
-            metrics_writer=metrics_writer,
             unified_memory_level=0,  # unit tests currently broken with UVM
         )
 
@@ -203,7 +201,7 @@ def test_engine_logging_step_interval_zero(self):
         mock_wandb.__name__ = "wandb"
         mock_wandb.log = Mock()
 
-        dynamic_context = self._get_dynamic_context(metrics_writer=mock_wandb)
+        dynamic_context = self._get_dynamic_context()
 
         # Create mock controller with proper spec to pass isinstance checks
         mock_controller = create_autospec(TextGenerationController, instance=True)
@@ -218,6 +216,7 @@ def test_engine_logging_step_interval_zero(self):
             context=dynamic_context,
             random_seed=123,
             inference_logging_step_interval=0,  # Disabled
+            metrics_writer=mock_wandb,
         )
 
         # Verify log was never called
@@ -259,7 +258,7 @@ def test_paused_requests_in_stats(self):
     @pytest.mark.internal
     def test_metrics_writer_none_handling(self):
         """Test that engine handles None metrics_writer gracefully."""
-        dynamic_context = self._get_dynamic_context(metrics_writer=None)
+        dynamic_context = self._get_dynamic_context()
 
         # Create mock controller with proper spec to pass isinstance checks
         mock_controller = create_autospec(TextGenerationController, instance=True)
@@ -279,4 +278,4 @@ def test_metrics_writer_none_handling(self):
 
         # Verify engine was created successfully
         assert engine.inference_logging_step_interval == 10
-        assert engine.context.metrics_writer is None
+        assert engine.metrics_writer is None

From 2dea9b3b4876748a5a560f4d0c641e0a035f8007 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Thu, 22 Jan 2026 09:35:36 -0800
Subject: [PATCH 11/30] Address reviewer comments

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 megatron/core/inference/contexts/dynamic_context.py | 4 ++--
 megatron/core/inference/engines/dynamic_engine.py   | 2 +-
 megatron/core/inference/engines/static_engine.py    | 1 -
 3 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index ee24b5961fe..d78146c89fe 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -288,8 +288,8 @@ def __init__(
         ]
         if any(param is not None for param in deprecated_params):
             raise TypeError(
-                "Passing `TransformerConfig` arguments directly is deprecated. "
-                "Please pass `model_config` instead."
+                "Passing `TransformerConfig` arguments individually is deprecated. "
+                "Please pass the `TransformerConfig` directly to `model_config` instead."
             )
 
         self.cache_mla_latent = (
diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index 2447bd243fa..134b4db5e14 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -285,7 +285,7 @@ def from_model_and_args(
             controller,
             context,
             enable_cuda_graph=model.config.cuda_graph_impl == "local",
-            random_seed=amodel.config.seed,
+            random_seed=model.config.seed,
             track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events,
             enable_chunked_prefill=not args.disable_chunked_prefill,
             metrics_writer=metrics_writer,
diff --git a/megatron/core/inference/engines/static_engine.py b/megatron/core/inference/engines/static_engine.py
index 196fc1054da..1ac84ee5f79 100644
--- a/megatron/core/inference/engines/static_engine.py
+++ b/megatron/core/inference/engines/static_engine.py
@@ -118,7 +118,6 @@ def __init__(
                     random_seed=self.random_seed,
                 )
         except Exception as e:
-            torch.distributed.breakpoint(0)
             # Get exception details for better debugging
             exception_msg = str(e) if str(e) else f"{type(e).__name__}: {repr(e)}"
             warnings.warn(

From e5edb9ab94fc6fe24266c3ae1d5f7923e36d11cf Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Thu, 22 Jan 2026 10:15:19 -0800
Subject: [PATCH 12/30] Bug fixes

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 megatron/core/inference/engines/dynamic_engine.py              | 3 ++-
 .../text_generation_controllers/text_generation_controller.py  | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index 134b4db5e14..1c19547d49d 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -44,6 +44,7 @@
 from megatron.core.utils import (
     experimental_api,
     get_asyncio_loop,
+    get_attr_wrapped_model,
     get_pg_rank,
     get_pg_size,
     get_pg_src_rank,
@@ -285,7 +286,7 @@ def from_model_and_args(
             controller,
             context,
             enable_cuda_graph=model.config.cuda_graph_impl == "local",
-            random_seed=model.config.seed,
+            random_seed=args.seed,
             track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events,
             enable_chunked_prefill=not args.disable_chunked_prefill,
             metrics_writer=metrics_writer,
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index e9eefa573c0..e5c40131e0e 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -124,7 +124,7 @@ def from_model_and_args(
         """
         tokenizer = build_tokenizer(args)
         model = model_inference_wrapper_cls(model, context)
-        model.model_is_pipeline_parallel = model.config.pipeline_parallel_size > 1
+        model.model_is_pipeline_parallel = model.config.pipeline_model_parallel_size > 1
         return cls(model, tokenizer)
 
     def set_stop_word_finished_ids_callback(self, callback):

From 4b6cb9caa76db2d94ea3d57437b1b614729df21e Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Thu, 22 Jan 2026 18:30:52 -0800
Subject: [PATCH 13/30] Major refactor - introduce DynamicInferenceConfig

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 .../inference/gpt/gpt_dynamic_inference.py    | 126 ++-----
 examples/inference/gpt/utils.py               | 167 ----------
 megatron/core/inference/config.py             | 205 ++++++++++++
 .../attention_context/mamba_metadata.py       |  26 +-
 .../core/inference/contexts/base_context.py   |   8 +-
 .../inference/contexts/dynamic_context.py     | 296 +++--------------
 .../core/inference/contexts/static_context.py |   5 +-
 .../core/inference/engines/dynamic_engine.py  | 141 ++------
 .../core/inference/engines/static_engine.py   |  25 +-
 .../text_generation_controller.py             |  84 ++---
 megatron/core/models/gpt/gpt_model.py         |   7 +-
 megatron/core/models/mamba/mamba_model.py     |   7 +-
 megatron/core/utils.py                        |  19 --
 megatron/inference/__init__.py                |   0
 megatron/inference/utils.py                   | 311 ++++++++++++++++++
 megatron/rl/inference/megatron.py             |   3 +-
 megatron/training/arguments.py                |   6 +-
 .../contexts/test_dynamic_context.py          |  81 +++--
 .../inference/engines/test_dynamic_engine.py  |  48 ++-
 .../gpt/test_gpt_inference_wrapper.py         |   4 +-
 .../inference/test_wandb_logging.py           |  13 +-
 .../test_text_generation_controller.py        |  15 +-
 tests/unit_tests/models/test_gpt_model.py     |  11 +-
 .../models/test_gpt_model_batch_invariant.py  |  33 +-
 tests/unit_tests/models/test_mamba_model.py   |  12 +-
 tools/run_dynamic_text_generation_server.py   |  20 +-
 tools/run_inference_performance_test.py       |   4 +-
 27 files changed, 812 insertions(+), 865 deletions(-)
 create mode 100644 megatron/core/inference/config.py
 create mode 100644 megatron/inference/__init__.py
 create mode 100644 megatron/inference/utils.py

diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py
index 28f6a35a227..7daa095d845 100644
--- a/examples/inference/gpt/gpt_dynamic_inference.py
+++ b/examples/inference/gpt/gpt_dynamic_inference.py
@@ -1,16 +1,15 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+# pylint: disable=bad-builtin
+
 import hashlib
 import io
 import json
-import math
 import os
-import pickle
 import sys
 import warnings
-from argparse import ArgumentParser
 from collections import defaultdict
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional
 
 import torch
 from tqdm import tqdm
@@ -19,23 +18,14 @@
     os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
 )
 
-import megatron
 from examples.inference.gpt.utils import (
     Request,
-    add_common_inference_args,
     build_dynamic_engine_setup_prefix,
     build_requests,
     get_curr_time,
     get_global_peak_memory_stats_bytes,
-    get_model,
-)
-from megatron.core.inference.contexts.attention_context.mamba_metadata import (
-    MambaInferenceStateConfig,
-)
-from megatron.core.inference.contexts.dynamic_context import (
-    ContextOverflowError,
-    DynamicInferenceContext,
 )
+from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
 from megatron.core.inference.engines import DynamicInferenceEngine, EngineSuspendedError
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
     GPTInferenceWrapper,
@@ -45,87 +35,26 @@
     TextGenerationController,
 )
 from megatron.core.tokenizers.text.utils.build_tokenizer import build_tokenizer
-from megatron.core.transformer.module import MegatronModule
-from megatron.core.utils import get_mamba_inference_state_config_from_model
+from megatron.inference.utils import (
+    add_inference_args,
+    get_dynamic_inference_config_from_model_and_args,
+    get_model,
+)
 
 sys.path.append(
     os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
 )
 import logging
 
+import megatron
 from megatron.core.utils import configure_nvtx_profiling
-from megatron.training import get_args
-from megatron.training import get_model as _get_model
-from megatron.training import get_tokenizer, initialize_megatron
+from megatron.training import get_args, get_tokenizer, initialize_megatron
 
 torch.serialization.add_safe_globals([io.BytesIO])
 torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunState])
 torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunDiagnostic])
 
 
-def add_dynamic_inference_args(parser: ArgumentParser) -> ArgumentParser:
-    """Dynamic inference arguments."""
-
-    add_common_inference_args(parser)
-
-    group = parser.add_argument_group(title='Dynamic inference')
-    group.add_argument(
-        "--inference-ckpt-non-strict",
-        action="store_true",
-        help="Load checkpoint with `strict=False`.",
-    )
-    group.add_argument(
-        "--termination-id",
-        type=int,
-        default=None,
-        help="Termination ID that overrides `tokenizer.eod`.",
-    )
-    group.add_argument(
-        "--suspend-resume-interval",
-        type=int,
-        default=None,
-        help="Suspend and resume the dynamic engine every "
-        "`suspend_resume_interval` steps. This is used to tet the suspend/resume "
-        "system.",
-    )
-    group.add_argument(
-        "--inference-repeat-n",
-        type=int,
-        default=1,
-        help="Repeat inference iterations N times for benchmarking.",
-    )
-    group.add_argument(
-        "--throughput-check-only",
-        action='store_true',
-        default=False,
-        help="If true, only run throughput check without verifying outputs.",
-    )
-
-    return parser
-
-
-def get_inference_context(
-    model,
-    requests: List[Request],
-    sampling_params: Optional[SamplingParams] = None,
-    calculate_max_sequence_length_from_requests: bool = True,
-    mamba_inference_state_config: Optional[MambaInferenceStateConfig] = None,
-):
-    """The inference context manages the KV cache and other inference state."""
-
-    args = get_args()
-
-    overrides = None
-
-    # Max sequence length.
-    if calculate_max_sequence_length_from_requests:
-        max_gen_length = sampling_params.num_tokens_to_generate
-        max_context_length = max(len(r.prompt_tokens) for r in requests)
-        max_sequence_length = max_context_length + max_gen_length
-        overrides = {"max_sequence_length": max_sequence_length}
-
-    return DynamicInferenceContext.from_model_and_args(model, args, overrides)
-
 def run_inference(
     requests: List[Request],
     engine: DynamicInferenceEngine,
@@ -210,7 +139,8 @@ def _add_request():
             pass  # ignore error in order to call 'engine.resume()' below.
         attempted_step_count += 1
 
-        # After step, we lost track of last iteration's is_decode_only, so we need to get it from the engine
+        # After step, we lost track of last iteration's is_decode_only,
+        # so we need to get it from the engine
         is_decode_only = engine.is_decode_only
 
         # Test suspending and resuming engine.
@@ -309,10 +239,10 @@ def _add_request():
 
 @torch.inference_mode()
 def main():
-
+    """Run dynamic inference."""
     # Initialize Megatron.
     initialize_megatron(
-        extra_args_provider=add_dynamic_inference_args,
+        extra_args_provider=add_inference_args,
         args_defaults={'no_load_rng': True, 'no_load_optim': True},
     )
 
@@ -353,11 +283,18 @@ def main():
 
     # Requests, context, controller.
     requests = build_requests(args, tokenizer, sampling_params)
-    context = get_inference_context(model, requests=requests, sampling_params=sampling_params)
-    controller = TextGenerationController.from_model_and_args(model, args, context)
+    inference_config = get_dynamic_inference_config_from_model_and_args(model, args)
+
+    # Calculate max_sequence_length from requests
+    max_gen_length = sampling_params.num_tokens_to_generate
+    max_context_length = max(len(r.prompt_tokens) for r in requests)
+    inference_config.max_sequence_length = max_context_length + max_gen_length
+    context = DynamicInferenceContext(model.config, inference_config)
+    wrapped_model = GPTInferenceWrapper(model, context)
+    controller = TextGenerationController(wrapped_model, tokenizer)
 
     # Validate all context_length's <= max_tokens.
-    if args.disable_chunked_prefill:
+    if not args.enable_chunked_prefill:
         invalid_prompt_length_map = {}
         for request_idx, request in enumerate(requests):
             if len(request.prompt_tokens) > context.max_tokens:
@@ -369,7 +306,7 @@ def main():
         )
 
     # Inference engine.
-    engine = DynamicInferenceEngine.from_model_and_args(model, args, controller, context)
+    engine = DynamicInferenceEngine(controller, context)
 
     setup_prefix = build_dynamic_engine_setup_prefix(args, model, context, requests)
     print("~~~")
@@ -424,7 +361,10 @@ def escape_str(s):
             # ---- Prompt summary line ----
             prompt_len = len(requests[request_idxs[0]].prompt_tokens)
             escaped_prompt_text = escape_str(prompt_text)
-            print(f"\n{unique_idx+1}/{len(unique_prompt_map)} [n {len(request_idxs)}, l {prompt_len}] {escaped_prompt_text}")
+            print(
+                f"\n{unique_idx+1}/{len(unique_prompt_map)}"
+                f"[n {len(request_idxs)}, l {prompt_len}] {escaped_prompt_text}"
+            )
 
             # ---- Group all outputs for this prompt ----
             output_map = defaultdict(list)
@@ -451,8 +391,10 @@ def escape_str(s):
                     o_hash = "--"
                     o_len = 0
                     escaped_output_text = "--"
-                print(f"  >>>> [n {len(output_request_idxs)}, {o_len} tokens, hash {o_hash}"
-                f"{', <evicted>' if evicted else ''}] {escaped_output_text}")
+                print(
+                    f"  >>>> [n {len(output_request_idxs)}, {o_len} tokens, hash {o_hash}"
+                    f"{', <evicted>' if evicted else ''}] {escaped_output_text}"
+                )
                 text_hashes.append(o_hash)
 
         # Write results to JSON. Primarily used for functional testing.
diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py
index fc87fb764d2..6f50f6a4b57 100644
--- a/examples/inference/gpt/utils.py
+++ b/examples/inference/gpt/utils.py
@@ -12,8 +12,6 @@
 import torch
 from tqdm import tqdm
 
-from gpt_builders import gpt_builder
-from mamba_builders import mamba_builder
 from megatron.core.inference.contexts import DynamicInferenceContext
 from megatron.core.inference.contexts.dynamic_context import get_mem_size_str
 from megatron.core.inference.inference_request import DynamicInferenceRequest
@@ -21,171 +19,6 @@
 from megatron.core.transformer.module import MegatronModule
 from megatron.training import get_args
 from megatron.training import get_model as _get_model
-from megatron.training.checkpointing import load_checkpoint
-from model_provider import model_provider
-
-
-def get_model() -> MegatronModule:
-    """Initialize model and load checkpoint."""
-
-    args = get_args()
-
-    if args.model_provider == "gpt":
-        model_builder = gpt_builder
-    elif args.model_provider == "mamba":
-        model_builder = mamba_builder
-    else:
-        raise ValueError(f"Invalid model provider {args.model_provider}")
-
-    # Build model.
-    model = _get_model(partial(model_provider, model_builder), wrap_with_ddp=False)
-
-    # Load checkpoint.
-    assert args.load is not None
-    args.exit_on_missing_checkpoint = True
-    load_checkpoint(
-        ddp_model=model,
-        optimizer=None,
-        opt_param_scheduler=None,
-        strict=not args.inference_ckpt_non_strict,
-    )
-
-    # No virtual PP.
-    assert len(model) == 1, "Above condition should have caught this"
-    model = model[0]
-
-    # Eval mode.
-    model.eval()
-
-    return model
-
-
-def add_common_inference_args(parser: ArgumentParser) -> ArgumentParser:
-    """Common inference arguments."""
-
-    group = parser.add_argument_group(title='Common inference')
-
-    group.add_argument("--temperature", type=float, default=1.0, help='Sampling temperature.')
-    group.add_argument("--top_k", type=int, default=1, help='Top k sampling.')
-    group.add_argument("--top_p", type=float, default=0.0, help='Top p sampling.')
-    group.add_argument(
-        "--return-log-probs",
-        action='store_true',
-        default=False,
-        help='Return the log probabilities of the final output tokens',
-    )
-    group.add_argument(
-        "--prompts",
-        metavar='N',
-        type=str,
-        nargs='+',
-        help='Input prompts with each prompt within quotes and seperated by space',
-    )
-    group.add_argument(
-        "--num-tokens-to-prompt",
-        type=int,
-        nargs="+",
-        default=[64, 1024],
-        help='Number of tokens to use for simulated prompts. This should be a '
-        'space-separated pair of integers, and the generated prompt lengths will '
-        'be uniformly sampled within this range.',
-    )
-    group.add_argument(
-        "--num-tokens-to-generate",
-        type=int,
-        default=30,
-        help='Number of tokens to generate for each prompt',
-    )
-    group.add_argument(
-        "--num-tokens-from-file",
-        action='store_true',
-        default=False,
-        help='Use per-prompt num_tokens_to_generate from prompt file',
-    )
-    group.add_argument(
-        "--top-n-logprobs",
-        type=int,
-        default=0,
-        help='Return the top n logprobs for the generated tokens and their corresponding token as a dictionary',
-    )
-    group.add_argument(
-        "--incoming-requests-per-step",
-        type=int,
-        default=None,
-        help="Add a deterministic number of requests per step. This arg is "
-        "prioritized over `--incoming-requests-per-sec` below (which is non-"
-        "deterministic). Note that the number of requests added per step is "
-        "additionally limited by the inference context's `max_requests`, "
-        "`max_tokens`, and KV buffer size.",
-    )
-    group.add_argument(
-        "--incoming-requests-per-sec",
-        type=float,
-        default=100.0,
-        help="Simulated number of requests per second. Set to -1 to add all requests together.",
-    )
-    group.add_argument(
-        "--incoming-requests-duration",
-        type=float,
-        default=10.0,
-        help="Total amount of time to simulate that requests are "
-        "arriving. Multiply this value with "
-        "`--incoming-requests-per-sec` to get the approximate "
-        "total number of requests. Set to -1 to add all requests together.",
-    )
-    group.add_argument(
-        "--model-provider", choices=["mamba", "gpt"], default="gpt", help="Model provider"
-    )
-    group.add_argument(
-        "--skip-prompt-log-probs", action='store_true', default=False, help='Skip prompt log probs.'
-    )
-    group.add_argument(
-        "--stop-words",
-        metavar='WORD',
-        type=str,
-        nargs='+',
-        default=None,
-        help='Stop words to terminate generation. Each word should be quoted and '
-        'separated by space. Example: --stop-words "\\n\\n" "END" "###"',
-    )
-    group.add_argument(
-        "--output-path", type=str, default=None, help="Path to save generations as JSON"
-    )
-    group.add_argument(
-        "--output-every-n-results",
-        type=int,
-        default=1,
-        help="To minimize the output file size of larger runs, only write the "
-        "results of every `n` requests.",
-    )
-    group.add_argument(
-        "--prompt-file",
-        help='Jsonl file containing input prompts, where each item (i.e., line) '
-        'contains the field \'text\' where the value is the prompt. All other '
-        'fields within each item are ignored, and may be customized for each '
-        'application.',
-    )
-    group.add_argument(
-        "--prompt-file-num-truncate",
-        type=int,
-        help='Number of samples to use from the loaded prompt file (see '
-        '`--prompt-file` above). The first `--prompt-file-num-truncate` samples '
-        'will be used, in order.',
-    )
-    group.add_argument(
-        "--use-flashinfer-fused-rope",
-        action='store_true',
-        default=False,
-        help='Use flashinfer fused rope implementation.',
-    )
-    group.add_argument(
-        "--no-record-throughput",
-        action='store_false',
-        dest="record_throughput",
-        help="Disable throughput recording in --output-file",
-    )
-
-    return parser
 
 
 def get_default_sampling_params(termination_id: int = None):
diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py
new file mode 100644
index 00000000000..dae39532ba0
--- /dev/null
+++ b/megatron/core/inference/config.py
@@ -0,0 +1,205 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+import abc
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+import torch
+
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.utils import get_attr_wrapped_model
+
+
+@dataclass
+class MambaInferenceStateConfig:
+    """
+    Config for initializing Mamba model inference state tensors.
+
+    Note that we maintain separate metadata for decode, regular prefill, and
+    chunked prefill requests because the Mamba kernels do not yet support mixing
+    these. Once the kernels have been updated we can simplify this code.
+    """
+
+    layer_type_list: List[str]
+    """
+    A list of strings that indicates the layer type (Mamba / Attention / MLP) for each layer.
+    See `megatron/core/ssm/mamba_hybrid_layer_allocation.py` for the list of symbols.
+    """
+
+    mamba_conv_states_shape: Tuple[int]
+    """Mamba conv states shape per request."""
+
+    mamba_ssm_states_shape: Tuple[int]
+    """Mamba ssm states shape per request."""
+
+    @classmethod
+    def from_model(cls, model: MegatronModule) -> Optional["MambaInferenceStateConfig"]:
+        """Returns Mamba inference state config from the model if it is a hybrid model."""
+        from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols
+
+        decoder = get_attr_wrapped_model(model, "decoder")
+        layer_type_list = getattr(decoder, "layer_type_list", None)
+        if layer_type_list is not None and Symbols.MAMBA in layer_type_list:
+            (mamba_conv_states_shape, mamba_ssm_states_shape) = (
+                decoder.mamba_state_shapes_per_request()
+            )
+            return cls(
+                layer_type_list=layer_type_list,
+                mamba_conv_states_shape=mamba_conv_states_shape,
+                mamba_ssm_states_shape=mamba_ssm_states_shape,
+            )
+        return None
+
+
+@dataclass
+class BaseInferenceConfig(abc.ABC):
+    """Base inference config."""
+
+    pass
+
+
+@dataclass
+class DynamicInferenceConfig(BaseInferenceConfig):
+    """
+    Config for dynamic inference.
+
+    Used to initialize `DynamicInferenceContext`, `TextGenerationController`, and
+    `DynamicInferenceEngine`.
+
+    NOTE: Must remain mutually exclusive with the `TransformerConfig`.
+    """
+
+    # =================================
+    # KV cache config
+    # =================================
+    block_size_tokens: int = 256
+    """Size of KV cache block size."""
+
+    buffer_size_gb: int = 20
+    """
+    Buffer size reserved on the GPU for the KV cache.
+    If `unified_memory_level` >= 1, then CPU memory is additionally utilized, resulting in a total
+    buffer size of `buffer_size_gb + paused_buffer_size_gb`.
+    """
+
+    paused_buffer_size_gb: Optional[int] = None
+    """
+    Portion of buffer reserved for paused requests. Active requests are paused when there are not
+    enough active blocks available to continue generating a request. The total buffer size
+    (active + paused) depends on `unified_memory_level` (uvm):
+        - uvm 0: buffer_size_gb (paused buffer is inclusive)
+        - uvm 1: buffer_size_gb + paused_buffer_size_gb
+    """
+
+    max_requests: Optional[int] = None
+    """
+    Max number of active requests to use for decode-only forward passes.
+    This is primarily limited by the combination of `buffer_size_gb` and `max_sequence_length`.
+    """
+
+    max_tokens: Optional[int] = None
+    """
+    Max number of tokens to use for forward passes. This is primarily limited by prefill activation
+    memory usage. (Defaults to 16384).
+    """
+
+    unified_memory_level: int = 0
+    """
+    Sets unified memory usage within the dynamic inference context.
+    The levels are:
+        0) no unified memory (default)
+        1) allocate `memory_buffer` in unified memory.
+    Eventually, additional levels will be included to control other tensors within the context.
+    """
+
+    # =================================
+    # CUDA graph config
+    # =================================
+    num_cuda_graphs: Optional[int] = None
+    """
+    Maximum number of cuda graphs to capture, where the cuda graph batch sizes range from 1 to
+    `max_requests`. Due to rounding, the actual number of cuda graphs may not equal this argument.
+    """
+
+    cuda_graph_mixed_prefill_count: Optional[int] = 16
+    """ 
+    The number of mixed prefill graphs to capture if mixed prefill/decode graphs are enabled.
+    """
+
+    use_cuda_graphs_for_non_decode_steps: bool = True
+    """
+    Whether to use CUDA graphs for non-decode steps.
+    """
+
+    persist_cuda_graphs: bool = False
+    """
+    Whether to persist CUDA graphs when the engine is suspended.
+    If False and `unified_memory_level` is 0, CUDA graphs are deleted on `suspend()`
+    and re-captured on `resume()` to save memory.
+    """
+
+    # =================================
+    # Model config
+    # =================================
+    max_sequence_length: int = 2560
+    """Max possible sequence length (prompt + output) that will occur."""
+
+    mamba_inference_state_config: Optional[MambaInferenceStateConfig] = None
+    """The Mamba inference state config if the model is a hybrid model."""
+
+    pg_collection: Optional[ProcessGroupCollection] = None
+    """A `ProcessGroupCollection` for distributed execution."""
+
+    use_flashinfer_fused_rope: Optional[bool] = False
+    """
+    If True, use flashinfer's fused rope implementation.
+    If None, defaults to using flash-infer if available.
+    """
+
+    materialize_only_last_token_logits: bool = True
+    """
+    Whether to only materialize logits for the last token. This should be set to False
+    if returning log probs.
+    """
+
+    # =================================
+    # Engine config
+    # =================================
+    enable_chunked_prefill: bool = False
+    """Whether to enable chunked prefill."""
+
+    # =================================
+    # Logging config
+    # =================================
+    track_paused_request_events: bool = False
+    """
+    Whether to track paused request events. If True, `add_event_pause()` is called on
+    requests when they are paused during bookkeeping.
+    """
+
+    metrics_writer: Optional["WandbModule"] = None
+    """Wandb module for writing metrics."""
+
+    logging_step_interval: int = 0
+    """
+    The step interval at which to log inference metrics to wandb.
+    Defaults to 0, which means no logging.
+    """
+
+    request_metadata_types: Optional[List[Tuple[str, torch.dtype, bool]]] = None
+    """
+    A list of the per-request metadata types to track. Each entry is a tuple consisting of the string
+    label, the target dtype, and whether to store the data on GPU.
+    """
+
+
+@dataclass
+class StaticInferenceConfig(DynamicInferenceConfig):
+    """
+    Static inference config. For now, exactly mimic the dynamic config.
+
+    TODO(ksanthanam): Remove when deprecating static inference.
+    """
+
+    pass
diff --git a/megatron/core/inference/contexts/attention_context/mamba_metadata.py b/megatron/core/inference/contexts/attention_context/mamba_metadata.py
index 6cf45aeb9e1..13179483f59 100644
--- a/megatron/core/inference/contexts/attention_context/mamba_metadata.py
+++ b/megatron/core/inference/contexts/attention_context/mamba_metadata.py
@@ -1,36 +1,12 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
-from dataclasses import dataclass
-from typing import List, Optional, Tuple
+from typing import Optional
 
 import torch
 
 from megatron.core.inference.batch_dimensions_utils import InferenceBatchDimensions
 
 
-@dataclass
-class MambaInferenceStateConfig:
-    """
-    Config for initializing Mamba model inference state tensors.
-
-    Note that we maintain separate metadata for decode, regular prefill, and
-    chunked prefill requests because the Mamba kernels do not yet support mixing
-    these. Once the kernels have been updated we can simplify this code.
-    """
-
-    layer_type_list: List[str]
-    """
-    A list of strings that indicates the layer type (Mamba / Attention / MLP) for each layer.
-    See `megatron/core/ssm/mamba_hybrid_layer_allocation.py` for the list of symbols.
-    """
-
-    mamba_conv_states_shape: Tuple[int]
-    """Mamba conv states shape per request."""
-
-    mamba_ssm_states_shape: Tuple[int]
-    """Mamba ssm states shape per request."""
-
-
 class MambaMetadata:
     """Manages the metadata tensors required for Mamba layers during inference."""
 
diff --git a/megatron/core/inference/contexts/base_context.py b/megatron/core/inference/contexts/base_context.py
index 3dfec6de3ad..738e8568bc8 100644
--- a/megatron/core/inference/contexts/base_context.py
+++ b/megatron/core/inference/contexts/base_context.py
@@ -2,6 +2,8 @@
 
 import abc
 
+from megatron.core.inference.config import BaseInferenceConfig
+
 
 class BaseInferenceContext(abc.ABC):
     """Base class for inference contexts.
@@ -10,13 +12,11 @@ class BaseInferenceContext(abc.ABC):
     Extend this class for any future contexts types.
     """
 
-    def __init__(self, materialize_only_last_token_logits: bool):
+    def __init__(self, inference_config: BaseInferenceConfig):
         """
         Args:
-            materialize_only_last_token_logits (bool):
-                If True, only the last-token logits will be extracted during decode
         """
-        self.materialize_only_last_token_logits = materialize_only_last_token_logits
+        self.inference_config = inference_config
 
     @abc.abstractmethod
     def is_static_batching(self) -> bool:
diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index d78146c89fe..030d116f35b 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -4,18 +4,18 @@
 import math
 import warnings
 from contextlib import nullcontext
-from typing import Any, Dict, List, Optional, Sequence, Tuple
+from typing import List, Optional, Sequence, Tuple
 
-import torch
-import torch.nn.functional as F
-from packaging.version import Version as PkgVersion
-from torch import Tensor
+import torch  # type: ignore
+import torch.nn.functional as F  # type: ignore
+from torch import Tensor  # type: ignore
 
 from megatron.core import parallel_state
 from megatron.core.inference.batch_dimensions_utils import (
     CUDAGraphBatchDimensionBuilder,
     InferenceBatchDimensions,
 )
+from megatron.core.inference.config import DynamicInferenceConfig
 from megatron.core.inference.inference_request import DynamicInferenceRequest
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.unified_memory import (
@@ -25,18 +25,12 @@
 from megatron.core.inference.utils import tensor_swap
 from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb
 from megatron.core.package_info import __version__ as mcore_version
-from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.ssm.mamba_hybrid_layer_allocation import get_layer_maps_from_layer_type_list
 from megatron.core.transformer import MLATransformerConfig, TransformerConfig
 from megatron.core.utils import divide as core_divide
-from megatron.core.utils import (
-    get_attr_wrapped_model,
-    get_mamba_inference_state_config_from_model,
-    get_pg_size,
-    internal_api,
-)
+from megatron.core.utils import get_pg_size, internal_api
 
-from .attention_context.mamba_metadata import MambaInferenceStateConfig, MambaMetadata
+from .attention_context.mamba_metadata import MambaMetadata
 from .attention_context.mha_metadata import GraphedMHAMetadata, NonGraphedMHAMetadata
 from .base_context import BaseInferenceContext
 from .dynamic_block_allocator import BlockAllocator
@@ -47,14 +41,7 @@
     triton_append_key_value_cache = None
 
 try:
-    from packaging.version import Version as PkgVersion
-
-    HAVE_PACKAGING = True
-except:
-    HAVE_PACKAGING = False
-
-try:
-    import flashinfer  # pylint: disable=unused-import
+    import flashinfer  # type: ignore # pylint: disable=unused-import
 
     HAVE_FLASHINFER = True
 except ImportError:
@@ -197,121 +184,28 @@ class DynamicInferenceContext(BaseInferenceContext):
 
     Args:
         model_config (TransformerConfig): Model config.
-        max_sequence_length (int): Max possible sequence length (prompt + output)
-            that will occur.
-        buffer_size_gb (float): Buffer size reserved on the GPU for the KV cache.
-            if `unified_memory_level` >= 1, then CPU memory is additionally
-            utilized, resulting in a total buffer size of `buffer_size_gb +
-            paused_buffer_size_gb`.
-        paused_buffer_size_gb (float | None): Portion of buffer reserved for
-            paused requests. Active requests are paused when there are not enough
-            active blocks available to continue generating a request. The total
-            buffer size (active + paused) depends on `unified_memory_level` (uvm):
-            - uvm 0: buffer_size_gb (paused buffer is inclusive)
-            - uvm 1: buffer_size_gb + paused_buffer_size_gb
-        mamba_inference_state_config (Optional[MambaInferenceStateConfig]): The Mamba
-            inference state config if the model is a hybrid model.
-        max_requests (int): Max number of active requests to use for
-            decode-only forward passes. This value is primarily limited by the
-            combination of `buffer_size_gb` and `max_sequence_length`.
-        max_tokens (int): Max number of tokens to use for forward passes. This is
-            primarily limited by prefill activation memory usage. (Defaults to
-            16384).
-        block_size_tokens (int): Size of KV cache block size.
-        num_cuda_graphs (Optional[int]): Maximum number of cuda graphs to capture,
-            where the cuda graph batch sizes range from 1 to `max_requests`
-            (as computed below). Due to rounding, the actual number of cuda graphs
-            may not equal this argument.
-        materialize_only_last_token_logits (Optional[bool]): Whether to only
-            materialize logits for the last token. This should be set to False
-            if returning log probs.
-        use_cuda_graphs_for_non_decode_steps (bool): If True, use cuda graphs for non-decode
-            engine steps.
-        unified_memory_level (Optional[int]): Set unified memory usage within the
-            dynamic inference context. The levels are: 0) no unified memory, 1)
-            allocate `memory_buffer` in unified memory. Eventually, additional
-            levels will be included to control other tensors within the context.
-        use_flashinfer_fused_rope (bool): If True, use flashinfer's fused rope implementation.
-            If None, defaults to using flash-infer if available.
-        request_metadata_types (Optional[List[Tuple[str, torch.dtype, bool]]]): A list of the
-            per-request metadata types to track. Each entry is a tuple consisting of the string
-            label, the target dtype, and whether to store the data on GPU.
+        inference_config (DynamicInferenceConfig): Inference config.
     """
 
     DEFAULT_MAX_TOKENS = 16384
     TOKEN_ROUNDER = 64
     REQUEST_ROUNDER = 4
 
-    def __init__(
-        self,
-        *,
-        model_config: TransformerConfig,
-        max_sequence_length: int,
-        buffer_size_gb: float,
-        paused_buffer_size_gb: float | None = None,
-        mamba_inference_state_config: Optional[MambaInferenceStateConfig] = None,
-        pg_collection: Optional[ProcessGroupCollection] = None,
-        max_requests: Optional[int] = None,
-        max_tokens: int = DEFAULT_MAX_TOKENS,
-        block_size_tokens: int = 256,
-        num_cuda_graphs: Optional[int] = None,
-        materialize_only_last_token_logits: Optional[bool] = True,
-        use_cuda_graphs_for_non_decode_steps: bool = True,
-        use_flashinfer_fused_rope: bool = False,
-        unified_memory_level: Optional[int] = 0,
-        cuda_graph_max_tokens: Optional[int] = None,
-        cuda_graph_mixed_prefill_count: Optional[int] = 16,
-        request_metadata_types: Optional[List[Tuple[str, torch.dtype, bool]]] = None,
-        persist_cuda_graphs: Optional[bool] = False,
-        params_dtype: Optional[Any] = None,  # Deprecated
-        num_layers: Optional[Any] = None,  # Deprecated
-        kv_channels: Optional[Any] = None,  # Deprecated
-        num_attention_heads: Optional[Any] = None,  # Deprecated
-        tensor_model_parallel_size: Optional[Any] = None,  # Deprecated
-        pipeline_model_parallel_size: Optional[Any] = None,  # Deprecated
-        cache_mla_latent: Optional[Any] = None,  # Deprecated
-        kv_lora_rank: Optional[Any] = None,  # Deprecate
-        qk_pos_emb_head_dim: Optional[Any] = None,  # Deprecated
-    ):
-        super().__init__(materialize_only_last_token_logits=materialize_only_last_token_logits)
-
-        deprecated_params = [
-            params_dtype,
-            num_layers,
-            kv_channels,
-            num_attention_heads,
-            tensor_model_parallel_size,
-            pipeline_model_parallel_size,
-            cache_mla_latent,
-            kv_lora_rank,
-            qk_pos_emb_head_dim,
-        ]
-        if any(param is not None for param in deprecated_params):
-            raise TypeError(
-                "Passing `TransformerConfig` arguments individually is deprecated. "
-                "Please pass the `TransformerConfig` directly to `model_config` instead."
-            )
+    def __init__(self, model_config: TransformerConfig, inference_config: DynamicInferenceConfig):
+        super().__init__(inference_config=inference_config)
 
         self.cache_mla_latent = (
             isinstance(model_config, MLATransformerConfig) and model_config.cache_mla_latents
         )
         if self.cache_mla_latent:
             assert (
-                block_size_tokens == 64
+                inference_config.block_size_tokens == 64
             ), "Flash MLA requires a block size of 64. Set --inference-dynamic-batching-block-size 64 to fix this assert"
 
-        # give deprecated args warning for cuda_graph_max_tokens
-        if cuda_graph_max_tokens is not None:
-            warnings.warn(
-                "`cuda_graph_max_tokens` is deprecated and will be removed in a future release. "
-                "The context now automatically sets the max tokens for cuda graphs based on "
-                "`max_requests`.",
-                DeprecationWarning,
-            )
-
         # Per partition num heads and hidden size.
         num_attention_heads = model_config.num_query_groups or model_config.num_attention_heads
         projection_size = model_config.kv_channels * num_attention_heads
+        pg_collection = inference_config.pg_collection
         if pg_collection is not None:
             tp_size = get_pg_size(pg_collection.tp)
         else:
@@ -337,6 +231,7 @@ def __init__(
             self.pipeline_parallel_group = None
 
         # Mamba states.
+        mamba_inference_state_config = inference_config.mamba_inference_state_config
         self.is_hybrid_model = mamba_inference_state_config is not None
         if self.is_hybrid_model:
             mamba_conv_states_shape = mamba_inference_state_config.mamba_conv_states_shape
@@ -373,10 +268,10 @@ def __init__(
 
         # Block size tokens, bytes.
         dtype_size_bytes = model_config.params_dtype.itemsize
-        self.block_size_tokens = block_size_tokens
+        self.block_size_tokens = inference_config.block_size_tokens
         if self.cache_mla_latent:
             #   one vector  c_t  (rank)  +  optional RoPE phase slice
-            self.kv_reduced_dim = kv_lora_rank + qk_pos_emb_head_dim
+            self.kv_reduced_dim = model_config.kv_lora_rank + model_config.qk_pos_emb_head_dim
             self.block_size_bytes = (
                 dtype_size_bytes
                 * self.num_attention_layers
@@ -402,9 +297,9 @@ def __init__(
             mamba_states_memory_per_request *= dtype_size_bytes
 
         # Unified memory.
-        self.unified_memory_level = unified_memory_level
-        self.persist_cuda_graphs = persist_cuda_graphs
-        if unified_memory_level > 0:
+        self.unified_memory_level = inference_config.unified_memory_level
+        self.persist_cuda_graphs = inference_config.persist_cuda_graphs
+        if self.unified_memory_level > 0:
             try:
                 self.unified_memory_mempool = create_unified_mempool()
             except UnifiedMemoryUnsupportedError:
@@ -415,9 +310,11 @@ def __init__(
                 self.unified_memory_level = 0
 
         # Initialize block allocator.
-        buffer_size_bytes = int(buffer_size_gb * 1024**3)
+        buffer_size_bytes = int(inference_config.buffer_size_gb * 1024**3)
         paused_buffer_size_bytes = (
-            0 if paused_buffer_size_gb is None else int(paused_buffer_size_gb * 1024**3)
+            0
+            if inference_config.paused_buffer_size_gb is None
+            else int(inference_config.paused_buffer_size_gb * 1024**3)
         )
         # TODO: Add parameter to control fraction of memory assigned to KV cache
         # versus Mamba state.
@@ -451,13 +348,14 @@ def __init__(
         )
 
         # Track request metadata.
+        request_metadata_types = inference_config.request_metadata_types
         if request_metadata_types is None:
             request_metadata_types = DynamicInferenceRequest.get_metadata_types()
         self.request_metadata_types = request_metadata_types
 
         # Initialize context state.
         self.params_dtype = model_config.params_dtype
-        self.max_sequence_length = max_sequence_length
+        self.max_sequence_length = inference_config.max_sequence_length
 
         # Request and token counts.
         self.total_request_count = 0
@@ -477,16 +375,16 @@ def __init__(
         self.max_kv_block_count = math.ceil(self.max_sequence_length / self.block_size_tokens)
 
         # Set max_requests, max_tokens.
-        if max_requests is None:
+        if inference_config.max_requests is None:
             # Maximize compute utilization by defaulting to 1 block per request.
             self.max_requests = self.block_allocator.total_count - 1  # -1 for dummy block
             self.max_requests = self.max_requests // tp_size * tp_size
             self.max_requests = self.max_requests // self.REQUEST_ROUNDER * self.REQUEST_ROUNDER
         else:
             # User can control request overflow via max_requests.
-            self.max_requests = max_requests
+            self.max_requests = inference_config.max_requests
 
-        self.max_tokens = max_tokens or self.DEFAULT_MAX_TOKENS
+        self.max_tokens = inference_config.max_tokens or self.DEFAULT_MAX_TOKENS
 
         assert self.max_tokens >= self.max_requests, (
             f"max_tokens ({self.max_tokens}) must be >= "
@@ -518,32 +416,33 @@ def __init__(
         )
 
         # CUDA graph config list
-        is_expert_parallel = parallel_state.get_expert_model_parallel_world_size() > 1
+        self.use_cuda_graphs_for_non_decode_steps = (
+            inference_config.use_cuda_graphs_for_non_decode_steps
+        )
         self.cuda_graph_batch_dimensions_list, self.cuda_graph_token_counts = (
             CUDAGraphBatchDimensionBuilder.generate_cuda_graph_batch_dimensions_list(
                 tp_size=tp_size,
-                num_cuda_graphs=num_cuda_graphs,
+                num_cuda_graphs=inference_config.num_cuda_graphs,
                 cuda_graph_max_tokens=self.max_requests,
-                cuda_graph_mixed_prefill_count=cuda_graph_mixed_prefill_count,
+                cuda_graph_mixed_prefill_count=inference_config.cuda_graph_mixed_prefill_count,
                 max_requests=self.max_requests,
                 max_tokens=self.max_tokens,
                 max_sequence_length=self.max_sequence_length,
-                use_cuda_graphs_for_non_decode_steps=use_cuda_graphs_for_non_decode_steps,
+                use_cuda_graphs_for_non_decode_steps=self.use_cuda_graphs_for_non_decode_steps,
             )
         )
 
         self._using_cuda_graph_this_step = False
-        self.use_cuda_graphs_for_non_decode_steps = use_cuda_graphs_for_non_decode_steps
         # Deal with chunked prefill
         self.chunked_prefill_request_id = -1
         self.has_explicit_chunked_prefill_req = False
 
         # FlashInfer.
-        if use_flashinfer_fused_rope is True:
+        if inference_config.use_flashinfer_fused_rope is True:
             assert HAVE_FLASHINFER, "flashinfer is not installed"
-        elif use_flashinfer_fused_rope is None:
-            use_flashinfer_fused_rope = HAVE_FLASHINFER
-        self.use_flashinfer_fused_rope = use_flashinfer_fused_rope
+        elif inference_config.use_flashinfer_fused_rope is None:
+            inference_config.use_flashinfer_fused_rope = HAVE_FLASHINFER
+        self.use_flashinfer_fused_rope = inference_config.use_flashinfer_fused_rope
 
         # Allocate GPU state.
         self.is_tensor_state_allocated = False
@@ -724,14 +623,7 @@ def deallocate_all_tensors(self):
 
     @classmethod
     def round_up_tokens(cls, value, tp_size=None):
-        """Round up to nearest multiple of `TOKEN_ROUNDER` (above) that is also divisible by tensor model parallel size."""
-        if not HAVE_PACKAGING:
-            raise ImportError(
-                "`packaging` is required for this functionality, please install it with `pip install packaging`"
-            )
-        if PkgVersion(mcore_version) < PkgVersion("0.13"):
-            return cls.round_up(value)
-
+        """Round up to nearest multiple of `TOKEN_ROUNDER` that is also divisible by tensor model parallel size."""
         # Make sure divisible by TP size
         if tp_size is None:
             # Check if parallel state is initialized before trying to get TP size
@@ -743,86 +635,9 @@ def round_up_tokens(cls, value, tp_size=None):
 
         return token_rounder * int(math.ceil(int(value) / token_rounder))
 
-    @classmethod
-    def from_model_and_args(cls, model, args, overrides: Optional[Dict[str, Any]] = None):
-        """
-        Instantiate a `DynamicInferenceContext` from the model and args.
-
-        Args:
-            model: The Megatron model instance.
-            args: The arguments object.
-            overrides (Optional[Dict[str, Any]]): A dictionary of values to override
-                the default arguments derived from `model` and `args`.
-
-        Returns:
-            DynamicInferenceContext: The initialized inference context.
-        """
-        config = model.config
-
-        # Max sequence length.
-        position_embedding_type = get_attr_wrapped_model(model, "position_embedding_type")
-        model_max_seq_len = get_attr_wrapped_model(model, "max_sequence_length")
-        inf_max_seq_len = args.inference_max_seq_length
-
-        if position_embedding_type == "learned_absolute":
-            # When using absolute position embeddings, it is critical that the
-            # context's `max_sequence_length` is less than or equal to the model's
-            # `max_sequence_length`. Otherwise, the context's `position_ids` will
-            # contain ids greater than the dimension of the position embedding
-            # tensor, which will result in an index error.
-            if inf_max_seq_len:
-                max_sequence_length = min(model_max_seq_len, inf_max_seq_len)
-            else:
-                max_sequence_length = model_max_seq_len
-            assert max_batch_size <= model_max_seq_len
-        else:
-            max_sequence_length = inf_max_seq_len
-        if args.inference_dynamic_batching_max_requests is not None:
-            max_sequence_length = max(
-                max_sequence_length, args.inference_dynamic_batching_max_requests
-            )
-
-        mamba_inference_state_config = get_mamba_inference_state_config_from_model(model)
-        pg_collection = get_attr_wrapped_model(model, "pg_collection")
-
-        kwargs = {
-            "model_config": config,
-            "max_sequence_length": max_sequence_length,
-            "mamba_inference_state_config": mamba_inference_state_config,
-            "pg_collection": pg_collection,
-            "num_cuda_graphs": (
-                args.inference_dynamic_batching_num_cuda_graphs
-                if args.cuda_graph_impl == "local"
-                else None
-            ),
-            "block_size_tokens": args.inference_dynamic_batching_block_size,
-            "buffer_size_gb": args.inference_dynamic_batching_buffer_size_gb,
-            "paused_buffer_size_gb": args.inference_dynamic_batching_paused_buffer_size_gb,
-            "max_requests": args.inference_dynamic_batching_max_requests,
-            "max_tokens": args.inference_dynamic_batching_max_tokens,
-            "materialize_only_last_token_logits": not args.return_log_probs,
-            "use_flashinfer_fused_rope": args.use_flashinfer_fused_rope,
-            "unified_memory_level": args.inference_dynamic_batching_unified_memory_level,
-            "cuda_graph_max_tokens": args.inference_dynamic_batching_cuda_graph_max_tokens,
-            "cuda_graph_mixed_prefill_count": args.inference_dynamic_batching_cuda_graph_mixed_prefill_count,
-            "persist_cuda_graphs": args.rl_training_cuda_graphs,
-        }
-
-        if overrides is not None:
-            kwargs.update(overrides)
-
-        return cls(**kwargs)
-
     @classmethod
     def round_up_requests(cls, value, tp_size=None):
-        """Round up to nearest multiple of `REQUEST_ROUNDER` (above) that is also divisible by tensor model parallel size."""
-        if not HAVE_PACKAGING:
-            raise ImportError(
-                "`packaging` is required for this functionality, please install it with `pip install packaging`"
-            )
-        if PkgVersion(mcore_version) < PkgVersion("0.13"):
-            return cls.round_up(value)
-
+        """Round up to nearest multiple of `REQUEST_ROUNDER` that is also divisible by tensor model parallel size."""
         # Make sure divisible by TP size
         if tp_size is None:
             # Check if parallel state is initialized before trying to get TP size
@@ -834,16 +649,6 @@ def round_up_requests(cls, value, tp_size=None):
 
         return request_rounder * int(math.ceil(int(value) / request_rounder))
 
-    @classmethod
-    def round_up(cls, value):
-        """Deprecated in favor of round_up_tokens and round_up_requests."""
-        warnings.warn(
-            "`round_up` is deprecated in favor of `round_up_tokens` or `round_up_requests` "
-            "and will be removed in `megatron-core` 0.14."
-        )
-        ROUNDER = getattr(cls, "ROUNDER", 64)
-        return ROUNDER * int(math.ceil(int(value) / ROUNDER))
-
     def is_static_batching(self) -> bool:
         """Is static batching? False."""
         return False
@@ -864,6 +669,7 @@ def has_unfinished_requests(self) -> bool:
 
     def cu_query_lengths(self) -> Tuple[Tensor, int]:
         """Cumulative query sequence lengths."""
+        assert self.active_attn_metadata is not None
         return (
             self.active_attn_metadata["mha_metadata"].state_data["cu_query_seq_lengths"],
             self.active_attn_metadata["mha_metadata"].state_data["max_seqlen_q"],
@@ -871,6 +677,7 @@ def cu_query_lengths(self) -> Tuple[Tensor, int]:
 
     def cu_kv_lengths(self) -> Tuple[Tensor, Tensor, int]:
         """Cumulative key/value sequence lengths."""
+        assert self.active_attn_metadata is not None
         return (
             self.active_attn_metadata["mha_metadata"].state_data["cu_kv_seq_lengths"],
             self.active_attn_metadata["mha_metadata"].state_data["kv_seq_lengths"],
@@ -940,18 +747,20 @@ def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor)
                 : self.padded_active_token_count
             ]
 
-    def key_value_cache(self, layer_number: int) -> Tuple[Tensor, Tensor]:
+    def key_value_cache(self, layer_number: int) -> Tuple[Tensor, Optional[Tensor], Tensor]:
         """Read from KV cache.
 
         Args:
             layer_number (int): Layer number.
 
         Return:
-            (Tuple[Tensor, Tensor]) The key and value pointer tensors that point
-            to blocks within the block-level memory buffer.
+            (Tuple[Tensor, Tensor, Tensor]) The key and value pointer tensors that point
+            to blocks within the block-level memory buffer as well as the block table.
         """
         attention_layer_number = self.layer_map[layer_number - 1]
 
+        assert self.active_attn_metadata is not None
+
         if self.cache_mla_latent:
             return (
                 self.memory_buffer[attention_layer_number],
@@ -1367,9 +1176,9 @@ def initialize_attention_state(
         ] = 0
 
         self.active_attn_metadata = (
-            self.graph_attn_metadata
+            self.graph_attn_metadata  # type: ignore[assignment]
             if self.using_cuda_graph_this_step()
-            else self.non_graph_attn_metadata
+            else self.non_graph_attn_metadata  # type: ignore[assignment]
         )
 
         # Update cu_query_seq_lengths, max_seqlen_q.
@@ -1394,6 +1203,7 @@ def initialize_attention_state(
                     has_explicit_chunked_prefill_req=False,
                 )
 
+        assert self.active_attn_metadata is not None
         self.active_attn_metadata["mha_metadata"].update(
             request_query_lengths=query_lengths_view,
             request_kv_length_offsets=request_kv_length_offsets_view,
@@ -1526,7 +1336,7 @@ def last_token_logits(self, logits: Tensor) -> Tensor:
 
         return last_token_logits
 
-    def check_availability(self, req: DynamicInferenceRequest) -> (bool, bool, bool):
+    def check_availability(self, req: DynamicInferenceRequest) -> Tuple[bool, bool, bool]:
         """
         Check if the request can be added to the context.
         """
@@ -1758,7 +1568,7 @@ def resume_paused_requests(
         active_request_count: int,
         newly_paused_request_ids: torch.Tensor,
         next_tokens: torch.Tensor,
-    ) -> tuple[int, int, torch.Tensor]:
+    ) -> tuple[int, torch.Tensor]:
         """Resume as many paused requests as we have space for in the active buffer.
 
         Args:
@@ -1837,7 +1647,7 @@ def resume_paused_requests(
 
     def evict_overflow_paused_requests(
         self, active_request_count: int, next_tokens: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> Optional[tuple[torch.Tensor, torch.Tensor]]:
         """Evict requests that overflow the paused buffer.
 
         Args:
diff --git a/megatron/core/inference/contexts/static_context.py b/megatron/core/inference/contexts/static_context.py
index 98ba8b5185d..5cfc510bf46 100644
--- a/megatron/core/inference/contexts/static_context.py
+++ b/megatron/core/inference/contexts/static_context.py
@@ -1,5 +1,7 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 
+from megatron.core.inference.config import StaticInferenceConfig
+
 from .base_context import BaseInferenceContext
 
 
@@ -15,7 +17,8 @@ class StaticInferenceContext(BaseInferenceContext):
     def __init__(
         self, max_batch_size: int, max_sequence_length: int, use_flashinfer_fused_rope: bool = None
     ):
-        super().__init__(materialize_only_last_token_logits=True)
+        inference_config = StaticInferenceConfig(materialize_only_last_token_logits=True)
+        super().__init__(inference_config=inference_config)
         self.max_sequence_length = max_sequence_length
         self.max_batch_size = max_batch_size
         self.sequence_len_offset = 0
diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index 1c19547d49d..07aae7b5528 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -44,12 +44,10 @@
 from megatron.core.utils import (
     experimental_api,
     get_asyncio_loop,
-    get_attr_wrapped_model,
     get_pg_rank,
     get_pg_size,
     get_pg_src_rank,
     internal_api,
-    log_single_rank,
     trace_async_exceptions,
 )
 
@@ -133,25 +131,9 @@ class DynamicInferenceEngine(AbstractEngine):
             outputs and detokenizer the output tokens.
         inference_context (DynamicInferenceContext): Context for managing in-flight
             batching and a dynamic block-level KV cache (similar to paged attention).
-        random_seed (Optional[int]): Use a random seed if you want deterministic
-            results. Defaults to None.
-        inference_logging_step_interval (int): The step interval at which to log
-        inference metrics to wandb. Defaults to 0, which means no logging.
     """
 
-    def __init__(
-        self,
-        controller: TextGenerationController,
-        context: DynamicInferenceContext,
-        enable_cuda_graph: Optional[bool] = None,
-        random_seed: Optional[int] = None,
-        *,
-        track_paused_request_events: bool = False,
-        enable_chunked_prefill: bool = True,
-        metrics_writer: Optional['WandbModule'] = None,
-        inference_logging_step_interval: int = 0,
-        pg_collection: Optional[ProcessGroupCollection] = None,
-    ):
+    def __init__(self, controller: TextGenerationController, context: DynamicInferenceContext):
 
         assert isinstance(
             controller, TextGenerationController
@@ -159,41 +141,28 @@ def __init__(
         assert isinstance(
             context, DynamicInferenceContext
         ), f"context must be a DynamicInferenceContext, got {type(context)}"
-        assert isinstance(random_seed, int), f"random_seed must be an int, got {type(random_seed)}"
-
-        # Deprecate `enable_cuda_graph`.
-        if enable_cuda_graph is not None:
-            warnings.warn(
-                "The `enable_cuda_graph` argument is deprecated and will be "
-                "removed in `megatron-core 0.15`. `enable_cuda_graph` is now "
-                "read directly from the transformer config object."
-            )
-            self.enable_cuda_graph = enable_cuda_graph
-        else:
-            self.enable_cuda_graph = (
-                controller.inference_wrapped_model.model.config.enable_cuda_graph
-            )
 
-        if pg_collection is not None:
-            self.pg_collection = pg_collection
+        model_config = controller.inference_wrapped_model.model.config
+        inference_config = context.inference_config
+
+        if inference_config.pg_collection is not None:
+            self.pg_collection = inference_config.pg_collection
         else:
             self.pg_collection = ProcessGroupCollection.use_mpu_process_groups()
 
         # Initialization options.
         self.controller = controller
         self.context = context
-        self.random_seed = random_seed
-        self.track_paused_request_events = track_paused_request_events
-        self.enable_chunked_prefill = enable_chunked_prefill
-        self.metrics_writer = metrics_writer
-        self.inference_logging_step_interval = inference_logging_step_interval
-        self.unified_memory_level = context.unified_memory_level
-        self.persist_cuda_graphs = context.persist_cuda_graphs
-
-        if enable_cuda_graph is not None:
-            self.cuda_graph_impl = "local" if enable_cuda_graph else "none"
-        else:
-            self.cuda_graph_impl = controller.inference_wrapped_model.model.config.cuda_graph_impl
+        self.track_paused_request_events = inference_config.track_paused_request_events
+        self.enable_chunked_prefill = inference_config.enable_chunked_prefill
+        self.metrics_writer = inference_config.metrics_writer
+        self.logging_step_interval = inference_config.logging_step_interval
+        self.unified_memory_level = inference_config.unified_memory_level
+        self.persist_cuda_graphs = inference_config.persist_cuda_graphs
+        self.materialize_only_last_token_logits = (
+            inference_config.materialize_only_last_token_logits
+        )
+        self.cuda_graph_impl = model_config.cuda_graph_impl
 
         # Initialize engine.
         self.reset()
@@ -204,7 +173,7 @@ def __init__(
         )
 
         # Configure wandb to use separate step counter for inference metrics (only once)
-        if self.inference_logging_step_interval > 0 and self.metrics_writer is not None:
+        if self.logging_step_interval > 0 and self.metrics_writer is not None:
             logging.info(
                 f"\033[1;93m[INFERENCE]\033[0m "
                 f"\033[1;95mLogging inference metrics to wandb (rank {self.rank})\033[0m"
@@ -231,69 +200,6 @@ def __init__(
         # Create cuda graphs.
         self.create_cuda_graphs()
 
-    @classmethod
-    def from_model_and_args(
-        cls,
-        model,
-        args,
-        controller: Optional[TextGenerationController] = None,
-        context: Optional[DynamicInferenceContext] = None,
-    ):
-        """
-        Initializes a `DynamicInferenceEngine` from the model and args.
-
-        Args:
-            model: The Megatron model instance.
-            args: The arguments object.
-            controller (Optional[TextGenerationController]): An optional existing
-                controller. If None, one is created from the model and args.
-            context (Optional[DynamicInferenceContext]): An optional existing
-                context. If None, one is created from the model and args.
-
-        Returns:
-            DynamicInferenceEngine: The initialized inference engine.
-        """
-        if context is None:
-            context = DynamicInferenceContext.from_model_and_args(model, args)
-        if controller is None:
-            controller = TextGenerationController.from_model_and_args(model, args, context)
-
-        # The model may have a custom ProcessGroupCollection with a different TP / PP size.
-        pg_collection = get_attr_wrapped_model(model, "pg_collection")
-
-        # Get inference logging configuration from args
-        log_inference_wandb = args.inference_wandb_logging
-        inference_logging_step_interval = args.inference_logging_step_interval
-
-        # Get metrics writer if logging is enabled and on the logging rank
-        # Use the same rank convention as training (last rank logs)
-        metrics_writer = None
-        if (
-            inference_logging_step_interval > 0
-            and log_inference_wandb
-            and args.rank == (args.world_size - 1)
-        ):
-            metrics_writer = get_wandb_writer()
-            if metrics_writer is None:
-                log_single_rank(
-                    logger,
-                    logging.WARNING,
-                    "WARNING: --rl-inference-logging-step-interval is set but no metrics writer "
-                    "wandb module is available. Inference logging will be disabled.",
-                )
-
-        return cls(
-            controller,
-            context,
-            enable_cuda_graph=model.config.cuda_graph_impl == "local",
-            random_seed=args.seed,
-            track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events,
-            enable_chunked_prefill=not args.disable_chunked_prefill,
-            metrics_writer=metrics_writer,
-            inference_logging_step_interval=args.inference_logging_step_interval,
-            pg_collection=pg_collection,
-        )
-
     def reset(self) -> None:
         """Reset by removing all requests and reset all state."""
 
@@ -755,7 +661,7 @@ def _add_request(
             request.sampling_params.return_log_probs
             and not request.sampling_params.skip_prompt_log_probs
         ):
-            assert not self.context.materialize_only_last_token_logits, (
+            assert not self.materialize_only_last_token_logits, (
                 "Prompt log probs cannot be calculated if only last token logits are materialized. "
                 "Set materialize_only_last_token_logits to False in DynamicInferenceContext "
                 "or skip_prompt_log_probs to True in SamplingParams."
@@ -944,7 +850,7 @@ def post_process_requests(
                 # For chunked prefill with materialize_only_last_token_logits, discard intermediate log probs
                 if (
                     request_id == self.context.chunked_prefill_request_id
-                    and self.context.materialize_only_last_token_logits
+                    and self.materialize_only_last_token_logits
                 ):
                     request.prompt_log_probs = []
                     request.generated_log_probs = []
@@ -1224,9 +1130,9 @@ async def async_forward(self) -> Tuple[Dict, Dict, float, int]:
         range_pop()
 
         if (
-            self.inference_logging_step_interval > 0
+            self.logging_step_interval > 0
             and self.step_count > 0
-            and self.step_count % self.inference_logging_step_interval == 0
+            and self.step_count % self.logging_step_interval == 0
             and self.metrics_writer is not None
         ):
             kvcache_util_stats = self.context.get_kvcache_utilization_stats()
@@ -1358,10 +1264,7 @@ async def async_bookkeep(
                 raise ValueError(f"Unsupported metrics writer type: {type(self.metrics_writer)}")
 
         # Print context state.
-        if (
-            self.inference_logging_step_interval > 0
-            and step_count % self.inference_logging_step_interval == 0
-        ):
+        if self.logging_step_interval > 0 and step_count % self.logging_step_interval == 0:
             mem = torch.cuda.memory_stats()
             step_type = "decode" if context_state["is_decode_only"] else "non-decode"
             output_str = (
diff --git a/megatron/core/inference/engines/static_engine.py b/megatron/core/inference/engines/static_engine.py
index 1ac84ee5f79..436b60c039f 100644
--- a/megatron/core/inference/engines/static_engine.py
+++ b/megatron/core/inference/engines/static_engine.py
@@ -8,6 +8,7 @@
 import torch
 
 from megatron.core.inference.async_stream import AsyncStream
+from megatron.core.inference.config import DynamicInferenceConfig, MambaInferenceStateConfig
 from megatron.core.inference.contexts import DynamicInferenceContext, StaticInferenceContext
 from megatron.core.inference.engines.abstract_engine import AbstractEngine
 from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine
@@ -17,7 +18,7 @@
 from megatron.core.inference.text_generation_controllers.text_generation_controller import (
     TextGenerationController,
 )
-from megatron.core.utils import get_asyncio_loop, get_mamba_inference_state_config_from_model
+from megatron.core.utils import get_asyncio_loop
 
 try:
     from tqdm import tqdm
@@ -91,7 +92,7 @@ def __init__(
 
         self.scheduler = Scheduler(max_batch_size=max_batch_size)
 
-        mamba_inference_state_config = get_mamba_inference_state_config_from_model(
+        mamba_inference_state_config = MambaInferenceStateConfig.from_model(
             self.inference_wrapped_model.model
         )
 
@@ -99,13 +100,15 @@ def __init__(
             if not legacy:
                 dynamic_context = DynamicInferenceContext(
                     model_config=self.config,
-                    max_sequence_length=original_context.max_sequence_length,
-                    buffer_size_gb=buffer_size_gb,
-                    mamba_inference_state_config=mamba_inference_state_config,
-                    max_requests=max_batch_size,
-                    num_cuda_graphs=1,
-                    block_size_tokens=256,
-                    unified_memory_level=0,
+                    inference_config=DynamicInferenceConfig(
+                        max_sequence_length=original_context.max_sequence_length,
+                        buffer_size_gb=buffer_size_gb,
+                        mamba_inference_state_config=mamba_inference_state_config,
+                        max_requests=max_batch_size,
+                        num_cuda_graphs=1,
+                        block_size_tokens=256,
+                        unified_memory_level=0,
+                    ),
                 )
 
                 self.controller.inference_wrapped_model.inference_context = dynamic_context
@@ -113,9 +116,7 @@ def __init__(
                 self.controller._init_dynamic_sampling_tensors()
 
                 self.dynamic_engine = DynamicInferenceEngine(
-                    controller=self.controller,
-                    context=dynamic_context,
-                    random_seed=self.random_seed,
+                    controller=self.controller, context=dynamic_context
                 )
         except Exception as e:
             # Get exception details for better debugging
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index e5c40131e0e..dc0c2af5b7b 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -11,37 +11,26 @@
 import torch
 import torch.nn.functional as F
 from torch import Tensor
-from torch.distributed import ProcessGroup
 
+from megatron.core import parallel_state
 from megatron.core.inference.async_stream import AsyncStream
 from megatron.core.inference.communication_utils import (
     broadcast_from_last_pipeline_stage,
-    is_pipeline_first_stage,
     is_pipeline_last_stage,
 )
-from megatron.core.inference.contexts.base_context import BaseInferenceContext
 from megatron.core.inference.contexts.dynamic_context import MaxSequenceLengthOverflowError
 from megatron.core.inference.contexts.static_context import StaticInferenceContext
 from megatron.core.inference.inference_request import InferenceRequest, Status
 from megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper import (
     AbstractModelInferenceWrapper,
 )
-from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
-    GPTInferenceWrapper,
-)
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.utils import get_attention_mask, set_decode_expert_padding
 from megatron.core.models.multimodal.llava_model import LLaVAModel
-from megatron.core.tokenizers.text.utils.build_tokenizer import build_tokenizer
 from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.transformer.moe.moe_layer import BaseMoELayer
 from megatron.core.transformer.utils import set_model_to_sequence_parallel
-from megatron.core.utils import (
-    get_asyncio_loop,
-    get_attr_wrapped_model,
-    get_model_config,
-    unwrap_model,
-)
+from megatron.core.utils import get_asyncio_loop, get_model_config, unwrap_model
 
 try:
     import transformer_engine as te  # pylint: disable=unused-import
@@ -64,36 +53,29 @@ class TextGenerationController:
         inference_wrapped_model (AbstractModelInferenceWrapper): A model that
             is wrapped using the specs given in the abstract_model_inference_wrapper.py
         tokenizer (_type_): Tokenizer used for tokenizing and detokenizing the prompts
-        pp_group (ProcessGroup): Process group for pipeline parallelism
     """
 
-    def __init__(
-        self,
-        inference_wrapped_model: AbstractModelInferenceWrapper,
-        tokenizer,
-        pp_group: ProcessGroup = None,
-    ):
+    def __init__(self, inference_wrapped_model: AbstractModelInferenceWrapper, tokenizer):
         self.inference_wrapped_model = inference_wrapped_model
         self.model_config = self.inference_wrapped_model.model.config
+        self.inference_config = self.inference_wrapped_model.inference_context.inference_config
         self.tokenizer = tokenizer
 
-        self.pp_group = pp_group
+        pg_collection = self.inference_config.pg_collection
+        if pg_collection is not None:
+            self.pp_group = pg_collection.pp
+        else:
+            self.pp_group = parallel_state.get_pipeline_model_parallel_group()
 
-        # For models without pipeline parallelism, is_first_stage and is_last_stage returns True
-        self.model_is_pipeline_parallel = not (
-            is_pipeline_first_stage(self.pp_group) and is_pipeline_last_stage(self.pp_group)
-        )
+        self.model_is_pipeline_parallel = self.model_config.pipeline_model_parallel_size > 1
 
         # Use padded vocab size because tokenizer vocab size might pad to nearest power of 2.
-        if isinstance(self.inference_wrapped_model.model, LLaVAModel):
-            # TODO(ksanthanam): Consider deprecating this check if LLaVAModel is no longer used
-            self.vocab_size = get_attr_wrapped_model(
-                self.inference_wrapped_model.model, "language_model"
-            ).vocab_size
+        # TODO(ksanthanam): Consider deprecating this check if LLaVAModel is no longer used
+        unwrapped_model = unwrap_model(self.inference_wrapped_model.model)
+        if isinstance(unwrapped_model, LLaVAModel):
+            self.vocab_size = unwrapped_model.language_model.vocab_size
         else:
-            self.vocab_size = get_attr_wrapped_model(
-                self.inference_wrapped_model.model, "vocab_size"
-            )
+            self.vocab_size = unwrapped_model.vocab_size
 
         self.sampling_rng = torch.Generator(device=torch.cuda.current_device())
         self.sampling_rng.manual_seed(self.model_config.inference_sampling_seed)
@@ -101,32 +83,6 @@ def __init__(
         if self.inference_wrapped_model.inference_context.is_dynamic_batching():
             self._init_dynamic_sampling_tensors()
 
-    @classmethod
-    def from_model_and_args(
-        cls,
-        model,
-        args,
-        context: BaseInferenceContext,
-        model_inference_wrapper_cls: type[AbstractModelInferenceWrapper] = GPTInferenceWrapper,
-    ):
-        """
-        Initializes a `TextGenerationController` from the model and args.
-
-        Args:
-            model: The Megatron model instance.
-            args: The arguments object.
-            context (BaseInferenceContext): The inference context.
-            model_inference_wrapper_cls (type[AbstractModelInferenceWrapper]): The class
-                used to wrap the model for inference. Defaults to GPTInferenceWrapper.
-
-        Returns:
-            TextGenerationController: The initialized text generation controller.
-        """
-        tokenizer = build_tokenizer(args)
-        model = model_inference_wrapper_cls(model, context)
-        model.model_is_pipeline_parallel = model.config.pipeline_model_parallel_size > 1
-        return cls(model, tokenizer)
-
     def set_stop_word_finished_ids_callback(self, callback):
         """Set a callback to get request IDs that should be marked as finished due to stop words.
 
@@ -626,7 +582,7 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
         if self.model_is_pipeline_parallel:
             logits_seq_len = (
                 active_request_count
-                if context.materialize_only_last_token_logits
+                if context.inference_config.materialize_only_last_token_logits
                 else input_ids.shape[1]
             )
             logits_shape = [1, logits_seq_len, self.vocab_size]
@@ -682,7 +638,7 @@ def _dynamic_step_sample_logits(self, logits: Tensor):
 
         # Last token logits.
         context = self.inference_wrapped_model.inference_context
-        if context.materialize_only_last_token_logits:
+        if context.inference_config.materialize_only_last_token_logits:
             # When materialize_only_last_token_logits is true, last_token_logits is
             # already called in the forward pass of GPT.
             last_token_logits = logits.squeeze(0)
@@ -727,7 +683,7 @@ def _dynamic_step_calculate_log_probs(self, logits: Tensor) -> Optional[Tensor]:
         return context.calculate_log_probs(
             logits,
             self._sampled_tokens_cuda[:active_request_count],
-            only_last_token_logits=context.materialize_only_last_token_logits,
+            only_last_token_logits=context.inference_config.materialize_only_last_token_logits,
         )
 
     def _dynamic_step_calculate_top_n_logprobs(
@@ -755,7 +711,7 @@ def _dynamic_step_calculate_top_n_logprobs(
         active_request_slice = slice(context.paused_request_count, context.total_request_count)
 
         # Handle decode-only mode (only last token)
-        if context.materialize_only_last_token_logits or context.is_decode_only():
+        if context.inference_config.materialize_only_last_token_logits or context.is_decode_only():
             # In decode mode or when only last token logits are materialized,
             # logits already represent only the last tokens
             log_probs = log_probs_tensor[:active_request_count]
@@ -1231,7 +1187,7 @@ def generate_all_output_tokens_static_batch(
                     or not (sampling_params.return_log_probs or sampling_params.top_n_logprobs > 0)
                 )
                 inference_context = self.inference_wrapped_model.inference_context
-                inference_context.materialize_only_last_token_logits = (
+                inference_context.inference_config.materialize_only_last_token_logits = (
                     materialize_only_last_token_logits
                 )
 
diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index e70221d2cfa..e4d7a879c30 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -635,7 +635,10 @@ def _postprocess(
                     )
         sequence_parallel_override = False
 
-        if in_inference_mode and inference_context.materialize_only_last_token_logits:
+        if (
+            in_inference_mode
+            and inference_context.inference_config.materialize_only_last_token_logits
+        ):
             if inference_context.is_static_batching():
                 hidden_states = hidden_states[-1:, :, :]
             else:
@@ -665,7 +668,7 @@ def _postprocess(
             assert (
                 in_inference_mode
                 and inference_context.is_dynamic_batching()
-                and inference_context.materialize_only_last_token_logits
+                and inference_context.inference_config.materialize_only_last_token_logits
             )
             self.output_layer.sequence_parallel = True
 
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index 0d71ead4b0f..baf8875735c 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -265,7 +265,10 @@ def forward(
             output_weight = self.shared_embedding_or_output_weight()
 
         sequence_parallel_override = False
-        if in_inference_mode and inference_context.materialize_only_last_token_logits:
+        if (
+            in_inference_mode
+            and inference_context.inference_config.materialize_only_last_token_logits
+        ):
             if inference_context.is_static_batching():
                 hidden_states = hidden_states[-1:, :, :]
             else:
@@ -295,7 +298,7 @@ def forward(
             assert (
                 in_inference_mode
                 and inference_context.is_dynamic_batching()
-                and inference_context.materialize_only_last_token_logits
+                and inference_context.inference_config.materialize_only_last_token_logits
             )
             self.output_layer.sequence_parallel = True
 
diff --git a/megatron/core/utils.py b/megatron/core/utils.py
index ee5ae0ec92a..64f20fec43d 100644
--- a/megatron/core/utils.py
+++ b/megatron/core/utils.py
@@ -2419,25 +2419,6 @@ async def wrapper(*args, **kwargs):
     return _decorate if func is None else _decorate(func)
 
 
-def get_mamba_inference_state_config_from_model(model) -> Optional["MambaInferenceStateConfig"]:
-    """Returns Mamba inference state config from the model if it is a hybrid model."""
-    from megatron.core.inference.contexts.attention_context.mamba_metadata import (
-        MambaInferenceStateConfig,
-    )
-    from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols
-
-    decoder = get_attr_wrapped_model(model, "decoder")
-    layer_type_list = getattr(decoder, "layer_type_list", None)
-    if layer_type_list is not None and Symbols.MAMBA in layer_type_list:
-        (mamba_conv_states_shape, mamba_ssm_states_shape) = decoder.mamba_state_shapes_per_request()
-        return MambaInferenceStateConfig(
-            layer_type_list=layer_type_list,
-            mamba_conv_states_shape=mamba_conv_states_shape,
-            mamba_ssm_states_shape=mamba_ssm_states_shape,
-        )
-    return None
-
-
 # ============================================================================
 # Backward Compatibility Decorators
 # ============================================================================
diff --git a/megatron/inference/__init__.py b/megatron/inference/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/megatron/inference/utils.py b/megatron/inference/utils.py
new file mode 100644
index 00000000000..fc3b64198f4
--- /dev/null
+++ b/megatron/inference/utils.py
@@ -0,0 +1,311 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+from argparse import ArgumentParser
+from functools import partial
+from typing import Optional
+
+from gpt_builders import gpt_builder
+from mamba_builders import mamba_builder
+from megatron.core.inference.config import DynamicInferenceConfig, MambaInferenceStateConfig
+from megatron.core.inference.contexts import DynamicInferenceContext
+from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
+    GPTInferenceWrapper,
+)
+from megatron.core.inference.text_generation_controllers.text_generation_controller import (
+    TextGenerationController,
+)
+from megatron.core.inference.engines import DynamicInferenceEngine
+from megatron.core.tokenizers.text.utils.build_tokenizer import build_tokenizer
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.utils import get_attr_wrapped_model
+from megatron.training import get_args, get_tokenizer
+from megatron.training import get_model as _get_model
+from megatron.training.checkpointing import load_checkpoint
+from model_provider import model_provider
+
+
+def get_model() -> MegatronModule:
+    """Initialize model and load checkpoint."""
+
+    args = get_args()
+
+    if args.model_provider == "gpt":
+        model_builder = gpt_builder
+    elif args.model_provider == "mamba":
+        model_builder = mamba_builder
+    else:
+        raise ValueError(f"Invalid model provider {args.model_provider}")
+
+    # Build model.
+    model = _get_model(partial(model_provider, model_builder), wrap_with_ddp=False)
+
+    # Load checkpoint.
+    assert args.load is not None
+    args.exit_on_missing_checkpoint = True
+    load_checkpoint(
+        ddp_model=model,
+        optimizer=None,
+        opt_param_scheduler=None,
+        strict=not args.inference_ckpt_non_strict,
+    )
+
+    # No virtual PP.
+    assert len(model) == 1, "Above condition should have caught this"
+    model = model[0]
+
+    # Eval mode.
+    model.eval()
+
+    return model
+
+
+def add_inference_args(parser: ArgumentParser) -> ArgumentParser:
+    """Add inference command line arguments to the parser."""
+
+    group = parser.add_argument_group(title='Inference')
+
+    group.add_argument("--temperature", type=float, default=1.0, help='Sampling temperature.')
+    group.add_argument("--top_k", type=int, default=1, help='Top k sampling.')
+    group.add_argument("--top_p", type=float, default=0.0, help='Top p sampling.')
+    group.add_argument(
+        "--return-log-probs",
+        action='store_true',
+        default=False,
+        help='Return the log probabilities of the final output tokens',
+    )
+    group.add_argument(
+        "--prompts",
+        metavar='N',
+        type=str,
+        nargs='+',
+        help='Input prompts with each prompt within quotes and seperated by space',
+    )
+    group.add_argument(
+        "--num-tokens-to-prompt",
+        type=int,
+        nargs="+",
+        default=[64, 1024],
+        help='Number of tokens to use for simulated prompts. This should be a '
+        'space-separated pair of integers, and the generated prompt lengths will '
+        'be uniformly sampled within this range.',
+    )
+    group.add_argument(
+        "--num-tokens-to-generate",
+        type=int,
+        default=30,
+        help='Number of tokens to generate for each prompt',
+    )
+    group.add_argument(
+        "--num-tokens-from-file",
+        action='store_true',
+        default=False,
+        help='Use per-prompt num_tokens_to_generate from prompt file',
+    )
+    group.add_argument(
+        "--top-n-logprobs",
+        type=int,
+        default=0,
+        help='Return the top n logprobs for the generated tokens and their corresponding token as a dictionary',
+    )
+    group.add_argument(
+        "--incoming-requests-per-step",
+        type=int,
+        default=None,
+        help="Add a deterministic number of requests per step. This arg is "
+        "prioritized over `--incoming-requests-per-sec` below (which is non-"
+        "deterministic). Note that the number of requests added per step is "
+        "additionally limited by the inference context's `max_requests`, "
+        "`max_tokens`, and KV buffer size.",
+    )
+    group.add_argument(
+        "--incoming-requests-per-sec",
+        type=float,
+        default=100.0,
+        help="Simulated number of requests per second. Set to -1 to add all requests together.",
+    )
+    group.add_argument(
+        "--incoming-requests-duration",
+        type=float,
+        default=10.0,
+        help="Total amount of time to simulate that requests are "
+        "arriving. Multiply this value with "
+        "`--incoming-requests-per-sec` to get the approximate "
+        "total number of requests. Set to -1 to add all requests together.",
+    )
+    group.add_argument(
+        "--model-provider", choices=["mamba", "gpt"], default="gpt", help="Model provider"
+    )
+    group.add_argument(
+        "--skip-prompt-log-probs", action='store_true', default=False, help='Skip prompt log probs.'
+    )
+    group.add_argument(
+        "--stop-words",
+        metavar='WORD',
+        type=str,
+        nargs='+',
+        default=None,
+        help='Stop words to terminate generation. Each word should be quoted and '
+        'separated by space. Example: --stop-words "\\n\\n" "END" "###"',
+    )
+    group.add_argument(
+        "--output-path", type=str, default=None, help="Path to save generations as JSON"
+    )
+    group.add_argument(
+        "--output-every-n-results",
+        type=int,
+        default=1,
+        help="To minimize the output file size of larger runs, only write the "
+        "results of every `n` requests.",
+    )
+    group.add_argument(
+        "--prompt-file",
+        help='Jsonl file containing input prompts, where each item (i.e., line) '
+        'contains the field \'text\' where the value is the prompt. All other '
+        'fields within each item are ignored, and may be customized for each '
+        'application.',
+    )
+    group.add_argument(
+        "--prompt-file-num-truncate",
+        type=int,
+        help='Number of samples to use from the loaded prompt file (see '
+        '`--prompt-file` above). The first `--prompt-file-num-truncate` samples '
+        'will be used, in order.',
+    )
+    group.add_argument(
+        "--use-flashinfer-fused-rope",
+        action='store_true',
+        default=False,
+        help='Use flashinfer fused rope implementation.',
+    )
+    group.add_argument(
+        "--no-record-throughput",
+        action='store_false',
+        dest="record_throughput",
+        help="Disable throughput recording in --output-file",
+    )
+    group.add_argument(
+        "--inference-ckpt-non-strict",
+        action="store_true",
+        help="Load checkpoint with `strict=False`.",
+    )
+    group.add_argument(
+        "--termination-id",
+        type=int,
+        default=None,
+        help="Termination ID that overrides `tokenizer.eod`.",
+    )
+    group.add_argument(
+        "--suspend-resume-interval",
+        type=int,
+        default=None,
+        help="Suspend and resume the dynamic engine every "
+        "`suspend_resume_interval` steps. This is used to tet the suspend/resume "
+        "system.",
+    )
+    group.add_argument(
+        "--inference-repeat-n",
+        type=int,
+        default=1,
+        help="Repeat inference iterations N times for benchmarking.",
+    )
+    group.add_argument(
+        "--throughput-check-only",
+        action='store_true',
+        default=False,
+        help="If true, only run throughput check without verifying outputs.",
+    )
+
+    return parser
+
+
+def get_dynamic_inference_config_from_model_and_args(model: MegatronModule, args):
+    """Returns a `DynamicInferenceConfig` constructed from the model and command line arguments."""
+
+    # Max sequence length.
+    position_embedding_type = get_attr_wrapped_model(model, "position_embedding_type")
+    model_max_seq_len = get_attr_wrapped_model(model, "max_sequence_length")
+    inf_max_seq_len = args.inference_max_seq_length
+
+    if position_embedding_type == "learned_absolute":
+        # When using absolute position embeddings, it is critical that the
+        # context's `max_sequence_length` is less than or equal to the model's
+        # `max_sequence_length`. Otherwise, the context's `position_ids` will
+        # contain ids greater than the dimension of the position embedding
+        # tensor, which will result in an index error.
+        if inf_max_seq_len:
+            max_sequence_length = min(model_max_seq_len, inf_max_seq_len)
+        else:
+            max_sequence_length = model_max_seq_len
+        assert max_batch_size <= model_max_seq_len
+    else:
+        max_sequence_length = inf_max_seq_len
+    if args.inference_dynamic_batching_max_requests is not None:
+        max_sequence_length = max(max_sequence_length, args.inference_dynamic_batching_max_requests)
+
+    mamba_inference_state_config = MambaInferenceStateConfig.from_model(model)
+    pg_collection = get_attr_wrapped_model(model, "pg_collection")
+
+    # Get inference logging configuration from args
+    log_inference_wandb = args.inference_wandb_logging
+    inference_logging_step_interval = args.inference_logging_step_interval
+
+    # Get metrics writer if logging is enabled and on the logging rank
+    # Use the same rank convention as training (last rank logs)
+    metrics_writer = None
+    if (
+        inference_logging_step_interval > 0
+        and log_inference_wandb
+        and args.rank == (args.world_size - 1)
+    ):
+        metrics_writer = get_wandb_writer()
+        if metrics_writer is None:
+            log_single_rank(
+                logger,
+                logging.WARNING,
+                "WARNING: --rl-inference-logging-step-interval is set but no metrics writer "
+                "wandb module is available. Inference logging will be disabled.",
+            )
+
+    return DynamicInferenceConfig(
+        block_size_tokens=args.inference_dynamic_batching_block_size,
+        buffer_size_gb=args.inference_dynamic_batching_buffer_size_gb,
+        paused_buffer_size_gb=args.inference_dynamic_batching_paused_buffer_size_gb,
+        num_cuda_graphs=(
+            args.inference_dynamic_batching_num_cuda_graphs
+            if args.cuda_graph_impl == "local"
+            else None
+        ),
+        max_requests=args.inference_dynamic_batching_max_requests,
+        max_tokens=args.inference_dynamic_batching_max_tokens,
+        unified_memory_level=args.inference_dynamic_batching_unified_memory_level,
+        cuda_graph_mixed_prefill_count=args.inference_dynamic_batching_cuda_graph_mixed_prefill_count,
+        use_cuda_graphs_for_non_decode_steps=not args.decode_only_cuda_graphs,
+        persist_cuda_graphs=args.rl_training_cuda_graphs,
+        max_sequence_length=max_sequence_length,
+        mamba_inference_state_config=mamba_inference_state_config,
+        pg_collection=pg_collection,
+        use_flashinfer_fused_rope=args.use_flashinfer_fused_rope,
+        materialize_only_last_token_logits=not args.return_log_probs,
+        track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events,
+        enable_chunked_prefill=args.enable_chunked_prefill,
+        metrics_writer=metrics_writer,
+        logging_step_interval=args.inference_logging_step_interval,
+    )
+
+
+def get_dynamic_inference_engine(model: Optional[MegatronModule] = None) -> DynamicInferenceEngine:
+    """Builds a `DynamicInferenceEngine`."""
+    args = get_args()
+    if model is None:
+        model = get_model()
+    if args.legacy_tokenizer:
+        tokenizer = get_tokenizer()
+    else:
+        tokenizer = build_tokenizer(args)
+
+    inference_config = get_dynamic_inference_config_from_model_and_args(model, args)
+    context = DynamicInferenceContext(model.config, inference_config)
+    inference_wrapped_model = GPTInferenceWrapper(model, context)
+    controller = TextGenerationController(inference_wrapped_model, tokenizer)
+    engine = DynamicInferenceEngine(controller, context)
+    return engine
diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py
index 6f10f504768..f0a984ceb40 100644
--- a/megatron/rl/inference/megatron.py
+++ b/megatron/rl/inference/megatron.py
@@ -22,6 +22,7 @@
 from megatron.core.models.gpt.gpt_model import GPTModel
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.utils import get_attr_wrapped_model, log_single_rank
+from megatron.inference import get_dynamic_inference_engine
 from megatron.training import get_wandb_writer
 from megatron.training.global_vars import get_args, get_tokenizer
 
@@ -130,7 +131,7 @@ async def launch(cls, model: GPTModel, **kwargs):
                 "WARNING: Tokenizer has no BOS token so prompt will not have BOS token",
             )
 
-        inference_engine: DynamicInferenceEngine = DynamicInferenceEngine.from_model_and_args(model, args)
+        inference_engine: DynamicInferenceEngine = get_dynamic_inference_engine()
         await inference_engine.start_listening_to_data_parallel_coordinator(
             inference_coordinator_port=41521, launch_inference_coordinator=True
         )
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 1881fcbf1e2..458fcfc72f5 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1685,11 +1685,9 @@ def _add_inference_args(parser):
                        help='Number of chunks along sequence dimension for MLP '
                        'computation during prefill')
     # TODO(ksanthanam): Clean this up in future PR
-    group.add_argument('--enable-chunked-prefill', dest='disable_chunked_prefill',
-                       action='store_false', default=True,
+    group.add_argument('--enable-chunked-prefill', dest='enable_chunked_prefill',
+                       action='store_true', default=False,
                        help="Enable chunked prefill (disabled by default)")
-    group.add_argument('--disable-chunked-prefill', dest='disable_chunked_prefill',
-                       action='store_true', help=argparse.SUPPRESS)
     group.add_argument('--inference-dynamic-batching-cuda-graph-max-tokens',
                        type=int, default=16384,
                        help='Maximum number of tokens to capture in a cuda graph.')
diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py
index ee2f26c6853..246551b19e6 100644
--- a/tests/unit_tests/inference/contexts/test_dynamic_context.py
+++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py
@@ -1,14 +1,13 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 
+import contextlib
 import math
 
 import pytest
 import torch
 
 from megatron.core import parallel_state
-from megatron.core.inference.contexts.attention_context.mamba_metadata import (
-    MambaInferenceStateConfig,
-)
+from megatron.core.inference.config import DynamicInferenceConfig, MambaInferenceStateConfig
 from megatron.core.inference.contexts.dynamic_context import (
     DynamicInferenceContext,
     RequestOverflowError,
@@ -22,11 +21,17 @@
 from tests.unit_tests.test_utilities import Utils
 
 
-def set_rounder(value):
-    """Utility function to set the DynamicInferenceContext rounder."""
-    DynamicInferenceContext.ROUNDER = value  # For backwards compatibility
-    DynamicInferenceContext.TOKEN_ROUNDER = value
-    DynamicInferenceContext.REQUEST_ROUNDER = value
+@contextlib.contextmanager
+def rounder_override(n):
+    original_token_rounder = DynamicInferenceContext.TOKEN_ROUNDER
+    original_request_rounder = DynamicInferenceContext.REQUEST_ROUNDER
+    try:
+        DynamicInferenceContext.TOKEN_ROUNDER = n
+        DynamicInferenceContext.REQUEST_ROUNDER = n
+        yield
+    finally:
+        DynamicInferenceContext.TOKEN_ROUNDER = original_token_rounder
+        DynamicInferenceContext.REQUEST_ROUNDER = original_request_rounder
 
 
 class TestDynamicContext:
@@ -53,11 +58,8 @@ def _get_dynamic_context(
         max_tokens,
         is_hybrid_model=False,
         layer_type_list=None,
-        rounder=64,
         paused_buffer_size_gb=None,
     ):
-        set_rounder(rounder)
-
         if is_hybrid_model:
             if layer_type_list is None:
                 layer_type_list = [Symbols.MAMBA, Symbols.MLP, Symbols.ATTENTION, Symbols.MLP]
@@ -76,19 +78,21 @@ def _get_dynamic_context(
                 kv_channels=kv_channels,
                 num_attention_heads=num_attention_heads,
             ),
-            max_sequence_length=max_sequence_length,
-            num_cuda_graphs=None,
-            use_cuda_graphs_for_non_decode_steps=True,
-            buffer_size_gb=buffer_size_gb,
-            paused_buffer_size_gb=(
-                0.2 * buffer_size_gb if paused_buffer_size_gb is None else paused_buffer_size_gb
+            inference_config=DynamicInferenceConfig(
+                max_sequence_length=max_sequence_length,
+                num_cuda_graphs=None,
+                use_cuda_graphs_for_non_decode_steps=True,
+                buffer_size_gb=buffer_size_gb,
+                paused_buffer_size_gb=(
+                    0.2 * buffer_size_gb if paused_buffer_size_gb is None else paused_buffer_size_gb
+                ),
+                block_size_tokens=block_size_tokens,
+                max_tokens=max_tokens,
+                mamba_inference_state_config=mamba_inference_state_config,
+                use_flashinfer_fused_rope=None,  # default to using flash-infer if available
+                # this is for compatibility with the LTS environment
+                unified_memory_level=0,  # unit tests currently broken with UVM
             ),
-            block_size_tokens=block_size_tokens,
-            max_tokens=max_tokens,
-            mamba_inference_state_config=mamba_inference_state_config,
-            use_flashinfer_fused_rope=None,  # default to using flash-infer if available
-            # this is for compatibility with the LTS environment
-            unified_memory_level=0,  # unit tests currently broken with UVM
         )
         return dynamic_context
 
@@ -96,6 +100,7 @@ def teardown_method(self, method):
         Utils.destroy_model_parallel()
 
     @pytest.mark.internal
+    @rounder_override(64)
     @pytest.mark.parametrize("is_hybrid_model", [False, True])
     def test_initialize_dynamic_context(self, is_hybrid_model: bool):
         self._setup_model_parallel_group(1, 1)
@@ -110,7 +115,6 @@ def test_initialize_dynamic_context(self, is_hybrid_model: bool):
             block_size_tokens=128,
             max_tokens=None,
             is_hybrid_model=is_hybrid_model,
-            rounder=64,
         )
 
         if not is_hybrid_model:
@@ -148,6 +152,7 @@ def test_is_static_batching(self):
         assert not dynamic_context.is_static_batching()
 
     @pytest.mark.internal
+    @rounder_override(64)
     @pytest.mark.parametrize("is_hybrid_model", [False, True])
     def test_is_memory_available(self, is_hybrid_model):
         self._setup_model_parallel_group(1, 1)
@@ -171,6 +176,7 @@ def test_is_memory_available(self, is_hybrid_model):
         assert not dynamic_context.block_allocator.is_memory_available(1)
 
     @pytest.mark.internal
+    @rounder_override(1)
     @pytest.mark.parametrize("is_hybrid_model", [False, True])
     def test_request_overflow(self, is_hybrid_model: bool):
         self._setup_model_parallel_group(1, 1)
@@ -184,7 +190,6 @@ def test_request_overflow(self, is_hybrid_model: bool):
             buffer_size_gb=0.01,
             block_size_tokens=32,
             max_tokens=None,
-            rounder=1,
             is_hybrid_model=is_hybrid_model,
         )
         dynamic_context.max_requests //= 2
@@ -201,6 +206,7 @@ def test_request_overflow(self, is_hybrid_model: bool):
                 )  # Adding more than allowed requests
 
     @pytest.mark.internal
+    @rounder_override(1)
     @pytest.mark.parametrize("is_hybrid_model", [False, True])
     def test_token_overflow_error(self, is_hybrid_model: bool):
         self._setup_model_parallel_group(1, 1)
@@ -214,7 +220,6 @@ def test_token_overflow_error(self, is_hybrid_model: bool):
             buffer_size_gb=0.1,
             block_size_tokens=128,
             max_tokens=200,  # setting low, but >= context.max_requests.
-            rounder=1,
             is_hybrid_model=is_hybrid_model,
         )
 
@@ -230,6 +235,7 @@ def test_token_overflow_error(self, is_hybrid_model: bool):
             )  # Exceeding max token count
 
     @pytest.mark.internal
+    @rounder_override(64)
     @pytest.mark.parametrize("is_hybrid_model", [False, True])
     def test_reset(self, is_hybrid_model: bool):
         self._setup_model_parallel_group(1, 1)
@@ -304,6 +310,7 @@ def test_reset(self, is_hybrid_model: bool):
             assert torch.all(dynamic_context.mamba_metadata.request_to_mamba_state_idx == -1)
 
     @pytest.mark.internal
+    @rounder_override(64)
     @pytest.mark.parametrize("is_hybrid_model", [False, True])
     def test_allocate_and_release_memory_blocks(self, is_hybrid_model):
         self._setup_model_parallel_group(1, 1)
@@ -352,6 +359,7 @@ def test_allocate_and_release_memory_blocks(self, is_hybrid_model):
         )
 
     @pytest.mark.internal
+    @rounder_override(64)
     @pytest.mark.parametrize("is_hybrid_model", [False, True])
     def test_add_request(self, is_hybrid_model: bool):
         self._setup_model_parallel_group(1, 1)
@@ -425,6 +433,7 @@ def test_add_request(self, is_hybrid_model: bool):
         )
 
     @pytest.mark.internal
+    @rounder_override(64)
     def test_add_dummy_requests_parallel_populates_state(self):
         self._setup_model_parallel_group(1, 1)
 
@@ -523,6 +532,7 @@ def test_add_dummy_requests_parallel_populates_state(self):
         )
 
     @pytest.mark.internal
+    @rounder_override(64)
     def test_add_dummy_requests_parallel_hybrid_allocates_mamba(self):
         self._setup_model_parallel_group(1, 1)
 
@@ -553,6 +563,7 @@ def test_add_dummy_requests_parallel_hybrid_allocates_mamba(self):
         assert torch.all(dynamic_context.mamba_ssm_states[:, mamba_idx] == 0)
 
     @pytest.mark.internal
+    @rounder_override(64)
     def test_add_dummy_requests_parallel_decode_does_not_count_as_prefill(self):
         self._setup_model_parallel_group(1, 1)
 
@@ -578,6 +589,7 @@ def test_add_dummy_requests_parallel_decode_does_not_count_as_prefill(self):
         assert dynamic_context.num_prefill_requests == 0
 
     @pytest.mark.internal
+    @rounder_override(64)
     @pytest.mark.parametrize("is_hybrid_model", [False, True])
     def test_update_request(self, is_hybrid_model: bool):
         self._setup_model_parallel_group(1, 1)
@@ -777,6 +789,7 @@ def test_update_request(self, is_hybrid_model: bool):
             )
 
     @pytest.mark.internal
+    @rounder_override(64)
     @pytest.mark.parametrize("is_hybrid_model", [False, True])
     def test_release_memory_blocks_for_finished_requests(self, is_hybrid_model):
         """Test that memory blocks are correctly released for finished requests."""
@@ -849,6 +862,7 @@ def test_release_memory_blocks_for_finished_requests(self, is_hybrid_model):
             assert mamba_idx[4] == -1
 
     @pytest.mark.internal
+    @rounder_override(64)
     @pytest.mark.parametrize("is_hybrid_model", [False, True])
     def test_finished_requests_with_multiple_blocks(self, is_hybrid_model):
         """Test that all memory blocks are correctly released for finished requests that use multiple blocks."""
@@ -916,6 +930,7 @@ def test_finished_requests_with_multiple_blocks(self, is_hybrid_model):
         assert dynamic_context.block_allocator.total_avail == initial_available_blocks + 6
 
     @pytest.mark.internal
+    @rounder_override(64)
     @pytest.mark.parametrize("is_hybrid_model", [False, True])
     def test_mamba_states_cache(self, is_hybrid_model: bool):
         self._setup_model_parallel_group(1, 1)
@@ -991,6 +1006,7 @@ def test_mamba_states_cache(self, is_hybrid_model: bool):
         assert torch.all(ssm_state_layer3 == 40.0)
 
     @pytest.mark.internal
+    @rounder_override(64)
     def test_calculate_and_store_log_probs(self):
         self._setup_model_parallel_group(1, 1)
         dynamic_context = self._get_dynamic_context(
@@ -1208,6 +1224,7 @@ def test_calculate_and_store_log_probs(self):
                 current_global_token_offset += expected_len
 
     @pytest.mark.internal
+    @rounder_override(64)
     def test_pipeline_parallel_uneven_layers(self):
         """
         Test that DynamicInferenceContext synchronizes the total block count across
@@ -1244,11 +1261,13 @@ def test_pipeline_parallel_uneven_layers(self):
                 tensor_model_parallel_size=1,
                 pipeline_dtype=torch.float32,
             ),
-            max_sequence_length=128,
-            buffer_size_gb=0.1,
-            block_size_tokens=16,
-            max_tokens=1024,
-            unified_memory_level=0,
+            inference_config=DynamicInferenceConfig(
+                max_sequence_length=128,
+                buffer_size_gb=0.1,
+                block_size_tokens=16,
+                max_tokens=1024,
+                unified_memory_level=0,
+            ),
         )
 
         # Collect the total block counts on each rank
diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py
index b8529c97784..81859d897b1 100644
--- a/tests/unit_tests/inference/engines/test_dynamic_engine.py
+++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py
@@ -13,9 +13,7 @@
 from transformer_engine.pytorch.fp8 import check_fp8_support
 
 from megatron.core import parallel_state
-from megatron.core.inference.contexts.attention_context.mamba_metadata import (
-    MambaInferenceStateConfig,
-)
+from megatron.core.inference.config import DynamicInferenceConfig, MambaInferenceStateConfig
 from megatron.core.inference.contexts.dynamic_context import (
     ActiveRequestCountOverflowError,
     BlockOverflowError,
@@ -45,11 +43,7 @@
 from megatron.core.transformer.cuda_graphs import CudaGraphManager, _CudagraphGlobalRecord
 from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.transformer.transformer_config import TransformerConfig
-from megatron.core.utils import (
-    get_mamba_inference_state_config_from_model,
-    is_fa_min_version,
-    is_te_min_version,
-)
+from megatron.core.utils import is_fa_min_version, is_te_min_version
 from tests.unit_tests.test_utilities import Utils
 
 
@@ -221,19 +215,21 @@ def _build_inference_context(
         # Inference context.
         context = DynamicInferenceContext(
             model_config=transformer_config,
-            max_sequence_length=test_config.max_sequence_length,
-            num_cuda_graphs=test_config.num_cuda_graphs,
-            use_cuda_graphs_for_non_decode_steps=True,
-            buffer_size_gb=test_config.context_buffer_size_gb,
-            paused_buffer_size_gb=test_config.context_paused_buffer_size_gb,
-            block_size_tokens=test_config.context_block_size_tokens,
-            max_requests=test_config.context_max_requests,
-            max_tokens=test_config.context_max_tokens,
-            mamba_inference_state_config=mamba_inference_state_config,
-            materialize_only_last_token_logits=test_config.materialize_only_last_token_logits,
-            use_flashinfer_fused_rope=None,  # default to using flash-infer if available
-            # this is for compatibility with the LTS environment
-            unified_memory_level=0,  # unit tests currently broken with UVM
+            inference_config=DynamicInferenceConfig(
+                max_sequence_length=test_config.max_sequence_length,
+                num_cuda_graphs=test_config.num_cuda_graphs,
+                use_cuda_graphs_for_non_decode_steps=True,
+                buffer_size_gb=test_config.context_buffer_size_gb,
+                paused_buffer_size_gb=test_config.context_paused_buffer_size_gb,
+                block_size_tokens=test_config.context_block_size_tokens,
+                max_requests=test_config.context_max_requests,
+                max_tokens=test_config.context_max_tokens,
+                mamba_inference_state_config=mamba_inference_state_config,
+                materialize_only_last_token_logits=test_config.materialize_only_last_token_logits,
+                use_flashinfer_fused_rope=None,  # default to using flash-infer if available
+                # this is for compatibility with the LTS environment
+                unified_memory_level=0,  # unit tests currently broken with UVM
+            ),
         )
 
         return context
@@ -373,7 +369,7 @@ def _build_test_env(cls, test_config):
 
         model.eval()
 
-        mamba_inference_state_config = get_mamba_inference_state_config_from_model(model)
+        mamba_inference_state_config = MambaInferenceStateConfig.from_model(model)
 
         # Inference context.
         inference_context = cls._build_inference_context(
@@ -405,13 +401,7 @@ def _build_test_env(cls, test_config):
         CudaGraphManager.global_mempool = None
 
         # Inference engine.
-        engine = DynamicInferenceEngine(
-            text_generation_controller,
-            inference_context,
-            random_seed=test_config.random_seed,
-            enable_cuda_graph=transformer_config.cuda_graph_impl == "local",
-            enable_chunked_prefill=test_config.enable_chunked_prefill,
-        )
+        engine = DynamicInferenceEngine(text_generation_controller, inference_context)
 
         # Test env.
         env = DynamicEngineTestEnv(config=test_config, requests=requests, engine=engine)
diff --git a/tests/unit_tests/inference/model_inference_wrappers/gpt/test_gpt_inference_wrapper.py b/tests/unit_tests/inference/model_inference_wrappers/gpt/test_gpt_inference_wrapper.py
index 086fc45118c..86fa525b389 100644
--- a/tests/unit_tests/inference/model_inference_wrappers/gpt/test_gpt_inference_wrapper.py
+++ b/tests/unit_tests/inference/model_inference_wrappers/gpt/test_gpt_inference_wrapper.py
@@ -67,7 +67,7 @@ def test_inference_pipeline_parallel(self, materialize_only_last_token_logits):
             .cuda()
         )
         self.inference_wrapped_model.prep_model_for_inference()
-        self.inference_wrapped_model.inference_context.materialize_only_last_token_logits = (
+        self.inference_wrapped_model.inference_context.inference_config.materialize_only_last_token_logits = (
             materialize_only_last_token_logits
         )
 
@@ -102,7 +102,7 @@ def test_inference_only_tensor_parallel(self, materialize_only_last_token_logits
             .cuda()
         )
         self.inference_wrapped_model.prep_model_for_inference()
-        self.inference_wrapped_model.inference_context.materialize_only_last_token_logits = (
+        self.inference_wrapped_model.inference_context.inference_config.materialize_only_last_token_logits = (
             materialize_only_last_token_logits
         )
 
diff --git a/tests/unit_tests/inference/test_wandb_logging.py b/tests/unit_tests/inference/test_wandb_logging.py
index 7d111ac1d93..99c16449616 100644
--- a/tests/unit_tests/inference/test_wandb_logging.py
+++ b/tests/unit_tests/inference/test_wandb_logging.py
@@ -7,6 +7,7 @@
 import pytest
 import torch
 
+from megatron.core.inference.config import DynamicInferenceConfig
 from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
 from megatron.core.inference.engines import DynamicInferenceEngine
 from megatron.core.inference.inference_request import DynamicInferenceRequest
@@ -60,11 +61,13 @@ def _get_dynamic_context(
                 kv_channels=kv_channels,
                 num_attention_heads=num_attention_heads,
             ),
-            max_sequence_length=max_sequence_length,
-            num_cuda_graphs=None,
-            buffer_size_gb=buffer_size_gb,
-            block_size_tokens=block_size_tokens,
-            unified_memory_level=0,  # unit tests currently broken with UVM
+            inference_config=DynamicInferenceConfig(
+                max_sequence_length=max_sequence_length,
+                num_cuda_graphs=None,
+                buffer_size_gb=buffer_size_gb,
+                block_size_tokens=block_size_tokens,
+                unified_memory_level=0,  # unit tests currently broken with UVM
+            ),
         )
 
     @pytest.mark.internal
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index 61ea73f823b..09c4a83acd7 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -14,6 +14,7 @@
 from transformer_engine.pytorch.fp8 import check_fp8_support
 
 from megatron.core import parallel_state
+from megatron.core.inference.config import DynamicInferenceConfig
 from megatron.core.inference.contexts import DynamicInferenceContext, StaticInferenceContext
 from megatron.core.inference.contexts.dynamic_context import MaxSequenceLengthOverflowError
 from megatron.core.inference.inference_request import (
@@ -104,12 +105,14 @@ def setup_model(
         else:
             inference_context = DynamicInferenceContext(
                 model_config=transformer_config,
-                max_sequence_length=2048,
-                buffer_size_gb=0.2,
-                materialize_only_last_token_logits=False,
-                use_flashinfer_fused_rope=None,  # default to using flash-infer if available
-                # this is for compatibility with the LTS environment
-                unified_memory_level=0,  # unit tests currently broken with UVM
+                inference_config=DynamicInferenceConfig(
+                    max_sequence_length=2048,
+                    buffer_size_gb=0.2,
+                    materialize_only_last_token_logits=False,
+                    use_flashinfer_fused_rope=None,  # default to using flash-infer if available
+                    # this is for compatibility with the LTS environment
+                    unified_memory_level=0,  # unit tests currently broken with UVM
+                ),
             )
 
         inference_wrapped_model = GPTInferenceWrapper(gpt_model, inference_context)
diff --git a/tests/unit_tests/models/test_gpt_model.py b/tests/unit_tests/models/test_gpt_model.py
index 8c8b1be638f..6c51f401fda 100644
--- a/tests/unit_tests/models/test_gpt_model.py
+++ b/tests/unit_tests/models/test_gpt_model.py
@@ -12,6 +12,7 @@
 
 from megatron.core import parallel_state
 from megatron.core.hyper_comm_grid import HyperCommGrid
+from megatron.core.inference.config import DynamicInferenceConfig
 from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
 from megatron.core.inference.inference_request import DynamicInferenceRequest
 from megatron.core.inference.sampling_params import SamplingParams
@@ -398,10 +399,12 @@ def test_dynamic_inference_padding_with_fp8(self):
                 kv_channels=config.hidden_size // config.num_attention_heads,
                 num_attention_heads=config.num_attention_heads,
             ),
-            max_sequence_length=self.gpt_model.module.max_sequence_length,
-            buffer_size_gb=1.0,
-            block_size_tokens=256,
-            materialize_only_last_token_logits=False,
+            inference_config=DynamicInferenceConfig(
+                max_sequence_length=self.gpt_model.module.max_sequence_length,
+                buffer_size_gb=1.0,
+                block_size_tokens=256,
+                materialize_only_last_token_logits=False,
+            ),
         )
 
         # Add a request with 10 tokens. Since 10 is not a multiple of 64,
diff --git a/tests/unit_tests/models/test_gpt_model_batch_invariant.py b/tests/unit_tests/models/test_gpt_model_batch_invariant.py
index 9284e3a6acc..98383e76346 100644
--- a/tests/unit_tests/models/test_gpt_model_batch_invariant.py
+++ b/tests/unit_tests/models/test_gpt_model_batch_invariant.py
@@ -5,6 +5,7 @@
 import torch
 import torch.distributed as dist
 
+from megatron.core.inference.config import DynamicInferenceConfig
 from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
 from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
@@ -184,13 +185,15 @@ def test_dynamic_engine_matches_batched_forward_rl(self):
 
         ctx = DynamicInferenceContext(
             model_config=base_model.config,
-            max_sequence_length=seq_len,
-            buffer_size_gb=0.125,
-            block_size_tokens=16,
-            num_cuda_graphs=None,
-            materialize_only_last_token_logits=False,
-            use_cuda_graphs_for_non_decode_steps=False,
-            unified_memory_level=0,
+            inference_config=DynamicInferenceConfig(
+                max_sequence_length=seq_len,
+                buffer_size_gb=0.125,
+                block_size_tokens=16,
+                num_cuda_graphs=None,
+                materialize_only_last_token_logits=False,
+                use_cuda_graphs_for_non_decode_steps=False,
+                unified_memory_level=0,
+            ),
         )
 
         wrapper = GPTInferenceWrapper(inference_model, ctx)
@@ -260,13 +263,15 @@ def test_dynamic_engine_is_batch_invariant(self):
         def _run_engine_with_order(order):
             ctx = DynamicInferenceContext(
                 model_config=based_model.config,
-                max_sequence_length=seq_len,
-                buffer_size_gb=0.125,
-                block_size_tokens=16,
-                num_cuda_graphs=None,
-                materialize_only_last_token_logits=False,
-                use_cuda_graphs_for_non_decode_steps=False,
-                unified_memory_level=0,
+                inference_config=DynamicInferenceConfig(
+                    max_sequence_length=seq_len,
+                    buffer_size_gb=0.125,
+                    block_size_tokens=16,
+                    num_cuda_graphs=None,
+                    materialize_only_last_token_logits=False,
+                    use_cuda_graphs_for_non_decode_steps=False,
+                    unified_memory_level=0,
+                ),
             )
 
             wrapper = GPTInferenceWrapper(inference_model, ctx)
diff --git a/tests/unit_tests/models/test_mamba_model.py b/tests/unit_tests/models/test_mamba_model.py
index 849581b2d3d..c720a6f4eb4 100644
--- a/tests/unit_tests/models/test_mamba_model.py
+++ b/tests/unit_tests/models/test_mamba_model.py
@@ -350,11 +350,13 @@ def test_dynamic_inference_padding_with_fp8(self):
 
         inference_context = DynamicInferenceContext(
             model_config=self.model.config,
-            max_sequence_length=self.model.module.max_sequence_length,
-            buffer_size_gb=1.0,
-            block_size_tokens=256,
-            materialize_only_last_token_logits=False,
-            mamba_inference_state_config=mamba_inference_state_config,
+            inference_config=DynamicInferenceConfig(
+                max_sequence_length=self.model.module.max_sequence_length,
+                buffer_size_gb=1.0,
+                block_size_tokens=256,
+                materialize_only_last_token_logits=False,
+                mamba_inference_state_config=mamba_inference_state_config,
+            ),
         )
 
         # Add a request with 10 tokens. Since 10 is not a multiple of 64 (TOKEN_ROUNDER),
diff --git a/tools/run_dynamic_text_generation_server.py b/tools/run_dynamic_text_generation_server.py
index 31382841fbf..e9474b21abd 100644
--- a/tools/run_dynamic_text_generation_server.py
+++ b/tools/run_dynamic_text_generation_server.py
@@ -5,23 +5,19 @@
 
 import torch
 
-from examples.inference.gpt.gpt_dynamic_inference import (
-    add_dynamic_inference_args,
-    get_model,
-)
-from megatron.core.inference.engines import DynamicInferenceEngine
 from megatron.core.inference.text_generation_server.dynamic_text_gen_server import run_flask_server
 from megatron.core.tokenizers.text.utils.build_tokenizer import build_tokenizer
-from megatron.core.utils import get_mamba_inference_state_config_from_model, trace_async_exceptions
+from megatron.core.utils import trace_async_exceptions
+from megatron.inference.utils import add_inference_args, get_dynamic_inference_engine
+from megatron.core.inference.engines import DynamicInferenceEngine
 from megatron.post_training.arguments import add_modelopt_args
 from megatron.training import get_args, get_tokenizer
 from megatron.training.initialize import initialize_megatron
 
-
 def add_text_generation_server_args(parser: argparse.ArgumentParser):
     """Adds the required command line arguments for running the text generation server."""
     parser = add_modelopt_args(parser)
-    parser = add_dynamic_inference_args(parser)
+    parser = add_inference_args(parser)
     parser.add_argument("--port", type=int, default=5000, help="Port for Flask server to run on")
     return parser
 
@@ -72,14 +68,12 @@ async def run_text_generation_server(
             args_defaults={'no_load_rng': True, 'no_load_optim': True},
         )
 
-        args = get_args()
-        model = get_model()
-
         # Enable return_log_probs to allow prompt logprobs computation for echo=True requests
         # This sets materialize_only_last_token_logits=False in the inference context,
         # which is required for lm-eval compatibility (loglikelihood evaluation tasks)
+        args = get_args()
         args.return_log_probs = True
 
-        engine = DynamicInferenceEngine.from_model_and_args(model, args)
-
+        engine = get_dynamic_inference_engine()
+        
         asyncio.run(run_text_generation_server(engine, args.inference_coordinator_port, args.port))
diff --git a/tools/run_inference_performance_test.py b/tools/run_inference_performance_test.py
index ff08cd13dbd..0607c12ee90 100644
--- a/tools/run_inference_performance_test.py
+++ b/tools/run_inference_performance_test.py
@@ -10,6 +10,7 @@
 
 from gpt_builders import gpt_builder
 from mamba_builders import mamba_builder
+from megatron.inference.utils import get_dynamic_inference_engine
 from megatron.core.inference.contexts import DynamicInferenceContext, StaticInferenceContext
 from megatron.core.inference.engines import DynamicInferenceEngine, StaticInferenceEngine
 from megatron.core.inference.engines.abstract_engine import AbstractEngine
@@ -85,7 +86,8 @@ def get_inference_engine(args: argparse.Namespace, model: MegatronModule) -> Abs
         )
         return StaticInferenceEngine(text_generation_controller=text_generation_controller)
     elif args.engine_type == "dynamic":
-        return DynamicInferenceEngine.from_model_and_args(model, args)
+        return get_dynamic_inference_engine(model=model)
+
 
 def get_random_prompt_tokens(tokenizer, num_input_tokens) -> List[int]:
     # Get the set of special token IDs to exclude

From bf66e636572ed0b7ef8a07dbe297b7b544b73d66 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Thu, 22 Jan 2026 18:54:03 -0800
Subject: [PATCH 14/30] Add inference config unit test

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 .../inference/test_inference_config.py          | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)
 create mode 100644 tests/unit_tests/inference/test_inference_config.py

diff --git a/tests/unit_tests/inference/test_inference_config.py b/tests/unit_tests/inference/test_inference_config.py
new file mode 100644
index 00000000000..075e5a0cd56
--- /dev/null
+++ b/tests/unit_tests/inference/test_inference_config.py
@@ -0,0 +1,17 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+import dataclasses
+
+from megatron.core.inference.config import DynamicInferenceConfig
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+class TestDynamicInferenceConfig:
+    def test_mutual_exclusivity_with_transformer_config(self):
+        """
+        Ensure mutual exclusivity between fields in `DynamicInferenceConfig` and
+        `TransformerConfig`.
+        """
+        dynamic_inference_config_fields = set(dataclasses.fields(DynamicInferenceConfig))
+        transformer_config_fields = set(dataclasses.fields(TransformerConfig))
+        assert len(dynamic_inference_config_fields.intersection(transformer_config_fields)) == 0

From 188dcf2b897a4907c27dc9cd088ca0fa52bb5f55 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Thu, 22 Jan 2026 20:47:22 -0800
Subject: [PATCH 15/30] Address reviewer feedback

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 .../inference/gpt/gpt_dynamic_inference.py    |  4 +--
 megatron/core/inference/config.py             | 30 +++----------------
 .../core/inference/contexts/base_context.py   |  6 ++--
 .../inference/contexts/dynamic_context.py     |  6 ++--
 .../core/inference/contexts/static_context.py |  6 ++--
 .../core/inference/engines/dynamic_engine.py  |  2 +-
 .../core/inference/engines/static_engine.py   |  4 +--
 .../abstract_model_inference_wrapper.py       | 12 ++++++++
 .../text_generation_controller.py             | 14 ++++-----
 megatron/core/models/gpt/gpt_model.py         |  7 ++---
 megatron/core/models/mamba/mamba_model.py     |  7 ++---
 megatron/inference/utils.py                   | 30 ++++++++++++-------
 .../contexts/test_dynamic_context.py          |  6 ++--
 .../inference/engines/test_dynamic_engine.py  |  4 +--
 .../gpt/test_gpt_inference_wrapper.py         |  4 +--
 .../inference/test_inference_config.py        |  8 ++---
 .../inference/test_wandb_logging.py           |  4 +--
 .../test_text_generation_controller.py        |  4 +--
 tests/unit_tests/models/test_gpt_model.py     |  4 +--
 .../models/test_gpt_model_batch_invariant.py  |  6 ++--
 tests/unit_tests/models/test_mamba_model.py   |  2 +-
 21 files changed, 81 insertions(+), 89 deletions(-)

diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py
index 7daa095d845..ce4a3c6f8cc 100644
--- a/examples/inference/gpt/gpt_dynamic_inference.py
+++ b/examples/inference/gpt/gpt_dynamic_inference.py
@@ -37,7 +37,7 @@
 from megatron.core.tokenizers.text.utils.build_tokenizer import build_tokenizer
 from megatron.inference.utils import (
     add_inference_args,
-    get_dynamic_inference_config_from_model_and_args,
+    get_inference_config_from_model_and_args,
     get_model,
 )
 
@@ -283,7 +283,7 @@ def main():
 
     # Requests, context, controller.
     requests = build_requests(args, tokenizer, sampling_params)
-    inference_config = get_dynamic_inference_config_from_model_and_args(model, args)
+    inference_config = get_inference_config_from_model_and_args(model, args)
 
     # Calculate max_sequence_length from requests
     max_gen_length = sampling_params.num_tokens_to_generate
diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py
index dae39532ba0..2d734a33aa6 100644
--- a/megatron/core/inference/config.py
+++ b/megatron/core/inference/config.py
@@ -1,6 +1,5 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
-import abc
 from dataclasses import dataclass
 from typing import List, Optional, Tuple
 
@@ -53,19 +52,9 @@ def from_model(cls, model: MegatronModule) -> Optional["MambaInferenceStateConfi
 
 
 @dataclass
-class BaseInferenceConfig(abc.ABC):
-    """Base inference config."""
-
-    pass
-
-
-@dataclass
-class DynamicInferenceConfig(BaseInferenceConfig):
+class InferenceConfig:
     """
-    Config for dynamic inference.
-
-    Used to initialize `DynamicInferenceContext`, `TextGenerationController`, and
-    `DynamicInferenceEngine`.
+    Config for inference.
 
     NOTE: Must remain mutually exclusive with the `TransformerConfig`.
     """
@@ -189,17 +178,6 @@ class DynamicInferenceConfig(BaseInferenceConfig):
 
     request_metadata_types: Optional[List[Tuple[str, torch.dtype, bool]]] = None
     """
-    A list of the per-request metadata types to track. Each entry is a tuple consisting of the string
-    label, the target dtype, and whether to store the data on GPU.
-    """
-
-
-@dataclass
-class StaticInferenceConfig(DynamicInferenceConfig):
+    A list of the per-request metadata types to track. Each entry is a tuple
+    consisting of the string label, the target dtype, and whether to store the data on GPU.
     """
-    Static inference config. For now, exactly mimic the dynamic config.
-
-    TODO(ksanthanam): Remove when deprecating static inference.
-    """
-
-    pass
diff --git a/megatron/core/inference/contexts/base_context.py b/megatron/core/inference/contexts/base_context.py
index 738e8568bc8..4f03726fe3d 100644
--- a/megatron/core/inference/contexts/base_context.py
+++ b/megatron/core/inference/contexts/base_context.py
@@ -2,7 +2,7 @@
 
 import abc
 
-from megatron.core.inference.config import BaseInferenceConfig
+from megatron.core.inference.config import InferenceConfig
 
 
 class BaseInferenceContext(abc.ABC):
@@ -12,11 +12,11 @@ class BaseInferenceContext(abc.ABC):
     Extend this class for any future contexts types.
     """
 
-    def __init__(self, inference_config: BaseInferenceConfig):
+    def __init__(self, inference_config: InferenceConfig):
         """
         Args:
         """
-        self.inference_config = inference_config
+        self.config = inference_config
 
     @abc.abstractmethod
     def is_static_batching(self) -> bool:
diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 5fcb52e71f0..9a4573d7236 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -15,7 +15,7 @@
     CUDAGraphBatchDimensionBuilder,
     InferenceBatchDimensions,
 )
-from megatron.core.inference.config import DynamicInferenceConfig
+from megatron.core.inference.config import InferenceConfig
 from megatron.core.inference.inference_request import DynamicInferenceRequest
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.inference.unified_memory import (
@@ -184,14 +184,14 @@ class DynamicInferenceContext(BaseInferenceContext):
 
     Args:
         model_config (TransformerConfig): Model config.
-        inference_config (DynamicInferenceConfig): Inference config.
+        inference_config (InferenceConfig): Inference config.
     """
 
     DEFAULT_MAX_TOKENS = 16384
     TOKEN_ROUNDER = 64
     REQUEST_ROUNDER = 4
 
-    def __init__(self, model_config: TransformerConfig, inference_config: DynamicInferenceConfig):
+    def __init__(self, model_config: TransformerConfig, inference_config: InferenceConfig):
         super().__init__(inference_config=inference_config)
 
         self.cache_mla_latent = (
diff --git a/megatron/core/inference/contexts/static_context.py b/megatron/core/inference/contexts/static_context.py
index 5cfc510bf46..a15b33c414a 100644
--- a/megatron/core/inference/contexts/static_context.py
+++ b/megatron/core/inference/contexts/static_context.py
@@ -1,6 +1,6 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 
-from megatron.core.inference.config import StaticInferenceConfig
+from megatron.core.inference.config import InferenceConfig
 
 from .base_context import BaseInferenceContext
 
@@ -17,8 +17,8 @@ class StaticInferenceContext(BaseInferenceContext):
     def __init__(
         self, max_batch_size: int, max_sequence_length: int, use_flashinfer_fused_rope: bool = None
     ):
-        inference_config = StaticInferenceConfig(materialize_only_last_token_logits=True)
-        super().__init__(inference_config=inference_config)
+        config = InferenceConfig(materialize_only_last_token_logits=True)
+        super().__init__(inference_config=config)
         self.max_sequence_length = max_sequence_length
         self.max_batch_size = max_batch_size
         self.sequence_len_offset = 0
diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index 290c24315d5..41ec1086dfa 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -144,7 +144,7 @@ def __init__(self, controller: TextGenerationController, context: DynamicInferen
         ), f"context must be a DynamicInferenceContext, got {type(context)}"
 
         model_config = controller.inference_wrapped_model.model.config
-        inference_config = context.inference_config
+        inference_config = context.config
 
         if inference_config.pg_collection is not None:
             self.pg_collection = inference_config.pg_collection
diff --git a/megatron/core/inference/engines/static_engine.py b/megatron/core/inference/engines/static_engine.py
index 436b60c039f..fc381848268 100644
--- a/megatron/core/inference/engines/static_engine.py
+++ b/megatron/core/inference/engines/static_engine.py
@@ -8,7 +8,7 @@
 import torch
 
 from megatron.core.inference.async_stream import AsyncStream
-from megatron.core.inference.config import DynamicInferenceConfig, MambaInferenceStateConfig
+from megatron.core.inference.config import InferenceConfig, MambaInferenceStateConfig
 from megatron.core.inference.contexts import DynamicInferenceContext, StaticInferenceContext
 from megatron.core.inference.engines.abstract_engine import AbstractEngine
 from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine
@@ -100,7 +100,7 @@ def __init__(
             if not legacy:
                 dynamic_context = DynamicInferenceContext(
                     model_config=self.config,
-                    inference_config=DynamicInferenceConfig(
+                    inference_config=InferenceConfig(
                         max_sequence_length=original_context.max_sequence_length,
                         buffer_size_gb=buffer_size_gb,
                         mamba_inference_state_config=mamba_inference_state_config,
diff --git a/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py b/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py
index ae8e8742fad..6ef5ac3a2e5 100644
--- a/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py
+++ b/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py
@@ -80,6 +80,18 @@ def prep_model_for_inference(self):
 
         self.inference_context.reset()
 
+    @abc.abstractmethod
+    def prep_inference_input(self, prompt_tokens) -> Dict[str, Any]:
+        """Prepares the inference input data.
+
+        Args:
+            prompts_tokens (torch.Tensor): A tensor of shape [batch_size, max_seq_len]
+
+        Returns:
+            A dict with all the inference input needed for the batch.
+        """
+        raise NotImplementedError()
+
     @abc.abstractmethod
     def get_batch_for_context_window(self, *args, **kwargs) -> Dict[str, Any]:
         """Returns the input data for inference
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index dc0c2af5b7b..617883414d4 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -58,10 +58,10 @@ class TextGenerationController:
     def __init__(self, inference_wrapped_model: AbstractModelInferenceWrapper, tokenizer):
         self.inference_wrapped_model = inference_wrapped_model
         self.model_config = self.inference_wrapped_model.model.config
-        self.inference_config = self.inference_wrapped_model.inference_context.inference_config
+        inference_config = self.inference_wrapped_model.inference_context.config
         self.tokenizer = tokenizer
 
-        pg_collection = self.inference_config.pg_collection
+        pg_collection = inference_config.pg_collection
         if pg_collection is not None:
             self.pp_group = pg_collection.pp
         else:
@@ -582,7 +582,7 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
         if self.model_is_pipeline_parallel:
             logits_seq_len = (
                 active_request_count
-                if context.inference_config.materialize_only_last_token_logits
+                if context.config.materialize_only_last_token_logits
                 else input_ids.shape[1]
             )
             logits_shape = [1, logits_seq_len, self.vocab_size]
@@ -638,7 +638,7 @@ def _dynamic_step_sample_logits(self, logits: Tensor):
 
         # Last token logits.
         context = self.inference_wrapped_model.inference_context
-        if context.inference_config.materialize_only_last_token_logits:
+        if context.config.materialize_only_last_token_logits:
             # When materialize_only_last_token_logits is true, last_token_logits is
             # already called in the forward pass of GPT.
             last_token_logits = logits.squeeze(0)
@@ -683,7 +683,7 @@ def _dynamic_step_calculate_log_probs(self, logits: Tensor) -> Optional[Tensor]:
         return context.calculate_log_probs(
             logits,
             self._sampled_tokens_cuda[:active_request_count],
-            only_last_token_logits=context.inference_config.materialize_only_last_token_logits,
+            only_last_token_logits=context.config.materialize_only_last_token_logits,
         )
 
     def _dynamic_step_calculate_top_n_logprobs(
@@ -711,7 +711,7 @@ def _dynamic_step_calculate_top_n_logprobs(
         active_request_slice = slice(context.paused_request_count, context.total_request_count)
 
         # Handle decode-only mode (only last token)
-        if context.inference_config.materialize_only_last_token_logits or context.is_decode_only():
+        if context.config.materialize_only_last_token_logits or context.is_decode_only():
             # In decode mode or when only last token logits are materialized,
             # logits already represent only the last tokens
             log_probs = log_probs_tensor[:active_request_count]
@@ -1187,7 +1187,7 @@ def generate_all_output_tokens_static_batch(
                     or not (sampling_params.return_log_probs or sampling_params.top_n_logprobs > 0)
                 )
                 inference_context = self.inference_wrapped_model.inference_context
-                inference_context.inference_config.materialize_only_last_token_logits = (
+                inference_context.config.materialize_only_last_token_logits = (
                     materialize_only_last_token_logits
                 )
 
diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index e4d7a879c30..f22bac3492e 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -635,10 +635,7 @@ def _postprocess(
                     )
         sequence_parallel_override = False
 
-        if (
-            in_inference_mode
-            and inference_context.inference_config.materialize_only_last_token_logits
-        ):
+        if in_inference_mode and inference_context.config.materialize_only_last_token_logits:
             if inference_context.is_static_batching():
                 hidden_states = hidden_states[-1:, :, :]
             else:
@@ -668,7 +665,7 @@ def _postprocess(
             assert (
                 in_inference_mode
                 and inference_context.is_dynamic_batching()
-                and inference_context.inference_config.materialize_only_last_token_logits
+                and inference_context.config.materialize_only_last_token_logits
             )
             self.output_layer.sequence_parallel = True
 
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index baf8875735c..02ec8a16c44 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -265,10 +265,7 @@ def forward(
             output_weight = self.shared_embedding_or_output_weight()
 
         sequence_parallel_override = False
-        if (
-            in_inference_mode
-            and inference_context.inference_config.materialize_only_last_token_logits
-        ):
+        if in_inference_mode and inference_context.config.materialize_only_last_token_logits:
             if inference_context.is_static_batching():
                 hidden_states = hidden_states[-1:, :, :]
             else:
@@ -298,7 +295,7 @@ def forward(
             assert (
                 in_inference_mode
                 and inference_context.is_dynamic_batching()
-                and inference_context.inference_config.materialize_only_last_token_logits
+                and inference_context.config.materialize_only_last_token_logits
             )
             self.output_layer.sequence_parallel = True
 
diff --git a/megatron/inference/utils.py b/megatron/inference/utils.py
index fc3b64198f4..f5776faa133 100644
--- a/megatron/inference/utils.py
+++ b/megatron/inference/utils.py
@@ -1,28 +1,32 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import logging
 from argparse import ArgumentParser
 from functools import partial
 from typing import Optional
 
 from gpt_builders import gpt_builder
 from mamba_builders import mamba_builder
-from megatron.core.inference.config import DynamicInferenceConfig, MambaInferenceStateConfig
+from megatron.core.inference.config import InferenceConfig, MambaInferenceStateConfig
 from megatron.core.inference.contexts import DynamicInferenceContext
+from megatron.core.inference.engines import DynamicInferenceEngine
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
     GPTInferenceWrapper,
 )
 from megatron.core.inference.text_generation_controllers.text_generation_controller import (
     TextGenerationController,
 )
-from megatron.core.inference.engines import DynamicInferenceEngine
 from megatron.core.tokenizers.text.utils.build_tokenizer import build_tokenizer
 from megatron.core.transformer.module import MegatronModule
-from megatron.core.utils import get_attr_wrapped_model
-from megatron.training import get_args, get_tokenizer
+from megatron.core.utils import get_attr_wrapped_model, log_single_rank
+from megatron.training import get_args
 from megatron.training import get_model as _get_model
+from megatron.training import get_tokenizer, get_wandb_writer
 from megatron.training.checkpointing import load_checkpoint
 from model_provider import model_provider
 
+logger = logging.getLogger(__name__)
+
 
 def get_model() -> MegatronModule:
     """Initialize model and load checkpoint."""
@@ -105,7 +109,10 @@ def add_inference_args(parser: ArgumentParser) -> ArgumentParser:
         "--top-n-logprobs",
         type=int,
         default=0,
-        help='Return the top n logprobs for the generated tokens and their corresponding token as a dictionary',
+        help=(
+            "Return the top n logprobs for the generated tokens and their "
+            "corresponding token as a dictionary"
+        ),
     )
     group.add_argument(
         "--incoming-requests-per-step",
@@ -218,13 +225,14 @@ def add_inference_args(parser: ArgumentParser) -> ArgumentParser:
     return parser
 
 
-def get_dynamic_inference_config_from_model_and_args(model: MegatronModule, args):
-    """Returns a `DynamicInferenceConfig` constructed from the model and command line arguments."""
+def get_inference_config_from_model_and_args(model: MegatronModule, args):
+    """Returns a `InferenceConfig` constructed from the model and command line arguments."""
 
     # Max sequence length.
     position_embedding_type = get_attr_wrapped_model(model, "position_embedding_type")
     model_max_seq_len = get_attr_wrapped_model(model, "max_sequence_length")
     inf_max_seq_len = args.inference_max_seq_length
+    max_batch_size = args.inference_dynamic_batching_max_requests
 
     if position_embedding_type == "learned_absolute":
         # When using absolute position embeddings, it is critical that the
@@ -236,11 +244,11 @@ def get_dynamic_inference_config_from_model_and_args(model: MegatronModule, args
             max_sequence_length = min(model_max_seq_len, inf_max_seq_len)
         else:
             max_sequence_length = model_max_seq_len
-        assert max_batch_size <= model_max_seq_len
+        assert max_batch_size is None or max_batch_size <= model_max_seq_len
     else:
         max_sequence_length = inf_max_seq_len
     if args.inference_dynamic_batching_max_requests is not None:
-        max_sequence_length = max(max_sequence_length, args.inference_dynamic_batching_max_requests)
+        max_sequence_length = max(max_sequence_length, max_batch_size)
 
     mamba_inference_state_config = MambaInferenceStateConfig.from_model(model)
     pg_collection = get_attr_wrapped_model(model, "pg_collection")
@@ -266,7 +274,7 @@ def get_dynamic_inference_config_from_model_and_args(model: MegatronModule, args
                 "wandb module is available. Inference logging will be disabled.",
             )
 
-    return DynamicInferenceConfig(
+    return InferenceConfig(
         block_size_tokens=args.inference_dynamic_batching_block_size,
         buffer_size_gb=args.inference_dynamic_batching_buffer_size_gb,
         paused_buffer_size_gb=args.inference_dynamic_batching_paused_buffer_size_gb,
@@ -278,7 +286,7 @@ def get_dynamic_inference_config_from_model_and_args(model: MegatronModule, args
         max_requests=args.inference_dynamic_batching_max_requests,
         max_tokens=args.inference_dynamic_batching_max_tokens,
         unified_memory_level=args.inference_dynamic_batching_unified_memory_level,
-        cuda_graph_mixed_prefill_count=args.inference_dynamic_batching_cuda_graph_mixed_prefill_count,
+        cuda_graph_mixed_prefill_count=args.inference_dynamic_batching_cuda_graph_mixed_prefill_count,  # pylint: disable=line-too-long
         use_cuda_graphs_for_non_decode_steps=not args.decode_only_cuda_graphs,
         persist_cuda_graphs=args.rl_training_cuda_graphs,
         max_sequence_length=max_sequence_length,
diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py
index 246551b19e6..f3ef0910f58 100644
--- a/tests/unit_tests/inference/contexts/test_dynamic_context.py
+++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py
@@ -7,7 +7,7 @@
 import torch
 
 from megatron.core import parallel_state
-from megatron.core.inference.config import DynamicInferenceConfig, MambaInferenceStateConfig
+from megatron.core.inference.config import InferenceConfig, MambaInferenceStateConfig
 from megatron.core.inference.contexts.dynamic_context import (
     DynamicInferenceContext,
     RequestOverflowError,
@@ -78,7 +78,7 @@ def _get_dynamic_context(
                 kv_channels=kv_channels,
                 num_attention_heads=num_attention_heads,
             ),
-            inference_config=DynamicInferenceConfig(
+            inference_config=InferenceConfig(
                 max_sequence_length=max_sequence_length,
                 num_cuda_graphs=None,
                 use_cuda_graphs_for_non_decode_steps=True,
@@ -1261,7 +1261,7 @@ def test_pipeline_parallel_uneven_layers(self):
                 tensor_model_parallel_size=1,
                 pipeline_dtype=torch.float32,
             ),
-            inference_config=DynamicInferenceConfig(
+            inference_config=InferenceConfig(
                 max_sequence_length=128,
                 buffer_size_gb=0.1,
                 block_size_tokens=16,
diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py
index 81859d897b1..2e935cab4bd 100644
--- a/tests/unit_tests/inference/engines/test_dynamic_engine.py
+++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py
@@ -13,7 +13,7 @@
 from transformer_engine.pytorch.fp8 import check_fp8_support
 
 from megatron.core import parallel_state
-from megatron.core.inference.config import DynamicInferenceConfig, MambaInferenceStateConfig
+from megatron.core.inference.config import InferenceConfig, MambaInferenceStateConfig
 from megatron.core.inference.contexts.dynamic_context import (
     ActiveRequestCountOverflowError,
     BlockOverflowError,
@@ -215,7 +215,7 @@ def _build_inference_context(
         # Inference context.
         context = DynamicInferenceContext(
             model_config=transformer_config,
-            inference_config=DynamicInferenceConfig(
+            inference_config=InferenceConfig(
                 max_sequence_length=test_config.max_sequence_length,
                 num_cuda_graphs=test_config.num_cuda_graphs,
                 use_cuda_graphs_for_non_decode_steps=True,
diff --git a/tests/unit_tests/inference/model_inference_wrappers/gpt/test_gpt_inference_wrapper.py b/tests/unit_tests/inference/model_inference_wrappers/gpt/test_gpt_inference_wrapper.py
index 86fa525b389..d7ddaa1e680 100644
--- a/tests/unit_tests/inference/model_inference_wrappers/gpt/test_gpt_inference_wrapper.py
+++ b/tests/unit_tests/inference/model_inference_wrappers/gpt/test_gpt_inference_wrapper.py
@@ -67,7 +67,7 @@ def test_inference_pipeline_parallel(self, materialize_only_last_token_logits):
             .cuda()
         )
         self.inference_wrapped_model.prep_model_for_inference()
-        self.inference_wrapped_model.inference_context.inference_config.materialize_only_last_token_logits = (
+        self.inference_wrapped_model.inference_context.config.materialize_only_last_token_logits = (
             materialize_only_last_token_logits
         )
 
@@ -102,7 +102,7 @@ def test_inference_only_tensor_parallel(self, materialize_only_last_token_logits
             .cuda()
         )
         self.inference_wrapped_model.prep_model_for_inference()
-        self.inference_wrapped_model.inference_context.inference_config.materialize_only_last_token_logits = (
+        self.inference_wrapped_model.inference_context.config.materialize_only_last_token_logits = (
             materialize_only_last_token_logits
         )
 
diff --git a/tests/unit_tests/inference/test_inference_config.py b/tests/unit_tests/inference/test_inference_config.py
index 075e5a0cd56..6d58328dade 100644
--- a/tests/unit_tests/inference/test_inference_config.py
+++ b/tests/unit_tests/inference/test_inference_config.py
@@ -2,16 +2,16 @@
 
 import dataclasses
 
-from megatron.core.inference.config import DynamicInferenceConfig
+from megatron.core.inference.config import InferenceConfig
 from megatron.core.transformer.transformer_config import TransformerConfig
 
 
-class TestDynamicInferenceConfig:
+class TestInferenceConfig:
     def test_mutual_exclusivity_with_transformer_config(self):
         """
-        Ensure mutual exclusivity between fields in `DynamicInferenceConfig` and
+        Ensure mutual exclusivity between fields in `InferenceConfig` and
         `TransformerConfig`.
         """
-        dynamic_inference_config_fields = set(dataclasses.fields(DynamicInferenceConfig))
+        dynamic_inference_config_fields = set(dataclasses.fields(InferenceConfig))
         transformer_config_fields = set(dataclasses.fields(TransformerConfig))
         assert len(dynamic_inference_config_fields.intersection(transformer_config_fields)) == 0
diff --git a/tests/unit_tests/inference/test_wandb_logging.py b/tests/unit_tests/inference/test_wandb_logging.py
index 99c16449616..cc4eb6fc70b 100644
--- a/tests/unit_tests/inference/test_wandb_logging.py
+++ b/tests/unit_tests/inference/test_wandb_logging.py
@@ -7,7 +7,7 @@
 import pytest
 import torch
 
-from megatron.core.inference.config import DynamicInferenceConfig
+from megatron.core.inference.config import InferenceConfig
 from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
 from megatron.core.inference.engines import DynamicInferenceEngine
 from megatron.core.inference.inference_request import DynamicInferenceRequest
@@ -61,7 +61,7 @@ def _get_dynamic_context(
                 kv_channels=kv_channels,
                 num_attention_heads=num_attention_heads,
             ),
-            inference_config=DynamicInferenceConfig(
+            inference_config=InferenceConfig(
                 max_sequence_length=max_sequence_length,
                 num_cuda_graphs=None,
                 buffer_size_gb=buffer_size_gb,
diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
index 09c4a83acd7..bdf95c2d9bf 100644
--- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
+++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py
@@ -14,7 +14,7 @@
 from transformer_engine.pytorch.fp8 import check_fp8_support
 
 from megatron.core import parallel_state
-from megatron.core.inference.config import DynamicInferenceConfig
+from megatron.core.inference.config import InferenceConfig
 from megatron.core.inference.contexts import DynamicInferenceContext, StaticInferenceContext
 from megatron.core.inference.contexts.dynamic_context import MaxSequenceLengthOverflowError
 from megatron.core.inference.inference_request import (
@@ -105,7 +105,7 @@ def setup_model(
         else:
             inference_context = DynamicInferenceContext(
                 model_config=transformer_config,
-                inference_config=DynamicInferenceConfig(
+                inference_config=InferenceConfig(
                     max_sequence_length=2048,
                     buffer_size_gb=0.2,
                     materialize_only_last_token_logits=False,
diff --git a/tests/unit_tests/models/test_gpt_model.py b/tests/unit_tests/models/test_gpt_model.py
index 6c51f401fda..87aba9c6ed9 100644
--- a/tests/unit_tests/models/test_gpt_model.py
+++ b/tests/unit_tests/models/test_gpt_model.py
@@ -12,7 +12,7 @@
 
 from megatron.core import parallel_state
 from megatron.core.hyper_comm_grid import HyperCommGrid
-from megatron.core.inference.config import DynamicInferenceConfig
+from megatron.core.inference.config import InferenceConfig
 from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
 from megatron.core.inference.inference_request import DynamicInferenceRequest
 from megatron.core.inference.sampling_params import SamplingParams
@@ -399,7 +399,7 @@ def test_dynamic_inference_padding_with_fp8(self):
                 kv_channels=config.hidden_size // config.num_attention_heads,
                 num_attention_heads=config.num_attention_heads,
             ),
-            inference_config=DynamicInferenceConfig(
+            inference_config=InferenceConfig(
                 max_sequence_length=self.gpt_model.module.max_sequence_length,
                 buffer_size_gb=1.0,
                 block_size_tokens=256,
diff --git a/tests/unit_tests/models/test_gpt_model_batch_invariant.py b/tests/unit_tests/models/test_gpt_model_batch_invariant.py
index 98383e76346..9ab7e445c0d 100644
--- a/tests/unit_tests/models/test_gpt_model_batch_invariant.py
+++ b/tests/unit_tests/models/test_gpt_model_batch_invariant.py
@@ -5,7 +5,7 @@
 import torch
 import torch.distributed as dist
 
-from megatron.core.inference.config import DynamicInferenceConfig
+from megatron.core.inference.config import InferenceConfig
 from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
 from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
@@ -185,7 +185,7 @@ def test_dynamic_engine_matches_batched_forward_rl(self):
 
         ctx = DynamicInferenceContext(
             model_config=base_model.config,
-            inference_config=DynamicInferenceConfig(
+            inference_config=InferenceConfig(
                 max_sequence_length=seq_len,
                 buffer_size_gb=0.125,
                 block_size_tokens=16,
@@ -263,7 +263,7 @@ def test_dynamic_engine_is_batch_invariant(self):
         def _run_engine_with_order(order):
             ctx = DynamicInferenceContext(
                 model_config=based_model.config,
-                inference_config=DynamicInferenceConfig(
+                inference_config=InferenceConfig(
                     max_sequence_length=seq_len,
                     buffer_size_gb=0.125,
                     block_size_tokens=16,
diff --git a/tests/unit_tests/models/test_mamba_model.py b/tests/unit_tests/models/test_mamba_model.py
index c720a6f4eb4..453a42626b9 100644
--- a/tests/unit_tests/models/test_mamba_model.py
+++ b/tests/unit_tests/models/test_mamba_model.py
@@ -350,7 +350,7 @@ def test_dynamic_inference_padding_with_fp8(self):
 
         inference_context = DynamicInferenceContext(
             model_config=self.model.config,
-            inference_config=DynamicInferenceConfig(
+            inference_config=InferenceConfig(
                 max_sequence_length=self.model.module.max_sequence_length,
                 buffer_size_gb=1.0,
                 block_size_tokens=256,

From 92b1a8d11cc23088382d59f6ed5e00be25c8d86e Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Thu, 22 Jan 2026 20:56:26 -0800
Subject: [PATCH 16/30] Fix copyright

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 megatron/inference/__init__.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/megatron/inference/__init__.py b/megatron/inference/__init__.py
index e69de29bb2d..26496bfed70 100644
--- a/megatron/inference/__init__.py
+++ b/megatron/inference/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.

From c76d953eb00250adb25f432862ef49b2b0c8dd9e Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Fri, 23 Jan 2026 09:54:06 -0800
Subject: [PATCH 17/30] Fix import

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 megatron/rl/inference/megatron.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py
index f0a984ceb40..07ee70d96ea 100644
--- a/megatron/rl/inference/megatron.py
+++ b/megatron/rl/inference/megatron.py
@@ -22,7 +22,7 @@
 from megatron.core.models.gpt.gpt_model import GPTModel
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.utils import get_attr_wrapped_model, log_single_rank
-from megatron.inference import get_dynamic_inference_engine
+from megatron.inference.utils import get_dynamic_inference_engine
 from megatron.training import get_wandb_writer
 from megatron.training.global_vars import get_args, get_tokenizer
 

From 952dd38c61491529a8c650a9572cfbba36ddfa57 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Fri, 23 Jan 2026 10:06:23 -0800
Subject: [PATCH 18/30] Fix mamba test

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 tests/unit_tests/models/test_mamba_model.py | 12 +++---------
 1 file changed, 3 insertions(+), 9 deletions(-)

diff --git a/tests/unit_tests/models/test_mamba_model.py b/tests/unit_tests/models/test_mamba_model.py
index 453a42626b9..29e3630d7bb 100644
--- a/tests/unit_tests/models/test_mamba_model.py
+++ b/tests/unit_tests/models/test_mamba_model.py
@@ -10,6 +10,7 @@
 
 from megatron.core import parallel_state
 from megatron.core.hyper_comm_grid import HyperCommGrid
+from megatron.core.inference.config import InferenceConfig, MambaInferenceStateConfig
 from megatron.core.inference.contexts import BaseInferenceContext, StaticInferenceContext
 from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
 from megatron.core.inference.inference_request import DynamicInferenceRequest
@@ -21,12 +22,7 @@
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.enums import AttnBackend
 from megatron.core.transformer.module import Float16Module
-from megatron.core.utils import (
-    divide,
-    get_mamba_inference_state_config_from_model,
-    is_fa_min_version,
-    is_torch_min_version,
-)
+from megatron.core.utils import divide, is_fa_min_version, is_torch_min_version
 from tests.unit_tests.test_utilities import Utils
 
 
@@ -344,9 +340,7 @@ def test_dynamic_inference_padding_with_fp8(self):
         self.model.eval()
         config = self.model.config
 
-        mamba_inference_state_config = get_mamba_inference_state_config_from_model(
-            self.model.module
-        )
+        mamba_inference_state_config = MambaInferenceStateConfig.from_model(self.model.module)
 
         inference_context = DynamicInferenceContext(
             model_config=self.model.config,

From 84c61237ed3adb1907768703f979f1ae562b6e73 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Fri, 23 Jan 2026 10:49:09 -0800
Subject: [PATCH 19/30] Fix unit tests

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 .../inference/test_wandb_logging.py           | 41 +++++++++----------
 .../unit_tests/models/test_mamba_moe_model.py |  2 +
 2 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/tests/unit_tests/inference/test_wandb_logging.py b/tests/unit_tests/inference/test_wandb_logging.py
index cc4eb6fc70b..1417926f13b 100644
--- a/tests/unit_tests/inference/test_wandb_logging.py
+++ b/tests/unit_tests/inference/test_wandb_logging.py
@@ -52,6 +52,8 @@ def _get_dynamic_context(
         max_sequence_length=512,
         buffer_size_gb=0.03,
         block_size_tokens=128,
+        logging_step_interval=0,
+        metrics_writer=None,
     ):
         """Helper to create a DynamicInferenceContext."""
         return DynamicInferenceContext(
@@ -67,6 +69,8 @@ def _get_dynamic_context(
                 buffer_size_gb=buffer_size_gb,
                 block_size_tokens=block_size_tokens,
                 unified_memory_level=0,  # unit tests currently broken with UVM
+                logging_step_interval=logging_step_interval,
+                metrics_writer=metrics_writer,
             ),
         )
 
@@ -199,12 +203,14 @@ def test_kvcache_utilization_stats_types(self):
     @pytest.mark.internal
     @patch('megatron.core.inference.engines.dynamic_engine.HAVE_WANDB', True)
     def test_engine_logging_step_interval_zero(self):
-        """Test that no logging occurs when inference_logging_step_interval is 0."""
+        """Test that no logging occurs when logging_step_interval is 0."""
         mock_wandb = Mock()
         mock_wandb.__name__ = "wandb"
         mock_wandb.log = Mock()
 
-        dynamic_context = self._get_dynamic_context()
+        dynamic_context = self._get_dynamic_context(
+            logging_step_interval=0, metrics_writer=mock_wandb
+        )
 
         # Create mock controller with proper spec to pass isinstance checks
         mock_controller = create_autospec(TextGenerationController, instance=True)
@@ -214,13 +220,7 @@ def test_engine_logging_step_interval_zero(self):
         mock_controller.inference_wrapped_model.model.config = Mock()
         mock_controller.inference_wrapped_model.model.config.cuda_graph_impl = "none"
 
-        engine = DynamicInferenceEngine(
-            controller=mock_controller,
-            context=dynamic_context,
-            random_seed=123,
-            inference_logging_step_interval=0,  # Disabled
-            metrics_writer=mock_wandb,
-        )
+        engine = DynamicInferenceEngine(controller=mock_controller, context=dynamic_context)
 
         # Verify log was never called
         mock_wandb.log.assert_not_called()
@@ -233,11 +233,13 @@ def test_paused_requests_in_stats(self):
             model_config=TransformerConfig(
                 params_dtype=torch.float32, num_layers=2, kv_channels=64, num_attention_heads=8
             ),
-            max_sequence_length=128,
-            num_cuda_graphs=None,
-            buffer_size_gb=0.01,  # Small buffer to force pausing
-            block_size_tokens=32,
-            unified_memory_level=0,  # unit tests currently broken with UVM
+            inference_config=InferenceConfig(
+                max_sequence_length=128,
+                num_cuda_graphs=None,
+                buffer_size_gb=0.01,  # Small buffer to force pausing
+                block_size_tokens=32,
+                unified_memory_level=0,  # unit tests currently broken with UVM
+            ),
         )
 
         # Add multiple requests to potentially trigger pausing
@@ -261,7 +263,7 @@ def test_paused_requests_in_stats(self):
     @pytest.mark.internal
     def test_metrics_writer_none_handling(self):
         """Test that engine handles None metrics_writer gracefully."""
-        dynamic_context = self._get_dynamic_context()
+        dynamic_context = self._get_dynamic_context(logging_step_interval=10, metrics_writer=None)
 
         # Create mock controller with proper spec to pass isinstance checks
         mock_controller = create_autospec(TextGenerationController, instance=True)
@@ -272,13 +274,8 @@ def test_metrics_writer_none_handling(self):
         mock_controller.inference_wrapped_model.model.config.cuda_graph_impl = "none"
 
         # Should not raise error even with logging interval set
-        engine = DynamicInferenceEngine(
-            controller=mock_controller,
-            context=dynamic_context,
-            random_seed=123,
-            inference_logging_step_interval=10,
-        )
+        engine = DynamicInferenceEngine(controller=mock_controller, context=dynamic_context)
 
         # Verify engine was created successfully
-        assert engine.inference_logging_step_interval == 10
+        assert engine.logging_step_interval == 10
         assert engine.metrics_writer is None
diff --git a/tests/unit_tests/models/test_mamba_moe_model.py b/tests/unit_tests/models/test_mamba_moe_model.py
index 5680751f63f..4f1eeba52b4 100644
--- a/tests/unit_tests/models/test_mamba_moe_model.py
+++ b/tests/unit_tests/models/test_mamba_moe_model.py
@@ -166,6 +166,7 @@
     "moe_layer_freq": 1,
     "moe_layer_recompute": False,
     "moe_pad_expert_input_to_capacity": False,
+    "moe_pad_experts_for_cuda_graph_inference": False,
     "moe_per_layer_logging": False,
     "moe_permute_fusion": False,
     "moe_router_bias_update_rate": 0.001,
@@ -196,6 +197,7 @@
     "mtp_num_layers": None,
     "mtp_standalone": False,
     "multi_latent_attention": False,
+    "nccl_all_reduce_for_prefill": False,
     "no_rope_freq": None,
     "no_sync_func": None,
     "normalization": "RMSNorm",

From ef508920376c6471ac3badbdf1b1dba1a18edbab Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Fri, 23 Jan 2026 14:31:45 -0800
Subject: [PATCH 20/30] Fix example scripts

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 .../inference/gpt/gpt_dynamic_inference.py    |  4 +--
 .../gpt_dynamic_inference_with_coordinator.py | 28 +++++++----------
 .../inference/gpt/gpt_static_inference.py     | 31 ++++---------------
 examples/inference/gpt/utils.py               |  1 -
 megatron/inference/utils.py                   |  8 ++---
 tools/run_dynamic_text_generation_server.py   |  8 ++---
 6 files changed, 27 insertions(+), 53 deletions(-)

diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py
index ce4a3c6f8cc..7fcac70c11a 100644
--- a/examples/inference/gpt/gpt_dynamic_inference.py
+++ b/examples/inference/gpt/gpt_dynamic_inference.py
@@ -38,7 +38,7 @@
 from megatron.inference.utils import (
     add_inference_args,
     get_inference_config_from_model_and_args,
-    get_model,
+    get_model_for_inference,
 )
 
 sys.path.append(
@@ -279,7 +279,7 @@ def main():
         stop_words=args.stop_words,
     )
 
-    model = get_model()
+    model = get_model_for_inference()
 
     # Requests, context, controller.
     requests = build_requests(args, tokenizer, sampling_params)
diff --git a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
index 97117b9dbe1..2925ce1cdb1 100644
--- a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
+++ b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
@@ -11,24 +11,18 @@
 
 import torch
 import torch.distributed as dist
-from tqdm import tqdm
-
-from examples.inference.gpt.utils import (
-    Request,
-    add_dynamic_inference_args,
-    add_common_inference_args,
-    build_dynamic_engine_setup_prefix,
-    build_requests,
-    get_model
-)
-from megatron.core import parallel_state
+
+from examples.inference.gpt.utils import Request, build_dynamic_engine_setup_prefix, build_requests
 from megatron.core.inference.engines import DynamicInferenceEngine
 from megatron.core.inference.inference_client import InferenceClient
 from megatron.core.inference.inference_request import DynamicInferenceRequestRecord
 from megatron.core.inference.sampling_params import SamplingParams
-from megatron.core.utils import get_mamba_inference_state_config_from_model
+from megatron.inference.utils import (
+    add_inference_args,
+    get_dynamic_inference_engine,
+    get_model_for_inference,
+)
 from megatron.training import get_args, get_tokenizer, initialize_megatron
-from megatron.training.arguments import parse_args
 
 # pylint: disable=line-too-long
 
@@ -190,7 +184,7 @@ async def main(
     # check for it.
     with torch.inference_mode():
         initialize_megatron(
-            extra_args_provider=add_dynamic_inference_args,
+            extra_args_provider=add_inference_args,
             args_defaults={'no_load_rng': True, 'no_load_optim': True},
         )
 
@@ -209,16 +203,16 @@ async def main(
             ),
         )
 
-        model = get_model()
+        model = get_model_for_inference()
 
         requests = (
             build_requests(args, tokenizer, sampling_params) if dist.get_rank() == 0 else None
         )
 
-        engine = DynamicInferenceEngine.from_model_and_args(model, args)
+        engine = get_dynamic_inference_engine(model=model)
 
         if dist.get_rank() == 0:
-            setup_prefix = build_dynamic_engine_setup_prefix(args, model, context, requests)
+            setup_prefix = build_dynamic_engine_setup_prefix(args, model, engine.context, requests)
             print("~~~")
             print(setup_prefix)
             print("~~~")
diff --git a/examples/inference/gpt/gpt_static_inference.py b/examples/inference/gpt/gpt_static_inference.py
index 906748d5d2e..4cd584cb517 100644
--- a/examples/inference/gpt/gpt_static_inference.py
+++ b/examples/inference/gpt/gpt_static_inference.py
@@ -1,18 +1,11 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 
 import os
-from model_provider import model_provider
-from gpt_builders import gpt_builder
-from mamba_builders import mamba_builder
-import torch
 import sys
 import time
-import warnings
-from functools import partial
 from argparse import Namespace
 
 import torch
-import tqdm
 
 from megatron.core.inference.contexts import StaticInferenceContext
 from megatron.core.inference.engines import StaticInferenceEngine
@@ -26,8 +19,6 @@
 )
 from megatron.core.tokenizers.text.utils.build_tokenizer import build_tokenizer
 from megatron.core.transformer.module import MegatronModule
-from pretrain_gpt import model_provider as gpt_model_provider
-from pretrain_mamba import model_provider as mamba_model_provider
 
 sys.path.append(
     os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
@@ -35,19 +26,18 @@
 
 import asyncio
 import json
-from typing import Any, AsyncIterator, List
+from typing import List
 
-from examples.inference.gpt.utils import add_common_inference_args, build_requests
-from megatron.core import mpu
-from megatron.training import get_args, get_model, get_tokenizer, print_rank_0
-from megatron.training.checkpointing import load_checkpoint
+from examples.inference.gpt.utils import build_requests
+from megatron.inference.utils import add_inference_args, get_model_for_inference
+from megatron.training import get_args, get_tokenizer, print_rank_0
 from megatron.training.initialize import initialize_megatron
 
 
 def add_static_inference_args(parser):
     """Static inference arguments."""
 
-    add_common_inference_args(parser)
+    add_inference_args(parser)
 
     group = parser.add_argument_group(title='Static inference')
     group.add_argument(
@@ -146,16 +136,7 @@ def main():
 
     args = get_args()
 
-    # Set up model and load checkpoint
-    if args.model_provider == "gpt":
-        model_builder = gpt_builder
-    elif args.model_provider == "mamba":
-        model_builder = mamba_builder
-    else:
-        raise ValueError(f"Invalid model provider {args.model_provider}")
-    model = get_model(partial(model_provider, model_builder), wrap_with_ddp=False)
-    load_checkpoint(model, None, None, strict=False)
-    model = model[0]
+    model = get_model_for_inference()
 
     inference_engine = get_inference_engine(args, model)
 
diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py
index 6f50f6a4b57..b7a3977605c 100644
--- a/examples/inference/gpt/utils.py
+++ b/examples/inference/gpt/utils.py
@@ -18,7 +18,6 @@
 from megatron.core.inference.sampling_params import SamplingParams
 from megatron.core.transformer.module import MegatronModule
 from megatron.training import get_args
-from megatron.training import get_model as _get_model
 
 
 def get_default_sampling_params(termination_id: int = None):
diff --git a/megatron/inference/utils.py b/megatron/inference/utils.py
index f5776faa133..f4fd04e8f83 100644
--- a/megatron/inference/utils.py
+++ b/megatron/inference/utils.py
@@ -28,8 +28,8 @@
 logger = logging.getLogger(__name__)
 
 
-def get_model() -> MegatronModule:
-    """Initialize model and load checkpoint."""
+def get_model_for_inference() -> MegatronModule:
+    """Initialize model and load checkpoint for inference."""
 
     args = get_args()
 
@@ -305,13 +305,13 @@ def get_dynamic_inference_engine(model: Optional[MegatronModule] = None) -> Dyna
     """Builds a `DynamicInferenceEngine`."""
     args = get_args()
     if model is None:
-        model = get_model()
+        model = get_model_for_inference()
     if args.legacy_tokenizer:
         tokenizer = get_tokenizer()
     else:
         tokenizer = build_tokenizer(args)
 
-    inference_config = get_dynamic_inference_config_from_model_and_args(model, args)
+    inference_config = get_inference_config_from_model_and_args(model, args)
     context = DynamicInferenceContext(model.config, inference_config)
     inference_wrapped_model = GPTInferenceWrapper(model, context)
     controller = TextGenerationController(inference_wrapped_model, tokenizer)
diff --git a/tools/run_dynamic_text_generation_server.py b/tools/run_dynamic_text_generation_server.py
index e9474b21abd..74f1e69679e 100644
--- a/tools/run_dynamic_text_generation_server.py
+++ b/tools/run_dynamic_text_generation_server.py
@@ -5,15 +5,15 @@
 
 import torch
 
+from megatron.core.inference.engines import DynamicInferenceEngine
 from megatron.core.inference.text_generation_server.dynamic_text_gen_server import run_flask_server
-from megatron.core.tokenizers.text.utils.build_tokenizer import build_tokenizer
 from megatron.core.utils import trace_async_exceptions
 from megatron.inference.utils import add_inference_args, get_dynamic_inference_engine
-from megatron.core.inference.engines import DynamicInferenceEngine
 from megatron.post_training.arguments import add_modelopt_args
-from megatron.training import get_args, get_tokenizer
+from megatron.training import get_args
 from megatron.training.initialize import initialize_megatron
 
+
 def add_text_generation_server_args(parser: argparse.ArgumentParser):
     """Adds the required command line arguments for running the text generation server."""
     parser = add_modelopt_args(parser)
@@ -75,5 +75,5 @@ async def run_text_generation_server(
         args.return_log_probs = True
 
         engine = get_dynamic_inference_engine()
-        
+
         asyncio.run(run_text_generation_server(engine, args.inference_coordinator_port, args.port))

From 39c92c63a206d411ed8a822294479fa4c5e11790 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Sat, 24 Jan 2026 15:28:38 -0800
Subject: [PATCH 21/30] Update tools/run_inference_performance_test.py

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 tools/run_inference_performance_test.py | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/tools/run_inference_performance_test.py b/tools/run_inference_performance_test.py
index 0607c12ee90..fb7b97c07ae 100644
--- a/tools/run_inference_performance_test.py
+++ b/tools/run_inference_performance_test.py
@@ -10,8 +10,7 @@
 
 from gpt_builders import gpt_builder
 from mamba_builders import mamba_builder
-from megatron.inference.utils import get_dynamic_inference_engine
-from megatron.core.inference.contexts import DynamicInferenceContext, StaticInferenceContext
+from megatron.core.inference.contexts import StaticInferenceContext
 from megatron.core.inference.engines import DynamicInferenceEngine, StaticInferenceEngine
 from megatron.core.inference.engines.abstract_engine import AbstractEngine
 from megatron.core.inference.inference_request import InferenceRequest
@@ -23,16 +22,15 @@
     TextGenerationController,
 )
 from megatron.core.transformer.module import MegatronModule
-from megatron.core.utils import get_mamba_inference_state_config_from_model
+from megatron.inference.utils import get_dynamic_inference_engine
 from model_provider import model_provider
 
 sys.path.append(
     os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir))
 )
 
-import asyncio
 from functools import partial
-from typing import List, Union
+from typing import List
 
 from examples.inference.gpt.utils import add_common_inference_args
 from megatron.core import mpu
@@ -71,9 +69,9 @@ def get_inference_engine(args: argparse.Namespace, model: MegatronModule) -> Abs
     Returns:
         AbstractBackend: The chosen backend
     """
-    tokenizer = get_tokenizer()
 
     if args.engine_type == "static":
+        tokenizer = get_tokenizer()
         context = StaticInferenceContext(
             args.inference_max_requests, args.inference_max_sequence_length
         )
@@ -150,8 +148,6 @@ def generate_dynamic(
 def main():
     """Main program."""
 
-    # Note: The default args passed here can be overwritten by using appropriate params (check arguments.py file)
-    # Micro batch size is not needed to be set by user. (It is calculated based on inference-batch-times-seqlen-threshold argument)
     initialize_megatron(
         extra_args_provider=add_inference_benchmarking_args,
         args_defaults={

From bdde0605259aa52dcf2b0b63e1c5b7c71da9afe2 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Mon, 26 Jan 2026 11:21:59 -0800
Subject: [PATCH 22/30] Add explicit deprecation error

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 .../inference/contexts/dynamic_context.py     | 39 ++++++++++++++
 .../core/inference/engines/dynamic_engine.py  | 14 +++++
 megatron/core/utils.py                        | 51 +++++++++++++++----
 3 files changed, 93 insertions(+), 11 deletions(-)

diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 9a4573d7236..c8b0b993843 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -27,6 +27,7 @@
 from megatron.core.package_info import __version__ as mcore_version
 from megatron.core.ssm.mamba_hybrid_layer_allocation import get_layer_maps_from_layer_type_list
 from megatron.core.transformer import MLATransformerConfig, TransformerConfig
+from megatron.core.utils import deprecate_args
 from megatron.core.utils import divide as core_divide
 from megatron.core.utils import get_pg_size, internal_api
 
@@ -48,6 +49,37 @@
     HAVE_FLASHINFER = False
 
 
+DEPRECATED_ARGS = [
+    "params_dtype",
+    "num_layers",
+    "kv_channels",
+    "num_attention_heads",
+    "max_sequence_length",
+    "buffer_size_gb",
+    "paused_buffer_size_gb",
+    "max_requests",
+    "max_tokens",
+    "block_size_tokens",
+    "tensor_model_parallel_size",
+    "pipeline_model_parallel_size",
+    "pg_collection",
+    "cache_mla_latent",
+    "kv_lora_rank",
+    "qk_pos_emb_head_dim",
+    "num_cuda_graphs",
+    "materialize_only_last_token_logits",
+    "mamba_inference_state_config",
+    "use_cuda_graphs_for_non_decode_steps",
+    "use_flashinfer_fused_rope",
+    "unified_memory_level",
+    "cuda_graph_max_tokens",
+    "cuda_graph_mixed_prefill_count",
+    "metrics_writer",
+    "request_metadata_types",
+    "persist_cuda_graphs",
+]
+
+
 class ContextOverflowError(Exception):
     """Base exception for when a new request does not fit.
 
@@ -191,6 +223,13 @@ class DynamicInferenceContext(BaseInferenceContext):
     TOKEN_ROUNDER = 64
     REQUEST_ROUNDER = 4
 
+    @deprecate_args(
+        *DEPRECATED_ARGS,
+        message=(
+            "Argument `{name}` has been deprecated. "
+            "Only pass `model_config` and `inference_config`"
+        ),
+    )
     def __init__(self, model_config: TransformerConfig, inference_config: InferenceConfig):
         super().__init__(inference_config=inference_config)
 
diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index 8675500e7dd..bb7fc2a5e9d 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -43,6 +43,7 @@
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.cuda_graphs import delete_cuda_graphs
 from megatron.core.utils import (
+    deprecate_args,
     experimental_api,
     get_asyncio_loop,
     get_pg_rank,
@@ -90,6 +91,15 @@
 except ImportError:
     HAVE_PSUTIL = False
 
+DEPRECATED_ARGS = [
+    "enable_cuda_graph",
+    "random_seed",
+    "track_paused_request_events",
+    "enable_chunked_prefill",
+    "inference_logging_step_interval",
+    "pg_collection",
+]
+
 
 class EngineSuspendedError(Exception):
     """Engine is currently suspended and not performing steps."""
@@ -134,6 +144,10 @@ class DynamicInferenceEngine(AbstractEngine):
             batching and a dynamic block-level KV cache (similar to paged attention).
     """
 
+    @deprecate_args(
+        *DEPRECATED_ARGS,
+        message="Argument `{name}` has been deprecated. Only pass `controller` and `context`",
+    )
     def __init__(self, controller: TextGenerationController, context: DynamicInferenceContext):
 
         assert isinstance(
diff --git a/megatron/core/utils.py b/megatron/core/utils.py
index 64f20fec43d..0e1e704e826 100644
--- a/megatron/core/utils.py
+++ b/megatron/core/utils.py
@@ -491,17 +491,6 @@ def divide(numerator, denominator):
     return numerator // denominator
 
 
-def deprecate_inference_params(inference_context, inference_params):
-    """Print warning for deprecated `inference_params`."""
-    if inference_context is None and inference_params is not None:
-        warnings.warn(
-            "`inference_params` renamed to `inference_context`, and will be "
-            "removed in `megatron-core` 0.13."
-        )
-        return inference_params
-    return inference_context
-
-
 def get_tensor_model_parallel_group_if_none(tp_group, is_expert=False, check_initialized=True):
     """Issue a deprecation warning if tp_group is None and return the default tp group."""
     # TODO(zijiey): remove this function later.
@@ -2553,3 +2542,43 @@ class ExperimentalModel:
     """
     func._experimental_api = True
     return func
+
+
+def deprecate_args(
+    *deprecated_keys, message="Argument '{name}' has been deprecated and should not be used."
+):
+    """
+    Intercepts specific keyword arguments to raise a custom TypeError.
+
+    Args:
+        *deprecated_keys: Strings representing the argument names to block.
+        message: Custom error message string. Use {name} as a placeholder.
+    """
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            # Check if any deprecated key is present in kwargs
+            found_deprecated = set(deprecated_keys) & set(kwargs.keys())
+
+            if found_deprecated:
+                bad_key = list(found_deprecated)[0]
+                raise TypeError(message.format(name=bad_key))
+
+            # Send args to the real function
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def deprecate_inference_params(inference_context, inference_params):
+    """Print warning for deprecated `inference_params`."""
+    if inference_context is None and inference_params is not None:
+        warnings.warn(
+            "`inference_params` renamed to `inference_context`, and will be "
+            "removed in `megatron-core` 0.13."
+        )
+        return inference_params
+    return inference_context

From 964302c097c18e2ceb52a182acb4888352d6edd3 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Mon, 26 Jan 2026 16:18:52 -0800
Subject: [PATCH 23/30] Remove --inference-max-batch-size

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 examples/rl/README.md                             | 2 +-
 examples/rl/model_configs/llama3p1_8b_instruct.sh | 2 +-
 examples/rl/model_configs/nemotron5_56b.sh        | 2 +-
 examples/rl/model_configs/nemotron5_8b.sh         | 2 +-
 examples/rl/model_configs/nemotron5p5_12b_H.sh    | 2 +-
 examples/rl/model_configs/nemotron6_3b_moe.sh     | 2 +-
 examples/rl/model_configs/qwen3_30b_a3b_moe.sh    | 2 +-
 examples/rl/model_configs/qwen3_32b.sh            | 2 +-
 examples/rl/model_configs/qwen3_4b.sh             | 2 +-
 examples/rl/model_configs/qwen3_8b.sh             | 2 +-
 examples/rl/model_configs/qwen_2p5_32b.sh         | 2 +-
 examples/rl/model_configs/qwen_2p5_3b.sh          | 2 +-
 examples/rl/model_configs/qwen_2p5_distill_7b.sh  | 2 +-
 examples/rl/model_configs/qwen_2p5_math_7b.sh     | 2 +-
 14 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/examples/rl/README.md b/examples/rl/README.md
index 34b6fafa517..9c2de3ec088 100644
--- a/examples/rl/README.md
+++ b/examples/rl/README.md
@@ -94,7 +94,7 @@ MODEL_OPTIONS="\
   --ckpt-format torch \
   --seq-length $MAX_SEQ_LENGTH \
   --inference-max-seq-length $MAX_SEQ_LENGTH \
-  --inference-max-batch-size $MAX_INFERENCE_BS \
+  --inference-max-requests $MAX_INFERENCE_BS \
   --pretrained-checkpoint $CHECKPOINT \
   --untie-embeddings-and-output-weights \
   --disable-bias-linear \
diff --git a/examples/rl/model_configs/llama3p1_8b_instruct.sh b/examples/rl/model_configs/llama3p1_8b_instruct.sh
index 24d285a6cf7..3b3a36452da 100644
--- a/examples/rl/model_configs/llama3p1_8b_instruct.sh
+++ b/examples/rl/model_configs/llama3p1_8b_instruct.sh
@@ -77,7 +77,7 @@ MODEL_OPTIONS="\
   --ckpt-format torch_dist \
   --seq-length $MAX_SEQ_LENGTH \
   --inference-max-seq-length $MAX_SEQ_LENGTH \
-  --inference-max-batch-size $MAX_INFERENCE_BS \
+  --inference-max-requests $MAX_INFERENCE_BS \
   --pretrained-checkpoint $CHECKPOINT \
   --add-qkv-bias \
   --normalization RMSNorm \
diff --git a/examples/rl/model_configs/nemotron5_56b.sh b/examples/rl/model_configs/nemotron5_56b.sh
index fd2cc4f7212..741cd054b73 100644
--- a/examples/rl/model_configs/nemotron5_56b.sh
+++ b/examples/rl/model_configs/nemotron5_56b.sh
@@ -58,7 +58,7 @@ MODEL_OPTIONS="\
   --calculate-per-token-loss \
   --seq-length $MAX_SEQ_LENGTH \
   --inference-max-seq-length $MAX_SEQ_LENGTH \
-  --inference-max-batch-size $MAX_INFERENCE_BS \
+  --inference-max-requests $MAX_INFERENCE_BS \
   --pretrained-checkpoint $CHECKPOINT \
  --fp8-format hybrid \
     --fp8-amax-history-len 1 \
diff --git a/examples/rl/model_configs/nemotron5_8b.sh b/examples/rl/model_configs/nemotron5_8b.sh
index 7b8947ae763..753d4e493a2 100644
--- a/examples/rl/model_configs/nemotron5_8b.sh
+++ b/examples/rl/model_configs/nemotron5_8b.sh
@@ -58,7 +58,7 @@ MODEL_OPTIONS="\
   --calculate-per-token-loss \
   --seq-length $MAX_SEQ_LENGTH \
   --inference-max-seq-length $MAX_SEQ_LENGTH \
-  --inference-max-batch-size $MAX_INFERENCE_BS \
+  --inference-max-requests $MAX_INFERENCE_BS \
   --pretrained-checkpoint $CHECKPOINT \
   --hybrid-override-pattern M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M- \
     --spec megatron.core.models.mamba.mamba_layer_specs mamba_stack_spec \
diff --git a/examples/rl/model_configs/nemotron5p5_12b_H.sh b/examples/rl/model_configs/nemotron5p5_12b_H.sh
index 9e97051e087..adbcc8d03f0 100644
--- a/examples/rl/model_configs/nemotron5p5_12b_H.sh
+++ b/examples/rl/model_configs/nemotron5p5_12b_H.sh
@@ -65,7 +65,7 @@ MODEL_OPTIONS="\
   --calculate-per-token-loss \
   --seq-length $MAX_SEQ_LENGTH \
   --inference-max-seq-length $MAX_SEQ_LENGTH \
-  --inference-max-batch-size $MAX_INFERENCE_BS \
+  --inference-max-requests $MAX_INFERENCE_BS \
   --pretrained-checkpoint $CHECKPOINT \
   --fp8-recipe blockwise \
   --fp8-format e4m3 \
diff --git a/examples/rl/model_configs/nemotron6_3b_moe.sh b/examples/rl/model_configs/nemotron6_3b_moe.sh
index 8efe0b2debb..f7308ab942e 100644
--- a/examples/rl/model_configs/nemotron6_3b_moe.sh
+++ b/examples/rl/model_configs/nemotron6_3b_moe.sh
@@ -91,7 +91,7 @@ MODEL_OPTIONS="\
   --rl-importance-sampling-truncation-coef 10.0 \
   --seq-length $MAX_SEQ_LENGTH \
   --inference-max-seq-length $MAX_SEQ_LENGTH \
-  --inference-max-batch-size $MAX_INFERENCE_BS \
+  --inference-max-requests $MAX_INFERENCE_BS \
   --pretrained-checkpoint $CHECKPOINT \
   --distributed-timeout-minutes 60 \
   --use-mcore-models \
diff --git a/examples/rl/model_configs/qwen3_30b_a3b_moe.sh b/examples/rl/model_configs/qwen3_30b_a3b_moe.sh
index 775a9587ba4..eb55ba35cc6 100644
--- a/examples/rl/model_configs/qwen3_30b_a3b_moe.sh
+++ b/examples/rl/model_configs/qwen3_30b_a3b_moe.sh
@@ -37,7 +37,7 @@ ENV_DEPENDENT="\
 MODEL_OPTIONS="
 --seq-length $MAX_SEQ_LENGTH \
 --inference-max-seq-length $MAX_SEQ_LENGTH \
---inference-max-batch-size $MAX_INFERENCE_BS \
+--inference-max-requests $MAX_INFERENCE_BS \
 --pretrained-checkpoint $CHECKPOINT \
 --no-use-tokenizer-model-from-checkpoint-args \
 --seq-length 8192 \
diff --git a/examples/rl/model_configs/qwen3_32b.sh b/examples/rl/model_configs/qwen3_32b.sh
index cd153a04f3c..c06c5f55b53 100644
--- a/examples/rl/model_configs/qwen3_32b.sh
+++ b/examples/rl/model_configs/qwen3_32b.sh
@@ -38,7 +38,7 @@ MODEL_OPTIONS="\
   --ckpt-format torch_dist \
   --seq-length $MAX_SEQ_LENGTH \
   --inference-max-seq-length $MAX_SEQ_LENGTH \
-  --inference-max-batch-size $MAX_INFERENCE_BS \
+  --inference-max-requests $MAX_INFERENCE_BS \
   --pretrained-checkpoint $CHECKPOINT \
   --untie-embeddings-and-output-weights \
   --num-layers 64 \
diff --git a/examples/rl/model_configs/qwen3_4b.sh b/examples/rl/model_configs/qwen3_4b.sh
index da238511fd3..6f6c6b6bf57 100644
--- a/examples/rl/model_configs/qwen3_4b.sh
+++ b/examples/rl/model_configs/qwen3_4b.sh
@@ -38,7 +38,7 @@ MODEL_OPTIONS="\
   --ckpt-format torch_dist \
   --seq-length $MAX_SEQ_LENGTH \
   --inference-max-seq-length $MAX_SEQ_LENGTH \
-  --inference-max-batch-size $MAX_INFERENCE_BS \
+  --inference-max-requests $MAX_INFERENCE_BS \
   --pretrained-checkpoint $CHECKPOINT \
   --num-layers 36 \
   --hidden-size 2560 \
diff --git a/examples/rl/model_configs/qwen3_8b.sh b/examples/rl/model_configs/qwen3_8b.sh
index 6758cd84c3d..54ff7385331 100644
--- a/examples/rl/model_configs/qwen3_8b.sh
+++ b/examples/rl/model_configs/qwen3_8b.sh
@@ -38,7 +38,7 @@ MODEL_OPTIONS="\
   --ckpt-format torch_dist \
   --seq-length $MAX_SEQ_LENGTH \
   --inference-max-seq-length $MAX_SEQ_LENGTH \
-  --inference-max-batch-size $MAX_INFERENCE_BS \
+  --inference-max-requests $MAX_INFERENCE_BS \
   --pretrained-checkpoint $CHECKPOINT \
   --untie-embeddings-and-output-weights \
   --num-layers 36 \
diff --git a/examples/rl/model_configs/qwen_2p5_32b.sh b/examples/rl/model_configs/qwen_2p5_32b.sh
index d82972ba477..2a2a9ae2420 100644
--- a/examples/rl/model_configs/qwen_2p5_32b.sh
+++ b/examples/rl/model_configs/qwen_2p5_32b.sh
@@ -59,7 +59,7 @@ MODEL_OPTIONS="\
   --ckpt-format torch_dist \
   --seq-length $MAX_SEQ_LENGTH \
   --inference-max-seq-length $MAX_SEQ_LENGTH \
-  --inference-max-batch-size $MAX_INFERENCE_BS \
+  --inference-max-requests $MAX_INFERENCE_BS \
   --pretrained-checkpoint $CHECKPOINT \
   --untie-embeddings-and-output-weights \
   --disable-bias-linear \
diff --git a/examples/rl/model_configs/qwen_2p5_3b.sh b/examples/rl/model_configs/qwen_2p5_3b.sh
index 246afae6ad2..f3250f39ecc 100644
--- a/examples/rl/model_configs/qwen_2p5_3b.sh
+++ b/examples/rl/model_configs/qwen_2p5_3b.sh
@@ -62,7 +62,7 @@ MODEL_OPTIONS="\
   --ckpt-format torch_dist \
   --seq-length $MAX_SEQ_LENGTH \
   --inference-max-seq-length $MAX_SEQ_LENGTH \
-  --inference-max-batch-size $MAX_INFERENCE_BS \
+  --inference-max-requests $MAX_INFERENCE_BS \
   --pretrained-checkpoint $CHECKPOINT \
   --disable-bias-linear \
   --add-qkv-bias \
diff --git a/examples/rl/model_configs/qwen_2p5_distill_7b.sh b/examples/rl/model_configs/qwen_2p5_distill_7b.sh
index 149ac77965f..1438bca0726 100644
--- a/examples/rl/model_configs/qwen_2p5_distill_7b.sh
+++ b/examples/rl/model_configs/qwen_2p5_distill_7b.sh
@@ -44,7 +44,7 @@ MODEL_OPTIONS="\
   --ckpt-format torch \
   --seq-length $MAX_SEQ_LENGTH \
   --inference-max-seq-length $MAX_SEQ_LENGTH \
-  --inference-max-batch-size $MAX_INFERENCE_BS \
+  --inference-max-requests $MAX_INFERENCE_BS \
   --pretrained-checkpoint $CHECKPOINT \
   --untie-embeddings-and-output-weights \
   --disable-bias-linear \
diff --git a/examples/rl/model_configs/qwen_2p5_math_7b.sh b/examples/rl/model_configs/qwen_2p5_math_7b.sh
index 1d631fa80a5..b598bb127bd 100644
--- a/examples/rl/model_configs/qwen_2p5_math_7b.sh
+++ b/examples/rl/model_configs/qwen_2p5_math_7b.sh
@@ -58,7 +58,7 @@ MODEL_OPTIONS="\
   --ckpt-format torch \
   --seq-length $MAX_SEQ_LENGTH \
   --inference-max-seq-length $MAX_SEQ_LENGTH \
-  --inference-max-batch-size $MAX_INFERENCE_BS \
+  --inference-max-requests $MAX_INFERENCE_BS \
   --pretrained-checkpoint $CHECKPOINT \
   --untie-embeddings-and-output-weights \
   --disable-bias-linear \

From 97cb2d61760bf9d44a5ef507dcffb91dc4a8d5e4 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Tue, 27 Jan 2026 11:21:09 -0800
Subject: [PATCH 24/30] RL fixes

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 megatron/rl/inference/megatron.py | 7 +++++--
 train_rl.py                       | 3 +++
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py
index 07ee70d96ea..fe05808c81f 100644
--- a/megatron/rl/inference/megatron.py
+++ b/megatron/rl/inference/megatron.py
@@ -22,7 +22,6 @@
 from megatron.core.models.gpt.gpt_model import GPTModel
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.utils import get_attr_wrapped_model, log_single_rank
-from megatron.inference.utils import get_dynamic_inference_engine
 from megatron.training import get_wandb_writer
 from megatron.training.global_vars import get_args, get_tokenizer
 
@@ -69,6 +68,7 @@ def get_static_inference_engine(args: Namespace, model: MegatronModule) -> Abstr
         ),
     )
 
+
 class MegatronLocal(InferenceServer, ReturnsTokens, ReturnsRaw):
     """Interface to use MCoreEngine directly as an inference engine."""
 
@@ -121,6 +121,9 @@ async def base_generate(self, request: InferenceRequest):
 
     @classmethod
     async def launch(cls, model: GPTModel, **kwargs):
+        # Import here to avoid circular imports
+        from megatron.inference.utils import get_dynamic_inference_engine
+
         args = get_args()
         tokenizer = get_tokenizer()
 
@@ -131,7 +134,7 @@ async def launch(cls, model: GPTModel, **kwargs):
                 "WARNING: Tokenizer has no BOS token so prompt will not have BOS token",
             )
 
-        inference_engine: DynamicInferenceEngine = get_dynamic_inference_engine()
+        inference_engine: DynamicInferenceEngine = get_dynamic_inference_engine(model=model)
         await inference_engine.start_listening_to_data_parallel_coordinator(
             inference_coordinator_port=41521, launch_inference_coordinator=True
         )
diff --git a/train_rl.py b/train_rl.py
index 299843bcff3..cb8bd8fccee 100644
--- a/train_rl.py
+++ b/train_rl.py
@@ -369,6 +369,8 @@ def __getitem__(self, idx):
 
 if __name__ == "__main__":
 
+    from megatron.inference.utils import add_inference_args
+
     # Temporary for transition to core datasets
     train_valid_test_datasets_provider.is_distributed = True
 
@@ -400,4 +402,5 @@ def _model_builder(
         ModelType.encoder_or_decoder,
         forward_step,
         args_defaults={},
+        extra_args_provider=add_inference_args,
     )

From 70e4e3a7371b38e3550fc45d0749933123156683 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Tue, 27 Jan 2026 11:26:31 -0800
Subject: [PATCH 25/30] Fix perf test

---
 tools/run_inference_performance_test.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/tools/run_inference_performance_test.py b/tools/run_inference_performance_test.py
index fb7b97c07ae..abe22113dc6 100644
--- a/tools/run_inference_performance_test.py
+++ b/tools/run_inference_performance_test.py
@@ -22,7 +22,7 @@
     TextGenerationController,
 )
 from megatron.core.transformer.module import MegatronModule
-from megatron.inference.utils import get_dynamic_inference_engine
+from megatron.inference.utils import add_inference_args, get_dynamic_inference_engine
 from model_provider import model_provider
 
 sys.path.append(
@@ -32,7 +32,6 @@
 from functools import partial
 from typing import List
 
-from examples.inference.gpt.utils import add_common_inference_args
 from megatron.core import mpu
 from megatron.training import get_args, get_model, get_tokenizer
 from megatron.training.checkpointing import load_checkpoint
@@ -43,7 +42,7 @@
 
 def add_inference_benchmarking_args(parser):
     """Inference benchmarking arguments."""
-    parser = add_common_inference_args(parser)
+    parser = add_inference_args(parser)
 
     group = parser.add_argument_group(title='inference_benchmarking')
 
@@ -187,6 +186,7 @@ def main():
         return_log_probs=args.return_log_probs,
         top_n_logprobs=args.top_n_logprobs,
         num_tokens_to_generate=args.num_tokens_to_generate,
+        termination_id=-1,
     )
     sampling_params.add_attributes({"no_early_termination": True})
 
@@ -216,10 +216,10 @@ def main():
                 )
             )
 
-    if args.cuda_graph_impl == "local" and args.engine_type == "static":
-        print(f"Running warmup for CUDA graphs...")
-        warmup_sampling_params = SamplingParams(num_tokens_to_generate=10)
-        warmup_sampling_params.add_attributes({"no_early_termination": True})
+    # TODO(ksanthanam): Use a command line argument for warmup iterations
+    for i in range(3):
+        print(f"Running warmup iteration {i+1}...")
+        warmup_sampling_params = SamplingParams(num_tokens_to_generate=10, termination_id=-1)
         inference_engine.generate(prompts=["warmup"], sampling_params=warmup_sampling_params)
 
     if args.benchmark_profile:
@@ -261,6 +261,10 @@ def main():
                 result_dict['generated_output'] = tokenizer.detokenize(result.generated_tokens)
             print(result_dict)
 
+    total_output_tokens = args.num_tokens_to_generate * args.inference_max_requests
+    throughput = total_output_tokens / latency
+    print(f"Throughput: {throughput} output tokens / second")
+
 
 if __name__ == "__main__":
     main()

From c3e10a25e8979d672a87131b23613af8f03bbaf3 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Thu, 29 Jan 2026 18:54:53 -0800
Subject: [PATCH 26/30] Fix static functional test

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 .../model_config.yaml                                            | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_logitsmatch/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_logitsmatch/model_config.yaml
index be00e4b3ce7..1c78b466b1e 100644
--- a/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_logitsmatch/model_config.yaml
+++ b/tests/functional_tests/test_cases/gpt/gpt_static_inference_tp1_pp1_583m_logitsmatch/model_config.yaml
@@ -44,6 +44,7 @@ MODEL_ARGS:
   --flash-decode: true
   --dist-ckpt-strictness: log_unexpected
   --output-path: ${INFERENCE_OUTPUT_PATH}
+  --use-legacy-static-engine: true
   --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies."
   --incoming-requests-per-sec: -1 # all requests arrive up front.
 METRICS:

From a57af29c6ab9534c918d71c918d060d2000ca4da Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Fri, 30 Jan 2026 10:36:39 -0800
Subject: [PATCH 27/30] Try removing arguments that are now in
 TransformerConfig

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 megatron/training/arguments.py | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 1768ebbd6e9..aefa89e08ec 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1505,9 +1505,6 @@ def _add_inference_args(parser):
                        '1) allocate `memory_buffer` in unified memory. '
                        'Eventually, additional levels will be included to '
                        'control other tensors within the context.')
-    group.add_argument('--nccl-all-reduce-for-prefill',
-                       action='store_true', default=False,
-                       help='When using symmeric all reduce kernels this will use regular nccl kernels for prefill. This can be more effecient when prefill is large as the nccl kernels can be more bandwith optimized')
     # TODO(ksanthanam): Clean this up in future PR
     group.add_argument('--enable-chunked-prefill', dest='enable_chunked_prefill',
                        action='store_true', default=False,
@@ -2709,10 +2706,6 @@ def _add_moe_args(parser):
     group.add_argument('--moe-upcycling-granularity', type=int, default=1,
                        help='This param sepecifics how many times smaller is the expert hidden size compared with the original dense FFN hidden size. '
                        'For using granular upcycling strategy, please set this param as a positive integer. If this param is set to 1, it means using the default upcycling strategy.')
-    group.add_argument('--moe-pad-experts-for-cuda-graph-inference', action='store_true',
-                       help="some MoE routers have a D2H sync that will break cuda graphs.  If this flag is set the router will switch" \
-                       " to dropping and padding during decode time which does not have a D2H sync. The capacity factor is set to the" \
-                       " max that an expert could see during inference so no tokens are actually dropped.")
     return parser
 
 def _add_mla_args(parser):

From dbb861461bf1f89991fb83ae5beabbd05fdedfb5 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Fri, 30 Jan 2026 15:28:02 -0800
Subject: [PATCH 28/30] Bug fixes

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 examples/inference/gpt/gpt_static_inference.py    |  7 +++----
 examples/rl/model_configs/llama3p1_8b_instruct.sh |  1 +
 megatron/core/inference/engines/static_engine.py  |  7 +++++--
 tools/run_inference_performance_test.py           | 11 +++++++----
 4 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/examples/inference/gpt/gpt_static_inference.py b/examples/inference/gpt/gpt_static_inference.py
index 4cd584cb517..298ebfebd86 100644
--- a/examples/inference/gpt/gpt_static_inference.py
+++ b/examples/inference/gpt/gpt_static_inference.py
@@ -202,15 +202,14 @@ def main():
         from collections import defaultdict
 
         unique_prompt_map = defaultdict(list)
-        for result_idx, record in enumerate(results):
-            result = record.requests[0]
+        for result_idx, result in enumerate(results):
             unique_prompt_map[result.prompt].append(result_idx)
 
         # Print unique prompts + outputs.
         for unique_idx, (prompt_text, result_idxs) in enumerate(unique_prompt_map.items()):
             result_idx = result_idxs[0]
-            record = results[result_idx]
-            generated_text = record.requests[0].generated_text.replace("\n", "\\n")
+            result = results[result_idx]
+            generated_text = result.generated_text.replace("\n", "\\n")
             print(
                 f"{unique_idx}/{len(unique_prompt_map)} [{len(result_idxs)}]. {prompt_text} "
                 f"... {generated_text}"
diff --git a/examples/rl/model_configs/llama3p1_8b_instruct.sh b/examples/rl/model_configs/llama3p1_8b_instruct.sh
index 3b3a36452da..5398dad1a4e 100644
--- a/examples/rl/model_configs/llama3p1_8b_instruct.sh
+++ b/examples/rl/model_configs/llama3p1_8b_instruct.sh
@@ -101,6 +101,7 @@ MODEL_OPTIONS="\
   --max-position-embeddings 131072  \
   --tokenizer-type HuggingFaceTokenizer \
   --tokenizer-model unsloth/Meta-Llama-3.1-8B-Instruct \
+  --legacy-tokenizer \
   --langrl-inference-server-type "inplace_megatron_chat" \
   --langrl-inference-server-conversation-template "unsloth/Meta-Llama-3.1-8B-Instruct" \
   --lr 3e-7 \
diff --git a/megatron/core/inference/engines/static_engine.py b/megatron/core/inference/engines/static_engine.py
index fc381848268..103558902f6 100644
--- a/megatron/core/inference/engines/static_engine.py
+++ b/megatron/core/inference/engines/static_engine.py
@@ -230,13 +230,16 @@ def generate_using_dynamic_engine(
         if prompts:
             if add_BOS:
                 sampling_params.add_BOS = True
-            return self.dynamic_engine.generate(prompts=prompts, sampling_params=sampling_params)
+            request_records = self.dynamic_engine.generate(prompts=prompts, sampling_params=sampling_params)
         elif inference_requests:
             prompts = [request.prompt for request in inference_requests]
             sampling_params = inference_requests[0].sampling_params
             if add_BOS:
                 sampling_params.add_BOS = True
-            return self.dynamic_engine.generate(prompts=prompts, sampling_params=sampling_params)
+            request_records = self.dynamic_engine.generate(prompts=prompts, sampling_params=sampling_params)
+
+        # Return the underlying `InferenceRequest` objects from the `DynamicInferenceRequestRecord`s.
+        return [record.results[0] for record in request_records]
 
     def generate_using_legacy_static_engine(
         self,
diff --git a/tools/run_inference_performance_test.py b/tools/run_inference_performance_test.py
index abe22113dc6..e10bed635e4 100644
--- a/tools/run_inference_performance_test.py
+++ b/tools/run_inference_performance_test.py
@@ -13,7 +13,10 @@
 from megatron.core.inference.contexts import StaticInferenceContext
 from megatron.core.inference.engines import DynamicInferenceEngine, StaticInferenceEngine
 from megatron.core.inference.engines.abstract_engine import AbstractEngine
-from megatron.core.inference.inference_request import InferenceRequest
+from megatron.core.inference.inference_request import (
+    DynamicInferenceRequestRecord,
+    InferenceRequest,
+)
 from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
     GPTInferenceWrapper,
 )
@@ -232,9 +235,10 @@ def main():
         )
     else:
         prompts = [request.prompt_tokens for request in requests]
-        results: List[InferenceRequest] = inference_engine.generate(
+        records: List[DynamicInferenceRequestRecord] = inference_engine.generate(
             prompts=prompts, sampling_params=sampling_params
         )
+        results: List[InferenceRequest] = [record.requests[0] for record in records]
 
     end_time = time.perf_counter()
     latency = end_time - start_time
@@ -245,8 +249,7 @@ def main():
         torch.cuda.cudart().cudaProfilerStop()
 
     if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
-        for idx, record in enumerate(results):
-            result = record.requests[0]
+        for idx, result in enumerate(results):
             print(f' \n------------- RESULT FOR PROMPT {idx} --------------- ')
             generated_log_probs = result.generated_log_probs
             result_dict = {

From ca27708bac987fb9df34163df8f00dc64471e644 Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Fri, 30 Jan 2026 15:31:32 -0800
Subject: [PATCH 29/30] Fix formatting

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 megatron/core/inference/engines/static_engine.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/megatron/core/inference/engines/static_engine.py b/megatron/core/inference/engines/static_engine.py
index 103558902f6..578549e5a09 100644
--- a/megatron/core/inference/engines/static_engine.py
+++ b/megatron/core/inference/engines/static_engine.py
@@ -230,13 +230,17 @@ def generate_using_dynamic_engine(
         if prompts:
             if add_BOS:
                 sampling_params.add_BOS = True
-            request_records = self.dynamic_engine.generate(prompts=prompts, sampling_params=sampling_params)
+            request_records = self.dynamic_engine.generate(
+                prompts=prompts, sampling_params=sampling_params
+            )
         elif inference_requests:
             prompts = [request.prompt for request in inference_requests]
             sampling_params = inference_requests[0].sampling_params
             if add_BOS:
                 sampling_params.add_BOS = True
-            request_records = self.dynamic_engine.generate(prompts=prompts, sampling_params=sampling_params)
+            request_records = self.dynamic_engine.generate(
+                prompts=prompts, sampling_params=sampling_params
+            )
 
         # Return the underlying `InferenceRequest` objects from the `DynamicInferenceRequestRecord`s.
         return [record.results[0] for record in request_records]

From 80ddc9a051e863c5a56a9278a18976aa812e31bf Mon Sep 17 00:00:00 2001
From: Keshav Santhanam <ksanthanam@nvidia.com>
Date: Fri, 30 Jan 2026 15:37:57 -0800
Subject: [PATCH 30/30] More bug fixes

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
---
 megatron/core/inference/engines/static_engine.py         | 2 +-
 tests/unit_tests/inference/engines/test_static_engine.py | 2 --
 tools/run_inference_performance_test.py                  | 2 +-
 3 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/megatron/core/inference/engines/static_engine.py b/megatron/core/inference/engines/static_engine.py
index 578549e5a09..5ae37d5967e 100644
--- a/megatron/core/inference/engines/static_engine.py
+++ b/megatron/core/inference/engines/static_engine.py
@@ -243,7 +243,7 @@ def generate_using_dynamic_engine(
             )
 
         # Return the underlying `InferenceRequest` objects from the `DynamicInferenceRequestRecord`s.
-        return [record.results[0] for record in request_records]
+        return [record.merge() for record in request_records]
 
     def generate_using_legacy_static_engine(
         self,
diff --git a/tests/unit_tests/inference/engines/test_static_engine.py b/tests/unit_tests/inference/engines/test_static_engine.py
index c7d7f223d6e..483a21d13bd 100644
--- a/tests/unit_tests/inference/engines/test_static_engine.py
+++ b/tests/unit_tests/inference/engines/test_static_engine.py
@@ -188,8 +188,6 @@ def test_generate_dynamic(self, batch_size: int, num_trials: int, empty_prompt:
 
             assert len(results) == batch_size
             for result in results:
-                if isinstance(result, DynamicInferenceRequestRecord):
-                    result = result.merge()
                 assert isinstance(result, InferenceRequest), (
                     "expected <InferenceRequest>; found <%s>." % type(result).__name__
                 )
diff --git a/tools/run_inference_performance_test.py b/tools/run_inference_performance_test.py
index e10bed635e4..430bb7ebb9a 100644
--- a/tools/run_inference_performance_test.py
+++ b/tools/run_inference_performance_test.py
@@ -238,7 +238,7 @@ def main():
         records: List[DynamicInferenceRequestRecord] = inference_engine.generate(
             prompts=prompts, sampling_params=sampling_params
         )
-        results: List[InferenceRequest] = [record.requests[0] for record in records]
+        results: List[InferenceRequest] = [record.merge() for record in records]
 
     end_time = time.perf_counter()
     latency = end_time - start_time