diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index 251aa100cba..fc2d17c1ebd 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -132,6 +132,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 @@ -159,11 +160,8 @@ 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, + max_tokens=args.inference_dynamic_batching_max_tokens, tensor_model_parallel_size=args.tensor_model_parallel_size, materialize_only_last_token_logits=not args.return_log_probs, layer_type_list=layer_type_list, @@ -546,14 +544,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. diff --git a/examples/inference/gpt/gpt_dynamic_inference_12b.sh b/examples/inference/gpt/gpt_dynamic_inference_12b.sh index a16fe5176d5..341bd1a8ac0 100644 --- a/examples/inference/gpt/gpt_dynamic_inference_12b.sh +++ b/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} \ " diff --git a/examples/inference/gpt/gpt_dynamic_inference_357m.sh b/examples/inference/gpt/gpt_dynamic_inference_357m.sh index c095371714f..c0e24e1a77c 100644 --- a/examples/inference/gpt/gpt_dynamic_inference_357m.sh +++ b/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} \ " diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py index 0ea1f5a3df0..e80c9b546e1 100644 --- a/examples/inference/gpt/utils.py +++ b/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( @@ -262,10 +263,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 @@ -342,7 +360,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. @@ -379,17 +397,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 = [ @@ -399,7 +410,6 @@ def build_dynamic_engine_setup_prefix( uvm_str, request_str, buffer_limits_str, - guaranteed_fraction_str, ] return " | ".join(parts) diff --git a/megatron/core/inference/contexts/dynamic_block_allocator.py b/megatron/core/inference/contexts/dynamic_block_allocator.py index 4baa3f5212c..8b4cc70cfb8 100644 --- a/megatron/core/inference/contexts/dynamic_block_allocator.py +++ b/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,8 +106,8 @@ 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. @@ -89,4 +115,4 @@ def reset(self) -> None: 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 diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 000b58200f8..22a4b1b1423 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1,5 +1,6 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import logging import math import warnings from contextlib import nullcontext @@ -113,7 +114,7 @@ class BlockOverflowError(ContextOverflowError): class ActiveRequestCountOverflowError(ContextOverflowError): '''Used when `initialize_attention_state()` is called with - `num_warmup_requests > max_requests.''' + `num_warmup_requests > max_active_requests.''' def __init__(self, max_request_count, active_request_count): assert active_request_count > max_request_count @@ -175,6 +176,15 @@ class WarmupEngineMode(Enum): NON_DECODE = "non_decode" +def get_mem_size_str(n_bytes: int) -> str: + """Convert number of bytes to human-readable string.""" + for exp, suffix in ((4, "TB"), (3, "GB"), (2, "MB"), (3, "KB"), (0, "bytes")): + nquery = int(1024**exp) + if round(n_bytes / nquery) >= 1: + return "%.3g %s" % (n_bytes / nquery, suffix) + raise Exception(f"something went wrong, n_bytes={n_bytes}.") + + # pylint: disable=line-too-long class DynamicInferenceContext(BaseInferenceContext): """Inference context that is passed to the main model in order @@ -185,9 +195,9 @@ class DynamicInferenceContext(BaseInferenceContext): arbitrary sequence length may be added, paused, or removed from the context at any step. The only constraint is the maximum number of requests or tokens that the context is defined to support. For the block-level KV cache, a memory - buffer is allocated up front (size `buffer_size_gb`), that is divided into - blocks and dynamically assigned to requests. At any given step, any unassigned - blocks equate to unused space. + buffer is allocated up front (size `2 * active_buffer_size_gb`), that is + divided into blocks and dynamically assigned to requests. At any given step, + any unassigned blocks equate to unused space. Additionally, a fraction of the memory buffer (`gtd_request_fraction`, i.e., the 'guaranteed' request fraction) is reserved for guaranteeing that a @@ -212,26 +222,19 @@ class DynamicInferenceContext(BaseInferenceContext): 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): Total buffer size (GB), shared by main and - fallback contexts. + active_buffer_size_gb (float): Buffer size reserved for active requests + that live on the GPU. The total buffer size (stored in unified memory) + is 2x this value, with the the other half of the buffer reserved for + paused requests that live on the CPU. + 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. - buffer_guaranteed_fraction (float): Fraction of the memory buffer that is - reserved to guarantee that one or more active requests are able to - run to completion. Without reserving this memory, paused requests are - able to fill the memory buffer and block execution of any requests. - buffer_overflow_factor (Optional[float]): Scaling factor over the buffer - size for auto computing `max_requests` and `max_tokens`. This scaling - factor is used for fitting more requests and tokens in the memory - buffer than it can safely hold, which in turn increases throughput. - max_requests_override (Optional[int]): If set, overrides value computed - from `buffer_overflow_factor`. - max_tokens_override (Optional[int]): If set, overrides value computed - from `buffer_overflow_factor`. 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. + where the cuda graph batch sizes range from 1 to `max_active_requests` + (as computed below). Due to rounding, the actual number of cuda graphs + 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. @@ -254,6 +257,10 @@ class DynamicInferenceContext(BaseInferenceContext): metrics_writer (Optional['WandbModule']): Wandb module for writing metrics. """ + DEFAULT_MAX_TOKENS = 16384 + TOKEN_ROUNDER = 64 + REQUEST_ROUNDER = 4 + def __init__( self, *, @@ -262,12 +269,9 @@ def __init__( kv_channels: int, num_attention_heads: int, max_sequence_length: int, - buffer_size_gb: float, - buffer_guaranteed_fraction: float, + active_buffer_size_gb: float, + max_tokens: int = DEFAULT_MAX_TOKENS, block_size_tokens: int = 256, - buffer_overflow_factor: Optional[float] = None, - max_requests_override: Optional[int] = None, - max_tokens_override: Optional[int] = None, tensor_model_parallel_size: Optional[int] = None, cache_mla_latent: bool = False, kv_lora_rank: Optional[int] = None, @@ -279,7 +283,7 @@ def __init__( mamba_ssm_states_shape: Optional[Tuple[int]] = None, use_cuda_graphs_for_non_decode_steps: bool = True, use_flashinfer_fused_rope: bool = False, - unified_memory_level: Optional[int] = 0, + unified_memory_level: Optional[int] = 1, metrics_writer: Optional['WandbModule'] = None, ): super().__init__(materialize_only_last_token_logits=materialize_only_last_token_logits) @@ -356,11 +360,6 @@ def __init__( ) assert self.block_size_bytes > 0 - # Adjust buffer to be a multiple of block size. - buffer_size_bytes = int(buffer_size_gb * 1024**3) - buffer_size_bytes_rem = buffer_size_bytes % self.block_size_bytes - buffer_size_bytes = buffer_size_bytes - buffer_size_bytes_rem - mamba_states_memory_per_request = 0 if self.is_hybrid_model: mamba_states_memory_per_request += math.prod(mamba_conv_states_shape) @@ -368,40 +367,25 @@ def __init__( mamba_states_memory_per_request *= self.num_mamba_layers mamba_states_memory_per_request *= dtype_size_bytes - # Compute max_requets, max_tokens from buffer size, overflow factor, and Mamba state size. - def bytes_to_max_requests_and_tokens(n_bytes): - bytes_per_token = self.block_size_bytes / self.block_size_tokens - cost_per_request_bytes = ( - mamba_states_memory_per_request + max_sequence_length * bytes_per_token - ) - # TODO(ksanthanam): Leave room for an extra request in the event of padding - # for non-decode CUDA graphs - n_requests = n_bytes / cost_per_request_bytes - n_tokens = n_requests * max_sequence_length - n_requests = self.round_up_requests(int(n_requests), tp_size=tp_size) - n_tokens = self.round_up_tokens(int(n_tokens), tp_size=tp_size) - return n_requests, n_tokens - - self.max_requests, self.max_tokens = bytes_to_max_requests_and_tokens(buffer_size_bytes) - if buffer_overflow_factor is not None: - self.max_requests = self.round_up_requests( - int(self.max_requests * buffer_overflow_factor), tp_size=tp_size - ) - self.max_tokens = self.round_up_tokens( - int(self.max_tokens * buffer_overflow_factor / 50.0), tp_size=tp_size - ) - - if max_requests_override is not None: - self.max_requests = ( - max_requests_override - if max_requests_override < self.REQUEST_ROUNDER - else self.round_up_requests(max_requests_override, tp_size=tp_size) - ) - - if max_tokens_override is not None: - self.max_tokens = self.round_up_tokens(max_tokens_override, tp_size=tp_size) - - self.max_requests = min(self.max_requests, self.max_tokens) # e.g., decode only. + # Initialize block allocator. + active_buffer_size_bytes = int(active_buffer_size_gb * 1024**3) + active_block_count_total = active_buffer_size_bytes // ( + self.block_size_bytes + mamba_states_memory_per_request + ) + self.block_allocator = BlockAllocator(context=self, active_count=active_block_count_total) + del active_block_count_total # use self.block_allocator.active_count + active_buffer_size_bytes = self.block_allocator.active_count * self.block_size_bytes + + # Set max_total_requests, max_active_requests, max_tokens. + self.max_total_requests = self.block_allocator.total_count - 1 # -1 for dummy block + self.max_active_requests = self.block_allocator.active_count + self.max_tokens = max_tokens or self.DEFAULT_MAX_TOKENS + + assert self.max_tokens >= self.max_active_requests, ( + f"max_tokens ({self.max_tokens}) must be >= " + f"max_active_requests ({self.max_active_requests}), " + "to have consistency between cuda graph sizes and the block table size." + ) # Initialize context state. self.params_dtype = params_dtype @@ -429,7 +413,7 @@ def bytes_to_max_requests_and_tokens(n_bytes): # Per-request state. self.request_ids = torch.full( - (self.max_requests,), -1, dtype=torch.int32, device=torch.cuda.current_device() + (self.max_total_requests,), -1, dtype=torch.int32, device=torch.cuda.current_device() ) # request_query_lengths is the input prompt tokens length during prefill phase (1st step) and then 1 for the decode phase (i.e During generation) self.request_query_lengths = torch.empty_like(self.request_ids) @@ -455,24 +439,15 @@ def bytes_to_max_requests_and_tokens(n_bytes): self.token_to_position_in_request = torch.empty_like(self.token_to_input_ids) self.token_to_local_position_within_kv_block = torch.empty_like(self.token_to_input_ids) - # Calculate the total number of chunks available in the buffer - total_mamba_states_memory = mamba_states_memory_per_request * self.max_requests - block_count_total = ( - max(0, buffer_size_bytes - total_mamba_states_memory) // self.block_size_bytes - ) - # Memory buffer. - ctx_manager = ( - torch.cuda.use_mem_pool(self.unified_memory_mempool) - if self.unified_memory_level > 0 - else nullcontext() - ) - with ctx_manager: + def allocate_memory_buffer(): + """Allocate the memory buffer. This function is called below within + `with ctx_manager:`.""" if cache_mla_latent: self.memory_buffer = torch.full( ( self.num_attention_layers, - block_count_total, + self.block_allocator.total_count, self.block_size_tokens, kv_reduced_dim, ), @@ -485,7 +460,7 @@ def bytes_to_max_requests_and_tokens(n_bytes): ( 2, # key and value self.num_attention_layers, - block_count_total, + self.block_allocator.total_count, self.block_size_tokens, num_attention_heads_per_partition, hidden_size_per_attention_head, @@ -498,7 +473,7 @@ def bytes_to_max_requests_and_tokens(n_bytes): # Block ids. self.max_kv_block_count = math.ceil(self.max_sequence_length / self.block_size_tokens) self.request_to_kv_block_ids = torch.full( - (self.max_requests, self.max_kv_block_count), + (self.max_total_requests, self.max_kv_block_count), -1, dtype=torch.int, device=torch.cuda.current_device(), @@ -509,11 +484,11 @@ def bytes_to_max_requests_and_tokens(n_bytes): if num_cuda_graphs is not None: # Ensure valid num_cuda_graphs. - num_cuda_graphs = min(max(num_cuda_graphs, 1), self.max_requests) + num_cuda_graphs = min(max(num_cuda_graphs, 1), self.max_active_requests) # Cuda graph step size. cuda_graph_rounder = 8 - self.cuda_graph_step_size = self.max_requests / num_cuda_graphs + self.cuda_graph_step_size = self.max_active_requests / num_cuda_graphs self.cuda_graph_step_size = ( math.ceil(self.cuda_graph_step_size / cuda_graph_rounder) * cuda_graph_rounder ) @@ -522,13 +497,17 @@ def bytes_to_max_requests_and_tokens(n_bytes): # Cuda graph token counts. if num_cuda_graphs == 1: - self.cuda_graph_token_counts = [self.max_requests] + self.cuda_graph_token_counts = [self.max_active_requests] else: self.cuda_graph_token_counts = list( - range(self.cuda_graph_step_size, self.max_requests, self.cuda_graph_step_size) + range( + self.cuda_graph_step_size, + self.max_active_requests, + self.cuda_graph_step_size, + ) ) - if self.cuda_graph_token_counts[-1] != self.max_requests: - self.cuda_graph_token_counts.append(self.max_requests) + if self.cuda_graph_token_counts[-1] != self.max_active_requests: + self.cuda_graph_token_counts.append(self.max_active_requests) self.cuda_graph_token_counts.reverse() # Set used for validating active cuda graph token count. @@ -550,66 +529,51 @@ def bytes_to_max_requests_and_tokens(n_bytes): self.active_attn_metadata = None self.graph_attn_metadata["mha_metadata"] = GraphedMHAMetadata( - block_count_total=block_count_total, + block_count_total=self.block_allocator.total_count, max_kv_block_count=self.max_kv_block_count, - max_requests=self.max_requests, + max_requests=self.max_total_requests, block_size_tokens=self.block_size_tokens, max_seqlen=self.max_sequence_length, ) self.non_graph_attn_metadata["mha_metadata"] = NonGraphedMHAMetadata( - block_count_total=block_count_total, + block_count_total=self.block_allocator.total_count, max_kv_block_count=self.max_kv_block_count, - max_requests=self.max_requests, + max_requests=self.max_total_requests, block_size_tokens=self.block_size_tokens, max_seqlen=self.max_sequence_length, ) - # Guaranteed active requests. - # * See details in the class docstring above. `gtd_request_fraction` is - # the fraction of blocks in the memory buffer that are reserved for - # guaranteeing that some number of active requests can always proceed - # with their generations. The number of blocks defined by - # `buffer_guaranteed_fraction * block_count_total` is converted to a - # number of requests that this reserved space can safely handle - # (`gtd_request_count`). - # * Note: computing the size of this guaranteed space from blocks rather - # than bytes is safer due to the non-linear impacts of a large - # `block_size_tokens` or `max_kv_block_count`. When computing from - # blocks, this space will always be less than `block_count_total`. When - # computing from bytes, this space can unexpectedly be much larger than - # `block_count_total`, resulting in stalled generations. - gtd_block_count = int(buffer_guaranteed_fraction * block_count_total) - gtd_block_count = min(gtd_block_count, block_count_total) - self.gtd_request_count = max(1, gtd_block_count // self.max_kv_block_count) - self.gtd_block_count = self.gtd_request_count * self.max_kv_block_count - - # Initialize allocator for KV memory blocks - self.block_allocator = BlockAllocator( - block_count_total=block_count_total, gtd_block_count=self.gtd_block_count - ) - # Optional state tensors for hybrid models - if self.is_hybrid_model: - self.mamba_metadata = MambaMetadata(max_requests=self.max_requests) - - with ctx_manager: + def allocate_mamba_states(): + """Allocate Mamba states. This function is called below within + `with ctx_manager:`.""" + if self.is_hybrid_model: + self.mamba_metadata = MambaMetadata(max_requests=self.max_total_requests) self.mamba_conv_states = torch.zeros( - (self.num_mamba_layers, self.max_requests) + mamba_conv_states_shape, + (self.num_mamba_layers, self.max_total_requests) + mamba_conv_states_shape, dtype=self.params_dtype, device=torch.cuda.current_device(), ) self.mamba_ssm_states = torch.zeros( - (self.num_mamba_layers, self.max_requests) + mamba_ssm_states_shape, + (self.num_mamba_layers, self.max_total_requests) + mamba_ssm_states_shape, dtype=self.params_dtype, device=torch.cuda.current_device(), ) - else: - self.mamba_metadata = None + else: + self.mamba_metadata = None - # Store the dummy block idx reference for convenience - self.dummy_block_idx = self.block_allocator.dummy_block_idx + # Allocate `ctx_manager`-managed buffers. (For currently unknown reasons, + # `ctx_manager` can only be used once.) + ctx_manager = ( + torch.cuda.use_mem_pool(self.unified_memory_mempool) + if self.unified_memory_level > 0 + else nullcontext() + ) + with ctx_manager: + allocate_memory_buffer() + allocate_mamba_states() # Deal with chunked prefill self.chunked_prefill_request_id = -1 @@ -618,14 +582,18 @@ def bytes_to_max_requests_and_tokens(n_bytes): self.reset_attention_state() self.reset_mamba_state() + # FlashInfer. if 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 - TOKEN_ROUNDER = 64 - REQUEST_ROUNDER = 4 + # Print info. + logging.info( + "DynamicInferenceContext: allocated context with active buffer size %s (%d blocks)." + % (get_mem_size_str(active_buffer_size_bytes), self.block_allocator.active_count) + ) @classmethod def round_up_tokens(cls, value, tp_size=None): @@ -654,7 +622,7 @@ def from_config( inference_config: InferenceWrapperConfig, model, max_batch_size: int, - buffer_size_gb: float = 40, + active_buffer_size_gb: float = 40, num_cuda_graphs: int = None, ): """ @@ -662,7 +630,6 @@ def from_config( """ # TODO: Add other necessary configs from inference_config - buffer_guaranteed_fraction = 0.1 model_config = model.config max_sequence_length = ( inference_config.inference_max_seq_length or model_config.max_sequence_length @@ -674,10 +641,8 @@ def from_config( kv_channels=model_config.kv_channels, num_attention_heads=model_config.num_query_groups, max_sequence_length=inference_config.inference_max_seq_length, - buffer_size_gb=buffer_size_gb, - buffer_guaranteed_fraction=buffer_guaranteed_fraction, + active_buffer_size_gb=active_buffer_size_gb, materialize_only_last_token_logits=False, - max_requests_override=max_batch_size, num_cuda_graphs=num_cuda_graphs, use_flashinfer_fused_rope=None, ) @@ -988,7 +953,7 @@ def initialize_attention_state( Args: num_warmup_tokens (Optional[int]): Number of tokens to use for warming up cuda graphs. Must be less than or equal to - `max_requests`. + `max_active_requests`. warmup_engine_mode (WarmupEngineMode): Denote whether to setup for a decode or a non-decode cuda-graph warmup. num_warmup_requests (Optional[int]): [DEPRECATED] Use num_warmup_tokens instead. @@ -1008,8 +973,8 @@ def initialize_attention_state( # warmup both decode and non-decode engine steps if num_warmup_tokens is not None: - if num_warmup_tokens > self.max_requests: - raise ActiveRequestCountOverflowError(self.max_requests, num_warmup_tokens) + if num_warmup_tokens > self.max_active_requests: + raise ActiveRequestCountOverflowError(self.max_active_requests, num_warmup_tokens) if warmup_engine_mode == WarmupEngineMode.NON_DECODE: assert self.non_decode_cuda_graphs, "Set non-decode cuda graphs to True" @@ -1028,7 +993,9 @@ def initialize_attention_state( math.ceil(active_token_count / self.cuda_graph_step_size) * self.cuda_graph_step_size ) - self.padded_active_token_count = min(self.padded_active_token_count, self.max_requests) + self.padded_active_token_count = min( + self.padded_active_token_count, self.max_active_requests + ) assert ( self.padded_active_token_count in self.cuda_graph_token_counts_set ), f"padded_active_token_count: {self.padded_active_token_count} not in cuda_graph_token_counts_set: {self.cuda_graph_token_counts_set}" @@ -1038,7 +1005,7 @@ def initialize_attention_state( if self.is_decode_only(): # For decode-only, the padded active token count cannot exceed max-requests. self.padded_active_token_count = min( - self.padded_active_token_count, self.max_requests + self.padded_active_token_count, self.max_active_requests ) # How are we calculating the padded active request count? @@ -1056,7 +1023,7 @@ def initialize_attention_state( # Update token position indexes. self.token_to_block_idx[self.active_token_count : self.padded_active_token_count] = ( - self.dummy_block_idx + self.block_allocator.dummy_block_idx ) self.token_to_local_position_within_kv_block[ self.active_token_count : self.padded_active_token_count @@ -1198,20 +1165,20 @@ def last_token_logits(self, logits: Tensor) -> Tensor: return last_token_logits - def check_availability( - self, req: DynamicInferenceRequest, safe: bool = False - ) -> (bool, bool, bool): + def check_availability(self, req: DynamicInferenceRequest) -> (bool, bool, bool): """ Check if the request can be added to the context. """ - request_can_be_added = self.total_request_count < self.max_requests + request_can_be_added = ( + self.total_request_count - self.paused_request_count < self.max_active_requests + ) request_tokens_can_be_added = ( self.active_token_count + req.remaining_prompt_length <= self.max_tokens ) blocks = math.ceil( (req.remaining_prompt_length + req.finished_chunk_token_count) / self.block_size_tokens ) - math.ceil(req.finished_chunk_token_count / self.block_size_tokens) - kv_cache_available = self.block_allocator.is_memory_available(blocks, safe=safe) + kv_cache_available = self.block_allocator.is_memory_available(blocks) return request_can_be_added, request_tokens_can_be_added, kv_cache_available def add_request(self, req: DynamicInferenceRequest, chunk_length: Optional[int] = None) -> None: @@ -1251,9 +1218,7 @@ def add_request(self, req: DynamicInferenceRequest, chunk_length: Optional[int] num_blocks_needed = overall_required_blocks - already_allocated_blocks if num_blocks_needed > 0: - new_block_ids = self.block_allocator.allocate_memory_blocks( - num_blocks_needed, safe=not is_chunked_prefill - ) + new_block_ids = self.block_allocator.allocate_memory_blocks(num_blocks_needed) if new_block_ids is None or len(new_block_ids) != num_blocks_needed: raise BlockOverflowError(req.request_id) @@ -1271,7 +1236,7 @@ def add_request(self, req: DynamicInferenceRequest, chunk_length: Optional[int] else: current_id = self.total_request_count - if current_id >= self.max_requests: + if current_id >= self.max_active_requests: raise RequestOverflowError(req.request_id) if self.active_token_count + chunk_length > self.max_tokens: @@ -1389,7 +1354,7 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T between these request groups. - 0:paused_request_count -> paused requests - paused_request_count:total_request_count -> active requests - - total_request_count:max_requests -> completed requests are moved here. + - total_request_count:max_active_requests -> completed requests are moved here. The reason for maintaining contiguous tensors rather than multiple smaller (e.g., per-group or per-request) tensors is for both 1) speed (avoid unnecessary tensor allocations), and 2) compatibility with the @@ -1432,6 +1397,9 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T # Reset attention state. self.reset_attention_state() + # Update total_request_count. + self.total_request_count = active_request_count + self.paused_request_count + # 2. If no paused requests are present and no active requests we release memory and reset. if active_request_count + self.paused_request_count == 0: if finished_request_count > 0: @@ -1531,6 +1499,11 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T (active_requests_requiring_new_block == 1).sum().item() ) + if active_requests_requiring_new_block_count > 0: + newly_paused_request_ids = self.request_ids[ + torch.nonzero(active_requests_requiring_new_block) + self.paused_request_count + ] + # Swap unfinished active requests on the left side with paused requests on the right side # NOTE : We add paused request count because we concatenate # paused tokens to the left at the beginning of update requests @@ -1563,7 +1536,6 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T self._move_book_keeping_tensors( src_idxs=src_idxs, dst_idxs=dst_idxs, next_tokens=next_tokens ) - newly_paused_request_ids = self.request_ids[dst_idxs] self.paused_request_count += active_requests_requiring_new_block_count active_request_count -= active_requests_requiring_new_block_count @@ -1572,15 +1544,16 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T # We determine how many requests we can resume and resume them # Assign released blocks to paused requests. # todo: @shanmugamr, un-pause requests using FIFO, rather than LIFO. - num_non_gtd_blocks = max(0, self.block_allocator.block_count_avail - self.gtd_block_count) - if num_non_gtd_blocks: - # if we have non-gtd blocks, use them. Do not dip into the gtd-block pool - resume_request_count = min(num_non_gtd_blocks, self.paused_request_count) - else: - # only dip into the gtd-block pool if we have run out of non-gtd-blocks and the active - # request count has fallen below a certain threshold. + resume_request_count = 0 + if self.paused_request_count > 0: + active_block_count_avail = self.block_allocator.get_active_avail() + paused_block_counts = self.request_kv_block_counts[: self.paused_request_count] + paused_block_counts = paused_block_counts.flip(dims=[0]) + paused_block_counts += 1 # +1 for newly added block + paused_block_counts_cumsum = paused_block_counts.cumsum(dim=0) resume_request_count = min( - max(self.gtd_request_count - active_request_count, 0), self.paused_request_count + torch.nonzero(paused_block_counts_cumsum <= active_block_count_avail).numel(), + self.block_allocator.total_avail, ) self.paused_request_count -= resume_request_count @@ -1640,6 +1613,7 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T == 0 ), "The request_last_kv_block_offset should be 0 for the requests that just got resumed this step. " + assert resume_request_count <= self.block_allocator.total_avail block_ids = self.block_allocator.allocate_memory_blocks(resume_request_count) row_idx = torch.arange( self.paused_request_count, @@ -1761,11 +1735,11 @@ def get_kvcache_utilization_stats(self) -> dict: } """ # Total usable blocks exclude the reserved dummy block. - total_blocks = max(self.block_allocator.block_count_total - 1, 1) - block_count_avail = int(self.block_allocator.block_count_avail) + total_blocks = max(self.block_allocator.total_count - 1, 1) + block_count_avail = int(self.block_allocator.total_avail) # Overall allocated blocks in the buffer right now. - allocated_blocks = (self.block_allocator.block_count_total - 1) - block_count_avail + allocated_blocks = (self.block_allocator.total_count - 1) - block_count_avail allocated_blocks = int(max(0, allocated_blocks)) # Active unique blocks referenced by current active requests only. @@ -1787,7 +1761,6 @@ def get_kvcache_utilization_stats(self) -> dict: active_utilization = float(active_unique_blocks) / float(total_blocks) # Diagnostic helpers - num_non_gtd_blocks = max(0, block_count_avail - int(self.gtd_block_count)) total_request_count = int(self.total_request_count) return { 'total_blocks': int(total_blocks), @@ -1797,10 +1770,9 @@ def get_kvcache_utilization_stats(self) -> dict: 'active_utilization': active_utilization, 'active_request_count': int(self.get_active_request_count()), 'paused_request_count': int(self.paused_request_count), - 'gtd_block_count': int(self.gtd_block_count), 'block_count_avail': int(block_count_avail), - 'num_non_gtd_blocks': int(num_non_gtd_blocks), 'active_token_count': int(self.active_token_count), 'total_request_count': int(total_request_count), - 'max_requests': int(self.max_requests), + 'max_total_requests': int(self.max_total_requests), + 'max_active_requests': int(self.max_active_requests), } diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 4bff4f85fa8..d18f2d2233a 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -642,7 +642,7 @@ def schedule_non_chunked_prefill(self): while self.waiting_request_ids: req = self.requests[self.waiting_request_ids[0]] request_can_be_added, request_tokens_can_be_added, kv_cache_available = ( - self.context.check_availability(req, safe=True) + self.context.check_availability(req) ) if request_can_be_added and request_tokens_can_be_added and kv_cache_available: self.context.add_request(req) @@ -716,9 +716,7 @@ def schedule_chunked_prefill(self): self.context.active_token_count + remaining_len <= self.context.max_tokens ) token_partially_can_be_added = self.context.active_token_count < self.context.max_tokens - request_can_be_added, _, kv_cache_available = self.context.check_availability( - req, safe=not is_continuing_chunked_prefill - ) + request_can_be_added, _, kv_cache_available = self.context.check_availability(req) request_can_be_added = is_continuing_chunked_prefill or request_can_be_added if request_can_be_added and kv_cache_available: @@ -865,7 +863,8 @@ async def async_step( step_type = "decode" if is_decode_only else "non-decode" output_str = ( "* step %d | %s ... time: %.3f%s ... " - "reqs: %d [ gtd %d, active %d, paused %d, finished %d ] ... " + "reqs: a %d/%d, p %d/%d, w %d, f %d ... " + "blocks: a %d/%d, p %d/%d ... " "mem: tensors %d, alloc %.1f gb, res %.1f gb." % ( self.step_count, @@ -883,11 +882,16 @@ async def async_step( ), ) ), - prev_total_request_count, - context.gtd_request_count, prev_total_request_count - prev_paused_request_count, + context.block_allocator.active_count, prev_paused_request_count, + context.block_allocator.paused_count, + len(self.waiting_request_ids), self.finished_request_count, + context.block_allocator.get_active_used(), + context.block_allocator.active_count, + context.block_allocator.get_paused_used(), + context.block_allocator.paused_count, mem["allocation.all.current"], mem["allocated_bytes.all.current"] / (1024**3), mem["reserved_bytes.all.current"] / (1024**3), diff --git a/megatron/core/inference/engines/static_engine.py b/megatron/core/inference/engines/static_engine.py index dc86eb775f9..d9f9e9e051d 100644 --- a/megatron/core/inference/engines/static_engine.py +++ b/megatron/core/inference/engines/static_engine.py @@ -99,7 +99,7 @@ def __init__( 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, + active_buffer_size_gb=buffer_size_gb, num_cuda_graphs=1, ) self.controller.inference_wrapped_model.inference_context = dynamic_context diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 0b14140529a..cb2691c9d49 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1136,9 +1136,8 @@ def validate_args(args, defaults={}): ), "Pipeline-parallel microbatched inference is incompatible with CUDA graphs" if args.inference_dynamic_batching: - assert args.inference_dynamic_batching_buffer_size_gb is not None + assert args.inference_dynamic_batching_active_buffer_size_gb is not None assert args.inference_dynamic_batching_block_size % 256 == 0, "block size should be a multiple of 256" - assert args.inference_dynamic_batching_buffer_guaranteed_fraction is not None # MoE upcycling check if args.moe_use_upcycling: @@ -1435,36 +1434,19 @@ def _add_inference_args(parser): group.add_argument('--inference-dynamic-batching', action='store_true', default=False, help='Enable dynamic batching mode.') - group.add_argument('--inference-dynamic-batching-buffer-size-gb', + group.add_argument('--inference-dynamic-batching-active-buffer-size-gb', type=float, default=40., - help='Total buffer size (GB) allocated for the block-level KV ' - 'memory.') + help='Buffer size (GB) allocated for the active (on-GPU) ' + 'portion of the chunked KV memory. The total buffer size ' + 'is 2x this value, which includes the same-size on-CPU ' + 'paused buffer.') group.add_argument('--inference-dynamic-batching-block-size', type=int, default=256, help='KV cache block size. ' 'It should be a multiple of 256') - group.add_argument('--inference-dynamic-batching-buffer-guaranteed-fraction', - type=float, default=0.2, - help='Space is reserved within the inference context ' - 'memory buffer to guarantee that a minimum number of ' - 'active requests will always be able to run to ' - 'completion. This is to avoid the context being deadlocked ' - 'by paused requests.') - group.add_argument('--inference-dynamic-batching-buffer-overflow-factor', - type=float, default=None, - help='Scaling factor over the memory buffer size for auto ' - 'computing `max_requests` and `max_tokens`. This scaling ' - 'factor is used for fitting more requests and tokens in ' - 'the memory buffer than it can safely hold, which in turn ' - 'increases throughput.') - group.add_argument('--inference-dynamic-batching-max-requests-override', + group.add_argument('--inference-dynamic-batching-max-tokens', type=int, default=None, - help='If set, this overrides the max requests as computed ' - 'from `--inference-dynamic-batching-buffer-overflow-factor`.') - group.add_argument('--inference-dynamic-batching-max-tokens-override', - type=int, default=None, - help='If set, this overrides the max tokens as computed ' - 'from `--inference-dynamic-batching-buffer-overflow-factor`.') + help='Override the inference context\'s default `max_tokens`.') group.add_argument('--inference-dynamic-batching-num-cuda-graphs', type=int, default=16, help='Maximum number of cuda graphs to capture, where the ' @@ -1481,7 +1463,7 @@ def _add_inference_args(parser): action='store_true', default=False, help='Only use cuda graphs for decode-only steps, not prefill and mixed steps.') group.add_argument('--inference-dynamic-batching-unified-memory-level', - type=int, default=0, choices=[0, 1], + type=int, default=1, choices=[0, 1], help='Set unified memory usage within the dynamic ' 'inference context. The levels are: 0) no unified memory, ' '1) allocate `memory_buffer` in unified memory. ' diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 0674cdfcabd..00c0e2a9dab 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -41,12 +41,9 @@ def _get_dynamic_context( kv_channels, num_attention_heads, max_sequence_length, - buffer_size_gb, + active_buffer_size_gb, block_size_tokens, - buffer_guaranteed_fraction, - buffer_overflow_factor, - max_requests_override, - max_tokens_override, + max_tokens, is_hybrid_model=False, layer_type_list=None, rounder=64, @@ -64,12 +61,9 @@ def _get_dynamic_context( max_sequence_length=max_sequence_length, num_cuda_graphs=None, use_cuda_graphs_for_non_decode_steps=not is_hybrid_model, - buffer_size_gb=buffer_size_gb, - buffer_guaranteed_fraction=buffer_guaranteed_fraction, + active_buffer_size_gb=active_buffer_size_gb, block_size_tokens=block_size_tokens, - buffer_overflow_factor=buffer_overflow_factor, - max_requests_override=max_requests_override, - max_tokens_override=max_tokens_override, + max_tokens=max_tokens, layer_type_list=layer_type_list, mamba_conv_states_shape=(544, 4), mamba_ssm_states_shape=(8, 64, 16), @@ -92,29 +86,26 @@ def test_initialize_dynamic_context(self, is_hybrid_model: bool): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=0.03, block_size_tokens=128, - max_requests_override=None, - max_tokens_override=None, - buffer_overflow_factor=None, + max_tokens=None, is_hybrid_model=is_hybrid_model, ) if not is_hybrid_model: - assert dynamic_context.gtd_block_count == 48 - assert dynamic_context.gtd_request_count == 12 - assert dynamic_context.block_allocator.block_count_total == 491 - assert dynamic_context.max_requests == 128 - assert dynamic_context.max_tokens == 62848 + assert dynamic_context.block_allocator.total_count == 981 + assert dynamic_context.block_allocator.active_count == 490 + assert dynamic_context.max_total_requests == 980 + assert dynamic_context.max_active_requests == 490 + assert dynamic_context.max_tokens == 16384 assert dynamic_context.num_mamba_layers == 0 assert dynamic_context.mamba_metadata is None else: - assert dynamic_context.gtd_block_count == 112 - assert dynamic_context.gtd_request_count == 28 - assert dynamic_context.block_allocator.block_count_total == 1156 - assert dynamic_context.max_requests == 320 - assert dynamic_context.max_tokens == 154176 + assert dynamic_context.block_allocator.total_count == 1111 + assert dynamic_context.block_allocator.active_count == 555 + assert dynamic_context.max_total_requests == 1110 + assert dynamic_context.max_active_requests == 555 + assert dynamic_context.max_tokens == 16384 assert dynamic_context.num_mamba_layers == 1 assert dynamic_context.mamba_metadata is not None @@ -130,12 +121,9 @@ def test_is_static_batching(self): kv_channels=64, num_attention_heads=8, max_sequence_length=512, - buffer_size_gb=1.0, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=1.0, block_size_tokens=128, - max_requests_override=None, - max_tokens_override=None, - buffer_overflow_factor=None, + max_tokens=None, ) assert not dynamic_context.is_static_batching() @@ -149,27 +137,19 @@ def test_is_memory_available(self, is_hybrid_model): kv_channels=64, num_attention_heads=8, max_sequence_length=512, - buffer_size_gb=1.0, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=1.0, block_size_tokens=128, - max_requests_override=None, - max_tokens_override=None, - buffer_overflow_factor=None, + max_tokens=None, is_hybrid_model=is_hybrid_model, ) - dynamic_context.block_allocator.block_count_avail = 10 + dynamic_context.block_allocator.active_count = 10 assert dynamic_context.block_allocator.is_memory_available(10) assert not dynamic_context.block_allocator.is_memory_available(11) assert dynamic_context.block_allocator.is_memory_available(1) - dynamic_context.block_allocator.block_count_avail = 0 + dynamic_context.block_allocator.active_count = 0 assert not dynamic_context.block_allocator.is_memory_available(1) - dynamic_context.block_allocator.block_count_avail = 10 - dynamic_context.gtd_block_count = 5 - assert dynamic_context.block_allocator.is_memory_available(6) - assert not dynamic_context.block_allocator.is_memory_available(6, safe=True) - @pytest.mark.internal @pytest.mark.parametrize("is_hybrid_model", [False, True]) def test_request_overflow(self, is_hybrid_model: bool): @@ -181,17 +161,15 @@ def test_request_overflow(self, is_hybrid_model: bool): kv_channels=64, num_attention_heads=8, max_sequence_length=128, - buffer_size_gb=0.01, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=0.01, block_size_tokens=32, - max_requests_override=None, - max_tokens_override=None, - buffer_overflow_factor=None, + max_tokens=None, rounder=1, is_hybrid_model=is_hybrid_model, ) + dynamic_context.max_active_requests //= 2 with pytest.raises(RequestOverflowError): - for i in range(dynamic_context.max_requests + 1): + for i in range(dynamic_context.max_active_requests + 1): dynamic_context.add_request( DynamicInferenceRequest( request_id=i, @@ -213,12 +191,9 @@ def test_token_overflow_error(self, is_hybrid_model: bool): kv_channels=64, num_attention_heads=8, max_sequence_length=512, - buffer_size_gb=0.1, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=0.1, block_size_tokens=128, - buffer_overflow_factor=1.0, - max_requests_override=2, - max_tokens_override=20, # Setting a very low token limit + max_tokens=200, # setting low, but >= context.max_active_requests. rounder=1, is_hybrid_model=is_hybrid_model, ) @@ -227,7 +202,7 @@ def test_token_overflow_error(self, is_hybrid_model: bool): dynamic_context.add_request( DynamicInferenceRequest( request_id=1, - prompt_tokens=torch.arange(0, 25, device='cuda'), + prompt_tokens=torch.arange(0, 225, device='cuda'), sampling_params=SamplingParams( num_tokens_to_generate=dynamic_context.max_tokens - 25 ), @@ -245,12 +220,9 @@ def test_reset(self, is_hybrid_model: bool): kv_channels=64, num_attention_heads=8, max_sequence_length=128, - buffer_size_gb=1.0, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=1.0, block_size_tokens=128, - max_requests_override=None, - max_tokens_override=None, - buffer_overflow_factor=None, + max_tokens=None, is_hybrid_model=is_hybrid_model, ) @@ -273,7 +245,6 @@ def test_reset(self, is_hybrid_model: bool): dynamic_context.token_to_position_in_request.fill_(1) dynamic_context.token_to_block_idx.fill_(1) dynamic_context.token_to_local_position_within_kv_block.fill_(1) - dynamic_context.block_allocator.block_count_avail = 5 dynamic_context.memory_buffer.fill_(1) dynamic_context.request_to_kv_block_ids.fill_(1) if is_hybrid_model: @@ -303,8 +274,8 @@ def test_reset(self, is_hybrid_model: bool): assert torch.all(dynamic_context.token_to_block_idx == -1) assert torch.all(dynamic_context.token_to_local_position_within_kv_block == 0) assert ( - dynamic_context.block_allocator.block_count_avail - == dynamic_context.block_allocator.block_count_total - 1 + dynamic_context.block_allocator.active_count + == dynamic_context.block_allocator.total_count // 2 ) assert torch.all(dynamic_context.request_to_kv_block_ids == -1) if is_hybrid_model: @@ -322,19 +293,16 @@ def test_allocate_and_release_memory_blocks(self, is_hybrid_model): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=0.03, block_size_tokens=128, - max_requests_override=None, - max_tokens_override=None, - buffer_overflow_factor=None, + max_tokens=None, is_hybrid_model=is_hybrid_model, ) if is_hybrid_model: - expected_memory_blocks = [1151, 1152, 1153, 1154] + expected_memory_blocks = [1106, 1107, 1108, 1109] else: - expected_memory_blocks = [486, 487, 488, 489] + expected_memory_blocks = [976, 977, 978, 979] expected_block_count_avail = expected_memory_blocks[0] assert ( @@ -345,20 +313,20 @@ def test_allocate_and_release_memory_blocks(self, is_hybrid_model): .tolist() == expected_memory_blocks ) - assert dynamic_context.block_allocator.block_count_avail == expected_block_count_avail + assert dynamic_context.block_allocator.total_avail == expected_block_count_avail dynamic_context.block_allocator.release_memory_blocks( torch.tensor(expected_memory_blocks[-2:], device='cuda') ) - assert dynamic_context.block_allocator.block_count_avail == expected_block_count_avail + 2 + assert dynamic_context.block_allocator.total_avail == expected_block_count_avail + 2 assert ( dynamic_context.block_allocator.allocate_memory_blocks(1).item() == expected_memory_blocks[-1] ) - assert dynamic_context.block_allocator.block_count_avail == expected_block_count_avail + 1 + assert dynamic_context.block_allocator.total_avail == expected_block_count_avail + 1 # Should return None since we allocate more blocks than what we have. assert ( dynamic_context.block_allocator.allocate_memory_blocks( - dynamic_context.block_allocator.block_count_avail + 100 + dynamic_context.block_allocator.total_avail + 100 ) == None ) @@ -374,12 +342,9 @@ def test_add_request(self, is_hybrid_model: bool): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=0.03, block_size_tokens=128, - max_requests_override=None, - max_tokens_override=None, - buffer_overflow_factor=None, + max_tokens=None, is_hybrid_model=is_hybrid_model, ) assert dynamic_context.block_size_tokens == 128 @@ -401,7 +366,7 @@ def test_add_request(self, is_hybrid_model: bool): assert dynamic_context.request_kv_length_offsets[0] == 0 assert dynamic_context.request_kv_block_counts[0] == 2 assert dynamic_context.request_last_kv_block_id[0].item() == ( - 1154 if is_hybrid_model else 489 + 1109 if is_hybrid_model else 979 ) assert dynamic_context.request_last_kv_block_offset[0].item() == 15 assert torch.all( @@ -450,12 +415,9 @@ def test_update_request(self, is_hybrid_model: bool): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=0.03, block_size_tokens=128, - max_requests_override=None, - max_tokens_override=None, - buffer_overflow_factor=None, + max_tokens=None, is_hybrid_model=is_hybrid_model, ) @@ -464,7 +426,7 @@ def test_update_request(self, is_hybrid_model: bool): dynamic_context.paused_request_count = 0 dynamic_context.total_request_count = 3 dynamic_context.request_kv_block_counts[0:3] = 1 - new_block_ids = dynamic_context.block_allocator.allocate_memory_blocks(3, safe=True) + new_block_ids = dynamic_context.block_allocator.allocate_memory_blocks(3) dynamic_context.request_to_kv_block_ids[0:3, 0] = new_block_ids if is_hybrid_model: @@ -497,12 +459,9 @@ def test_update_request(self, is_hybrid_model: bool): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=0.03, block_size_tokens=128, - max_requests_override=None, - max_tokens_override=None, - buffer_overflow_factor=None, + max_tokens=None, is_hybrid_model=is_hybrid_model, ) @@ -520,18 +479,16 @@ def test_update_request(self, is_hybrid_model: bool): ) total_request_count = 10 - dynamic_context.block_allocator.block_count_avail -= 11 # We align 11 blocks to the 10 requests we have. 3rd request alone we setup like it requires 2 blocks + dynamic_context.block_allocator.total_avail -= 11 # We align 11 blocks to the 10 requests we have. 3rd request alone we setup like it requires 2 blocks dynamic_context.total_request_count = total_request_count dynamic_context.request_to_kv_block_ids[0:total_request_count, 0] = torch.arange( - dynamic_context.block_allocator.block_count_avail, - dynamic_context.block_allocator.block_count_avail + 10, + dynamic_context.block_allocator.total_avail, + dynamic_context.block_allocator.total_avail + 10, ) dynamic_context.request_to_kv_block_ids[3][ 1 - ] = ( - dynamic_context.block_allocator.block_count_avail - ) # Assign one extra block to request 3. + ] = dynamic_context.block_allocator.total_avail # Assign one extra block to request 3. dynamic_context.request_kv_length_offsets[0:total_request_count] = 10 # For 0, 1, 5, 6, the total number of tokens in last block is block size -1, so that they will all need extra blocks dynamic_context.request_kv_length_offsets[0:2] = dynamic_context.block_size_tokens - 1 @@ -617,13 +574,13 @@ def test_update_request(self, is_hybrid_model: bool): dynamic_context.request_to_kv_block_ids[0:10].cpu() == torch.tensor( [ - [1144, 1147, -1, -1], - [1145, 1144, -1, -1], - [1149, 1151, -1, -1], - [1150, 1152, -1, -1], - [1148, -1, -1, -1], - [1146, -1, -1, -1], - [1153, -1, -1, -1], + [1099, 1102, -1, -1], + [1100, 1099, -1, -1], + [1104, 1106, -1, -1], + [1105, 1107, -1, -1], + [1103, -1, -1, -1], + [1101, -1, -1, -1], + [1108, -1, -1, -1], [-1, -1, -1, -1], [-1, -1, -1, -1], [-1, -1, -1, -1], @@ -635,13 +592,13 @@ def test_update_request(self, is_hybrid_model: bool): dynamic_context.request_to_kv_block_ids[0:10].cpu() == torch.tensor( [ - [479, 482, -1, -1], - [480, 479, -1, -1], - [484, 486, -1, -1], - [485, 487, -1, -1], - [483, -1, -1, -1], - [481, -1, -1, -1], - [488, -1, -1, -1], + [969, 972, -1, -1], + [970, 969, -1, -1], + [974, 976, -1, -1], + [975, 977, -1, -1], + [973, -1, -1, -1], + [971, -1, -1, -1], + [978, -1, -1, -1], [-1, -1, -1, -1], [-1, -1, -1, -1], [-1, -1, -1, -1], @@ -661,23 +618,20 @@ def test_release_memory_blocks_for_finished_requests(self, is_hybrid_model): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=0.03, block_size_tokens=128, - max_requests_override=None, - max_tokens_override=None, - buffer_overflow_factor=None, + max_tokens=None, is_hybrid_model=is_hybrid_model, ) # Set up the initial state with 5 requests # Allocate 5 blocks for 5 requests - initial_blocks = dynamic_context.block_allocator.allocate_memory_blocks(5, safe=True) + initial_blocks = dynamic_context.block_allocator.allocate_memory_blocks(5) dynamic_context.total_request_count = 5 dynamic_context.paused_request_count = 0 # Record the available blocks before releasing memory - initial_available_blocks = dynamic_context.block_allocator.block_count_avail + initial_available_blocks = dynamic_context.block_allocator.total_avail # Assign blocks to the requests (one block per request) for i in range(5): @@ -708,7 +662,7 @@ def test_release_memory_blocks_for_finished_requests(self, is_hybrid_model): assert dynamic_context.active_token_count == 2 # Verify that 3 blocks were released by checking the available blocks - assert dynamic_context.block_allocator.block_count_avail == initial_available_blocks + 3 + assert dynamic_context.block_allocator.total_avail == initial_available_blocks + 3 if is_hybrid_model: # Request at position 3 now moves into finished request position 0 @@ -736,23 +690,20 @@ def test_finished_requests_with_multiple_blocks(self, is_hybrid_model): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=0.03, block_size_tokens=128, - max_requests_override=None, - max_tokens_override=None, - buffer_overflow_factor=None, + max_tokens=None, is_hybrid_model=is_hybrid_model, ) # Set up the initial state with 3 requests, where some use multiple blocks # Allocate 6 blocks in total for the requests - initial_blocks = dynamic_context.block_allocator.allocate_memory_blocks(6, safe=True) + initial_blocks = dynamic_context.block_allocator.allocate_memory_blocks(6) dynamic_context.total_request_count = 3 dynamic_context.paused_request_count = 0 # Record the available blocks before releasing memory - initial_available_blocks = dynamic_context.block_allocator.block_count_avail + initial_available_blocks = dynamic_context.block_allocator.total_avail # Assign blocks to the requests: # - Request 0: 1 block @@ -792,7 +743,7 @@ def test_finished_requests_with_multiple_blocks(self, is_hybrid_model): assert dynamic_context.active_token_count == 0 # Verify that all 6 blocks were released by checking the available blocks - assert dynamic_context.block_allocator.block_count_avail == initial_available_blocks + 6 + assert dynamic_context.block_allocator.total_avail == initial_available_blocks + 6 if is_hybrid_model: # All mamba states should be zeroed out @@ -812,12 +763,9 @@ def test_mamba_states_cache(self, is_hybrid_model: bool): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=0.03, block_size_tokens=128, - max_requests_override=None, - max_tokens_override=None, - buffer_overflow_factor=None, + max_tokens=None, is_hybrid_model=False, ) with pytest.raises(AssertionError) as error: @@ -830,12 +778,9 @@ def test_mamba_states_cache(self, is_hybrid_model: bool): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=0.03, block_size_tokens=128, - max_requests_override=None, - max_tokens_override=None, - buffer_overflow_factor=None, + max_tokens=None, is_hybrid_model=is_hybrid_model, layer_type_list=[Symbols.MAMBA, Symbols.ATTENTION, Symbols.MAMBA, Symbols.ATTENTION], ) @@ -889,12 +834,9 @@ def test_calculate_and_store_log_probs(self): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=0.03, block_size_tokens=128, - max_requests_override=None, - max_tokens_override=None, - buffer_overflow_factor=None, + max_tokens=None, ) # Add a few requests to the context @@ -1119,8 +1061,8 @@ def test_unified_memory(self): gpu_size_gb = ( torch.cuda.get_device_properties(torch.cuda.current_device()).total_memory / 1024**3 ) - buffer_size_gb = 20 - num_contexts = math.ceil(gpu_size_gb / buffer_size_gb) + 1 + active_buffer_size_gb = 20 + num_contexts = math.ceil(gpu_size_gb / active_buffer_size_gb) + 1 # Allocate enough contexts to fill GPU memory. def init_contexts(*, unified_memory_level): @@ -1129,13 +1071,11 @@ def init_contexts(*, unified_memory_level): contexts.append( DynamicInferenceContext( params_dtype=torch.float32, - num_layers=4, - kv_channels=8, - num_attention_heads=2, + num_layers=64, + kv_channels=16, + num_attention_heads=4, max_sequence_length=512, - buffer_size_gb=buffer_size_gb, - buffer_overflow_factor=1, - buffer_guaranteed_fraction=0, + active_buffer_size_gb=active_buffer_size_gb, unified_memory_level=unified_memory_level, ) ) diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index f3c05257552..457f1121f24 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -84,12 +84,9 @@ class DynamicEngineTestConfig: num_gap_steps: int = 2 - context_buffer_size_gb: float = 0.1 # enough room for all tokens. + context_active_buffer_size_gb: float = 0.1 # enough room for all tokens. context_block_size_tokens: int = 256 - context_buffer_guaranteed_fraction: float = 0.01 - context_buffer_overflow_factor: Optional[float] = None - context_max_requests_override: Optional[int] = None - context_max_tokens_override: Optional[int] = None + context_max_tokens: Optional[int] = None tensor_model_parallel_size: int = 1 pipeline_model_parallel_size: int = 1 expert_model_parallel_size: int = 1 @@ -125,17 +122,6 @@ def __post_init__(self): assert self.num_tokens_total is not None self.max_sequence_length = self.num_tokens_total - # Update overrides if not using overflow factor. - if self.context_buffer_overflow_factor is None: - - # Enough room for all requests. - if self.context_max_requests_override is None: - self.context_max_requests_override = self.num_requests - - # Enough room for all tokens. - if self.context_max_tokens_override is None: - self.context_max_tokens_override = self.num_requests * self.max_sequence_length - @dataclass class DynamicEngineTestEnv: @@ -227,12 +213,9 @@ def _build_inference_context( max_sequence_length=test_config.max_sequence_length, num_cuda_graphs=test_config.num_cuda_graphs, use_cuda_graphs_for_non_decode_steps=not test_config.model_provider == "mamba", - buffer_size_gb=test_config.context_buffer_size_gb, - buffer_guaranteed_fraction=test_config.context_buffer_guaranteed_fraction, + active_buffer_size_gb=test_config.context_active_buffer_size_gb, block_size_tokens=test_config.context_block_size_tokens, - buffer_overflow_factor=test_config.context_buffer_overflow_factor, - max_requests_override=test_config.context_max_requests_override, - max_tokens_override=test_config.context_max_tokens_override, + max_tokens=test_config.context_max_tokens, tensor_model_parallel_size=transformer_config.tensor_model_parallel_size, layer_type_list=layer_type_list, mamba_conv_states_shape=mamba_conv_states_shape, @@ -520,14 +503,12 @@ def test_simple(self, model_provider, num_cuda_graphs, cuda_graph_scope) -> None env = self._run_test( model_provider=model_provider, num_cuda_graphs=num_cuda_graphs, - context_max_requests_override=32, cuda_graph_scope=cuda_graph_scope, force_build_cuda_graphs=True, ) # Validate max_requests, max_tokens. - assert env.engine.context.max_requests == 32 - assert env.engine.context.max_tokens == 160 + assert env.engine.context.max_tokens == DynamicInferenceContext.DEFAULT_MAX_TOKENS # Validate output tokens. gpt_expected_generated_tokens = [ @@ -568,41 +549,6 @@ def test_simple(self, model_provider, num_cuda_graphs, cuda_graph_scope) -> None f"expected ({expected_generated_tokens})." ) - @pytest.mark.internal - @pytest.mark.skipif( - not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" - ) - def test_overflow_factor(self, model_provider: str = "gpt") -> None: - """Test overflow factor arg.""" - skip_if_mamba_sequence_packing_not_available(model_provider) - - # Run test. - env = self._run_test( - context_buffer_overflow_factor=0.1, - context_max_requests_override=None, - context_max_tokens_override=None, - model_provider=model_provider, - ) - - # Validate max_requests, max_tokens. - if model_provider == "gpt": - assert env.engine.context.max_requests == 420 - assert env.engine.context.max_tokens == 420 - elif model_provider == "mamba": - assert env.engine.context.max_requests == 16 - assert env.engine.context.max_tokens == 16 - - @pytest.mark.internal - @pytest.mark.skipif( - not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" - ) - @pytest.mark.parametrize("model_provider", ["gpt", "mamba"]) - def test_request_overflow(self, model_provider: str) -> None: - """Test request overflow.""" - skip_if_mamba_sequence_packing_not_available(model_provider) - - self._run_test(context_max_requests_override=4, model_provider=model_provider) - @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) @@ -610,7 +556,7 @@ def test_request_overflow(self, model_provider: str) -> None: def test_token_overflow_transient(self) -> None: """Test token overflow.""" test_config = DynamicEngineTestConfig( - num_requests=2, min_prompt_length=8, max_prompt_length=8, context_max_tokens_override=12 + num_requests=2, min_prompt_length=512, max_prompt_length=512, context_max_tokens=900 ) env = self._build_test_env(test_config) env.engine._add_request(env.requests[0]) @@ -629,7 +575,7 @@ def test_token_overflow_transient(self) -> None: ) def test_token_overflow_nontransient(self) -> None: """Test token overflow (non-transient).""" - test_config = DynamicEngineTestConfig(context_max_tokens_override=8) + test_config = DynamicEngineTestConfig(context_max_tokens=8) env = self._build_test_env(test_config) try: env.engine._add_request(env.requests[0]) @@ -649,9 +595,9 @@ def test_block_overflow(self, model_provider: str) -> None: env = self._build_test_env(DynamicEngineTestConfig(model_provider=model_provider)) context = env.engine.context block_size_bytes = context.block_size_bytes - buffer_size_gb = (block_size_bytes + 1) / 1024**3 + active_buffer_size_gb = (block_size_bytes + 1) / 1024**3 test_config = DynamicEngineTestConfig( - context_buffer_size_gb=buffer_size_gb, model_provider=model_provider + context_active_buffer_size_gb=active_buffer_size_gb, model_provider=model_provider ) env = self._build_test_env(test_config) env.engine._add_request(env.requests[0]) @@ -686,19 +632,21 @@ def test_cuda_graph_token_counts(self) -> None: # Test num_cuda_graphs. for num_cuda_graphs, expected_cuda_graph_token_counts in [ - (0, [64]), - (1, [64]), - (2, [64, 32]), - (4, [64, 48, 32, 16]), - (8, [64, 56, 48, 40, 32, 24, 16, 8]), - (16, [64, 56, 48, 40, 32, 24, 16, 8]), - (64, [64, 56, 48, 40, 32, 24, 16, 8]), - (1024, [64, 56, 48, 40, 32, 24, 16, 8]), + (0, [80]), + (1, [80]), + (2, [80, 40]), + (4, [80, 72, 48, 24]), + (8, [80, 64, 48, 32, 16]), + (16, [80, 72, 64, 56, 48, 40, 32, 24, 16, 8]), + (64, [80, 72, 64, 56, 48, 40, 32, 24, 16, 8]), + (1024, [80, 72, 64, 56, 48, 40, 32, 24, 16, 8]), ]: # Build cuda graphs (inside dynamic engine). env = self._build_test_env( - DynamicEngineTestConfig(num_requests=64, num_cuda_graphs=num_cuda_graphs) + DynamicEngineTestConfig( + context_active_buffer_size_gb=0.01, num_cuda_graphs=num_cuda_graphs + ) ) actual_cuda_graph_token_counts = env.engine.context.cuda_graph_token_counts assert ( @@ -745,7 +693,9 @@ def test_cuda_graph_warmup( # Initialize context. env = self._build_test_env( - DynamicEngineTestConfig(num_requests=32, num_cuda_graphs=8, num_tokens_to_generate=1) + DynamicEngineTestConfig( + context_active_buffer_size_gb=0.0041, num_cuda_graphs=8, num_tokens_to_generate=1 + ) ) context = env.engine.context @@ -1034,9 +984,8 @@ def test_events(self): num_requests=16, max_prompt_length=10, num_tokens_to_generate=32, - context_buffer_size_gb=0.001, # 0.001, # 8 blocks - context_max_requests_override=8, - context_max_tokens_override=8, + context_active_buffer_size_gb=0.001, # 0.001, # 8 blocks + context_max_tokens=8, num_gap_steps=1, ) @@ -1085,27 +1034,5 @@ def test_chunked_prefill(self, model_provider: str): materialize_only_last_token_logits=False, model_provider=model_provider, context_block_size_tokens=256, - context_max_tokens_override=300, + context_max_tokens=1000, ) - - -if __name__ == "__main__": - test = TestDynamicInferenceEngine() - test.test_simple(4) - test.test_overflow_factor() - test.test_request_overflow() - test.test_token_overflow_transient() - # test.test_token_overflow_nontransient() # uncomment in megatron-core 0.16 - test.test_block_overflow() - test.test_multi_add() - test.test_fixed_output_lengths() - test.test_cuda_graph_request_counts() - test.test_cuda_graph_warmup(WarmupEngineMode.DECODE, 1, 8) - test.test_generate_function() - asyncio.run(test.test_run_engine()) - test.test_return_log_probs() - test.test_parallel_inference() - # test.test_events() # uncomment in megatron-core 0.16 - test.teardown_method(None) - print("~~~") - print("success.") diff --git a/tests/unit_tests/inference/test_wandb_logging.py b/tests/unit_tests/inference/test_wandb_logging.py index 1512e805f9c..d7a45783e00 100644 --- a/tests/unit_tests/inference/test_wandb_logging.py +++ b/tests/unit_tests/inference/test_wandb_logging.py @@ -48,9 +48,8 @@ def _get_dynamic_context( kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, + active_buffer_size_gb=0.03, block_size_tokens=128, - buffer_guaranteed_fraction=0.1, metrics_writer=None, ): """Helper to create a DynamicInferenceContext.""" @@ -61,8 +60,7 @@ def _get_dynamic_context( num_attention_heads=num_attention_heads, max_sequence_length=max_sequence_length, num_cuda_graphs=None, - buffer_size_gb=buffer_size_gb, - buffer_guaranteed_fraction=buffer_guaranteed_fraction, + active_buffer_size_gb=active_buffer_size_gb, block_size_tokens=block_size_tokens, metrics_writer=metrics_writer, ) @@ -83,12 +81,11 @@ def test_get_kvcache_utilization_stats_with_requests(self): assert 'active_utilization' in stats assert 'active_request_count' in stats assert 'paused_request_count' in stats - assert 'gtd_block_count' in stats assert 'block_count_avail' in stats - assert 'num_non_gtd_blocks' in stats assert 'active_token_count' in stats assert 'total_request_count' in stats - assert 'max_requests' in stats + assert 'max_total_requests' in stats + assert 'max_active_requests' in stats # Verify values for empty context assert stats['allocated_blocks'] == 0 @@ -134,12 +131,11 @@ def test_get_kvcache_utilization_stats_with_requests(self): assert stats_after['total_blocks'] == stats['total_blocks'] assert stats_after['total_blocks'] > 0 - # Verify that gtd_block_count remains constant - assert stats_after['gtd_block_count'] == stats['gtd_block_count'] - # Verify that max_requests remains constant - assert stats_after['max_requests'] == stats['max_requests'] - assert stats_after['max_requests'] > 0 + assert stats_after['max_total_requests'] == stats['max_total_requests'] + assert stats_after['max_total_requests'] > 0 + assert stats_after['max_active_requests'] == stats['max_active_requests'] + assert stats_after['max_active_requests'] > 0 # Verify block availability decreased after allocation assert stats_after['block_count_avail'] < stats['block_count_avail'] @@ -147,7 +143,7 @@ def test_get_kvcache_utilization_stats_with_requests(self): # Verify relationship: allocated_blocks + block_count_avail + 1 (dummy) = total assert ( stats_after['allocated_blocks'] + stats_after['block_count_avail'] + 1 - == dynamic_context.block_allocator.block_count_total + == dynamic_context.block_allocator.total_count ) # Verify utilization bounds [0, 1] @@ -180,12 +176,11 @@ def test_kvcache_utilization_stats_types(self): 'active_unique_blocks', 'active_request_count', 'paused_request_count', - 'gtd_block_count', 'block_count_avail', - 'num_non_gtd_blocks', 'active_token_count', 'total_request_count', - 'max_requests', + 'max_total_requests', + 'max_active_requests', ] for field in int_fields: @@ -239,8 +234,7 @@ def test_paused_requests_in_stats(self): num_attention_heads=8, max_sequence_length=128, num_cuda_graphs=None, - buffer_size_gb=0.01, # Small buffer to force pausing - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=0.01, # Small buffer to force pausing block_size_tokens=32, ) 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_simple_text_generation_controller.py index 10ffe2fdd40..5fbdef762eb 100644 --- a/tests/unit_tests/inference/text_generation_controllers/test_simple_text_generation_controller.py +++ b/tests/unit_tests/inference/text_generation_controllers/test_simple_text_generation_controller.py @@ -116,8 +116,7 @@ def setup_model( kv_channels=transformer_config.kv_channels, num_attention_heads=transformer_config.num_attention_heads, max_sequence_length=2048, - buffer_size_gb=1, - buffer_guaranteed_fraction=0.1, + active_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