diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py
index 88b744b3ac0..7fcac70c11a 100644
--- a/examples/inference/gpt/gpt_dynamic_inference.py
+++ b/examples/inference/gpt/gpt_dynamic_inference.py
@@ -1,40 +1,31 @@
 # 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
-import torch
-from argparse import ArgumentParser
 from collections import defaultdict
-from functools import partial
+from typing import Dict, List, Optional
+
+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))
 )
 
-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,
 )
-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.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,
@@ -44,194 +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_inference_config_from_model_and_args,
+    get_model_for_inference,
+)
 
 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
 
+import megatron
 from megatron.core.utils import configure_nvtx_profiling
-import logging
+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_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(
-    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()
-
-    # 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
-    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,
-        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,
-        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,
-        offload_kv_cache=args.rl_offload_kv_cache_during_training
-    )
-
-    return context
-
-
-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
-
-
 def run_inference(
     requests: List[Request],
     engine: DynamicInferenceEngine,
@@ -284,11 +107,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
@@ -305,10 +124,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)
 
@@ -318,11 +136,12 @@ 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 
+        # 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.
         if args.suspend_resume_interval is not None:
@@ -335,9 +154,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()
@@ -349,7 +168,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"]
@@ -408,29 +229,29 @@ 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,
     }
 
 
 @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},
     )
 
     # 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)
@@ -456,42 +277,36 @@ 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)
+    model = get_model_for_inference()
 
     # 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)
+    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
+    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:
                 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(controller, context)
 
     setup_prefix = build_dynamic_engine_setup_prefix(args, model, context, requests)
     print("~~~")
@@ -522,14 +337,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")
 
@@ -547,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)
@@ -567,16 +384,17 @@ 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}{', <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.
@@ -592,14 +410,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
 
@@ -631,7 +451,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(
@@ -643,18 +463,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 cbb7a1aa745..ab84ee5bf5c 100644
--- a/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
+++ b/examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py
@@ -2,43 +2,33 @@
 
 import asyncio
 import json
+import logging
 import os
 import time
-import torch
-import torch.distributed as dist
+import warnings
 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, 
-    build_requests,
-    add_common_inference_args
-)
+import torch
+import torch.distributed as dist
 
-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
 
 logging.basicConfig(level=logging.INFO, force=True)
 
+
 async def main(
     engine: DynamicInferenceEngine,
     requests: List[Request],
@@ -51,12 +41,11 @@ 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.
     # leaving inference_coordinator_port as None will find a free port automatically.
-    
     dp_addr = await engine.start_listening_to_data_parallel_coordinator(
         inference_coordinator_port=port,
         launch_inference_coordinator=True,
@@ -69,14 +58,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()
@@ -98,7 +84,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
@@ -114,10 +103,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
@@ -135,7 +123,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)
 
@@ -170,16 +158,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()
@@ -190,11 +181,11 @@ 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(
-            extra_args_provider=add_dynamic_inference_args,
+            extra_args_provider=add_inference_args,
             args_defaults={'no_load_rng': True, 'no_load_optim': True},
         )
 
@@ -213,34 +204,16 @@ async def main(
             ),
         )
 
-        # Requests, context, conroller.
-        model = get_model()
-        mamba_inference_state_config = get_mamba_inference_state_config_from_model(model)
+        model = get_model_for_inference()
+
         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 = 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("~~~")
@@ -249,13 +222,7 @@ 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"):
diff --git a/examples/inference/gpt/gpt_static_inference.py b/examples/inference/gpt/gpt_static_inference.py
index 03a60927ab2..298ebfebd86 100644
--- a/examples/inference/gpt/gpt_static_inference.py
+++ b/examples/inference/gpt/gpt_static_inference.py
@@ -1,21 +1,11 @@
 # 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
-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
@@ -23,17 +13,12 @@
 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,
 )
 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))
@@ -41,18 +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(
@@ -83,30 +68,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,22 +136,7 @@ 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
-    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)
 
@@ -276,7 +232,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 +249,5 @@ def main():
         torch.distributed.destroy_process_group()
 
 
-
 if __name__ == "__main__":
     main()
diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py
index a04b856c0a6..b7a3977605c 100644
--- a/examples/inference/gpt/utils.py
+++ b/examples/inference/gpt/utils.py
@@ -1,158 +1,23 @@
 # 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 megatron.core.inference.contexts import DynamicInferenceContext
 from megatron.core.inference.contexts.dynamic_context import get_mem_size_str
-from megatron.core.transformer.module import MegatronModule
-
+from megatron.core.inference.inference_request import DynamicInferenceRequest
 from megatron.core.inference.sampling_params import SamplingParams
-
-
-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
+from megatron.core.transformer.module import MegatronModule
+from megatron.training import get_args
 
 
 def get_default_sampling_params(termination_id: int = None):
@@ -162,9 +27,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 +54,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 +70,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 +101,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 +117,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 +133,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 +145,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 +165,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 +212,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 +282,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 +306,7 @@ def build_dynamic_engine_setup_prefix(
         f"[r {context.max_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 +322,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/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..5398dad1a4e 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 \
@@ -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/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 eff4f6cf0b3..7d98f4eda63 100644
--- a/examples/rl/model_configs/nemotron6_3b_moe.sh
+++ b/examples/rl/model_configs/nemotron6_3b_moe.sh
@@ -85,7 +85,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 \
diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py
new file mode 100644
index 00000000000..5970b4f14f6
--- /dev/null
+++ b/megatron/core/inference/config.py
@@ -0,0 +1,186 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+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 InferenceConfig:
+    """
+    Config for inference.
+
+    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.
+    """
+
+    offload_kv_cache: bool = False
+    """If True, offload KV cache during RL training."""
+
+    # =================================
+    # 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.
+    """
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..4f03726fe3d 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 InferenceConfig
+
 
 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: InferenceConfig):
         """
         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.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 5dc2d503097..915180a5ca2 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -4,22 +4,19 @@
 import math
 import warnings
 from contextlib import nullcontext
-from typing import TYPE_CHECKING, 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 InferenceConfig
 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,
@@ -28,13 +25,13 @@
 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 TransformerConfig
+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_attr_wrapped_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
@@ -45,14 +42,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:
@@ -66,16 +56,36 @@
 except ImportError:
     HAVE_TORCH_MEMORY_SAVER = 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
+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",
+    "offload_kv_cache",
+]
 
 
 class ContextOverflowError(Exception):
@@ -213,130 +223,45 @@ class DynamicInferenceContext(BaseInferenceContext):
     given step, 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.
-        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
-        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.
-        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_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.
-        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
-            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.
-        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.
+        model_config (TransformerConfig): Model config.
+        inference_config (InferenceConfig): Inference config.
     """
 
     DEFAULT_MAX_TOKENS = 16384
     TOKEN_ROUNDER = 64
     REQUEST_ROUNDER = 4
 
-    def __init__(
-        self,
-        *,
-        params_dtype: torch.dtype,
-        num_layers: int,
-        kv_channels: int,
-        num_attention_heads: int,
-        max_sequence_length: int,
-        buffer_size_gb: float,
-        paused_buffer_size_gb: float | None = None,
-        max_requests: 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,
-        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,
-        offload_kv_cache: Optional[bool] = False,
-    ):
-        super().__init__(materialize_only_last_token_logits=materialize_only_last_token_logits)
-
-        self.cache_mla_latent = cache_mla_latent
+    @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)
+
+        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,
-            )
-
-        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:
-            tp_size = (
-                get_pg_size(pg_collection.tp)
-                if pg_collection is not None
-                else parallel_state.get_tensor_model_parallel_world_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)
+            pp_size = get_pg_size(pg_collection.pp)
         else:
-            tp_size = tensor_model_parallel_size
+            tp_size = model_config.tensor_model_parallel_size
+            pp_size = model_config.pipeline_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:
-            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
-
         # 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.
         # Otherwise:
@@ -357,6 +282,7 @@ def __init__(
             self.expert_model_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
@@ -381,7 +307,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)}
@@ -392,11 +318,11 @@ def __init__(
             )
 
         # Block size tokens, bytes.
-        dtype_size_bytes = params_dtype.itemsize
-        self.block_size_tokens = block_size_tokens
+        dtype_size_bytes = model_config.params_dtype.itemsize
+        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
@@ -422,9 +348,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:
@@ -435,9 +361,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.
@@ -471,13 +399,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 = params_dtype
-        self.max_sequence_length = max_sequence_length
+        self.params_dtype = model_config.params_dtype
+        self.max_sequence_length = inference_config.max_sequence_length
 
         # Request and token counts.
         self.total_request_count = 0
@@ -497,16 +426,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 >= "
@@ -538,37 +467,39 @@ def __init__(
         )
 
         # CUDA graph config list
+        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,
             )
         )
 
         # Whether to offload the KV cache. Determines where the KV cache is allocated within memory.
-        self.offload_kv_cache = offload_kv_cache
+        self.offload_kv_cache = inference_config.offload_kv_cache
         assert not (
             self.offload_kv_cache and self.unified_memory_level
         ), "The KV cache should not be instantiated in unified memory when it is offloaded during training."
 
         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
@@ -756,14 +687,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
@@ -775,72 +699,9 @@ 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,
-    ):
-        """
-        Instantiate a `DynamicInferenceContext` from a `TransformerConfig` and an `InferenceWrapperConfig`.
-        """
-        # TODO: Add other necessary configs from inference_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
-
-        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 = (
-                inference_config.inference_max_seq_length or model_config.max_sequence_length
-            )
-        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,
-        )
-
     @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
@@ -852,16 +713,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
@@ -882,6 +733,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"],
@@ -889,6 +741,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"],
@@ -958,18 +811,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],
@@ -1386,9 +1241,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.
@@ -1413,6 +1268,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,
@@ -1545,7 +1401,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.
         """
@@ -1784,7 +1640,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:
@@ -1863,7 +1719,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 8c83d2f09b3..a15b33c414a 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.inference.config import InferenceConfig
 
 from .base_context import BaseInferenceContext
 
@@ -19,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)
+        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
@@ -27,13 +26,6 @@ def __init__(
         self.key_value_memory_dict = {}
         self.decode_mode = False
 
-    @classmethod
-    def from_config(cls, config: InferenceWrapperConfig) -> "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 0a95e8f4a53..882db6b3a6a 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -42,6 +42,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,
@@ -89,6 +90,14 @@
 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",
+]
 from megatron.core.inference.contexts.dynamic_context import HAVE_TORCH_MEMORY_SAVER
 
 if HAVE_TORCH_MEMORY_SAVER:
@@ -136,24 +145,13 @@ 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,
-        inference_logging_step_interval: int = 0,
-        pg_collection: Optional[ProcessGroupCollection] = None,
-    ):
+    @deprecate_args(
+        *DEPRECATED_ARGS,
+        message="Argument `{name}` has been deprecated. Only pass `controller` and `context`",
+    )
+    def __init__(self, controller: TextGenerationController, context: DynamicInferenceContext):
 
         assert isinstance(
             controller, TextGenerationController
@@ -161,40 +159,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.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.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()
@@ -205,12 +191,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.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(
@@ -288,8 +274,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()
 
@@ -733,7 +717,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."
@@ -922,7 +906,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 = []
@@ -1202,10 +1186,10 @@ 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.context.metrics_writer is not None
+            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()
         else:
@@ -1338,18 +1322,13 @@ 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 (
-            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 d4c61965d2b..5ae37d5967e 100644
--- a/megatron/core/inference/engines/static_engine.py
+++ b/megatron/core/inference/engines/static_engine.py
@@ -8,7 +8,8 @@
 import torch
 
 from megatron.core.inference.async_stream import AsyncStream
-from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
+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
 from megatron.core.inference.inference_request import InferenceRequest
@@ -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
@@ -42,8 +43,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,53 +68,55 @@ 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
+        # 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 = text_generation_controller.inference_wrapped_model.inference_context
-
-        mamba_inference_state_config = get_mamba_inference_state_config_from_model(
-            text_generation_controller.inference_wrapped_model.model
+        mamba_inference_state_config = MambaInferenceStateConfig.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,
-                    buffer_size_gb=buffer_size_gb,
-                    num_cuda_graphs=1,
-                    mamba_inference_state_config=mamba_inference_state_config,
+                dynamic_context = DynamicInferenceContext(
+                    model_config=self.config,
+                    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,
+                        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,
+                    controller=self.controller, context=dynamic_context
                 )
         except Exception as e:
             # Get exception details for better debugging
@@ -229,13 +230,20 @@ 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.merge() for record in request_records]
 
     def generate_using_legacy_static_engine(
         self,
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..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
@@ -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
+        # TODO(ksanthanam): Add support for fp4
 
-    @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
-
-    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
@@ -137,7 +96,9 @@ def prep_inference_input(self, prompt_tokens) -> Dict[str, Any]:
     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 +144,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 +166,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 +176,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 +192,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 +234,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 +244,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 +263,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..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
@@ -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,13 @@ class GPTInferenceWrapper(AbstractModelInferenceWrapper):
     def __init__(
         self,
         model: GPTModel,
-        inference_wrapper_config: InferenceWrapperConfig,
         inference_context: Optional[BaseInferenceContext] = None,
         pg_collection: Optional[ProcessGroupCollection] = None,
+        inference_wrapper_config: Optional[Any] = None,  # Deprecated
     ):
-        super().__init__(model, inference_wrapper_config, inference_context, pg_collection)
+        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]:
         """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/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 a5233983ed0..617883414d4 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -11,21 +11,22 @@
 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.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.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.transformer.enums import CudaGraphScope
 from megatron.core.transformer.moe.moe_layer import BaseMoELayer
 from megatron.core.transformer.utils import set_model_to_sequence_parallel
@@ -52,28 +53,32 @@ 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
+        inference_config = self.inference_wrapped_model.inference_context.config
         self.tokenizer = tokenizer
 
-        self.pp_group = pp_group
+        pg_collection = 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()
+
+        self.model_is_pipeline_parallel = self.model_config.pipeline_model_parallel_size > 1
 
-        # 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)
-        )
+        # Use padded vocab size because tokenizer vocab size might pad to nearest power of 2.
+        # 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 = unwrapped_model.vocab_size
 
-        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()
@@ -98,9 +103,7 @@ 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
-        # 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
+        logits_dtype = self.inference_wrapped_model.config.params_dtype
 
         self._sampling_backend = "torch"
         self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
@@ -505,7 +508,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 +519,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 +571,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
 
@@ -582,18 +582,17 @@ 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.config.materialize_only_last_token_logits
                 else input_ids.shape[1]
             )
-            vocab_size = inference_wrapper_config.padded_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
 
             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,
             )
@@ -639,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.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)
@@ -684,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.config.materialize_only_last_token_logits,
         )
 
     def _dynamic_step_calculate_top_n_logprobs(
@@ -712,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.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]
@@ -1024,9 +1023,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_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
+        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(
@@ -1066,10 +1066,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 = inference_wrapper_config.padded_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
@@ -1130,14 +1126,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)
@@ -1191,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.config.materialize_only_last_token_logits = (
                     materialize_only_last_token_logits
                 )
 
@@ -1212,14 +1208,14 @@ 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],
-                        dtype=inference_wrapper_config.params_dtype,
+                        [batch_size, logits_seq_len, self.vocab_size],
+                        dtype=self.model_config.params_dtype,
                         tensor=logits,
                         pp_group=self.pp_group,
                     )
@@ -1248,7 +1244,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/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index e287344c13d..f44aed613e7 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -661,7 +661,7 @@ 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.config.materialize_only_last_token_logits:
             if inference_context.is_static_batching():
                 hidden_states = hidden_states[-1:, :, :]
             else:
@@ -691,7 +691,7 @@ def _postprocess(
             assert (
                 in_inference_mode
                 and inference_context.is_dynamic_batching()
-                and inference_context.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 8d45e1d0147..6d43f5583df 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -267,7 +267,7 @@ 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.config.materialize_only_last_token_logits:
             if inference_context.is_static_batching():
                 hidden_states = hidden_states[-1:, :, :]
             else:
@@ -297,7 +297,7 @@ def forward(
             assert (
                 in_inference_mode
                 and inference_context.is_dynamic_batching()
-                and inference_context.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/ssm/mamba_layer.py b/megatron/core/ssm/mamba_layer.py
index ac6e8b5bf40..0b4ef42457d 100644
--- a/megatron/core/ssm/mamba_layer.py
+++ b/megatron/core/ssm/mamba_layer.py
@@ -193,6 +193,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/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
index eaae585905e..48b04c35134 100644
--- a/megatron/core/transformer/transformer_config.py
+++ b/megatron/core/transformer/transformer_config.py
@@ -696,6 +696,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: Literal['probs', 'position'] = "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
@@ -830,6 +836,9 @@ class TransformerConfig(ModelParallelConfig):
     which is no use of symmetric memory.
     """
 
+    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/core/utils.py b/megatron/core/utils.py
index d7b702f25ec..cb2f7d34128 100644
--- a/megatron/core/utils.py
+++ b/megatron/core/utils.py
@@ -496,17 +496,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.
@@ -2405,25 +2394,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
 # ============================================================================
@@ -2558,3 +2528,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
diff --git a/megatron/inference/__init__.py b/megatron/inference/__init__.py
new file mode 100644
index 00000000000..26496bfed70
--- /dev/null
+++ b/megatron/inference/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/inference/utils.py b/megatron/inference/utils.py
new file mode 100644
index 00000000000..145af726c4f
--- /dev/null
+++ b/megatron/inference/utils.py
@@ -0,0 +1,320 @@
+# 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 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.tokenizers.text.utils.build_tokenizer import build_tokenizer
+from megatron.core.transformer.module import MegatronModule
+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_for_inference() -> MegatronModule:
+    """Initialize model and load checkpoint for inference."""
+
+    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_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
+        # 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 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, max_batch_size)
+
+    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 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,
+        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,
+        offload_kv_cache=args.rl_offload_kv_cache_during_training,
+        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,
+        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_for_inference()
+    if args.legacy_tokenizer:
+        tokenizer = get_tokenizer()
+    else:
+        tokenizer = build_tokenizer(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)
+    engine = DynamicInferenceEngine(controller, context)
+    return engine
diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py
index 4e9364b3ae9..602ff4f7450 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
@@ -16,23 +15,13 @@
 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_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,134 +55,20 @@ 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(
-        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,
         max_batch_size=(
-            args.inference_max_batch_size if args.inference_max_batch_size is not None else 1
+            args.inference_max_requests if args.inference_max_requests is not None else 1
         ),
     )
 
 
-## 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.
-
-    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:
-        AbstractBackend: The chosen backend
-    """
-    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,
-        persist_cuda_graphs=args.rl_training_cuda_graphs,
-        offload_kv_cache=args.rl_offload_kv_cache_during_training
-    )
-
-    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,
-    )
-
-
 class MegatronLocal(InferenceServer, ReturnsTokens, ReturnsRaw):
     """Interface to use MCoreEngine directly as an inference engine."""
 
@@ -246,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()
 
@@ -256,30 +134,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 = get_dynamic_inference_engine(model=model)
         dp_addr = 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 46f3c28b1da..f56bc6c5e2f 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -1446,13 +1446,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.')
@@ -1508,15 +1505,10 @@ 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='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.')
@@ -2714,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):
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:
diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py
index 05e0306bfd8..f3ef0910f58 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 InferenceConfig, MambaInferenceStateConfig
 from megatron.core.inference.contexts.dynamic_context import (
     DynamicInferenceContext,
     RequestOverflowError,
@@ -18,14 +17,21 @@
 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
 
 
-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:
@@ -52,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]
@@ -69,23 +72,27 @@ 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,
-            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
+            model_config=TransformerConfig(
+                params_dtype=params_dtype,
+                num_layers=num_layers,
+                kv_channels=kv_channels,
+                num_attention_heads=num_attention_heads,
+            ),
+            inference_config=InferenceConfig(
+                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
 
@@ -93,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)
@@ -107,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:
@@ -145,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)
@@ -168,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)
@@ -181,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
@@ -198,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)
@@ -211,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,
         )
 
@@ -227,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)
@@ -301,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)
@@ -349,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)
@@ -422,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)
 
@@ -520,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)
 
@@ -550,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)
 
@@ -575,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)
@@ -774,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."""
@@ -846,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."""
@@ -913,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)
@@ -988,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(
@@ -1205,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
@@ -1215,23 +1235,39 @@ 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,
-            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,
+            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,
+            ),
+            inference_config=InferenceConfig(
+                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 d5803b3638e..2e935cab4bd 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 InferenceConfig, MambaInferenceStateConfig
 from megatron.core.inference.contexts.dynamic_context import (
     ActiveRequestCountOverflowError,
     BlockOverflowError,
@@ -28,9 +26,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,
@@ -48,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
 
 
@@ -223,26 +214,22 @@ 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,
-            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,
-            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
-            # this is for compatibility with the LTS environment
-            unified_memory_level=0,  # unit tests currently broken with UVM
+            model_config=transformer_config,
+            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,
+                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
@@ -382,17 +369,7 @@ def _build_test_env(cls, test_config):
 
         model.eval()
 
-        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,
-        )
+        mamba_inference_state_config = MambaInferenceStateConfig.from_model(model)
 
         # Inference context.
         inference_context = cls._build_inference_context(
@@ -403,7 +380,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 (
@@ -424,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/engines/test_static_engine.py b/tests/unit_tests/inference/engines/test_static_engine.py
index 03b3712e39a..483a21d13bd 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
@@ -200,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/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..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
@@ -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 = (
@@ -82,7 +67,7 @@ def test_inference_pipeline_parallel_small_size(self, materialize_only_last_toke
             .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.config.materialize_only_last_token_logits = (
             materialize_only_last_token_logits
         )
 
@@ -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)
@@ -153,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.config.materialize_only_last_token_logits = (
             materialize_only_last_token_logits
         )
 
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..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
@@ -7,9 +9,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 +76,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.from_config(inference_wrapper_config)
+        inference_context = StaticInferenceContext(max_batch_size=8, max_sequence_length=2560)
 
-        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_inference_config.py b/tests/unit_tests/inference/test_inference_config.py
new file mode 100644
index 00000000000..6d58328dade
--- /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 InferenceConfig
+from megatron.core.transformer.transformer_config import TransformerConfig
+
+
+class TestInferenceConfig:
+    def test_mutual_exclusivity_with_transformer_config(self):
+        """
+        Ensure mutual exclusivity between fields in `InferenceConfig` and
+        `TransformerConfig`.
+        """
+        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 cab464af503..1417926f13b 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 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
@@ -15,6 +16,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
 
 
@@ -50,20 +52,26 @@ 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(
-            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,
-            block_size_tokens=block_size_tokens,
-            metrics_writer=metrics_writer,
-            unified_memory_level=0,  # unit tests currently broken with UVM
+            model_config=TransformerConfig(
+                params_dtype=params_dtype,
+                num_layers=num_layers,
+                kv_channels=kv_channels,
+                num_attention_heads=num_attention_heads,
+            ),
+            inference_config=InferenceConfig(
+                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
+                logging_step_interval=logging_step_interval,
+                metrics_writer=metrics_writer,
+            ),
         )
 
     @pytest.mark.internal
@@ -195,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(metrics_writer=mock_wandb)
+        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)
@@ -210,12 +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
-        )
+        engine = DynamicInferenceEngine(controller=mock_controller, context=dynamic_context)
 
         # Verify log was never called
         mock_wandb.log.assert_not_called()
@@ -225,15 +230,16 @@ 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,
-            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
+            model_config=TransformerConfig(
+                params_dtype=torch.float32, num_layers=2, kv_channels=64, num_attention_heads=8
+            ),
+            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
@@ -257,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(metrics_writer=None)
+        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)
@@ -268,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.context.metrics_writer is None
+        assert engine.logging_step_interval == 10
+        assert engine.metrics_writer is None
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..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
@@ -12,9 +14,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 +84,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.from_config(inference_wrapper_config)
+        inference_context = StaticInferenceContext(max_batch_size=8, max_sequence_length=2560)
 
-        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 96%
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..bdf95c2d9bf 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
@@ -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 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 (
@@ -24,9 +25,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,37 +98,24 @@ 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,
-                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
+                model_config=transformer_config,
+                inference_config=InferenceConfig(
+                    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_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..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
@@ -13,9 +15,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 +91,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.from_config(inference_wrapper_config)
+        inference_context = StaticInferenceContext(max_batch_size=8, max_sequence_length=2560)
 
-        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.py b/tests/unit_tests/models/test_gpt_model.py
index cf3bd40ee4b..87aba9c6ed9 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 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
@@ -392,14 +393,18 @@ 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,
-            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,
+            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,
+            ),
+            inference_config=InferenceConfig(
+                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 ead9125e5ec..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,17 +5,15 @@
 import torch
 import torch.distributed as dist
 
+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 (
     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 +89,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,32 +184,21 @@ 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,
-            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,
+            model_config=base_model.config,
+            inference_config=InferenceConfig(
+                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_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
         )
@@ -273,32 +262,21 @@ 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,
-                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,
+                model_config=based_model.config,
+                inference_config=InferenceConfig(
+                    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_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 9eb7b2dea9a..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,20 +340,17 @@ 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(
-            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.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,
+            model_config=self.model.config,
+            inference_config=InferenceConfig(
+                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/tests/unit_tests/models/test_mamba_moe_model.py b/tests/unit_tests/models/test_mamba_moe_model.py
index 3c7ae93a17c..a5590a0ffad 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,
@@ -197,6 +198,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",
diff --git a/tools/run_dynamic_text_generation_server.py b/tools/run_dynamic_text_generation_server.py
index 615073b8fd0..74f1e69679e 100644
--- a/tools/run_dynamic_text_generation_server.py
+++ b/tools/run_dynamic_text_generation_server.py
@@ -5,25 +5,19 @@
 
 import torch
 
-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
 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.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)
-    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
 
@@ -74,36 +68,12 @@ async def run_text_generation_server(
             args_defaults={'no_load_rng': True, 'no_load_optim': True},
         )
 
-        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 = get_args()
         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,
-            enable_chunked_prefill=not args.disable_chunked_prefill,
-        )
+        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 32d61444530..430bb7ebb9a 100644
--- a/tools/run_inference_performance_test.py
+++ b/tools/run_inference_performance_test.py
@@ -10,33 +10,31 @@
 
 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 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,
 )
-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,
 )
 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_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
 from megatron.training import get_args, get_model, get_tokenizer
 from megatron.training.checkpointing import load_checkpoint
@@ -47,7 +45,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')
 
@@ -60,7 +58,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
 
 
@@ -74,24 +71,13 @@ def get_inference_engine(args: argparse.Namespace, model: MegatronModule) -> Abs
     Returns:
         AbstractBackend: The chosen backend
     """
-    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)
+        tokenizer = get_tokenizer()
+        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()
         )
@@ -100,98 +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,
-        )
-
-
-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
+        return get_dynamic_inference_engine(model=model)
 
 
 def get_random_prompt_tokens(tokenizer, num_input_tokens) -> List[int]:
@@ -232,14 +127,12 @@ 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 = []
     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
@@ -257,8 +150,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={
@@ -298,13 +189,14 @@ 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})
 
     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(
@@ -327,33 +219,27 @@ def main():
                 )
             )
 
-    if args.cuda_graph_impl == "local":
-        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:
         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]
+        records: List[DynamicInferenceRequestRecord] = inference_engine.generate(
+            prompts=prompts, sampling_params=sampling_params
+        )
+        results: List[InferenceRequest] = [record.merge() for record in records]
+
     end_time = time.perf_counter()
     latency = end_time - start_time
 
@@ -378,6 +264,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()
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":
diff --git a/train_rl.py b/train_rl.py
index cfc010b3c04..4b5cec5fcc8 100644
--- a/train_rl.py
+++ b/train_rl.py
@@ -370,6 +370,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
 
@@ -401,4 +403,5 @@ def _model_builder(
         ModelType.encoder_or_decoder,
         forward_step,
         args_defaults={},
+        extra_args_provider=add_inference_args,
     )