Dynamic inference context | UVM only.

examples/inference/gpt/gpt_dynamic_inference.py

@@ -120,6 +120,7 @@ def get_inference_context(
     """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    
@@ -147,11 +148,7 @@ def get_inference_context(
             else None
         ),
         block_size_tokens=args.inference_dynamic_batching_block_size,
-        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,
+        active_buffer_size_gb=args.inference_dynamic_batching_active_buffer_size_gb,
         tensor_model_parallel_size=args.tensor_model_parallel_size,
         materialize_only_last_token_logits=not args.return_log_probs,
         cache_mla_latent=args.multi_latent_attention and args.cache_mla_latents,
@@ -514,14 +511,14 @@ def escape_str(s):
         if engine.capture_stats else
         "--"
     )
-    print(
-        f"{setup_prefix} … "
-        f"capture {capture_str} … "
-        f"mem {peak_alloc_gb:.1f}/{peak_resvd_gb:.1f} GB … "
-        f"total time: {total_time:.3f}s … "
-        f"steps: {engine.step_count:d} … "
-        f"throughput: {throughput:.3f} tok/s"
-    )
+    print(" … ".join((
+        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}",
+    )))
     print("~~~")
 
     # Stop Nsight profiler.

examples/inference/gpt/gpt_dynamic_inference_12b.sh

@@ -23,9 +23,7 @@ export CUDA_DEVICE_MAX_CONNECTIONS=1
 : ${INCOMING_REQUESTS_PER_SEC=100.}
 
 # Dynamic context.
-: ${BUFFER_SIZE_GB=50.}
-: ${BUFFER_OVERFLOW_FACTOR=1.}
-: ${BUFFER_GUARANTEED_FRACTION=0.05}
+: ${ACTIVE_BUFFER_SIZE_GB=50.}
 
 # Cuda graphs.
 : ${CUDA_GRAPH_IMPL=local}
@@ -78,9 +76,7 @@ ARGS=" \
     --inference-rng-tracker \
     \
     --inference-dynamic-batching \
-    --inference-dynamic-batching-buffer-size-gb ${BUFFER_SIZE_GB} \
-    --inference-dynamic-batching-buffer-overflow-factor ${BUFFER_OVERFLOW_FACTOR} \
-    --inference-dynamic-batching-buffer-guaranteed-fraction ${BUFFER_GUARANTEED_FRACTION} \
+    --inference-dynamic-batching-active-buffer-size-gb ${ACTIVE_BUFFER_SIZE_GB} \
     \
     ${EXTRA_ARGS} \
 "

examples/inference/gpt/gpt_dynamic_inference_357m.sh

@@ -24,9 +24,7 @@ export CUDA_DEVICE_MAX_CONNECTIONS=1
 : ${INCOMING_REQUESTS_PER_SEC=100.}
 
 # Dynamic context.
-: ${BUFFER_SIZE_GB=50.}
-: ${BUFFER_OVERFLOW_FACTOR=1.}
-: ${BUFFER_GUARANTEED_FRACTION=0.05}
+: ${ACTIVE_BUFFER_SIZE_GB=50.}
 
 # Cuda graphs.
 : ${CUDA_GRAPH_IMPL=local}
@@ -64,9 +62,7 @@ ARGS=" \
     --inference-rng-tracker \
     \
     --inference-dynamic-batching \
-    --inference-dynamic-batching-buffer-size-gb ${BUFFER_SIZE_GB} \
-    --inference-dynamic-batching-buffer-overflow-factor ${BUFFER_OVERFLOW_FACTOR} \
-    --inference-dynamic-batching-buffer-guaranteed-fraction ${BUFFER_GUARANTEED_FRACTION} \
+    --inference-dynamic-batching-active-buffer-size-gb ${ACTIVE_BUFFER_SIZE_GB} \
     \
     ${EXTRA_ARGS} \
 "

examples/inference/gpt/utils.py

@@ -11,6 +11,7 @@
 
 from megatron.core.inference.inference_request import DynamicInferenceRequest
 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.sampling_params import SamplingParams
@@ -65,7 +66,7 @@ def add_common_inference_args(parser: ArgumentParser) -> ArgumentParser:
         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`, "
+        "additionally limited by the inference context's `max_active_requests`, "
         "`max_tokens`, and KV buffer size.",
     )
     group.add_argument(
@@ -259,10 +260,27 @@ def get_synthetic_requests(
         int(args.incoming_requests_per_sec * args.incoming_requests_duration),
     )
 
+    # 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]
+    )
+    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 ]
+
     # Init requests.
+    assert len(prompt_texts) == len(time_offsets)
     requests = [
-        Request("hi " * random.randint(*args.num_tokens_to_prompt), t, tokenizer, sampling_params)
-        for t in time_offsets
+        Request(t, o, tokenizer, sampling_params=sampling_params)
+        for t, o in zip(prompt_texts, time_offsets)
     ]
 
     return requests
@@ -339,7 +357,7 @@ def build_dynamic_engine_setup_prefix(
 
     Args:
         args (Namespace): Command-line arguments for this run.
-        context (DynamicInferenceContext): Stores limits such as `max_requests`,
+        context (DynamicInferenceContext): Stores limits such as `max_active_requests`,
             `max_tokens`, and `gtd_request_count`.
         requests (List[DynamicInferenceRequest]): List of inference requests.
 
@@ -376,17 +394,10 @@ def build_dynamic_engine_setup_prefix(
     )
 
     # Buffer limits config
-    flw = args.inference_dynamic_batching_buffer_overflow_factor
-    flw_str = "no overflow" if flw is None else f"{flw:.1f}"
     buffer_limits_str = (
-        f"bf {args.inference_dynamic_batching_buffer_size_gb:.0f}, {flw_str} "
-        f"[r {context.max_requests}, t {context.max_tokens}]"
-    )
-
-    # Guaranteed request config
-    guaranteed_fraction_str = (
-        f"gtd {args.inference_dynamic_batching_buffer_guaranteed_fraction:.2f} "
-        f"[r {context.gtd_request_count}]"
+        f"bf: {get_mem_size_str(args.inference_dynamic_batching_active_buffer_size_gb*1024**3)}, "
+        f"{context.block_allocator.active_count} chunks "
+        f"[r {context.max_active_requests}, t {context.max_tokens}]"
     )
 
     parts = [
@@ -396,7 +407,6 @@ def build_dynamic_engine_setup_prefix(
         uvm_str,
         request_str,
         buffer_limits_str,
-        guaranteed_fraction_str,
     ]
 
     return " | ".join(parts)

megatron/core/inference/contexts/dynamic_block_allocator.py

@@ -13,60 +13,86 @@ class BlockAllocator:
     - Initializing a pool of block IDs
     - Allocating blocks from the pool
     - Releasing blocks back to the pool
-    - Managing the guaranteed block count for active requests
 
     Args:
-        block_count_total (int): Total number of blocks available in the buffer.
-        gtd_block_count (int): Number of blocks reserved for guaranteed requests.
+        context (DynamicInferenceContext): Dynamic inference context.
+        active_count (int): Total number of active blocks available in the buffer.
+            The full buffer size is 2*active_count, to accommodate an equal-size
+            space for paused requests that live on the CPU.
     """
 
-    def __init__(self, block_count_total: int, gtd_block_count: int):
-        self.block_count_total = block_count_total
-        self.gtd_block_count = gtd_block_count
+    def __init__(self, context: "DynamicInferenceContext", active_count: int):
 
-        # Reserve last block ID as dummy block for decode-only inference steps
-        self.block_count_avail = self.block_count_total - 1
-        self.dummy_block_idx = self.block_count_total - 1
+        self.context = context
+
+        active_count -= 1  # -1 for dummy_block_idx (see below)
+        active_count = max(1, active_count)  # need at least one block
+        self.total_count = 2 * active_count + 1  # +1 for dummy_block_idx
+        self.total_avail = self.total_count - 1  # -1 for dummy_block_idx
+        self.active_count = active_count
+        self.paused_count = self.total_count - self.active_count - 1  # -1 for dummy_block_idx
+        self.dummy_block_idx = self.total_count - 1
 
         # Initialize block pool as a "stack" data structure
         self.block_bag = torch.arange(
-            self.block_count_total, dtype=torch.int32, device=torch.cuda.current_device()
+            self.total_count, dtype=torch.int32, device=torch.cuda.current_device()
         )
 
-    def is_memory_available(self, num_blocks: int, safe: bool = False) -> bool:
-        """Check if memory blocks are available.
+    def __str__(self):
+        return (
+            f"total avail {self.total_avail} / {self.total_count - 1}"
+            f"; active {self.active_count}"
+        )
 
-        Use 'safe' to avoid all requests being deadlocked. A fraction of the KV cache
-        memory buffer is reserved to guarantee that a minimum number of active
-        requests can run on any given step.
+    def get_active_used(self):
+        """Compute number of active blocks used."""
+        return (
+            self.context.request_kv_block_counts[
+                self.context.paused_request_count : self.context.total_request_count
+            ]
+            .sum()
+            .item()
+        )
+
+    def get_paused_used(self):
+        """Compute number of paused blocks used."""
+        return (
+            self.context.request_kv_block_counts[: self.context.paused_request_count].sum().item()
+        )
+
+    def get_active_avail(self):
+        """Compute number of active blocks available."""
+        return self.active_count - self.get_active_used()
+
+    def get_paused_avail(self):
+        """Compute number of paused blocks available."""
+        return self.paused_count - self.get_paused_used()
+
+    def is_memory_available(self, num_blocks: int) -> bool:
+        """Check if memory blocks are available.
 
         Args:
             num_blocks (int): Number of blocks to check.
-            safe (bool): Include extra space for guaranteeing ability to run
-                requests to completion.
 
         Return:
             (bool) Is memory available?
         """
-        if safe:
-            return self.block_count_avail >= num_blocks + self.gtd_block_count
-        else:
-            return self.block_count_avail >= num_blocks
+        return self.get_active_avail() >= num_blocks
 
-    def allocate_memory_blocks(self, num_blocks: int = 1, safe: bool = False) -> Optional[Tensor]:
+    def allocate_memory_blocks(self, num_blocks: int) -> Optional[Tensor]:
         """Allocate memory blocks if available, else return None.
 
         Args:
             num_blocks (int): Number of blocks to allocate.
-            safe (bool): Include extra space for guaranteeing ability to run
-                requests to completion.
 
         Return:
             (Optional[Tensor]) Allocated block IDs.
         """
-        if self.is_memory_available(num_blocks, safe):
-            self.block_count_avail -= num_blocks
-            return self.block_bag[self.block_count_avail : (self.block_count_avail + num_blocks)]
+        if self.is_memory_available(num_blocks):
+            self.total_avail -= num_blocks
+            block_ids = self.block_bag[self.total_avail : (self.total_avail + num_blocks)]
+            assert num_blocks == block_ids.numel()
+            return block_ids
         else:
             return None
 
@@ -80,13 +106,13 @@ def release_memory_blocks(self, blocks: Tensor) -> None:
             None
         """
         num_blocks = blocks.size(dim=0)
-        self.block_bag[self.block_count_avail : (self.block_count_avail + num_blocks)] = blocks
-        self.block_count_avail += num_blocks
+        self.block_bag[self.total_avail : (self.total_avail + num_blocks)] = blocks
+        self.total_avail += num_blocks
 
     def reset(self) -> None:
         """Reset the allocator to initial state.
 
         This resets the available block count to the entire memory pool
         (except for the dummy block).
         """
-        self.block_count_avail = self.block_count_total - 1
+        self.total_avail = self.total_count - 1