Allow mixed-batch sampling in dynamic inference

examples/inference/gpt/gpt_dynamic_inference.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import hashlib
 import json
@@ -11,7 +11,7 @@
 from collections import defaultdict
 from functools import partial
 from tqdm import tqdm
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import torch
 from tqdm import tqdm
@@ -117,8 +117,11 @@ def get_model() -> MegatronModule:
     return model
 
 
-def get_inference_context(requests: List[Request], sampling_params: SamplingParams, 
-                          calculate_max_sequence_length_from_requests: bool =True):
+def get_inference_context(
+    requests: List[Request],
+    sampling_params: Optional[SamplingParams] = None,
+    calculate_max_sequence_length_from_requests: bool = True
+):
     """The inference context manages the KV cache and other inference state."""
 
     args = get_args()
@@ -199,19 +202,28 @@ def get_inference_controller(
 
 
 def run_inference(
-    requests: List[Request], sampling_params: SamplingParams, engine: DynamicInferenceEngine
+    requests: List[Request],
+    engine: DynamicInferenceEngine,
+    sampling_params: Optional[SamplingParams] = None,
 ) -> List[Dict[str, float]]:
     """Add requests to engine and generate tokens.
 
     Args:
         requests (List[Request]): Requests that are to be added and processed.
-        sampling_params (SamplingParams): Sampling params for the logits.
         engine (DynamicInferenceEngine): Inference engine that manages generating tokens.
+        sampling_params (SamplingParams): Deprecated as of megatron-core 0.16.
 
     Return:
         A dictionary of step times with `prefill` and `decode` keys.
     """
 
+    if sampling_params is not None and torch.distributed.get_rank() == 0:
+        warnings.warn(
+            "The `sampling_params` argument is deprecated. "
+            "Sampling parameters are specified per request.",
+            DeprecationWarning,
+        )
+
     args = get_args()
 
     # Initialize request arrival times.
@@ -244,7 +256,7 @@ def _add_request():
         engine.add_request(
             num_requests_added,
             _request.prompt_text,
-            sampling_params.num_tokens_to_generate,
+            _request.sampling_params,
         )
         _request.time_start = get_curr_time()
         _request.state = "started"
@@ -271,7 +283,7 @@ def _add_request():
 
         # Step inference engine (i.e., generate a token for each active request).
         # Before step, we haven't done the scheduling, so we cannot know the is_decode_only
-        result = engine.step_modern(sampling_params, verbose=True)
+        result = engine.step_modern(verbose=True)
         # 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 
         step_id += 1
@@ -301,7 +313,7 @@ def _add_request():
                 request.output_text = finished_request.generated_text
                 request.state = "finished"
                 request.request_id = finished_request.request_id
-                if sampling_params.return_log_probs:
+                if finished_request.sampling_params.return_log_probs:
                     request.log_probs = (
                         finished_request.prompt_log_probs + finished_request.generated_log_probs
                     )
@@ -349,11 +361,12 @@ def main():
         top_p=args.top_p,
         return_log_probs=args.return_log_probs,
         num_tokens_to_generate=args.num_tokens_to_generate,
+        termination_id=args.termination_id if args.termination_id is not None else tokenizer.eod,
     )
 
     # Requests, context, conroller.
     model = get_model()
-    requests = build_requests(args, tokenizer)
+    requests = build_requests(args, tokenizer, sampling_params)
     context = get_inference_context(requests, sampling_params)
     controller = get_inference_controller(model, context)
 
@@ -371,7 +384,6 @@ def main():
     engine = DynamicInferenceEngine(
         controller,
         context,
-        termination_id=args.termination_id if args.termination_id is not None else tokenizer.eod,
         enable_cuda_graph=args.cuda_graph_impl == "local",
         random_seed=args.seed,
         track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events,
@@ -387,7 +399,7 @@ def main():
     throughputs = []
     for _ in range(args.inference_repeat_n):
         t = get_curr_time()
-        result = run_inference(requests, sampling_params, engine)
+        result = run_inference(requests, engine)
         step_times = result["step_times"]
         add_times = result["add_times"]
         output_times = result["output_times"]
@@ -458,7 +470,7 @@ def escape_str(s):
                         "cuda_graph_request_count_map" : result["cuda_graph_request_count_map"],
                         "step_count" : engine.step_count,
                     }
-                    if sampling_params.return_log_probs:
+                    if req.sampling_params.return_log_probs:
                         response_logprobs = req.log_probs
                         result_dict["logprobs"] = response_logprobs
                     json_results[req.request_id] = result_dict

examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py

@@ -1,3 +1,5 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
 from megatron.core.inference.inference_client import InferenceClient
 from examples.inference.gpt.utils import add_common_inference_args
 import asyncio
@@ -18,13 +20,24 @@
 from megatron.training.arguments import parse_args
 from megatron.core import parallel_state
 
-async def main(engine: DynamicInferenceEngine, requests: List[Request], sampling_params: SamplingParams, port: int):
+async def main(
+    engine: DynamicInferenceEngine,
+    requests: List[Request],
+    port: int,
+    sampling_params: SamplingParams | None = None,
+):
+    if sampling_params is not None:
+        warnings.warn(
+            "The `sampling_params` argument is deprecated. "
+            "Sampling parameters are specified per request.",
+            DeprecationWarning,
+        )
     # once you call engine.start_listening_to_data_parallel_coordinator,
     # the engine will start accepting requests from the data parallel coordinator.
     # and processing them in an asyncio coroutine. 
-    await engine.start_listening_to_data_parallel_coordinator(sampling_params, 
-                                                inference_coordinator_port=port,
-                                                launch_inference_coordinator=True)   
+    await engine.start_listening_to_data_parallel_coordinator( 
+        inference_coordinator_port=port, launch_inference_coordinator=True
+    )
     # if you want to use your own inference coordinator - 
     # 1. set launch_inference_coordinator to False
     # 2. setup a router socket at tcp://MASTER_ADDR:PORT
@@ -50,8 +63,7 @@ async def main(engine: DynamicInferenceEngine, requests: List[Request], sampling
                 # 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
                 # request completion.
-                futures.append(client.add_request(request.prompt_text, 
-                                                        sampling_params))
+                futures.append(client.add_request(request.prompt_text, request.sampling_params))
                 num_requests_added += 1
                 #tbar.update(1)
             if num_requests_added == num_requests_total:
@@ -74,7 +86,7 @@ async def main(engine: DynamicInferenceEngine, requests: List[Request], sampling
                     "generated_tokens": req.generated_tokens,
                     "latency": req.latency, #InferenceClient populates this field in the returned future.
                 }
-                if sampling_params.return_log_probs:
+                if req.sampling_params["return_log_probs"]:
                     result_dict["logprobs"] = req.prompt_log_probs + req.generated_log_probs
                 json_results[req.request_id] = result_dict
             with open(args.output_path, "w") as fp:
@@ -115,13 +127,13 @@ async def main(engine: DynamicInferenceEngine, requests: List[Request], sampling
             top_p=args.top_p,
             return_log_probs=args.return_log_probs,
             num_tokens_to_generate=args.num_tokens_to_generate,
+            termination_id=args.termination_id if args.termination_id is not None else tokenizer.eod,
         )
 
         # Requests, context, conroller.
         model = get_model()
-        requests = build_requests(args, tokenizer) if dist.get_rank() == 0 else None
+        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)
@@ -132,7 +144,6 @@ async def main(engine: DynamicInferenceEngine, requests: List[Request], sampling
         engine = DynamicInferenceEngine(
             controller,
             context,
-            termination_id=tokenizer.eod,
             enable_cuda_graph=args.cuda_graph_impl == "local",
             random_seed=args.seed,
             enable_chunked_prefill=not args.disable_chunked_prefill
@@ -147,6 +158,5 @@ async def main(engine: DynamicInferenceEngine, requests: List[Request], sampling
 
         asyncio.run(main(engine, 
                         requests,
-                        sampling_params,
                         args.inference_coordinator_port))
 

examples/inference/gpt/utils.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import json
 import itertools
@@ -13,6 +13,8 @@
 from megatron.core.inference.contexts import DynamicInferenceContext
 from megatron.core.transformer.module import MegatronModule
 
+from megatron.core.inference.sampling_params import SamplingParams
+
 
 
 def add_common_inference_args(parser: ArgumentParser) -> ArgumentParser:
@@ -130,6 +132,16 @@ def add_common_inference_args(parser: ArgumentParser) -> ArgumentParser:
     return parser
 
 
+def get_default_sampling_params(termination_id: int = None):
+    return SamplingParams(
+        temperature=1.0,
+        top_k=1,
+        top_p=0.0,
+        return_log_probs=False,
+        num_tokens_to_generate=30,
+        termination_id = termination_id,
+    )
+
 def get_curr_time() -> float:
     """Get synchronized time across ranks."""
     curr_time = torch.cuda.LongTensor([time.time_ns()])
@@ -153,7 +165,7 @@ class Request:
         tokenizer (Any): Tokenizer for tokenizing the prompt.
     """
 
-    def __init__(self, prompt_text: str, time_offset: float, tokenizer: Any):
+    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
@@ -163,6 +175,7 @@ def __init__(self, prompt_text: str, time_offset: float, tokenizer: Any):
         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)
 
     def __str__(self) -> str:
         return "state '%s'; toffset %.1e; prompt len %d; output len %d; '%s'" % (
@@ -216,22 +229,26 @@ def arrival(r):
     return time_offsets
 
 
-def get_cli_requests(args: Namespace, tokenizer: Any) -> list[Request]:
+def get_cli_requests(
+        args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None
+) -> list[Request]:
 
     # Get time offsets.
-    time_offsets = get_time_offsets(
+    t_offsets = get_time_offsets(
         args.seed,
         args.incoming_requests_per_step,
         args.incoming_requests_per_sec,
         len(args.prompts),
     )
 
     # Init requests.
-    requests = [Request(p, t, tokenizer) for p,t in zip(args.prompts, time_offsets)]
+    requests = [Request(p, t, tokenizer, sampling_params) for p,t in zip(args.prompts, t_offsets)]
     return requests
 
 
-def get_synthetic_requests(args: Namespace, tokenizer: Any) -> list[Request]:
+def get_synthetic_requests(
+    args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None
+) -> list[Request]:
     """Get example requests."""
 
     # Get time offsets.
@@ -244,14 +261,16 @@ def get_synthetic_requests(args: Namespace, tokenizer: Any) -> list[Request]:
 
     # Init requests.
     requests = [
-        Request("hi " * random.randint(*args.num_tokens_to_prompt), t, tokenizer)
+        Request("hi " * random.randint(*args.num_tokens_to_prompt), t, tokenizer, sampling_params)
         for t in time_offsets
     ]
 
     return requests
 
 
-def get_requests_from_file(args: Namespace, tokenizer: Any) -> list[Request]:
+def get_requests_from_file(
+    args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None
+) -> list[Request]:
     """Get requests from a file."""
     if not args.prompt_file:
         raise ValueError("Prompt file is required to read requests from a file.")
@@ -275,23 +294,25 @@ def get_requests_from_file(args: Namespace, tokenizer: Any) -> list[Request]:
 
     # Init requests.
     requests = [
-        Request(p, t, tokenizer)
+        Request(p, t, tokenizer, sampling_params)
         for p, t in tqdm(zip(prompts, time_offsets), "init requests", total=len(prompts))
     ]
 
     return requests
 
 
-def build_requests(args: Namespace, tokenizer: Any) -> list[Request]:
+def build_requests(
+    args: Namespace, tokenizer: Any, sampling_params: Optional[SamplingParams] = None
+) -> list[Request]:
     # Check if we have any prompts (from command line or JSONL)
     if args.prompts:
         if args.prompt_file:
             raise ValueError("Cannot use both --prompts and --prompt-file")
-        return get_cli_requests(args, tokenizer)
+        return get_cli_requests(args, tokenizer, sampling_params)
     elif args.prompt_file:
-        return get_requests_from_file(args, tokenizer)
+        return get_requests_from_file(args, tokenizer, sampling_params)
     else:
-        return get_synthetic_requests(args, tokenizer)
+        return get_synthetic_requests(args, tokenizer, sampling_params)
 
 
 def get_model_size_str(model):