From 16609c299aa608b43ef5afd28ca4ac13d75fb664 Mon Sep 17 00:00:00 2001 From: Robert Kirby Date: Tue, 9 Sep 2025 04:46:29 -0700 Subject: [PATCH 01/69] copied in unified_memory.py. --- megatron/core/inference/unified_memory.py | 61 +++++++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 megatron/core/inference/unified_memory.py diff --git a/megatron/core/inference/unified_memory.py b/megatron/core/inference/unified_memory.py new file mode 100644 index 00000000000..27e2752923c --- /dev/null +++ b/megatron/core/inference/unified_memory.py @@ -0,0 +1,61 @@ +# build_and_use_managed_allocator.py +import os, pathlib +import torch +from torch.utils.cpp_extension import load_inline, CUDA_HOME +from torch.cuda.memory import CUDAPluggableAllocator, change_current_allocator + +# 1) Inline C++ source for a managed-memory allocator +src = r""" +#include +#include + +#define EXPORT extern "C" + +// malloc: allocate Unified (Managed) memory +EXPORT void* my_managed_malloc(ssize_t size, int device, void* stream) { + (void)stream; + int cur = -1; + cudaGetDevice(&cur); + if (device != cur && device >= 0) cudaSetDevice(device); + + void* ptr = nullptr; + cudaError_t err = cudaMallocManaged(&ptr, (size_t)size, cudaMemAttachGlobal); + if (err != cudaSuccess) return nullptr; + + // Optional hints to reduce first-touch page faults: + if (device >= 0) { + cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetPreferredLocation, device); + cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetAccessedBy, device); + } + return ptr; +} + +// free +EXPORT void my_managed_free(void* ptr, ssize_t size, int device, void* stream) { + (void)size; (void)device; (void)stream; + if (ptr) cudaFree(ptr); +} +""" + +# 2) JIT-compile a tiny extension that exports those C symbols. +# We link against cudart; on most systems, -lcudart is enough if CUDA is on LD_LIBRARY_PATH. +extra_ldflags = ["-lcudart"] +if CUDA_HOME: + # help linker find libcudart if not on the default path + cand = os.path.join(CUDA_HOME, "lib64") + if os.path.isdir(cand): + extra_ldflags = [f"-L{cand}", "-lcudart"] + +mod = load_inline( + name="managed_alloc_runtime", + cpp_sources=[src], + functions=[], # no pybind functions; we just want the .so with our exported symbols + with_cuda=True, # ensures CUDA include paths & NVCC toolchain are available + extra_ldflags=extra_ldflags, + verbose=False, +) + +# 3) Locate the built shared library and plug it into PyTorch +so_path = pathlib.Path(mod.__file__).as_posix() +alloc = CUDAPluggableAllocator(so_path, "my_managed_malloc", "my_managed_free").allocator() +unified_memory_mempool = torch.cuda.MemPool(allocator=alloc) \ No newline at end of file From f32fdf24593015487be10a50050aa5105f258999 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 9 Sep 2025 04:56:14 -0700 Subject: [PATCH 02/69] connected arg to context. --- .../inference/gpt/gpt_dynamic_inference.py | 1 + .../inference/contexts/dynamic_context.py | 53 +++++++++++-------- megatron/training/arguments.py | 7 +++ 3 files changed, 39 insertions(+), 22 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index 1a568157215..b904b03f94f 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -136,6 +136,7 @@ def get_inference_context(requests: List[Request], sampling_params: SamplingPara cache_mla_latent=args.multi_latent_attention and args.cache_mla_latents, kv_lora_rank=args.kv_lora_rank if args.multi_latent_attention else None, qk_pos_emb_head_dim=args.qk_pos_emb_head_dim, + unified_memory_level=args.inference_dynamic_batching_unified_memory_level, ) return context diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index e999df58a61..9f3a9161254 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -11,6 +11,7 @@ from torch import Tensor from megatron.core import parallel_state +from megatron.core.inference.unified_memory import unified_memory_mempool from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb from megatron.core.package_info import __version__ as mcore_version from megatron.core.transformer import TransformerConfig @@ -145,7 +146,11 @@ class DynamicInferenceContext(BaseInferenceContext): materialize_only_last_token_logits (bool): If True, only the last token logits are materialized in the context. use_cuda_graphs_for_non_decode_steps (bool): If True, use cuda graphs for non-decode - engine steps. + engine steps. + unified_memory_level (Optional[int]): Set unified memory usage within the + dynamic inference context. The levels are: 0) no unified memory, 1) + allocate `memory_buffer` in unified memory. Eventually, additional + levels will be included to control other tensors within the context. """ def __init__( @@ -169,6 +174,7 @@ def __init__( num_cuda_graphs: Optional[int] = None, materialize_only_last_token_logits: bool = True, use_cuda_graphs_for_non_decode_steps: bool = True, + unified_memory_level: Optional[int] = 0, ): super().__init__(materialize_only_last_token_logits=materialize_only_last_token_logits) @@ -240,6 +246,7 @@ def bytes_to_max_requests_and_tokens(n_bytes): self.params_dtype = params_dtype self.num_layers = num_layers self.max_sequence_length = max_sequence_length + self.unified_memory_level = unified_memory_level self.total_request_count = 0 self.active_token_count = 0 @@ -280,27 +287,29 @@ def bytes_to_max_requests_and_tokens(n_bytes): chunk_count_total = buffer_size_bytes // self.chunk_size_bytes # Memory buffer. - if cache_mla_latent: - self.memory_buffer = torch.full( - (self.num_layers, chunk_count_total, self.chunk_size_tokens, kv_reduced_dim), - -1, - dtype=self.params_dtype, - device=torch.cuda.current_device(), - ) - else: - self.memory_buffer = torch.full( - ( - 2, # key and value - self.num_layers, - chunk_count_total, - self.chunk_size_tokens, - num_attention_heads_per_partition, - hidden_size_per_attention_head, - ), - -1, - dtype=self.params_dtype, - device=torch.cuda.current_device(), - ) + ctx_manager = torch.cuda.use_mem_pool(unified_memory_mempool) if self.unified_memory_level > 0 else nullcontext() + with ctx_manager: + if cache_mla_latent: + self.memory_buffer = torch.full( + (self.num_layers, chunk_count_total, self.chunk_size_tokens, kv_reduced_dim), + -1, + dtype=self.params_dtype, + device=torch.cuda.current_device(), + ) + else: + self.memory_buffer = torch.full( + ( + 2, # key and value + self.num_layers, + chunk_count_total, + self.chunk_size_tokens, + num_attention_heads_per_partition, + hidden_size_per_attention_head, + ), + -1, + dtype=self.params_dtype, + device=torch.cuda.current_device(), + ) # Chunk ids. self.max_kv_chunk_count = math.ceil(self.max_sequence_length / self.chunk_size_tokens) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index d138e0cfa85..4158c1aba40 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1407,6 +1407,13 @@ def _add_inference_args(parser): '(See `dynamic_context.py` for details on how ' '`max_requests` is computed). Due to rounding, the actual ' 'number of cuda graphs may not equal this argument.') + group.add_argument('--inference-dynamic-batching-unified-memory-level', + type=int, default=0, + help='Set unified memory usage within the dynamic ' + 'inference context. The levels are: 0) no unified memory, ' + '1) allocate `memory_buffer` in unified memory. ' + 'Eventually, additional levels will be included to ' + 'control other tensors within the context.') group.add_argument('--symmetric-ar-type', type=str, default=None, choices=['two_shot', "one_shot", "multimem_all_reduce", None], help='What type of symmetric all reduce to use. The default is none which is no use of symetric memory') From beb7f7dbb2bb829dea4f85db3e56e56e628ac609 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 9 Sep 2025 05:08:29 -0700 Subject: [PATCH 03/69] minor fixes. --- megatron/core/inference/contexts/dynamic_context.py | 1 + megatron/training/arguments.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 9f3a9161254..eb1a845ff11 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -2,6 +2,7 @@ import math import warnings +from contextlib import nullcontext from enum import Enum from typing import List, Optional, Tuple diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 4158c1aba40..7a2ec6dda4a 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1408,7 +1408,7 @@ def _add_inference_args(parser): '`max_requests` is computed). Due to rounding, the actual ' 'number of cuda graphs may not equal this argument.') group.add_argument('--inference-dynamic-batching-unified-memory-level', - type=int, default=0, + type=int, default=0, 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. ' From c641430bbca05962d70cc44a0581685978edccbc Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 9 Sep 2025 05:54:06 -0700 Subject: [PATCH 04/69] added unit test. --- .../contexts/test_dynamic_context.py | 56 +++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 341ea78dc89..548f10ebcc7 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -1,3 +1,6 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +import math import pytest import torch @@ -830,3 +833,56 @@ def test_calculate_and_store_log_probs(self): ) current_global_token_offset += expected_len + + def test_unified_memory(self): + + self._setup_model_parallel_group(1, 1) + + 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 + + # >>> + # from lutil import pax + # pax("gpu_size_gb, buffer_size_gb, num_contexts") + # <<< + + # Allocate enough contexts to fill GPU memory. + def init_contexts(unified_memory_level): + contexts = [] + for i in range(num_contexts): + contexts.append(DynamicInferenceContext( + params_dtype=torch.float32, + num_layers=4, + kv_channels=8, + num_attention_heads=2, + max_sequence_length=512, + buffer_size_gb=buffer_size_gb, + buffer_overflow_factor=1, + buffer_guaranteed_fraction=0, + unified_memory_level=unified_memory_level, + )) + + # Pure GPU memory test should OOM. + try: + init_contexts(0) + except torch.OutOfMemoryError: + pass + else: + raise Exception("expected OOM.") + + # Unified memory test should succeed. + init_contexts(1) + + # >>> + from lutil import pax + pax() + # <<< + +# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> +if __name__ == "__main__": + test = TestDynamicContext() + test.test_unified_memory() + test.teardown_method(None) + raise Exception("hi.") +# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< From 2d053686b45acfadda349a731fc232102856c956 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 9 Sep 2025 05:58:34 -0700 Subject: [PATCH 05/69] clean up. --- .../inference/contexts/test_dynamic_context.py | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 548f10ebcc7..f778e6534ca 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -838,15 +838,11 @@ def test_unified_memory(self): self._setup_model_parallel_group(1, 1) + # Compute number of contexts needed to fill GPU memory. 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 - # >>> - # from lutil import pax - # pax("gpu_size_gb, buffer_size_gb, num_contexts") - # <<< - # Allocate enough contexts to fill GPU memory. def init_contexts(unified_memory_level): contexts = [] @@ -874,15 +870,10 @@ def init_contexts(unified_memory_level): # Unified memory test should succeed. init_contexts(1) - # >>> - from lutil import pax - pax() - # <<< -# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> if __name__ == "__main__": test = TestDynamicContext() test.test_unified_memory() test.teardown_method(None) - raise Exception("hi.") -# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< + print("~~~") + print("success.") From 088b3209452880326ba99452f7ae1b823dcdc2bd Mon Sep 17 00:00:00 2001 From: Mcore Bot Date: Tue, 9 Sep 2025 12:14:33 +0000 Subject: [PATCH 06/69] chore: Format files --- megatron/core/inference/contexts/dynamic_context.py | 6 +++++- megatron/core/inference/unified_memory.py | 12 +++++++----- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index eb1a845ff11..5b4945e127d 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -288,7 +288,11 @@ def bytes_to_max_requests_and_tokens(n_bytes): chunk_count_total = buffer_size_bytes // self.chunk_size_bytes # Memory buffer. - ctx_manager = torch.cuda.use_mem_pool(unified_memory_mempool) if self.unified_memory_level > 0 else nullcontext() + ctx_manager = ( + torch.cuda.use_mem_pool(unified_memory_mempool) + if self.unified_memory_level > 0 + else nullcontext() + ) with ctx_manager: if cache_mla_latent: self.memory_buffer = torch.full( diff --git a/megatron/core/inference/unified_memory.py b/megatron/core/inference/unified_memory.py index 27e2752923c..1a9220b1c69 100644 --- a/megatron/core/inference/unified_memory.py +++ b/megatron/core/inference/unified_memory.py @@ -1,8 +1,10 @@ # build_and_use_managed_allocator.py -import os, pathlib +import os +import pathlib + import torch -from torch.utils.cpp_extension import load_inline, CUDA_HOME from torch.cuda.memory import CUDAPluggableAllocator, change_current_allocator +from torch.utils.cpp_extension import CUDA_HOME, load_inline # 1) Inline C++ source for a managed-memory allocator src = r""" @@ -49,8 +51,8 @@ mod = load_inline( name="managed_alloc_runtime", cpp_sources=[src], - functions=[], # no pybind functions; we just want the .so with our exported symbols - with_cuda=True, # ensures CUDA include paths & NVCC toolchain are available + functions=[], # no pybind functions; we just want the .so with our exported symbols + with_cuda=True, # ensures CUDA include paths & NVCC toolchain are available extra_ldflags=extra_ldflags, verbose=False, ) @@ -58,4 +60,4 @@ # 3) Locate the built shared library and plug it into PyTorch so_path = pathlib.Path(mod.__file__).as_posix() alloc = CUDAPluggableAllocator(so_path, "my_managed_malloc", "my_managed_free").allocator() -unified_memory_mempool = torch.cuda.MemPool(allocator=alloc) \ No newline at end of file +unified_memory_mempool = torch.cuda.MemPool(allocator=alloc) From f322734a3377e34772b13eb94dfd7a96f697c742 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 9 Sep 2025 06:29:39 -0700 Subject: [PATCH 07/69] fix formatting and imports for ci. --- megatron/core/inference/unified_memory.py | 37 +++++++++++-------- .../contexts/test_dynamic_context.py | 6 +-- 2 files changed, 25 insertions(+), 18 deletions(-) diff --git a/megatron/core/inference/unified_memory.py b/megatron/core/inference/unified_memory.py index 1a9220b1c69..e1929ab8186 100644 --- a/megatron/core/inference/unified_memory.py +++ b/megatron/core/inference/unified_memory.py @@ -1,9 +1,11 @@ -# build_and_use_managed_allocator.py +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + import os import pathlib +import warnings import torch -from torch.cuda.memory import CUDAPluggableAllocator, change_current_allocator +from torch.cuda.memory import CUDAPluggableAllocator from torch.utils.cpp_extension import CUDA_HOME, load_inline # 1) Inline C++ source for a managed-memory allocator @@ -48,16 +50,21 @@ if os.path.isdir(cand): extra_ldflags = [f"-L{cand}", "-lcudart"] -mod = load_inline( - name="managed_alloc_runtime", - cpp_sources=[src], - functions=[], # no pybind functions; we just want the .so with our exported symbols - with_cuda=True, # ensures CUDA include paths & NVCC toolchain are available - extra_ldflags=extra_ldflags, - verbose=False, -) - -# 3) Locate the built shared library and plug it into PyTorch -so_path = pathlib.Path(mod.__file__).as_posix() -alloc = CUDAPluggableAllocator(so_path, "my_managed_malloc", "my_managed_free").allocator() -unified_memory_mempool = torch.cuda.MemPool(allocator=alloc) +try: + mod = load_inline( + name="managed_alloc_runtime", + cpp_sources=[src], + functions=[], # no pybind functions; we just want the .so with our exported symbols + with_cuda=True, # ensures CUDA include paths & NVCC toolchain are available + extra_ldflags=extra_ldflags, + verbose=False, + ) + + # 3) Locate the built shared library and plug it into PyTorch + so_path = pathlib.Path(mod.__file__).as_posix() + alloc = CUDAPluggableAllocator(so_path, "my_managed_malloc", "my_managed_free").allocator() + unified_memory_mempool = torch.cuda.MemPool(allocator=alloc) + +except RuntimeError: + warnings.warn("Ninja is required to load C++ extensions (pip install ninja to get it).") + unified_memory_mempool = None diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index f778e6534ca..c5f8f46fbb0 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -844,7 +844,7 @@ def test_unified_memory(self): num_contexts = math.ceil(gpu_size_gb / buffer_size_gb) + 1 # Allocate enough contexts to fill GPU memory. - def init_contexts(unified_memory_level): + def init_contexts(*, unified_memory_level): contexts = [] for i in range(num_contexts): contexts.append(DynamicInferenceContext( @@ -861,14 +861,14 @@ def init_contexts(unified_memory_level): # Pure GPU memory test should OOM. try: - init_contexts(0) + init_contexts(unified_memory_level=0) except torch.OutOfMemoryError: pass else: raise Exception("expected OOM.") # Unified memory test should succeed. - init_contexts(1) + init_contexts(unified_memory_level=1) if __name__ == "__main__": From 330c555efc70c850e48400faa7e7b2ca0ffa3e1f Mon Sep 17 00:00:00 2001 From: Mcore Bot Date: Tue, 9 Sep 2025 13:01:58 +0000 Subject: [PATCH 08/69] chore: Format files --- .../contexts/test_dynamic_context.py | 29 +++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index c5f8f46fbb0..552841ad2f6 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -1,6 +1,7 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. import math + import pytest import torch @@ -839,7 +840,9 @@ def test_unified_memory(self): self._setup_model_parallel_group(1, 1) # Compute number of contexts needed to fill GPU memory. - gpu_size_gb = torch.cuda.get_device_properties(torch.cuda.current_device()).total_memory / 1024**3 + 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 @@ -847,17 +850,19 @@ def test_unified_memory(self): def init_contexts(*, unified_memory_level): contexts = [] for i in range(num_contexts): - contexts.append(DynamicInferenceContext( - params_dtype=torch.float32, - num_layers=4, - kv_channels=8, - num_attention_heads=2, - max_sequence_length=512, - buffer_size_gb=buffer_size_gb, - buffer_overflow_factor=1, - buffer_guaranteed_fraction=0, - unified_memory_level=unified_memory_level, - )) + contexts.append( + DynamicInferenceContext( + params_dtype=torch.float32, + num_layers=4, + kv_channels=8, + num_attention_heads=2, + max_sequence_length=512, + buffer_size_gb=buffer_size_gb, + buffer_overflow_factor=1, + buffer_guaranteed_fraction=0, + unified_memory_level=unified_memory_level, + ) + ) # Pure GPU memory test should OOM. try: From 4e8148fcfec6fcd34e18270892c1a803648f26fc Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 9 Sep 2025 11:42:51 -0500 Subject: [PATCH 09/69] Clean up unified memory code --- .../inference/contexts/dynamic_context.py | 3 ++ megatron/core/inference/unified_memory.py | 31 +++++++------------ 2 files changed, 14 insertions(+), 20 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 5b4945e127d..d074ef58a56 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -248,6 +248,9 @@ def bytes_to_max_requests_and_tokens(n_bytes): self.num_layers = num_layers self.max_sequence_length = max_sequence_length self.unified_memory_level = unified_memory_level + if unified_memory_level > 0 and unified_memory_mempool is None: + warnings.warn("Unified memory requested but not available; defaulting to GPU memory.") + self.unified_memory_level = 0 self.total_request_count = 0 self.active_token_count = 0 diff --git a/megatron/core/inference/unified_memory.py b/megatron/core/inference/unified_memory.py index e1929ab8186..e3e9eb8c814 100644 --- a/megatron/core/inference/unified_memory.py +++ b/megatron/core/inference/unified_memory.py @@ -1,22 +1,20 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. import os -import pathlib import warnings +from pathlib import Path import torch from torch.cuda.memory import CUDAPluggableAllocator from torch.utils.cpp_extension import CUDA_HOME, load_inline -# 1) Inline C++ source for a managed-memory allocator src = r""" #include #include #define EXPORT extern "C" -// malloc: allocate Unified (Managed) memory -EXPORT void* my_managed_malloc(ssize_t size, int device, void* stream) { +EXPORT void* managed_malloc(ssize_t size, int device, void* stream) { (void)stream; int cur = -1; cudaGetDevice(&cur); @@ -26,7 +24,6 @@ cudaError_t err = cudaMallocManaged(&ptr, (size_t)size, cudaMemAttachGlobal); if (err != cudaSuccess) return nullptr; - // Optional hints to reduce first-touch page faults: if (device >= 0) { cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetPreferredLocation, device); cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetAccessedBy, device); @@ -34,37 +31,31 @@ return ptr; } -// free -EXPORT void my_managed_free(void* ptr, ssize_t size, int device, void* stream) { +EXPORT void managed_free(void* ptr, ssize_t size, int device, void* stream) { (void)size; (void)device; (void)stream; if (ptr) cudaFree(ptr); } """ -# 2) JIT-compile a tiny extension that exports those C symbols. -# We link against cudart; on most systems, -lcudart is enough if CUDA is on LD_LIBRARY_PATH. extra_ldflags = ["-lcudart"] if CUDA_HOME: - # help linker find libcudart if not on the default path - cand = os.path.join(CUDA_HOME, "lib64") - if os.path.isdir(cand): - extra_ldflags = [f"-L{cand}", "-lcudart"] + cuda_lib = os.path.join(CUDA_HOME, "lib64") + if os.path.isdir(cuda_lib): + extra_ldflags = [f"-L{cuda_lib}", "-lcudart"] try: mod = load_inline( name="managed_alloc_runtime", cpp_sources=[src], - functions=[], # no pybind functions; we just want the .so with our exported symbols - with_cuda=True, # ensures CUDA include paths & NVCC toolchain are available + functions=[] + with_cuda=True, extra_ldflags=extra_ldflags, verbose=False, ) - # 3) Locate the built shared library and plug it into PyTorch - so_path = pathlib.Path(mod.__file__).as_posix() - alloc = CUDAPluggableAllocator(so_path, "my_managed_malloc", "my_managed_free").allocator() + so_path = Path(mod.__file__).as_posix() + alloc = CUDAPluggableAllocator(so_path, "managed_malloc", "managed_free").allocator() unified_memory_mempool = torch.cuda.MemPool(allocator=alloc) - except RuntimeError: - warnings.warn("Ninja is required to load C++ extensions (pip install ninja to get it).") + warnings.warn("Unified memory mempool is not available.") unified_memory_mempool = None From 55a528e907399731598d957a08f000e72e3880c9 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 9 Sep 2025 12:02:19 -0500 Subject: [PATCH 10/69] Fix pipeline --- megatron/core/inference/unified_memory.py | 57 ++++++++++++++--------- 1 file changed, 34 insertions(+), 23 deletions(-) diff --git a/megatron/core/inference/unified_memory.py b/megatron/core/inference/unified_memory.py index e3e9eb8c814..8794fcb6af3 100644 --- a/megatron/core/inference/unified_memory.py +++ b/megatron/core/inference/unified_memory.py @@ -8,13 +8,24 @@ from torch.cuda.memory import CUDAPluggableAllocator from torch.utils.cpp_extension import CUDA_HOME, load_inline +from megatron.core.utils import is_torch_min_version + +try: + if is_torch_min_version("2.8.0"): + from torch.cuda.memory import MemPool + else: + from torch.cuda import MemPool + has_mem_pool = True +except ImportError: + has_mem_pool = False + src = r""" #include #include #define EXPORT extern "C" -EXPORT void* managed_malloc(ssize_t size, int device, void* stream) { +EXPORT void* managed_malloc(size_t size, int device, void* stream) { (void)stream; int cur = -1; cudaGetDevice(&cur); @@ -31,31 +42,31 @@ return ptr; } -EXPORT void managed_free(void* ptr, ssize_t size, int device, void* stream) { +EXPORT void managed_free(void* ptr, size_t size, int device, void* stream) { (void)size; (void)device; (void)stream; if (ptr) cudaFree(ptr); } """ -extra_ldflags = ["-lcudart"] -if CUDA_HOME: - cuda_lib = os.path.join(CUDA_HOME, "lib64") - if os.path.isdir(cuda_lib): - extra_ldflags = [f"-L{cuda_lib}", "-lcudart"] +unified_memory_mempool = None -try: - mod = load_inline( - name="managed_alloc_runtime", - cpp_sources=[src], - functions=[] - with_cuda=True, - extra_ldflags=extra_ldflags, - verbose=False, - ) - - so_path = Path(mod.__file__).as_posix() - alloc = CUDAPluggableAllocator(so_path, "managed_malloc", "managed_free").allocator() - unified_memory_mempool = torch.cuda.MemPool(allocator=alloc) -except RuntimeError: - warnings.warn("Unified memory mempool is not available.") - unified_memory_mempool = None +if has_mem_pool: + extra_ldflags = ["-lcudart"] + if CUDA_HOME: + cuda_lib = os.path.join(CUDA_HOME, "lib64") + if os.path.isdir(cuda_lib): + extra_ldflags = [f"-L{cuda_lib}", "-lcudart"] + try: + mod = load_inline( + name="managed_alloc_runtime", + cpp_sources=[src], + functions=[], + with_cuda=True, + extra_ldflags=extra_ldflags, + verbose=False, + ) + so_path = Path(mod.__file__).as_posix() + alloc = CUDAPluggableAllocator(so_path, "managed_malloc", "managed_free").allocator() + unified_memory_mempool = MemPool(allocator=alloc) + except (RuntimeError, ImportError): + warnings.warn("Unified memory mempool is not available.") From ca2cdd3c16961529c53ce1a2d81f1f968bf99df9 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 9 Sep 2025 12:31:05 -0500 Subject: [PATCH 11/69] lint --- megatron/core/inference/unified_memory.py | 1 - 1 file changed, 1 deletion(-) diff --git a/megatron/core/inference/unified_memory.py b/megatron/core/inference/unified_memory.py index 8794fcb6af3..f5fae1bc5a9 100644 --- a/megatron/core/inference/unified_memory.py +++ b/megatron/core/inference/unified_memory.py @@ -4,7 +4,6 @@ import warnings from pathlib import Path -import torch from torch.cuda.memory import CUDAPluggableAllocator from torch.utils.cpp_extension import CUDA_HOME, load_inline From 074e908e33141f3b0a9351d5377b6abe41b433a3 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 9 Sep 2025 11:52:14 -0700 Subject: [PATCH 12/69] updated prints/warnings. --- megatron/core/inference/contexts/dynamic_context.py | 5 ++++- .../unit_tests/inference/contexts/test_dynamic_context.py | 8 -------- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index d074ef58a56..b8c1828f17b 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -249,7 +249,10 @@ def bytes_to_max_requests_and_tokens(n_bytes): self.max_sequence_length = max_sequence_length self.unified_memory_level = unified_memory_level if unified_memory_level > 0 and unified_memory_mempool is None: - warnings.warn("Unified memory requested but not available; defaulting to GPU memory.") + if torch.distributed.get_rank() == 0: + warnings.warn( + "Unified memory requested but not available; defaulting to GPU memory." + ) self.unified_memory_level = 0 self.total_request_count = 0 diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 552841ad2f6..b9d24cf6673 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -874,11 +874,3 @@ def init_contexts(*, unified_memory_level): # Unified memory test should succeed. init_contexts(unified_memory_level=1) - - -if __name__ == "__main__": - test = TestDynamicContext() - test.test_unified_memory() - test.teardown_method(None) - print("~~~") - print("success.") From d1e9c4dd9956ff177796ed3a088a37d6eca50fa1 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 9 Sep 2025 14:24:55 -0500 Subject: [PATCH 13/69] Pytest skip if unified memory is not available --- tests/unit_tests/inference/contexts/test_dynamic_context.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index b9d24cf6673..46736cda35e 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -836,6 +836,9 @@ def test_calculate_and_store_log_probs(self): current_global_token_offset += expected_len def test_unified_memory(self): + from megatron.core.inference.unified_memory import unified_memory_mempool + if unified_memory_mempool is None: + pytest.skip("Unified memory not available due to bad environment.") self._setup_model_parallel_group(1, 1) From 8187df3062cca88d93c031d22c6ca497db38e761 Mon Sep 17 00:00:00 2001 From: Mcore Bot Date: Tue, 9 Sep 2025 19:28:48 +0000 Subject: [PATCH 14/69] chore: Format files --- tests/unit_tests/inference/contexts/test_dynamic_context.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 46736cda35e..bc304159a5b 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -837,6 +837,7 @@ def test_calculate_and_store_log_probs(self): def test_unified_memory(self): from megatron.core.inference.unified_memory import unified_memory_mempool + if unified_memory_mempool is None: pytest.skip("Unified memory not available due to bad environment.") From 18c589e064c9b2cf3ff71b3acc08c1bc3e392a88 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 9 Sep 2025 16:00:53 -0500 Subject: [PATCH 15/69] Address reviewer comments --- megatron/core/inference/unified_memory.py | 31 +++++++++++++++-------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/megatron/core/inference/unified_memory.py b/megatron/core/inference/unified_memory.py index f5fae1bc5a9..d9240ed99bf 100644 --- a/megatron/core/inference/unified_memory.py +++ b/megatron/core/inference/unified_memory.py @@ -18,7 +18,7 @@ except ImportError: has_mem_pool = False -src = r""" +_mempool_c_src = r""" #include #include @@ -30,18 +30,27 @@ cudaGetDevice(&cur); if (device != cur && device >= 0) cudaSetDevice(device); + // cudaMallocManaged allows for more memory to be allocated than the device memory size. + // The cudaMemAttachGlobal flag makes the memory accessible from both host and device. void* ptr = nullptr; cudaError_t err = cudaMallocManaged(&ptr, (size_t)size, cudaMemAttachGlobal); if (err != cudaSuccess) return nullptr; if (device >= 0) { + // cudaMemAdviseSetPreferredLocation sets the preferred location for the memory. + // This is a hint that tries to prevent data from being migrated away from the device. cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetPreferredLocation, device); + // cudaMemAdviseSetAccessedBy ensures the memory always lives in the device's page table. + // Even if the memory has to be migrated away from the device, it still does not page fault. + // The CUDA docs claim that cudaMemAdviseSetPreferredLocation completely overrides this flag, + // but there is no harm in adding this flag as well for future-proofing. cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetAccessedBy, device); } return ptr; } EXPORT void managed_free(void* ptr, size_t size, int device, void* stream) { + // Memory allocated with cudaMallocManaged should be released with cudaFree. (void)size; (void)device; (void)stream; if (ptr) cudaFree(ptr); } @@ -50,22 +59,22 @@ unified_memory_mempool = None if has_mem_pool: - extra_ldflags = ["-lcudart"] + _extra_ldflags = ["-lcudart"] if CUDA_HOME: - cuda_lib = os.path.join(CUDA_HOME, "lib64") - if os.path.isdir(cuda_lib): - extra_ldflags = [f"-L{cuda_lib}", "-lcudart"] + _cuda_lib = os.path.join(CUDA_HOME, "lib64") + if os.path.isdir(_cuda_lib): + _extra_ldflags = [f"-L{_cuda_lib}", "-lcudart"] try: - mod = load_inline( + _mod = load_inline( name="managed_alloc_runtime", - cpp_sources=[src], + cpp_sources=[_mempool_c_src], functions=[], with_cuda=True, - extra_ldflags=extra_ldflags, + extra_ldflags=_extra_ldflags, verbose=False, ) - so_path = Path(mod.__file__).as_posix() - alloc = CUDAPluggableAllocator(so_path, "managed_malloc", "managed_free").allocator() - unified_memory_mempool = MemPool(allocator=alloc) + _so_path = Path(_mod.__file__).as_posix() + _alloc = CUDAPluggableAllocator(_so_path, "managed_malloc", "managed_free").allocator() + unified_memory_mempool = MemPool(allocator=_alloc) except (RuntimeError, ImportError): warnings.warn("Unified memory mempool is not available.") From 618b3cffc33536b1d87ea4b4249a2d0e1a2df18f Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 22 Sep 2025 02:41:53 -0700 Subject: [PATCH 16/69] tested different compilation setups. --- megatron/core/inference/unified_memory.py | 110 +++++++++++++----- .../contexts/test_dynamic_context.py | 27 +++++ 2 files changed, 107 insertions(+), 30 deletions(-) diff --git a/megatron/core/inference/unified_memory.py b/megatron/core/inference/unified_memory.py index d9240ed99bf..6037838554b 100644 --- a/megatron/core/inference/unified_memory.py +++ b/megatron/core/inference/unified_memory.py @@ -18,6 +18,51 @@ except ImportError: has_mem_pool = False +# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> +# _mempool_c_src = r""" +# #include +# #include + +# #define EXPORT extern "C" + +# EXPORT void* managed_malloc(size_t size, int device, void* stream) { +# // >>> +# printf("******************************* inside managed_malloc.\n"); +# // <<< +# (void)stream; +# int cur = -1; +# cudaGetDevice(&cur); +# if (device != cur && device >= 0) cudaSetDevice(device); + +# // cudaMallocManaged allows for more memory to be allocated than the device memory size. +# // The cudaMemAttachGlobal flag makes the memory accessible from both host and device. +# void* ptr = nullptr; +# cudaError_t err = cudaMallocManaged(&ptr, (size_t)size, cudaMemAttachGlobal); +# if (err != cudaSuccess) return nullptr; + +# if (device >= 0) { +# // cudaMemAdviseSetPreferredLocation sets the preferred location for the memory. +# // This is a hint that tries to prevent data from being migrated away from the device. +# cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetPreferredLocation, device); +# // cudaMemAdviseSetAccessedBy ensures the memory always lives in the device's page table. +# // Even if the memory has to be migrated away from the device, it still does not page fault. +# // The CUDA docs claim that cudaMemAdviseSetPreferredLocation completely overrides this flag, +# // but there is no harm in adding this flag as well for future-proofing. +# cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetAccessedBy, device); +# } +# return ptr; +# } + +# EXPORT void managed_free(void* ptr, size_t size, int device, void* stream) { +# // >>> +# printf("******************************* inside managed_free.\n"); +# // <<< +# // Memory allocated with cudaMallocManaged should be released with cudaFree. +# (void)size; (void)device; (void)stream; +# if (ptr) cudaFree(ptr); +# } +# """ +# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ _mempool_c_src = r""" #include #include @@ -25,36 +70,22 @@ #define EXPORT extern "C" EXPORT void* managed_malloc(size_t size, int device, void* stream) { - (void)stream; - int cur = -1; - cudaGetDevice(&cur); - if (device != cur && device >= 0) cudaSetDevice(device); - - // cudaMallocManaged allows for more memory to be allocated than the device memory size. - // The cudaMemAttachGlobal flag makes the memory accessible from both host and device. + printf("******************************* inside managed_malloc.\n"); void* ptr = nullptr; - cudaError_t err = cudaMallocManaged(&ptr, (size_t)size, cudaMemAttachGlobal); - if (err != cudaSuccess) return nullptr; - - if (device >= 0) { - // cudaMemAdviseSetPreferredLocation sets the preferred location for the memory. - // This is a hint that tries to prevent data from being migrated away from the device. - cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetPreferredLocation, device); - // cudaMemAdviseSetAccessedBy ensures the memory always lives in the device's page table. - // Even if the memory has to be migrated away from the device, it still does not page fault. - // The CUDA docs claim that cudaMemAdviseSetPreferredLocation completely overrides this flag, - // but there is no harm in adding this flag as well for future-proofing. - cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetAccessedBy, device); + cudaError_t err = cudaMalloc(&ptr, size); + if (err != cudaSuccess) { + printf("******************************* ERROR.\n"); + return nullptr; } return ptr; } EXPORT void managed_free(void* ptr, size_t size, int device, void* stream) { - // Memory allocated with cudaMallocManaged should be released with cudaFree. - (void)size; (void)device; (void)stream; + printf("******************************* inside managed_free.\n"); if (ptr) cudaFree(ptr); } """ +# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< unified_memory_mempool = None @@ -65,16 +96,35 @@ if os.path.isdir(_cuda_lib): _extra_ldflags = [f"-L{_cuda_lib}", "-lcudart"] try: - _mod = load_inline( - name="managed_alloc_runtime", - cpp_sources=[_mempool_c_src], - functions=[], - with_cuda=True, - extra_ldflags=_extra_ldflags, - verbose=False, - ) - _so_path = Path(_mod.__file__).as_posix() + # >>> + # _mod = load_inline( + # name="managed_alloc_runtime", + # cpp_sources=[_mempool_c_src], + # functions=[], + # with_cuda=True, + # extra_ldflags=_extra_ldflags, + # verbose=False, + # ) + # _so_path = Path(_mod.__file__).as_posix() + # +++ + _so_path = "/lustre/fsw/portfolios/adlr/users/lmcafee/inference/megatrons/unified-memory/scripts_ignore_me/c/lib.so" + # <<< + # >>> _alloc = CUDAPluggableAllocator(_so_path, "managed_malloc", "managed_free").allocator() unified_memory_mempool = MemPool(allocator=_alloc) + # +++ + # _alloc = CUDAPluggableAllocator(_so_path, "managed_malloc", "managed_free") + # # unified_memory_mempool = MemPool(allocator=_alloc.allocator()) + # unified_memory_mempool = None + # +++ + # _alloc = CUDAPluggableAllocator(_so_path, "managed_malloc", "managed_free") + # unified_memory_mempool = MemPool(_alloc._allocator) + # <<< except (RuntimeError, ImportError): warnings.warn("Unified memory mempool is not available.") + +# >>> +def delete_mempool(): + global unified_memory_mempool + del unified_memory_mempool +# <<< diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index bc304159a5b..8d951fe4a1b 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -867,6 +867,12 @@ def init_contexts(*, unified_memory_level): unified_memory_level=unified_memory_level, ) ) + # >>> + # for context in contexts: + # print("++++ me, context.") + # del context.memory_buffer + # del context + # <<< # Pure GPU memory test should OOM. try: @@ -878,3 +884,24 @@ def init_contexts(*, unified_memory_level): # Unified memory test should succeed. init_contexts(unified_memory_level=1) + + # >>> + # # Destroy process group (otherwise, segfault will occur). + # torch.distributed.destroy_process_group() + # <<< + + # >>> + # torch.cuda.empty_cache() + # <<< + + # >>> + # from megatron.core.inference.unified_memory import delete_mempool + # delete_mempool() + # <<< + +# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> +if __name__ == "__main__": + test = TestDynamicContext() + test.test_unified_memory() + print("~~~\nsuccess.") +# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< From ea486976ec78010d01fc9ac5435f94d49bd5dcf0 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 22 Sep 2025 02:45:12 -0700 Subject: [PATCH 17/69] clean up. --- megatron/core/inference/unified_memory.py | 106 ++++++------------ .../contexts/test_dynamic_context.py | 27 ----- 2 files changed, 34 insertions(+), 99 deletions(-) diff --git a/megatron/core/inference/unified_memory.py b/megatron/core/inference/unified_memory.py index 6037838554b..8c6c335def8 100644 --- a/megatron/core/inference/unified_memory.py +++ b/megatron/core/inference/unified_memory.py @@ -18,51 +18,6 @@ except ImportError: has_mem_pool = False -# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> -# _mempool_c_src = r""" -# #include -# #include - -# #define EXPORT extern "C" - -# EXPORT void* managed_malloc(size_t size, int device, void* stream) { -# // >>> -# printf("******************************* inside managed_malloc.\n"); -# // <<< -# (void)stream; -# int cur = -1; -# cudaGetDevice(&cur); -# if (device != cur && device >= 0) cudaSetDevice(device); - -# // cudaMallocManaged allows for more memory to be allocated than the device memory size. -# // The cudaMemAttachGlobal flag makes the memory accessible from both host and device. -# void* ptr = nullptr; -# cudaError_t err = cudaMallocManaged(&ptr, (size_t)size, cudaMemAttachGlobal); -# if (err != cudaSuccess) return nullptr; - -# if (device >= 0) { -# // cudaMemAdviseSetPreferredLocation sets the preferred location for the memory. -# // This is a hint that tries to prevent data from being migrated away from the device. -# cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetPreferredLocation, device); -# // cudaMemAdviseSetAccessedBy ensures the memory always lives in the device's page table. -# // Even if the memory has to be migrated away from the device, it still does not page fault. -# // The CUDA docs claim that cudaMemAdviseSetPreferredLocation completely overrides this flag, -# // but there is no harm in adding this flag as well for future-proofing. -# cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetAccessedBy, device); -# } -# return ptr; -# } - -# EXPORT void managed_free(void* ptr, size_t size, int device, void* stream) { -# // >>> -# printf("******************************* inside managed_free.\n"); -# // <<< -# // Memory allocated with cudaMallocManaged should be released with cudaFree. -# (void)size; (void)device; (void)stream; -# if (ptr) cudaFree(ptr); -# } -# """ -# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ _mempool_c_src = r""" #include #include @@ -70,22 +25,42 @@ #define EXPORT extern "C" EXPORT void* managed_malloc(size_t size, int device, void* stream) { + // >>> printf("******************************* inside managed_malloc.\n"); + // <<< + (void)stream; + int cur = -1; + cudaGetDevice(&cur); + if (device != cur && device >= 0) cudaSetDevice(device); + + // cudaMallocManaged allows for more memory to be allocated than the device memory size. + // The cudaMemAttachGlobal flag makes the memory accessible from both host and device. void* ptr = nullptr; - cudaError_t err = cudaMalloc(&ptr, size); - if (err != cudaSuccess) { - printf("******************************* ERROR.\n"); - return nullptr; + cudaError_t err = cudaMallocManaged(&ptr, (size_t)size, cudaMemAttachGlobal); + if (err != cudaSuccess) return nullptr; + + if (device >= 0) { + // cudaMemAdviseSetPreferredLocation sets the preferred location for the memory. + // This is a hint that tries to prevent data from being migrated away from the device. + cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetPreferredLocation, device); + // cudaMemAdviseSetAccessedBy ensures the memory always lives in the device's page table. + // Even if the memory has to be migrated away from the device, it still does not page fault. + // The CUDA docs claim that cudaMemAdviseSetPreferredLocation completely overrides this flag, + // but there is no harm in adding this flag as well for future-proofing. + cudaMemAdvise(ptr, (size_t)size, cudaMemAdviseSetAccessedBy, device); } return ptr; } EXPORT void managed_free(void* ptr, size_t size, int device, void* stream) { + // >>> printf("******************************* inside managed_free.\n"); + // <<< + // Memory allocated with cudaMallocManaged should be released with cudaFree. + (void)size; (void)device; (void)stream; if (ptr) cudaFree(ptr); } """ -# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< unified_memory_mempool = None @@ -96,30 +71,17 @@ if os.path.isdir(_cuda_lib): _extra_ldflags = [f"-L{_cuda_lib}", "-lcudart"] try: - # >>> - # _mod = load_inline( - # name="managed_alloc_runtime", - # cpp_sources=[_mempool_c_src], - # functions=[], - # with_cuda=True, - # extra_ldflags=_extra_ldflags, - # verbose=False, - # ) - # _so_path = Path(_mod.__file__).as_posix() - # +++ - _so_path = "/lustre/fsw/portfolios/adlr/users/lmcafee/inference/megatrons/unified-memory/scripts_ignore_me/c/lib.so" - # <<< - # >>> + _mod = load_inline( + name="managed_alloc_runtime", + cpp_sources=[_mempool_c_src], + functions=[], + with_cuda=True, + extra_ldflags=_extra_ldflags, + verbose=False, + ) + _so_path = Path(_mod.__file__).as_posix() _alloc = CUDAPluggableAllocator(_so_path, "managed_malloc", "managed_free").allocator() unified_memory_mempool = MemPool(allocator=_alloc) - # +++ - # _alloc = CUDAPluggableAllocator(_so_path, "managed_malloc", "managed_free") - # # unified_memory_mempool = MemPool(allocator=_alloc.allocator()) - # unified_memory_mempool = None - # +++ - # _alloc = CUDAPluggableAllocator(_so_path, "managed_malloc", "managed_free") - # unified_memory_mempool = MemPool(_alloc._allocator) - # <<< except (RuntimeError, ImportError): warnings.warn("Unified memory mempool is not available.") diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 8d951fe4a1b..bc304159a5b 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -867,12 +867,6 @@ def init_contexts(*, unified_memory_level): unified_memory_level=unified_memory_level, ) ) - # >>> - # for context in contexts: - # print("++++ me, context.") - # del context.memory_buffer - # del context - # <<< # Pure GPU memory test should OOM. try: @@ -884,24 +878,3 @@ def init_contexts(*, unified_memory_level): # Unified memory test should succeed. init_contexts(unified_memory_level=1) - - # >>> - # # Destroy process group (otherwise, segfault will occur). - # torch.distributed.destroy_process_group() - # <<< - - # >>> - # torch.cuda.empty_cache() - # <<< - - # >>> - # from megatron.core.inference.unified_memory import delete_mempool - # delete_mempool() - # <<< - -# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> -if __name__ == "__main__": - test = TestDynamicContext() - test.test_unified_memory() - print("~~~\nsuccess.") -# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< From 2eed5a85ed08e92e6acb4399421e207d4c52d614 Mon Sep 17 00:00:00 2001 From: Mcore Bot Date: Mon, 22 Sep 2025 09:49:31 +0000 Subject: [PATCH 18/69] chore: Format files --- megatron/core/inference/unified_memory.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/megatron/core/inference/unified_memory.py b/megatron/core/inference/unified_memory.py index 8c6c335def8..1268a9c5a7b 100644 --- a/megatron/core/inference/unified_memory.py +++ b/megatron/core/inference/unified_memory.py @@ -85,8 +85,11 @@ except (RuntimeError, ImportError): warnings.warn("Unified memory mempool is not available.") + # >>> def delete_mempool(): global unified_memory_mempool del unified_memory_mempool + + # <<< From b14c607ebcc1d18ed75ff1c87e075d879bd9a4d6 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Mon, 22 Sep 2025 06:14:17 -0500 Subject: [PATCH 19/69] Import a mempool-builder instead of global mempool --- .../inference/contexts/dynamic_context.py | 12 ++++---- megatron/core/inference/unified_memory.py | 28 ++++++++++--------- 2 files changed, 22 insertions(+), 18 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index cb30de48515..c090ba43aca 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -12,7 +12,7 @@ from torch import Tensor from megatron.core import parallel_state -from megatron.core.inference.unified_memory import unified_memory_mempool +from megatron.core.inference.unified_memory import create_unified_mempool, has_unified_memory from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb from megatron.core.package_info import __version__ as mcore_version from megatron.core.transformer import TransformerConfig @@ -253,12 +253,14 @@ def bytes_to_max_requests_and_tokens(n_bytes): self.num_layers = num_layers self.max_sequence_length = max_sequence_length self.unified_memory_level = unified_memory_level - if unified_memory_level > 0 and unified_memory_mempool is None: - if torch.distributed.get_rank() == 0: + if unified_memory_level > 0: + if not has_unified_memory and torch.distributed.get_rank() == 0: warnings.warn( "Unified memory requested but not available; defaulting to GPU memory." ) - self.unified_memory_level = 0 + self.unified_memory_level = 0 + else: + self.unified_memory_mempool = create_unified_mempool() self.total_request_count = 0 self.active_token_count = 0 @@ -300,7 +302,7 @@ def bytes_to_max_requests_and_tokens(n_bytes): # Memory buffer. ctx_manager = ( - torch.cuda.use_mem_pool(unified_memory_mempool) + torch.cuda.use_mem_pool(self.unified_memory_mempool) if self.unified_memory_level > 0 else nullcontext() ) diff --git a/megatron/core/inference/unified_memory.py b/megatron/core/inference/unified_memory.py index 1268a9c5a7b..c4b34a16f67 100644 --- a/megatron/core/inference/unified_memory.py +++ b/megatron/core/inference/unified_memory.py @@ -5,6 +5,7 @@ from pathlib import Path from torch.cuda.memory import CUDAPluggableAllocator +from torch.distributed import get_rank from torch.utils.cpp_extension import CUDA_HOME, load_inline from megatron.core.utils import is_torch_min_version @@ -14,9 +15,9 @@ from torch.cuda.memory import MemPool else: from torch.cuda import MemPool - has_mem_pool = True + _has_mem_pool = True except ImportError: - has_mem_pool = False + _has_mem_pool = False _mempool_c_src = r""" #include @@ -62,9 +63,12 @@ } """ -unified_memory_mempool = None +# Avoid linting errors. +has_unified_memory = False +_alloc = None -if has_mem_pool: +# Build the .so upon import; this avoids issues. +if _has_mem_pool: _extra_ldflags = ["-lcudart"] if CUDA_HOME: _cuda_lib = os.path.join(CUDA_HOME, "lib64") @@ -81,15 +85,13 @@ ) _so_path = Path(_mod.__file__).as_posix() _alloc = CUDAPluggableAllocator(_so_path, "managed_malloc", "managed_free").allocator() - unified_memory_mempool = MemPool(allocator=_alloc) + has_unified_memory = True except (RuntimeError, ImportError): - warnings.warn("Unified memory mempool is not available.") + if get_rank() == 0: + warnings.warn("Failed to create unified memory mempool.") -# >>> -def delete_mempool(): - global unified_memory_mempool - del unified_memory_mempool - - -# <<< +def create_unified_mempool(): + """Create a unified memory mempool using CUDA managed memory.""" + assert has_unified_memory + return MemPool(allocator=_alloc) From df8ca830ccde18c6ad8e5f90d77f6a4b791833ea Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Mon, 22 Sep 2025 09:56:06 -0500 Subject: [PATCH 20/69] Fix import issue --- megatron/core/inference/unified_memory.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/megatron/core/inference/unified_memory.py b/megatron/core/inference/unified_memory.py index c4b34a16f67..9331fab2fc2 100644 --- a/megatron/core/inference/unified_memory.py +++ b/megatron/core/inference/unified_memory.py @@ -5,7 +5,6 @@ from pathlib import Path from torch.cuda.memory import CUDAPluggableAllocator -from torch.distributed import get_rank from torch.utils.cpp_extension import CUDA_HOME, load_inline from megatron.core.utils import is_torch_min_version @@ -87,8 +86,7 @@ _alloc = CUDAPluggableAllocator(_so_path, "managed_malloc", "managed_free").allocator() has_unified_memory = True except (RuntimeError, ImportError): - if get_rank() == 0: - warnings.warn("Failed to create unified memory mempool.") + warnings.warn("Failed to create unified memory mempool.") def create_unified_mempool(): From ff43bf416f4d2b44395812fa0cf316dc447d0d2d Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Mon, 22 Sep 2025 10:25:09 -0500 Subject: [PATCH 21/69] Clean-up, get ready for merge --- megatron/core/inference/unified_memory.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/megatron/core/inference/unified_memory.py b/megatron/core/inference/unified_memory.py index 9331fab2fc2..0a2ff13f112 100644 --- a/megatron/core/inference/unified_memory.py +++ b/megatron/core/inference/unified_memory.py @@ -25,9 +25,6 @@ #define EXPORT extern "C" EXPORT void* managed_malloc(size_t size, int device, void* stream) { - // >>> - printf("******************************* inside managed_malloc.\n"); - // <<< (void)stream; int cur = -1; cudaGetDevice(&cur); @@ -53,9 +50,6 @@ } EXPORT void managed_free(void* ptr, size_t size, int device, void* stream) { - // >>> - printf("******************************* inside managed_free.\n"); - // <<< // Memory allocated with cudaMallocManaged should be released with cudaFree. (void)size; (void)device; (void)stream; if (ptr) cudaFree(ptr); From 6922c76c21ef3266f68a3bf76b62149537b568a6 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Mon, 22 Sep 2025 11:05:36 -0500 Subject: [PATCH 22/69] Fix tests --- tests/unit_tests/inference/contexts/test_dynamic_context.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index bc304159a5b..2fd8db210a6 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -836,9 +836,9 @@ def test_calculate_and_store_log_probs(self): current_global_token_offset += expected_len def test_unified_memory(self): - from megatron.core.inference.unified_memory import unified_memory_mempool + from megatron.core.inference.unified_memory import has_unified_memory - if unified_memory_mempool is None: + if not has_unified_memory: pytest.skip("Unified memory not available due to bad environment.") self._setup_model_parallel_group(1, 1) From c05eb2a2dd8bd51b455060fff65f21c019ff28ea Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 30 Sep 2025 06:28:01 -0700 Subject: [PATCH 23/69] sep 30 alpha. --- .../inference/gpt/gpt_dynamic_inference.py | 163 +++++++++++++++++- .../gpt/gpt_dynamic_inference_12b.sh | 17 +- .../gpt/gpt_dynamic_inference_357m.sh | 17 +- .../inference/contexts/dynamic_context.py | 143 ++++++--------- megatron/training/arguments.py | 81 +++++---- 5 files changed, 275 insertions(+), 146 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index cfadb8a3c9e..f9046eed133 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -61,6 +61,14 @@ torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunState]) torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunDiagnostic]) +# >>> +from lutil import pax as _pax, get_mem_stats_str +import builtins +builtins.pax = _pax + +def print_mem(key): + print(f"+++++++++++ {key} ........... {get_mem_stats_str()}.") +# <<< def add_dynamic_inference_args(parser: ArgumentParser) -> ArgumentParser: @@ -113,11 +121,120 @@ 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_active_buffer_size_bytes(model: torch.nn.Module) -> int: + """Estimate memory remaining for context's `memory_buffer`. + + The following implementation is a placeholder, using only + `args.inference_dynamic_batching_active_buffer_size_gb`. Eventually, this will + be computed by first computing memory used by the model, activations, and the + KV cache. + """ + args = get_args() + active_buffer_size_gb = args.inference_dynamic_batching_active_buffer_size_gb + active_buffer_size_bytes = int(active_buffer_size_gb * 1024**3) + # >>> + # pax({ + # "active buffer gb" : get_args().inference_dynamic_batching_active_buffer_size_gb, + # "1024**3" : 1024**3, + # }, "active_buffer_size_bytes") + # <<< + return active_buffer_size_bytes + + +# # >>> +# def get_inference_context_max_tokens(model: torch.nn.Module) -> int: +# # def get_inference_context_max_tokens(controller: torch.nn.Module) -> int: +# # <<< +# """Estimate `max_tokens` for the context. + +# This function performs a binary search over the number of input tokens and +# analyzes the latency to find the 'knee of the curve'. +# """ + +# import time + +# # args = get_args() +# # model = GPTInferenceWrapper(model, args, inference_context="not-actually-set") # None) +# # pax("model") + +# num_tokens_step = 1000 +# num_tokens_min = num_tokens_step +# num_tokens_max = 24000 # 1000000 +# time_map = {} +# tbar = tqdm(range(num_tokens_min, num_tokens_max, num_tokens_step)) +# # tbar = tqdm((3000,)) +# # tbar = tqdm((3072,)) +# # tbar = tqdm((1000,)) +# # num_tokens_list = list( +# for num_tokens in tbar: +# tbar.set_description(f"num_tokens {num_tokens}") +# input_ids = torch.full( +# (num_tokens,), +# 0, +# dtype=torch.long, +# device=torch.cuda.current_device(), +# ).unsqueeze(0) +# # position_ids = torch.arange( +# # num_tokens, +# position_ids = torch.zeros( +# (num_tokens,), +# dtype=torch.long, +# device=torch.cuda.current_device(), +# ).unsqueeze(0) +# # pax("input_ids, position_ids") +# # result = model.run_one_forward_step( +# # { +# # "tokens": input_ids, +# # "position_ids": position_ids, +# # "attention_mask": None, +# # } +# # ) +# t = time.time() +# for _ in range(3): +# result = model( +# input_ids, +# position_ids, +# attention_mask=None, +# inference_context=None, +# runtime_gather_output=True, # Inference should always gather the logits +# ) +# t = time.time() - t +# time_map[num_tokens] = t +# # >>> +# # pax("result") +# # <<< +# print("~~ x ~~") +# [ print(x) for x in time_map.keys() ] +# print("~~ y ~~") +# [ print(y) for y in time_map.values() ] +# pax() +# pax("time_map") +# return max_tokens +def get_inference_context_max_tokens(model: torch.nn.Module) -> int: + """Estimate `max_tokens` for the context. + + The following implementation is a placeholder, simply returning a hard-coded + value for `max_tokens`. Eventually, this will be computed by analyzing the + correlation between `max_tokens` and request throughput. + """ + return 16384 +# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< + + +def get_inference_context( + requests: List[Request], + sampling_params: SamplingParams, + # >>> + # buffer_size_bytes: int, + model: torch.nn.Module, + # <<< + calculate_max_sequence_length_from_requests: bool =True, +): """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 @@ -126,6 +243,14 @@ def get_inference_context(requests: List[Request], sampling_params: SamplingPara else: max_sequence_length = args.inference_max_seq_length + # Active buffer size & max tokens. + active_buffer_size_bytes = get_inference_context_active_buffer_size_bytes(model) + max_tokens = get_inference_context_max_tokens(model) + + # >>> + # pax("active_buffer_size_bytes, max_tokens") + # <<< + # Inference context. context = DynamicInferenceContext( params_dtype=args.params_dtype, @@ -139,17 +264,28 @@ def get_inference_context(requests: List[Request], sampling_params: SamplingPara args.inference_dynamic_batching_num_cuda_graphs if args.enable_cuda_graph else None ), chunk_size_tokens=args.inference_dynamic_batching_chunk_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, + # >>> + # 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, + # +++ + # active_buffer_size_gb=active_buffer_size_gb, + active_buffer_size_bytes=active_buffer_size_bytes, + # model=model, + # <<< + # >>> + # max_tokens_override=args.inference_dynamic_batching_max_tokens_override, + max_tokens=max_tokens, + # <<< 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, kv_lora_rank=args.kv_lora_rank if args.multi_latent_attention else None, qk_pos_emb_head_dim=args.qk_pos_emb_head_dim, - unified_memory_level=args.inference_dynamic_batching_unified_memory_level, + # >>> + # unified_memory_level=args.inference_dynamic_batching_unified_memory_level, + # <<< ) return context @@ -320,6 +456,10 @@ def main(): args_defaults={'no_load_rng': True, 'no_load_optim': True}, ) + # >>> + print_mem("initialize_megatron()") + # <<< + # Start Nsight profiler. if os.environ.get("NSIGHT_PREFIX"): torch.cuda.cudart().cudaProfilerStart() @@ -342,7 +482,7 @@ def main(): # Requests, context, conroller. model = get_model() requests = build_requests(args, tokenizer) - context = get_inference_context(requests, sampling_params) + context = get_inference_context(requests, sampling_params, model) controller = get_inference_controller(model, context) # Validate all context_length's <= max_tokens. @@ -365,6 +505,11 @@ def main(): track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events, ) + # >>> + print_mem("engine") + exit() + # <<< + setup_prefix = build_dynamic_engine_setup_prefix(args, model, context, requests) print("~~~") print(setup_prefix) diff --git a/examples/inference/gpt/gpt_dynamic_inference_12b.sh b/examples/inference/gpt/gpt_dynamic_inference_12b.sh index b469adadea9..5ae70f22c89 100644 --- a/examples/inference/gpt/gpt_dynamic_inference_12b.sh +++ b/examples/inference/gpt/gpt_dynamic_inference_12b.sh @@ -23,9 +23,16 @@ 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} +# >>> +# : ${BUFFER_SIZE_GB=50.} +# : ${BUFFER_OVERFLOW_FACTOR=1.} +# : ${BUFFER_GUARANTEED_FRACTION=0.05} +# --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} \ +# +++ +: ${ACTIVE_BUFFER_SIZE_GB=50.} +# <<< # Cuda graphs. : ${ENABLE_CUDA_GRAPHS=1} @@ -77,9 +84,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 bbe60614f69..0383405257d 100644 --- a/examples/inference/gpt/gpt_dynamic_inference_357m.sh +++ b/examples/inference/gpt/gpt_dynamic_inference_357m.sh @@ -24,9 +24,16 @@ 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} +# >>> +# : ${BUFFER_SIZE_GB=50.} +# : ${BUFFER_OVERFLOW_FACTOR=1.} +# : ${BUFFER_GUARANTEED_FRACTION=0.05} +# --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} \ +# +++ +: ${ACTIVE_BUFFER_SIZE_GB=50.} +# <<< # Cuda graphs. : ${ENABLE_CUDA_GRAPHS=1} @@ -63,9 +70,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/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index c090ba43aca..a885e871cb4 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -129,21 +129,17 @@ 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_bytes (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. + # >>> + # model (torch.nn.Module): Model used for inference. This is needed for + # computing `max_tokens` below. + # <<< chunk_size_tokens (int): Size of KV cache chunk 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 @@ -153,10 +149,6 @@ class DynamicInferenceContext(BaseInferenceContext): are materialized in the context. use_cuda_graphs_for_non_decode_steps (bool): If True, use cuda graphs for non-decode engine steps. - unified_memory_level (Optional[int]): Set unified memory usage within the - dynamic inference context. The levels are: 0) no unified memory, 1) - allocate `memory_buffer` in unified memory. Eventually, additional - levels will be included to control other tensors within the context. """ def __init__( @@ -167,12 +159,12 @@ def __init__( kv_channels: int, num_attention_heads: int, max_sequence_length: int, - buffer_size_gb: float, - buffer_guaranteed_fraction: float, + active_buffer_size_bytes: float, + # >>> + # model: torch.nn.Module, + max_tokens: int, + # <<< chunk_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, @@ -180,7 +172,6 @@ def __init__( num_cuda_graphs: Optional[int] = None, materialize_only_last_token_logits: bool = True, use_cuda_graphs_for_non_decode_steps: bool = True, - unified_memory_level: Optional[int] = 0, ): super().__init__(materialize_only_last_token_logits=materialize_only_last_token_logits) @@ -218,49 +209,26 @@ def __init__( * hidden_size_per_attention_head ) - # Adjust buffer to be a multiple of chunk size. - buffer_size_bytes = int(buffer_size_gb * 1024**3) - buffer_size_bytes_rem = buffer_size_bytes % self.chunk_size_bytes - buffer_size_bytes = buffer_size_bytes - buffer_size_bytes_rem - - # Compute max_requets, max_tokens from buffer size and overflow factor. - def bytes_to_max_requests_and_tokens(n_bytes): - n_tokens = n_bytes / self.chunk_size_bytes * self.chunk_size_tokens - n_requests = n_tokens / max_sequence_length - return self.round_up_requests(int(n_requests), tp_size=tp_size), self.round_up_tokens( - int(n_tokens), tp_size=tp_size - ) - - 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 - ) + # Calculate the number of active chunks, and finalize the buffer size. + active_chunk_count_total = active_buffer_size_bytes // self.chunk_size_bytes + active_buffer_size_bytes = active_chunk_count_total * self.chunk_size_bytes + self.max_requests = active_chunk_count_total + self.max_tokens = max_tokens - if max_requests_override is not None: - self.max_requests = 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 chunk allocator + self.chunk_allocator = ChunkAllocator(active_count_total=active_chunk_count_total) + # >>> + pax({"chunk_allocator": self.chunk_allocator}) + # <<< # Initialize context state. self.params_dtype = params_dtype self.num_layers = num_layers self.max_sequence_length = max_sequence_length - self.unified_memory_level = unified_memory_level - if unified_memory_level > 0: - if not has_unified_memory and torch.distributed.get_rank() == 0: - warnings.warn( - "Unified memory requested but not available; defaulting to GPU memory." - ) - self.unified_memory_level = 0 - else: - self.unified_memory_mempool = create_unified_mempool() + assert has_unified_memory, ( + "CUDA unified memory must be available to use `DynamicInferenceContext`." + ) + self.unified_memory_mempool = create_unified_mempool() self.total_request_count = 0 self.active_token_count = 0 @@ -297,16 +265,8 @@ 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_chunk = torch.empty_like(self.token_to_input_ids) - # Calculate the total number of chunks available in the buffer - chunk_count_total = buffer_size_bytes // self.chunk_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: + with torch.cuda.use_mem_pool(self.unified_memory_mempool): if cache_mla_latent: self.memory_buffer = torch.full( (self.num_layers, chunk_count_total, self.chunk_size_tokens, kv_reduced_dim), @@ -403,29 +363,26 @@ def bytes_to_max_requests_and_tokens(n_bytes): device=torch.cuda.current_device(), ) - # Guaranteed active requests. - # * See details in the class docstring above. `gtd_request_fraction` is - # the fraction of chunks in the memory buffer that are reserved for - # guaranteeing that some number of active requests can always proceed - # with their generations. The number of chunks defined by - # `buffer_guaranteed_fraction * chunk_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 chunks rather - # than bytes is safer due to the non-linear impacts of a large - # `chunk_size_tokens` or `max_kv_chunk_count`. When computing from - # chunks, this space will always be less than `chunk_count_total`. When - # computing from bytes, this space can unexpectedly be much larger than - # `chunk_count_total`, resulting in stalled generations. - gtd_chunk_count = int(buffer_guaranteed_fraction * chunk_count_total) - gtd_chunk_count = min(gtd_chunk_count, chunk_count_total) - self.gtd_request_count = max(1, gtd_chunk_count // self.max_kv_chunk_count) - self.gtd_chunk_count = self.gtd_request_count * self.max_kv_chunk_count - - # Initialize chunk allocator - self.chunk_allocator = ChunkAllocator( - chunk_count_total=chunk_count_total, gtd_chunk_count=self.gtd_chunk_count - ) + # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> + # # Guaranteed active requests. + # # * See details in the class docstring above. `gtd_request_fraction` is + # # the fraction of chunks in the memory buffer that are reserved for + # # guaranteeing that some number of active requests can always proceed + # # with their generations. The number of chunks defined by + # # `buffer_guaranteed_fraction * chunk_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 chunks rather + # # than bytes is safer due to the non-linear impacts of a large + # # `chunk_size_tokens` or `max_kv_chunk_count`. When computing from + # # chunks, this space will always be less than `chunk_count_total`. When + # # computing from bytes, this space can unexpectedly be much larger than + # # `chunk_count_total`, resulting in stalled generations. + # gtd_chunk_count = int(buffer_guaranteed_fraction * chunk_count_total) + # gtd_chunk_count = min(gtd_chunk_count, chunk_count_total) + # self.gtd_request_count = max(1, gtd_chunk_count // self.max_kv_chunk_count) + # self.gtd_chunk_count = self.gtd_request_count * self.max_kv_chunk_count + # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< # Store the dummy chunk idx reference for convenience self.dummy_chunk_idx = self.chunk_allocator.dummy_chunk_idx diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 0c6cd68fe83..97305f45e9d 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1103,9 +1103,13 @@ def validate_args(args, defaults={}): assert not args.enable_cuda_graph, "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_buffer_size_gb is not None + # <<< assert args.inference_dynamic_batching_chunk_size % 256 == 0, "chunk size should be a multiple of 256" - assert args.inference_dynamic_batching_buffer_guaranteed_fraction is not None + # >>> + # assert args.inference_dynamic_batching_buffer_guaranteed_fraction is not None + # <<< # MoE upcycling check if args.moe_use_upcycling: @@ -1392,36 +1396,47 @@ 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-buffer-size-gb', + # type=float, default=40., + # help='Total buffer size (GB) allocated for the chunked KV ' + # 'memory.') + # +++ + group.add_argument('--inference-dynamic-batching-active-buffer-size-gb', type=float, default=40., help='Total buffer size (GB) allocated for the chunked KV ' 'memory.') + # <<< group.add_argument('--inference-dynamic-batching-chunk-size', type=int, default=256, help='KV cache chunk 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 blocked ' - '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', - 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`.') + # >>> + # 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 blocked ' + # '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', + # 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`.') + # <<< group.add_argument('--inference-dynamic-batching-num-cuda-graphs', type=int, default=16, help='Maximum number of cuda graphs to capture, where the ' @@ -1434,13 +1449,15 @@ def _add_inference_args(parser): help='Track paused request ids by adding \'paused\' events ' 'to each request\'s event history. This has a very minor ' 'impact on latency.') - group.add_argument('--inference-dynamic-batching-unified-memory-level', - type=int, default=0, 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. ' - 'Eventually, additional levels will be included to ' - 'control other tensors within the context.') + # >>> + # group.add_argument('--inference-dynamic-batching-unified-memory-level', + # type=int, default=0, 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. ' + # 'Eventually, additional levels will be included to ' + # 'control other tensors within the context.') + # <<< group.add_argument('--symmetric-ar-type', type=str, default=None, choices=['two_shot', "one_shot", "multimem_all_reduce", None], help='What type of symmetric all reduce to use. The default is none which is no use of symetric memory') From eba0a47055b3991b0bddbf658c43fffd69454058 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Wed, 1 Oct 2025 11:28:10 -0700 Subject: [PATCH 24/69] oct 1 alpha. --- .../inference/gpt/gpt_dynamic_inference.py | 2 +- examples/inference/gpt/utils.py | 36 +++++++--- .../contexts/dynamic_chunk_allocator.py | 72 +++++++++++-------- .../inference/contexts/dynamic_context.py | 44 ++++++++---- .../core/inference/engines/dynamic_engine.py | 3 +- 5 files changed, 100 insertions(+), 57 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index f9046eed133..b37ab50aa87 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -507,7 +507,7 @@ def main(): # >>> print_mem("engine") - exit() + # exit() # <<< setup_prefix = build_dynamic_engine_setup_prefix(args, model, context, requests) diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py index 8528970060c..9cb02bf92c2 100644 --- a/examples/inference/gpt/utils.py +++ b/examples/inference/gpt/utils.py @@ -344,27 +344,43 @@ def build_dynamic_engine_setup_prefix( f"r/step {args.incoming_requests_per_step}" ) + # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> + # # 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}]" + # ) + + # parts = [ + # get_model_size_str(model), + # "dynamic", + # cg_str, + # request_str, + # buffer_limits_str, + # guaranteed_fraction_str, + # ] + # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ # 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"bf {args.inference_dynamic_batching_active_buffer_size_gb:.0f} " 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}]" - ) - parts = [ get_model_size_str(model), "dynamic", cg_str, request_str, buffer_limits_str, - guaranteed_fraction_str, ] + # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< return " | ".join(parts) diff --git a/megatron/core/inference/contexts/dynamic_chunk_allocator.py b/megatron/core/inference/contexts/dynamic_chunk_allocator.py index cbc1127dd4b..d125c8d6dc4 100644 --- a/megatron/core/inference/contexts/dynamic_chunk_allocator.py +++ b/megatron/core/inference/contexts/dynamic_chunk_allocator.py @@ -13,60 +13,70 @@ class ChunkAllocator: - Initializing a pool of chunk IDs - Allocating chunks from the pool - Releasing chunks back to the pool - - Managing the guaranteed chunk count for active requests Args: - chunk_count_total (int): Total number of chunks available in the buffer. - gtd_chunk_count (int): Number of chunks reserved for guaranteed requests. + active_count (int): Total number of active chunks 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, chunk_count_total: int, gtd_chunk_count: int): - self.chunk_count_total = chunk_count_total - self.gtd_chunk_count = gtd_chunk_count - - # Reserve last chunk ID as dummy chunk for decode-only inference steps - self.chunk_count_avail = self.chunk_count_total - 1 - self.dummy_chunk_idx = self.chunk_count_total - 1 + # def __init__(self, active_count: int): + + # active_count -= 1 # subtract 1 for dummy_chunk_idx (see below) + # self.total_count = 2 * active_count + # self.active_count = active_count + # self.paused_count = active_count + # # >>> + # # self.active_avail = active_count - 1 # reserve last chunk for decode-only steps + # self.active_avail = active_count + # # <<< + # self.paused_avail = active_count + # self.dummy_chunk_idx = self.total_count - 1 + + # # Initialize chunk pool as a "stack" data structure + # self.chunk_bag = torch.arange( + # self.total_count, + # dtype=torch.int32, + # device=torch.cuda.current_device(), + # ) + def __init__(self, active_count: int): + + active_count -= 1 # -1 for dummy_chunk_idx (see below) + self.total_count = 2 * active_count + 1 + self.total_avail = self.total_count - 1 # -1 for dummy_chunk_idx + self.active_count = active_count + self.dummy_chunk_idx = self.total_count - 1 # Initialize chunk pool as a "stack" data structure self.chunk_bag = torch.arange( - self.chunk_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_chunks: int, safe: bool = False) -> bool: + def is_memory_available(self, num_chunks: int) -> bool: """Check if memory chunks are available. - Use 'safe' to avoid all requests being blocked. 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. - Args: num_chunks (int): Number of chunks to check. - safe (bool): Include extra space for guaranteeing ability to run - requests to completion. Return: (bool) Is memory available? """ - if safe: - return self.chunk_count_avail >= num_chunks + self.gtd_chunk_count - else: - return self.chunk_count_avail >= num_chunks + return self.total_avail >= num_chunks - def allocate_memory_chunks(self, num_chunks: int = 1, safe: bool = False) -> Optional[Tensor]: + def allocate_memory_chunks(self, num_chunks: int = 1) -> Optional[Tensor]: """Allocate memory chunks if available, else return None. Args: num_chunks (int): Number of chunks to allocate. - safe (bool): Include extra space for guaranteeing ability to run - requests to completion. Return: (Optional[Tensor]) Allocated chunk IDs. """ - if self.is_memory_available(num_chunks, safe): - self.chunk_count_avail -= num_chunks - return self.chunk_bag[self.chunk_count_avail : (self.chunk_count_avail + num_chunks)] + if self.is_memory_available(num_chunks): + self.total_avail -= num_chunks + return self.chunk_bag[self.total_avail : (self.total_avail + num_chunks)] else: return None @@ -80,8 +90,8 @@ def release_memory_chunks(self, chunks: Tensor) -> None: None """ num_chunks = chunks.size(dim=0) - self.chunk_bag[self.chunk_count_avail : (self.chunk_count_avail + num_chunks)] = chunks - self.chunk_count_avail += num_chunks + self.chunk_bag[self.total_avail : (self.total_avail + num_chunks)] = chunks + self.total_avail += num_chunks def reset(self) -> None: """Reset the allocator to initial state. @@ -89,4 +99,4 @@ def reset(self) -> None: This resets the available chunk count to the entire memory pool (except for the dummy chunk). """ - self.chunk_count_avail = self.chunk_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 a885e871cb4..d0fa28f77b3 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -216,9 +216,9 @@ def __init__( self.max_tokens = max_tokens # Initialize chunk allocator - self.chunk_allocator = ChunkAllocator(active_count_total=active_chunk_count_total) + self.chunk_allocator = ChunkAllocator(active_count=active_chunk_count_total) # >>> - pax({"chunk_allocator": self.chunk_allocator}) + # pax({"chunk_allocator": self.chunk_allocator}) # <<< # Initialize context state. @@ -269,7 +269,12 @@ def __init__( with torch.cuda.use_mem_pool(self.unified_memory_mempool): if cache_mla_latent: self.memory_buffer = torch.full( - (self.num_layers, chunk_count_total, self.chunk_size_tokens, kv_reduced_dim), + ( + self.num_layers, + self.chunk_allocator.total_count, + self.chunk_size_tokens, + kv_reduced_dim, + ), -1, dtype=self.params_dtype, device=torch.cuda.current_device(), @@ -279,7 +284,7 @@ def __init__( ( 2, # key and value self.num_layers, - chunk_count_total, + self.chunk_allocator.total_count, self.chunk_size_tokens, num_attention_heads_per_partition, hidden_size_per_attention_head, @@ -943,7 +948,7 @@ def add_request( # Preallocate chunks. num_chunks_needed = math.ceil(context_length / self.chunk_size_tokens) - new_chunk_ids = self.chunk_allocator.allocate_memory_chunks(num_chunks_needed, safe=True) + new_chunk_ids = self.chunk_allocator.allocate_memory_chunks(num_chunks_needed) if new_chunk_ids is None: raise ChunkOverflowError(request_id) @@ -1183,16 +1188,29 @@ 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 chunks to paused requests. # todo: @shanmugamr, un-pause requests using FIFO, rather than LIFO. - num_non_gtd_chunks = max(0, self.chunk_allocator.chunk_count_avail - self.gtd_chunk_count) - if num_non_gtd_chunks: - # if we have non-gtd chunks, use them. Do not dip into the gtd-chunk pool - resume_request_count = min(num_non_gtd_chunks, self.paused_request_count) - else: - # only dip into the gtd-chunk pool if we have run out of non-gtd-chunks and the active - # request count has fallen below a certain threshold. + # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> + # num_non_gtd_chunks = max(0, self.chunk_allocator.chunk_count_avail - self.gtd_chunk_count) + # if num_non_gtd_chunks: + # # if we have non-gtd chunks, use them. Do not dip into the gtd-chunk pool + # resume_request_count = min(num_non_gtd_chunks, self.paused_request_count) + # else: + # # only dip into the gtd-chunk pool if we have run out of non-gtd-chunks and the active + # # request count has fallen below a certain threshold. + # resume_request_count = min( + # max(self.gtd_request_count - active_request_count, 0), self.paused_request_count + # ) + # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ + if active_request_count < self.chunk_allocator.active_count: resume_request_count = min( - max(self.gtd_request_count - active_request_count, 0), self.paused_request_count + # >>> + # self.chunk_allocator.active_avail, + # <<< + self.chunk_allocator.active_count - active_request_count, + self.paused_request_count, ) + else: + resume_request_count = 0 + # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< self.paused_request_count -= resume_request_count active_request_count += resume_request_count diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index d72ffee2fd4..919a50075b2 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -622,7 +622,7 @@ 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: %d [ active %d, paused %d, finished %d ] ... " "mem: tensors %d, alloc %.1f gb, res %.1f gb." % ( self.step_count, @@ -641,7 +641,6 @@ async def async_step( ) ), prev_total_request_count, - context.gtd_request_count, prev_total_request_count - prev_paused_request_count, prev_paused_request_count, self.finished_request_count, From f604162a1cc3cb990bf3648bc566a7d6425e3ce3 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Wed, 1 Oct 2025 13:16:28 -0700 Subject: [PATCH 25/69] fixed max_requests. --- .../inference/gpt/gpt_dynamic_inference.py | 18 +--- .../contexts/dynamic_chunk_allocator.py | 8 +- .../inference/contexts/dynamic_context.py | 82 ++++++++++++++++--- 3 files changed, 82 insertions(+), 26 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index b37ab50aa87..a8c0f631889 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -263,29 +263,19 @@ def get_inference_context( num_cuda_graphs=( args.inference_dynamic_batching_num_cuda_graphs if args.enable_cuda_graph else None ), - chunk_size_tokens=args.inference_dynamic_batching_chunk_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, + # chunk_size_tokens=args.inference_dynamic_batching_chunk_size, + # active_buffer_size_bytes=active_buffer_size_bytes, # +++ - # active_buffer_size_gb=active_buffer_size_gb, - active_buffer_size_bytes=active_buffer_size_bytes, - # model=model, + chunk_size_tokens=4, # 256, + active_buffer_size_bytes=int(2 * 1024**2), # <<< - # >>> - # max_tokens_override=args.inference_dynamic_batching_max_tokens_override, max_tokens=max_tokens, - # <<< 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, kv_lora_rank=args.kv_lora_rank if args.multi_latent_attention else None, qk_pos_emb_head_dim=args.qk_pos_emb_head_dim, - # >>> - # unified_memory_level=args.inference_dynamic_batching_unified_memory_level, - # <<< ) return context diff --git a/megatron/core/inference/contexts/dynamic_chunk_allocator.py b/megatron/core/inference/contexts/dynamic_chunk_allocator.py index d125c8d6dc4..2a299d8d85d 100644 --- a/megatron/core/inference/contexts/dynamic_chunk_allocator.py +++ b/megatron/core/inference/contexts/dynamic_chunk_allocator.py @@ -42,7 +42,7 @@ class ChunkAllocator: def __init__(self, active_count: int): active_count -= 1 # -1 for dummy_chunk_idx (see below) - self.total_count = 2 * active_count + 1 + self.total_count = 2 * active_count + 1 # +1 for dummy_chunk_idx self.total_avail = self.total_count - 1 # -1 for dummy_chunk_idx self.active_count = active_count self.dummy_chunk_idx = self.total_count - 1 @@ -54,6 +54,12 @@ def __init__(self, active_count: int): device=torch.cuda.current_device(), ) + def __str__(self): + return ( + f"total avail {self.total_avail} / {self.total_count - 1}" + f"; active {self.active_count}" + ) + def is_memory_available(self, num_chunks: int) -> bool: """Check if memory chunks are available. diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index d0fa28f77b3..9e837b79c2c 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -92,6 +92,21 @@ 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 n_bytes >= nquery: + # return "%d %s" % (n_bytes // nquery, suffix) + # +++ + if round(n_bytes / nquery) >= 1: + # return "%d %s" % (round(n_bytes / nquery), suffix) + return "%.3g %s" % (n_bytes / nquery, suffix) + # <<< + raise Exception("something went wrong.") + + # pylint: disable=line-too-long class DynamicInferenceContext(BaseInferenceContext): """Inference context that is passed to the main model in order @@ -160,10 +175,7 @@ def __init__( num_attention_heads: int, max_sequence_length: int, active_buffer_size_bytes: float, - # >>> - # model: torch.nn.Module, max_tokens: int, - # <<< chunk_size_tokens: int = 256, tensor_model_parallel_size: Optional[int] = None, cache_mla_latent: bool = False, @@ -209,14 +221,38 @@ def __init__( * hidden_size_per_attention_head ) - # Calculate the number of active chunks, and finalize the buffer size. + # Initialize chunk allocator. + # >>> + # active_buffer_size_bytes_0 = active_buffer_size_bytes + # <<< active_chunk_count_total = active_buffer_size_bytes // self.chunk_size_bytes - active_buffer_size_bytes = active_chunk_count_total * self.chunk_size_bytes - self.max_requests = active_chunk_count_total + self.chunk_allocator = ChunkAllocator(active_count=active_chunk_count_total) + del active_chunk_count_total # no accidental use; use chunk_allocator + # >>> + # active_buffer_size_bytes = active_chunk_count_total * self.chunk_size_bytes + active_buffer_size_bytes = self.chunk_allocator.active_count * self.chunk_size_bytes + # <<< + + # >>> + # active_buffer_size_bytes_1 = active_buffer_size_bytes + # pax("active_chunk_count_total", { + # "chunk_allocator" : self.chunk_allocator, + # }, "active_buffer_size_bytes_0, active_buffer_size_bytes_1") + # <<< + + # Set max_requests, max_tokens. + # >>> + # self.max_requests = active_chunk_count_total + self.max_requests = self.chunk_allocator.active_count + # <<< + # >>> + # pax({ + # "chunk_allocator" : self.chunk_allocator, + # "max_requests" : self.max_requests, + # }) + # <<< self.max_tokens = max_tokens - # Initialize chunk allocator - self.chunk_allocator = ChunkAllocator(active_count=active_chunk_count_total) # >>> # pax({"chunk_allocator": self.chunk_allocator}) # <<< @@ -389,11 +425,26 @@ def __init__( # self.gtd_chunk_count = self.gtd_request_count * self.max_kv_chunk_count # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - # Store the dummy chunk idx reference for convenience - self.dummy_chunk_idx = self.chunk_allocator.dummy_chunk_idx + # >>> + # # Store the dummy chunk idx reference for convenience + # self.dummy_chunk_idx = self.chunk_allocator.dummy_chunk_idx + # <<< + # Reset attention state. self.reset_attention_state() + # >>> + # Print info. + print("DynamicInferenceContext: allocated context with active buffer size %s (%d chunks)." % ( + get_mem_size_str(active_buffer_size_bytes), + self.chunk_allocator.active_count, + )) + # <<< + + # >>> + # print("~~~\nexiting ..."); exit() + # <<< + TOKEN_ROUNDER = 64 REQUEST_ROUNDER = 4 @@ -715,7 +766,7 @@ def initialize_attention_state( # Update token position indexes. self.token_to_chunk_idx[self.active_token_count : self.padded_active_token_count] = ( - self.dummy_chunk_idx + self.chunk_allocator.dummy_chunk_idx ) self.token_to_local_position_within_kv_chunk[ self.active_token_count : self.padded_active_token_count @@ -998,6 +1049,11 @@ def add_request( self.total_request_count += 1 self.active_token_count += context_length + # >>> + print(f"++++++++ chunk allocator ... {self.chunk_allocator}.") + # pax({"max_requests": self.max_requests}) + # <<< + def _move_book_keeping_tensors(self, src_idxs, dst_idxs, next_tokens): """ Swaps all the relevent booking tensors with src idxs to dst idxs @@ -1223,6 +1279,10 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T # requests are newly paused during this update. if newly_paused_request_ids is not None and resume_request_count > 0: newly_paused_request_ids = newly_paused_request_ids[:-resume_request_count] + # >>> + if newly_paused_request_ids is not None and newly_paused_request_ids.numel() > 0: + pax("newly_paused_request_ids") + # <<< # 7. We make changes to the request book keeping tesnsors and setup the tokens for next iteration self.total_request_count = active_request_count + self.paused_request_count From 62ce81c1f16b4384eecaa8ea198ab200a0066c25 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Thu, 2 Oct 2025 09:08:17 -0700 Subject: [PATCH 26/69] fixed newly_paused_request_ids. --- .../inference/gpt/gpt_dynamic_inference.py | 8 +-- .../inference/contexts/dynamic_context.py | 69 ++++++++++++++++++- 2 files changed, 70 insertions(+), 7 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index a8c0f631889..aa338919b66 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -264,11 +264,11 @@ def get_inference_context( args.inference_dynamic_batching_num_cuda_graphs if args.enable_cuda_graph else None ), # >>> - # chunk_size_tokens=args.inference_dynamic_batching_chunk_size, - # active_buffer_size_bytes=active_buffer_size_bytes, + chunk_size_tokens=args.inference_dynamic_batching_chunk_size, + active_buffer_size_bytes=active_buffer_size_bytes, # +++ - chunk_size_tokens=4, # 256, - active_buffer_size_bytes=int(2 * 1024**2), + # chunk_size_tokens=4, # 256, + # active_buffer_size_bytes=int(2 * 1024**2), # <<< max_tokens=max_tokens, tensor_model_parallel_size=args.tensor_model_parallel_size, diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 9e837b79c2c..97763c0b5ec 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1115,6 +1115,10 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T # finished_request_count are requests that have reached the termination criterion active_request_count = (active_requests_mask == 1).sum().item() finished_request_count = (active_requests_mask == 0).sum().item() + # >>> + if finished_request_count > 0: + raise Exception("huh?") + # <<< assert ( active_request_count + finished_request_count + self.paused_request_count == self.total_request_count @@ -1202,6 +1206,21 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T active_requests_requiring_new_chunk_count = ( (active_requests_requiring_new_chunk == 1).sum().item() ) + # >>> + # if active_requests_requiring_new_chunk_count > 0: + # pax("active_requests_requiring_new_chunk") + # <<< + + # >>> [ new ] + if active_requests_requiring_new_chunk_count > 0: + newly_paused_request_ids = self.request_ids[ + torch.nonzero(active_requests_requiring_new_chunk) + + self.paused_request_count + ] + # >>> + # pax("newly_paused_request_ids") + # <<< + # <<< # Swap unfinished active requests on the left side with paused requests on the right side # NOTE : We add paused request count because we concatenate @@ -1232,10 +1251,18 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T ) dst_idxs = torch.cat((active_request_ids_on_left, paused_requests_idxs_on_right)) src_idxs = torch.cat((paused_requests_idxs_on_right, active_request_ids_on_left)) + # >>> [ new ] + # newly_paused_request_ids = self.request_ids[dst_idxs] + # <<< 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] + # >>> [ orig ] + # newly_paused_request_ids = self.request_ids[dst_idxs] + # <<< + # >>> + pax("active_request_ids_on_left, paused_requests_idxs_on_right") + # <<< self.paused_request_count += active_requests_requiring_new_chunk_count active_request_count -= active_requests_requiring_new_chunk_count @@ -1261,6 +1288,7 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T # >>> # self.chunk_allocator.active_avail, # <<< + self.chunk_allocator.total_avail, self.chunk_allocator.active_count - active_request_count, self.paused_request_count, ) @@ -1268,6 +1296,14 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T resume_request_count = 0 # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< + # >>> + if resume_request_count != 0 and self.chunk_allocator.total_avail == 0: + pax({ + "chunk_allocator" : self.chunk_allocator, + "paused_request_count" : self.paused_request_count, + }, "active_request_count, resume_request_count") + # <<< + self.paused_request_count -= resume_request_count active_request_count += resume_request_count assert active_request_count > 0, "active_request_count == %d." % active_request_count @@ -1280,8 +1316,8 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T if newly_paused_request_ids is not None and resume_request_count > 0: newly_paused_request_ids = newly_paused_request_ids[:-resume_request_count] # >>> - if newly_paused_request_ids is not None and newly_paused_request_ids.numel() > 0: - pax("newly_paused_request_ids") + # if newly_paused_request_ids is not None and newly_paused_request_ids.numel() > 0: + # pax("newly_paused_request_ids") # <<< # 7. We make changes to the request book keeping tesnsors and setup the tokens for next iteration @@ -1321,6 +1357,12 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T ), "The request_last_kv_chunk_offset should be 0 for the requests that just got resumed this step. " chunk_ids = self.chunk_allocator.allocate_memory_chunks(resume_request_count) + # >>> + if chunk_ids is None: + pax({ + "chunk_allocator" : self.chunk_allocator, + }, "resume_request_count, chunk_ids") + # <<< row_idx = torch.arange( self.paused_request_count, self.paused_request_count + resume_request_count, @@ -1352,6 +1394,27 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T self.request_last_kv_chunk_offset[self.paused_request_count : self.total_request_count] ) + # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> + def get_request_str(i): + r = self.request_ids[i].item() + q = self.request_query_lengths[i].item() + kv = self.request_kv_length_offsets[i].item() + # pax("request_id") + return "%d[%d:%d]" % (r, q, kv) + request_strs = [ + get_request_str(i) + for i in range(self.total_request_count) + ] + # pax("request_strs") + print("++++++++++ requests: %s --- %s ......... alloc: %d/%d [%d]" % ( + ", ".join(request_strs[:self.paused_request_count]), + ", ".join(request_strs[self.paused_request_count:self.total_request_count]), + self.chunk_allocator.total_avail, + self.chunk_allocator.total_count, + self.chunk_allocator.active_count, + )) + # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< + return newly_paused_request_ids def calculate_log_probs( From 1a74d5c341a98cb35de81fdda3cbf2a11ad13df2 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Thu, 2 Oct 2025 09:56:48 -0700 Subject: [PATCH 27/69] fixed resume_request_count; mostly vectorized. --- .../inference/contexts/dynamic_context.py | 58 +++++++++++++++---- .../core/inference/engines/dynamic_engine.py | 4 ++ 2 files changed, 50 insertions(+), 12 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 97763c0b5ec..9e38e9ad555 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1261,12 +1261,20 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T # newly_paused_request_ids = self.request_ids[dst_idxs] # <<< # >>> - pax("active_request_ids_on_left, paused_requests_idxs_on_right") + # pax("active_request_ids_on_left, paused_requests_idxs_on_right") # <<< self.paused_request_count += active_requests_requiring_new_chunk_count active_request_count -= active_requests_requiring_new_chunk_count + # >>> + # if active_requests_requiring_new_chunk_count > 0: + # pax({ + # "chunk_allocator" : self.chunk_allocator, + # "paused_request_count" : self.paused_request_count, + # }, "active_request_count") + # <<< + # 6. Now that we have the requests in following order [Paused, Active, Finished] # We determine how many requests we can resume and resume them # Assign released chunks to paused requests. @@ -1283,17 +1291,43 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T # max(self.gtd_request_count - active_request_count, 0), self.paused_request_count # ) # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - if active_request_count < self.chunk_allocator.active_count: - resume_request_count = min( - # >>> - # self.chunk_allocator.active_avail, - # <<< - self.chunk_allocator.total_avail, - self.chunk_allocator.active_count - active_request_count, - self.paused_request_count, - ) - else: - resume_request_count = 0 + # if active_request_count < self.chunk_allocator.active_count: + # resume_request_count = min( + # # >>> + # # self.chunk_allocator.active_avail, + # # <<< + # self.chunk_allocator.total_avail, + # self.chunk_allocator.active_count - active_request_count, + # self.paused_request_count, + # ) + # else: + # resume_request_count = 0 + # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ + resume_request_count = 0 + if self.paused_request_count > 0: + active_chunk_count_total = self.chunk_allocator.active_count + active_chunk_count_used = self.request_kv_chunk_counts[ + self.paused_request_count:self.total_request_count + ].sum().item() + active_chunk_count_avail = active_chunk_count_total - active_chunk_count_used + + paused_chunk_counts = self.request_kv_chunk_counts[:self.paused_request_count] + paused_chunk_counts = paused_chunk_counts.flip(dims=[0]) + paused_chunk_counts_cumsum = paused_chunk_counts.cumsum(dim=0) + paused_chunk_counts_cumsum += 1 # +1 for newly added chunk + resume_request_count = torch.nonzero( + paused_chunk_counts_cumsum <= active_chunk_count_avail + ).numel() + + # for i in range(self.paused_request_count - 1, -1, -1): + # current_chunk_count + # pax("i") + + # >>> + # pax("active_chunk_count_total, active_chunk_count_used, active_chunk_count_avail", + # "paused_chunk_counts, paused_chunk_counts_cumsum", + # "resume_request_count") + # <<< # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< # >>> diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 919a50075b2..c19cfd0d228 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -654,6 +654,10 @@ async def async_step( print(output_str) self.step_count += 1 + # >>> + if self.step_count >= 800: + raise Exception("step 800.") + # <<< range_pop() return { "active_requests": active_requests, From 9e0b2b3634ddbb68e4df2a9a2fd1c687d9834b10 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Thu, 2 Oct 2025 14:48:06 -0700 Subject: [PATCH 28/69] fixed paused_chunk_counts. --- .../inference/contexts/dynamic_context.py | 39 ++++++++++++++----- .../core/inference/engines/dynamic_engine.py | 4 +- 2 files changed, 32 insertions(+), 11 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 9e38e9ad555..5b438a2fe4a 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1116,8 +1116,8 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T active_request_count = (active_requests_mask == 1).sum().item() finished_request_count = (active_requests_mask == 0).sum().item() # >>> - if finished_request_count > 0: - raise Exception("huh?") + # if finished_request_count > 0: + # raise Exception("huh?") # <<< assert ( active_request_count + finished_request_count + self.paused_request_count @@ -1127,6 +1127,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: @@ -1307,14 +1310,20 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T if self.paused_request_count > 0: active_chunk_count_total = self.chunk_allocator.active_count active_chunk_count_used = self.request_kv_chunk_counts[ + # >>> self.paused_request_count:self.total_request_count + # self.paused_request_count:(self.paused_request_count + active_request_count) + # <<< ].sum().item() active_chunk_count_avail = active_chunk_count_total - active_chunk_count_used paused_chunk_counts = self.request_kv_chunk_counts[:self.paused_request_count] paused_chunk_counts = paused_chunk_counts.flip(dims=[0]) + paused_chunk_counts += 1 # +1 for newly added chunk paused_chunk_counts_cumsum = paused_chunk_counts.cumsum(dim=0) - paused_chunk_counts_cumsum += 1 # +1 for newly added chunk + # >>> + # paused_chunk_counts_cumsum += 1 # +1 for newly added chunk + # <<< resume_request_count = torch.nonzero( paused_chunk_counts_cumsum <= active_chunk_count_avail ).numel() @@ -1322,16 +1331,26 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T # for i in range(self.paused_request_count - 1, -1, -1): # current_chunk_count # pax("i") + # if resume_request_count != 1: + # pax("resume_request_count") # >>> - # pax("active_chunk_count_total, active_chunk_count_used, active_chunk_count_avail", - # "paused_chunk_counts, paused_chunk_counts_cumsum", - # "resume_request_count") + # if active_request_count + resume_request_count == 0: + # if active_request_count == 0: + if resume_request_count >= 2: + pax("active_chunk_count_total, active_chunk_count_used, active_chunk_count_avail", + "paused_chunk_counts, paused_chunk_counts_cumsum", + "active_request_count, resume_request_count") # <<< # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< # >>> - if resume_request_count != 0 and self.chunk_allocator.total_avail == 0: + newly_resumed_request_ids = self.request_ids[(self.paused_request_count - resume_request_count):self.paused_request_count] + # <<< + + # >>> + # if resume_request_count != 0 and self.chunk_allocator.total_avail == 0: + if active_request_count + resume_request_count == 0: pax({ "chunk_allocator" : self.chunk_allocator, "paused_request_count" : self.paused_request_count, @@ -1355,7 +1374,6 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T # <<< # 7. We make changes to the request book keeping tesnsors and setup the tokens for next iteration - self.total_request_count = active_request_count + self.paused_request_count # All these active requests are in decode phase, so they need only 1 token per request self.active_token_count = active_request_count # Always the first section of token input ids are only used. @@ -1440,9 +1458,12 @@ def get_request_str(i): for i in range(self.total_request_count) ] # pax("request_strs") - print("++++++++++ requests: %s --- %s ......... alloc: %d/%d [%d]" % ( + print("++++++++++ requests: %s --- %s ..... paused: %s, resumed: %s, finished: %s ..... alloc: %d/%d [%d]" % ( ", ".join(request_strs[:self.paused_request_count]), ", ".join(request_strs[self.paused_request_count:self.total_request_count]), + "--" if newly_paused_request_ids is None else newly_paused_request_ids.tolist(), + newly_resumed_request_ids.tolist(), + self.request_ids[self.total_request_count:(self.total_request_count + finished_request_count)].tolist(), self.chunk_allocator.total_avail, self.chunk_allocator.total_count, self.chunk_allocator.active_count, diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index c19cfd0d228..45850c141f0 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -655,8 +655,8 @@ async def async_step( self.step_count += 1 # >>> - if self.step_count >= 800: - raise Exception("step 800.") + if self.step_count >= 600: + raise Exception(f"step {self.step_count}.") # <<< range_pop() return { From db6dbdd4be013ebf54f2d62e19fe21c22ed7cbf5 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 6 Oct 2025 03:32:23 -0700 Subject: [PATCH 29/69] max_requests -> max_{total,active}_requests. --- examples/inference/gpt/utils.py | 6 +- .../inference/contexts/dynamic_context.py | 56 +++++++++++-------- .../core/inference/engines/dynamic_engine.py | 11 +++- 3 files changed, 46 insertions(+), 27 deletions(-) diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py index 9cb02bf92c2..3430db9f7b5 100644 --- a/examples/inference/gpt/utils.py +++ b/examples/inference/gpt/utils.py @@ -62,7 +62,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( @@ -311,7 +311,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. @@ -371,7 +371,7 @@ def build_dynamic_engine_setup_prefix( # Buffer limits config buffer_limits_str = ( f"bf {args.inference_dynamic_batching_active_buffer_size_gb:.0f} " - f"[r {context.max_requests}, t {context.max_tokens}]" + f"[r {context.max_active_requests}, t {context.max_tokens}]" ) parts = [ diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 5b438a2fe4a..a8cd95ae22c 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -243,7 +243,11 @@ def __init__( # Set max_requests, max_tokens. # >>> # self.max_requests = active_chunk_count_total - self.max_requests = self.chunk_allocator.active_count + # +++ + # self.max_requests = self.chunk_allocator.active_count + # +++ + self.max_total_requests = self.chunk_allocator.total_count - 1 # -1 for dummy chunk + self.max_active_requests = self.chunk_allocator.active_count # <<< # >>> # pax({ @@ -275,7 +279,7 @@ def __init__( # 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) @@ -333,7 +337,7 @@ def __init__( # Chunk ids. self.max_kv_chunk_count = math.ceil(self.max_sequence_length / self.chunk_size_tokens) self.request_to_kv_chunk_ids = torch.full( - (self.max_requests, self.max_kv_chunk_count), + (self.max_total_requests, self.max_kv_chunk_count), -1, dtype=torch.int, device=torch.cuda.current_device(), @@ -344,11 +348,11 @@ def __init__( 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 ) @@ -357,13 +361,13 @@ def __init__( # 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. @@ -385,20 +389,20 @@ def __init__( # corresponding tensors are used. self.query_seq_lengths_cudagraph_only = torch.full( - (self.max_requests,), 0, dtype=torch.int32, device=torch.cuda.current_device() + (self.max_total_requests,), 0, dtype=torch.int32, device=torch.cuda.current_device() ) self.cu_query_seq_lengths_cudagraph_only = torch.full( - (self.max_requests + 1,), 0, dtype=torch.int32, device=torch.cuda.current_device() + (self.max_total_requests + 1,), 0, dtype=torch.int32, device=torch.cuda.current_device() ) self.kv_seq_lengths_cudagraph_only = torch.full( - (self.max_requests,), 0, dtype=torch.int32, device=torch.cuda.current_device() + (self.max_total_requests,), 0, dtype=torch.int32, device=torch.cuda.current_device() ) self.cu_kv_seq_lengths_cudagraph_only = torch.full( - (self.max_requests + 1,), 0, dtype=torch.int32, device=torch.cuda.current_device() + (self.max_total_requests + 1,), 0, dtype=torch.int32, device=torch.cuda.current_device() ) self.request_to_kv_chunk_ids_cudagraph_only = torch.full( - (self.max_requests, self.max_kv_chunk_count), + (self.max_total_requests, self.max_kv_chunk_count), 0, dtype=torch.int, device=torch.cuda.current_device(), @@ -720,8 +724,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" @@ -740,7 +744,7 @@ 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 assert self.padded_active_token_count >= active_token_count else: @@ -748,7 +752,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? @@ -994,8 +998,13 @@ def add_request( # TODO : Should move this into some waiting queue if self.active_token_count + context_length > self.max_tokens: raise TokenOverflowError(request_id) - if self.total_request_count >= self.max_requests: + # >>> + # if self.total_request_count >= self.max_total_requests: + # raise RequestOverflowError(request_id) + # +++ + if self.total_request_count - self.paused_request_count >= self.max_active_requests: raise RequestOverflowError(request_id) + # <<< # Preallocate chunks. num_chunks_needed = math.ceil(context_length / self.chunk_size_tokens) @@ -1455,17 +1464,18 @@ def get_request_str(i): return "%d[%d:%d]" % (r, q, kv) request_strs = [ get_request_str(i) - for i in range(self.total_request_count) + for i in range(self.total_request_count + finished_request_count) ] # pax("request_strs") - print("++++++++++ requests: %s --- %s ..... paused: %s, resumed: %s, finished: %s ..... alloc: %d/%d [%d]" % ( + print("++++++++++ requests: %s -- %s -- %s ..... paused: %s, resumed: %s, finished: %s ..... alloc: %d/%d [%d]" % ( ", ".join(request_strs[:self.paused_request_count]), ", ".join(request_strs[self.paused_request_count:self.total_request_count]), + ", ".join(request_strs[self.total_request_count:(self.total_request_count + finished_request_count)]), "--" if newly_paused_request_ids is None else newly_paused_request_ids.tolist(), newly_resumed_request_ids.tolist(), self.request_ids[self.total_request_count:(self.total_request_count + finished_request_count)].tolist(), - self.chunk_allocator.total_avail, - self.chunk_allocator.total_count, + self.chunk_allocator.total_count - self.chunk_allocator.total_avail - 1, + self.chunk_allocator.total_count - 1, self.chunk_allocator.active_count, )) # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 45850c141f0..28d29759424 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -517,6 +517,9 @@ def post_process_requests( def schedule_waiting_requests(self): """Tries to schedule any requests in the waiting pool.""" + # >>> + # print(" > schedule_waiting_requests: %s." % self.waiting_request_ids) + # <<< for waiting_request_id in self.waiting_request_ids.copy(): waiting_request: DynamicInferenceRequest = self.requests[waiting_request_id] try: @@ -527,7 +530,13 @@ def schedule_waiting_requests(self): ) waiting_request.add_event_add() self.waiting_request_ids.popleft() + # >>> + print(" > schedule_waiting_requests: %s ... (success)." % waiting_request_id) + # <<< except Exception as e: + # >>> + print(" > schedule_waiting_requests: %s ... (%s)." % (waiting_request_id, e.__class__.__name__)) + # <<< break async def async_step( @@ -655,7 +664,7 @@ async def async_step( self.step_count += 1 # >>> - if self.step_count >= 600: + if self.step_count >= 302: # 600, 300 raise Exception(f"step {self.step_count}.") # <<< range_pop() From db36ad00b5d408ceb02a86439a1311ba035308db Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 6 Oct 2025 05:55:10 -0700 Subject: [PATCH 30/69] available memory determined by available active chunks. --- .../contexts/dynamic_chunk_allocator.py | 69 +++++++++++++++++-- .../inference/contexts/dynamic_context.py | 37 +++++++--- 2 files changed, 90 insertions(+), 16 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_chunk_allocator.py b/megatron/core/inference/contexts/dynamic_chunk_allocator.py index 2a299d8d85d..162dba9e86e 100644 --- a/megatron/core/inference/contexts/dynamic_chunk_allocator.py +++ b/megatron/core/inference/contexts/dynamic_chunk_allocator.py @@ -60,31 +60,90 @@ def __str__(self): f"; active {self.active_count}" ) - def is_memory_available(self, num_chunks: int) -> bool: + # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> + def get_active_avail(self, context: "DynamicInferenceContext"): + """Compute number of available active chunks. + + Args: + context (DynamicInferenceContext): Dynamic inference context. + """ + active_total = self.active_count + active_used = context.request_kv_chunk_counts[ + context.paused_request_count:context.total_request_count + ].sum().item() + active_avail = active_total - active_used + # >>> + # if active_used != 0: + # pax("active_total, active_used, active_avail") + # <<< + return active_avail + # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< + + # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> + # def is_memory_available(self, num_chunks: int) -> bool: + # """Check if memory chunks are available. + + # Args: + # num_chunks (int): Number of chunks to check. + + # Return: + # (bool) Is memory available? + # """ + # return self.total_avail >= num_chunks + # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ + def is_memory_available( + self, + context: "DynamicInferenceContext", + num_chunks: int, + ) -> bool: """Check if memory chunks are available. Args: + context (DynamicInferenceContext): Dynamic inference context. num_chunks (int): Number of chunks to check. Return: (bool) Is memory available? """ - return self.total_avail >= num_chunks - - def allocate_memory_chunks(self, num_chunks: int = 1) -> Optional[Tensor]: + return self.get_active_avail(context) >= num_chunks + # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< + + # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> + # def allocate_memory_chunks(self, num_chunks: int = 1) -> Optional[Tensor]: + # """Allocate memory chunks if available, else return None. + + # Args: + # num_chunks (int): Number of chunks to allocate. + + # Return: + # (Optional[Tensor]) Allocated chunk IDs. + # """ + # if self.is_memory_available(num_chunks): + # self.total_avail -= num_chunks + # return self.chunk_bag[self.total_avail : (self.total_avail + num_chunks)] + # else: + # return None + # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ + def allocate_memory_chunks( + self, + context: "DynamicInferenceContext", + num_chunks: int, + ) -> Optional[Tensor]: """Allocate memory chunks if available, else return None. Args: + context (DynamicInferenceContext): Dynamic inference context. num_chunks (int): Number of chunks to allocate. Return: (Optional[Tensor]) Allocated chunk IDs. """ - if self.is_memory_available(num_chunks): + if self.is_memory_available(context, num_chunks): self.total_avail -= num_chunks return self.chunk_bag[self.total_avail : (self.total_avail + num_chunks)] else: return None + # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< def release_memory_chunks(self, chunks: Tensor) -> None: """Release memory chunks. diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index a8cd95ae22c..b101b056260 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1002,13 +1002,30 @@ def add_request( # if self.total_request_count >= self.max_total_requests: # raise RequestOverflowError(request_id) # +++ - if self.total_request_count - self.paused_request_count >= self.max_active_requests: + if ( + self.total_request_count - self.paused_request_count >= self.max_active_requests + # or + # self.get_active_request_count() >= self.chunk_allocator.active_count + ): raise RequestOverflowError(request_id) # <<< + # >>> + if request_id == 4: + pax({ + "chunk_allocator" : self.chunk_allocator, + "total_request_count" : self.total_request_count, + "paused_request_count" : self.paused_request_count, + }) + # <<< + # Preallocate chunks. num_chunks_needed = math.ceil(context_length / self.chunk_size_tokens) - new_chunk_ids = self.chunk_allocator.allocate_memory_chunks(num_chunks_needed) + # >>> + # new_chunk_ids = self.chunk_allocator.allocate_memory_chunks(num_chunks_needed) + # +++ + new_chunk_ids = self.chunk_allocator.allocate_memory_chunks(self, num_chunks_needed) + # <<< if new_chunk_ids is None: raise ChunkOverflowError(request_id) @@ -1317,14 +1334,8 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ resume_request_count = 0 if self.paused_request_count > 0: - active_chunk_count_total = self.chunk_allocator.active_count - active_chunk_count_used = self.request_kv_chunk_counts[ - # >>> - self.paused_request_count:self.total_request_count - # self.paused_request_count:(self.paused_request_count + active_request_count) - # <<< - ].sum().item() - active_chunk_count_avail = active_chunk_count_total - active_chunk_count_used + active_chunk_count_avail = self.chunk_allocator.get_active_avail(self) + # pax("active_chunk_count_avail") paused_chunk_counts = self.request_kv_chunk_counts[:self.paused_request_count] paused_chunk_counts = paused_chunk_counts.flip(dims=[0]) @@ -1417,7 +1428,11 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T == 0 ), "The request_last_kv_chunk_offset should be 0 for the requests that just got resumed this step. " - chunk_ids = self.chunk_allocator.allocate_memory_chunks(resume_request_count) + # >>> + # chunk_ids = self.chunk_allocator.allocate_memory_chunks(resume_request_count) + # +++ + chunk_ids = self.chunk_allocator.allocate_memory_chunks(self, resume_request_count) + # <<< # >>> if chunk_ids is None: pax({ From 96c265620e880b0ae54440261f08bb4ed85896c6 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 6 Oct 2025 06:55:39 -0700 Subject: [PATCH 31/69] added methods to compute active/paused used/avail chunk counts. --- .../contexts/dynamic_chunk_allocator.py | 60 +++++++++++++++++-- .../inference/contexts/dynamic_context.py | 46 ++++++++++---- .../core/inference/engines/dynamic_engine.py | 8 +-- 3 files changed, 92 insertions(+), 22 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_chunk_allocator.py b/megatron/core/inference/contexts/dynamic_chunk_allocator.py index 162dba9e86e..65110aef5dc 100644 --- a/megatron/core/inference/contexts/dynamic_chunk_allocator.py +++ b/megatron/core/inference/contexts/dynamic_chunk_allocator.py @@ -61,22 +61,55 @@ def __str__(self): ) # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> - def get_active_avail(self, context: "DynamicInferenceContext"): - """Compute number of available active chunks. + def get_active_used(self, context: "DynamicInferenceContext"): + """Compute number of active chunks used. Args: context (DynamicInferenceContext): Dynamic inference context. """ - active_total = self.active_count - active_used = context.request_kv_chunk_counts[ + return context.request_kv_chunk_counts[ context.paused_request_count:context.total_request_count ].sum().item() + + def get_paused_used(self, context: "DynamicInferenceContext"): + """Compute number of paused chunks used. + + Args: + context (DynamicInferenceContext): Dynamic inference context. + """ + return context.request_kv_chunk_counts[ + :context.paused_request_count + ].sum().item() + + def get_active_avail(self, context: "DynamicInferenceContext"): + """Compute number of active chunks available. + + Args: + context (DynamicInferenceContext): Dynamic inference context. + """ + active_total = self.active_count + active_used = self.get_active_used(context) active_avail = active_total - active_used # >>> # if active_used != 0: # pax("active_total, active_used, active_avail") # <<< return active_avail + + def get_paused_avail(self, context: "DynamicInferenceContext"): + """Compute number of paused chunks available. + + Args: + context (DynamicInferenceContext): Dynamic inference context. + """ + paused_total = self.total_count - self.active_count - 1 # -1 for dummy chunk + paused_used = self.get_paused_used(context) + paused_avail = paused_total - paused_used + # >>> + # if paused_used != 0: + # pax("paused_total, paused_used, paused_avail") + # <<< + return paused_avail # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> @@ -105,6 +138,17 @@ def is_memory_available( Return: (bool) Is memory available? """ + # >>> + paused_avail = self.get_paused_avail(context) + active_avail = self.get_active_avail(context) + if active_avail >= num_chunks and self.total_avail == 0: + pax("self, num_chunks, paused_avail, active_avail", { + "paused_used" : self.get_paused_used(context), + "active_used" : self.get_active_used(context), + "total_request_count" : context.total_request_count, + "paused_request_count" : context.paused_request_count, + }) + # <<< return self.get_active_avail(context) >= num_chunks # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< @@ -140,7 +184,13 @@ def allocate_memory_chunks( """ if self.is_memory_available(context, num_chunks): self.total_avail -= num_chunks - return self.chunk_bag[self.total_avail : (self.total_avail + num_chunks)] + chunk_ids = self.chunk_bag[self.total_avail : (self.total_avail + num_chunks)] + # >>> + # assert num_chunks == chunk_ids.numel() + if num_chunks != chunk_ids.numel(): + pax("self, num_chunks") + # <<< + return chunk_ids else: return None # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index b101b056260..051d317b56c 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1010,15 +1010,6 @@ def add_request( raise RequestOverflowError(request_id) # <<< - # >>> - if request_id == 4: - pax({ - "chunk_allocator" : self.chunk_allocator, - "total_request_count" : self.total_request_count, - "paused_request_count" : self.paused_request_count, - }) - # <<< - # Preallocate chunks. num_chunks_needed = math.ceil(context_length / self.chunk_size_tokens) # >>> @@ -1026,6 +1017,14 @@ def add_request( # +++ new_chunk_ids = self.chunk_allocator.allocate_memory_chunks(self, num_chunks_needed) # <<< + # >>> + # if request_id == 4: + # pax({ + # "chunk_allocator" : self.chunk_allocator, + # "total_request_count" : self.total_request_count, + # "paused_request_count" : self.paused_request_count, + # }, "new_chunk_ids") + # <<< if new_chunk_ids is None: raise ChunkOverflowError(request_id) @@ -1076,7 +1075,7 @@ def add_request( self.active_token_count += context_length # >>> - print(f"++++++++ chunk allocator ... {self.chunk_allocator}.") + # print(f"++++++++ chunk allocator ... {self.chunk_allocator}.") # pax({"max_requests": self.max_requests}) # <<< @@ -1447,7 +1446,16 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T col_idx = self.request_kv_chunk_counts[ self.paused_request_count : (self.paused_request_count + resume_request_count) ] - self.request_to_kv_chunk_ids[row_idx, col_idx] = chunk_ids + # >>> + # self.request_to_kv_chunk_ids[row_idx, col_idx] = chunk_ids + # +++ + try: + self.request_to_kv_chunk_ids[row_idx, col_idx] = chunk_ids + except Exception as e: + pax({ + "request_to_kv_chunk_ids" : self.request_to_kv_chunk_ids, + }, "row_idx, col_idx, resume_request_count, chunk_ids") + # <<< self.request_kv_chunk_counts[ self.paused_request_count : (self.paused_request_count + resume_request_count) ] += 1 @@ -1482,17 +1490,29 @@ def get_request_str(i): for i in range(self.total_request_count + finished_request_count) ] # pax("request_strs") - print("++++++++++ requests: %s -- %s -- %s ..... paused: %s, resumed: %s, finished: %s ..... alloc: %d/%d [%d]" % ( + request_str = "requests: %s -- %s -- %s" % ( ", ".join(request_strs[:self.paused_request_count]), ", ".join(request_strs[self.paused_request_count:self.total_request_count]), ", ".join(request_strs[self.total_request_count:(self.total_request_count + finished_request_count)]), + ) + newly_str = "paused: %s, resumed: %s, finished: %s" % ( "--" if newly_paused_request_ids is None else newly_paused_request_ids.tolist(), newly_resumed_request_ids.tolist(), self.request_ids[self.total_request_count:(self.total_request_count + finished_request_count)].tolist(), + ) + alloc_str = "alloc: t %d/%d [ p %d/%d, a %d/%d ]" % ( self.chunk_allocator.total_count - self.chunk_allocator.total_avail - 1, self.chunk_allocator.total_count - 1, + self.chunk_allocator.get_paused_used(self), + self.chunk_allocator.total_count - self.chunk_allocator.active_count - 1, + self.chunk_allocator.get_active_used(self), self.chunk_allocator.active_count, - )) + ) + print("++++++++++ %s" % " ..... ".join(( + request_str, + # newly_str, + alloc_str, + ))) # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< return newly_paused_request_ids diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 28d29759424..391aa53cb9b 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -531,11 +531,11 @@ def schedule_waiting_requests(self): waiting_request.add_event_add() self.waiting_request_ids.popleft() # >>> - print(" > schedule_waiting_requests: %s ... (success)." % waiting_request_id) + # print(" > schedule_waiting_requests: %s ... (success)." % waiting_request_id) # <<< except Exception as e: # >>> - print(" > schedule_waiting_requests: %s ... (%s)." % (waiting_request_id, e.__class__.__name__)) + # print(" > schedule_waiting_requests: %s ... (%s)." % (waiting_request_id, e.__class__.__name__)) # <<< break @@ -664,8 +664,8 @@ async def async_step( self.step_count += 1 # >>> - if self.step_count >= 302: # 600, 300 - raise Exception(f"step {self.step_count}.") + # if self.step_count >= 602: # 600, 300 + # raise Exception(f"step {self.step_count}.") # <<< range_pop() return { From 42c0777e8944ebe6168ecb0e5f1e5df1187bd3b8 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 6 Oct 2025 08:54:52 -0700 Subject: [PATCH 32/69] minor clean up. --- .../inference/gpt/gpt_dynamic_inference.py | 25 +-------- examples/inference/gpt/utils.py | 25 --------- .../contexts/dynamic_chunk_allocator.py | 56 ++++++++++--------- .../inference/contexts/dynamic_context.py | 2 +- megatron/training/arguments.py | 55 ++---------------- 5 files changed, 38 insertions(+), 125 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index aa338919b66..837ee145b33 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -121,7 +121,6 @@ def get_model() -> MegatronModule: return model -# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> def get_inference_context_active_buffer_size_bytes(model: torch.nn.Module) -> int: """Estimate memory remaining for context's `memory_buffer`. @@ -133,12 +132,6 @@ def get_inference_context_active_buffer_size_bytes(model: torch.nn.Module) -> in args = get_args() active_buffer_size_gb = args.inference_dynamic_batching_active_buffer_size_gb active_buffer_size_bytes = int(active_buffer_size_gb * 1024**3) - # >>> - # pax({ - # "active buffer gb" : get_args().inference_dynamic_batching_active_buffer_size_gb, - # "1024**3" : 1024**3, - # }, "active_buffer_size_bytes") - # <<< return active_buffer_size_bytes @@ -156,7 +149,6 @@ def get_inference_context_active_buffer_size_bytes(model: torch.nn.Module) -> in # # args = get_args() # # model = GPTInferenceWrapper(model, args, inference_context="not-actually-set") # None) -# # pax("model") # num_tokens_step = 1000 # num_tokens_min = num_tokens_step @@ -182,7 +174,6 @@ def get_inference_context_active_buffer_size_bytes(model: torch.nn.Module) -> in # dtype=torch.long, # device=torch.cuda.current_device(), # ).unsqueeze(0) -# # pax("input_ids, position_ids") # # result = model.run_one_forward_step( # # { # # "tokens": input_ids, @@ -201,15 +192,11 @@ def get_inference_context_active_buffer_size_bytes(model: torch.nn.Module) -> in # ) # t = time.time() - t # time_map[num_tokens] = t -# # >>> -# # pax("result") -# # <<< # print("~~ x ~~") # [ print(x) for x in time_map.keys() ] # print("~~ y ~~") # [ print(y) for y in time_map.values() ] -# pax() -# pax("time_map") +# raise Exception("matplotlib.") # return max_tokens def get_inference_context_max_tokens(model: torch.nn.Module) -> int: """Estimate `max_tokens` for the context. @@ -219,7 +206,6 @@ def get_inference_context_max_tokens(model: torch.nn.Module) -> int: correlation between `max_tokens` and request throughput. """ return 16384 -# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< def get_inference_context( @@ -247,10 +233,6 @@ def get_inference_context( active_buffer_size_bytes = get_inference_context_active_buffer_size_bytes(model) max_tokens = get_inference_context_max_tokens(model) - # >>> - # pax("active_buffer_size_bytes, max_tokens") - # <<< - # Inference context. context = DynamicInferenceContext( params_dtype=args.params_dtype, @@ -263,13 +245,8 @@ def get_inference_context( num_cuda_graphs=( args.inference_dynamic_batching_num_cuda_graphs if args.enable_cuda_graph else None ), - # >>> chunk_size_tokens=args.inference_dynamic_batching_chunk_size, active_buffer_size_bytes=active_buffer_size_bytes, - # +++ - # chunk_size_tokens=4, # 256, - # active_buffer_size_bytes=int(2 * 1024**2), - # <<< max_tokens=max_tokens, tensor_model_parallel_size=args.tensor_model_parallel_size, materialize_only_last_token_logits=not args.return_log_probs, diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py index 3430db9f7b5..c81553f5164 100644 --- a/examples/inference/gpt/utils.py +++ b/examples/inference/gpt/utils.py @@ -344,30 +344,6 @@ def build_dynamic_engine_setup_prefix( f"r/step {args.incoming_requests_per_step}" ) - # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> - # # 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}]" - # ) - - # parts = [ - # get_model_size_str(model), - # "dynamic", - # cg_str, - # request_str, - # buffer_limits_str, - # guaranteed_fraction_str, - # ] - # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ # Buffer limits config buffer_limits_str = ( f"bf {args.inference_dynamic_batching_active_buffer_size_gb:.0f} " @@ -381,6 +357,5 @@ def build_dynamic_engine_setup_prefix( request_str, buffer_limits_str, ] - # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< return " | ".join(parts) diff --git a/megatron/core/inference/contexts/dynamic_chunk_allocator.py b/megatron/core/inference/contexts/dynamic_chunk_allocator.py index 65110aef5dc..410e8a84e81 100644 --- a/megatron/core/inference/contexts/dynamic_chunk_allocator.py +++ b/megatron/core/inference/contexts/dynamic_chunk_allocator.py @@ -45,6 +45,7 @@ def __init__(self, active_count: int): self.total_count = 2 * active_count + 1 # +1 for dummy_chunk_idx self.total_avail = self.total_count - 1 # -1 for dummy_chunk_idx self.active_count = active_count + self.paused_count = self.total_count - self.active_count - 1 # -1 for dummy_chunk_idx self.dummy_chunk_idx = self.total_count - 1 # Initialize chunk pool as a "stack" data structure @@ -87,14 +88,7 @@ def get_active_avail(self, context: "DynamicInferenceContext"): Args: context (DynamicInferenceContext): Dynamic inference context. """ - active_total = self.active_count - active_used = self.get_active_used(context) - active_avail = active_total - active_used - # >>> - # if active_used != 0: - # pax("active_total, active_used, active_avail") - # <<< - return active_avail + return self.active_count - self.get_active_used(context) def get_paused_avail(self, context: "DynamicInferenceContext"): """Compute number of paused chunks available. @@ -102,14 +96,24 @@ def get_paused_avail(self, context: "DynamicInferenceContext"): Args: context (DynamicInferenceContext): Dynamic inference context. """ - paused_total = self.total_count - self.active_count - 1 # -1 for dummy chunk - paused_used = self.get_paused_used(context) - paused_avail = paused_total - paused_used - # >>> - # if paused_used != 0: - # pax("paused_total, paused_used, paused_avail") - # <<< - return paused_avail + return self.paused_count - self.get_paused_used(context) + + def get_active_needed_addl(self, context: "DynamicInferenceContext"): + """Compute max number of additional chunks needed to complete currently + active requests. + + Args: + context (DynamicInferenceContext): Dynamic inference context. + """ + + active_request_count = context.total_request_count - context.paused_request_count + active_chunks_needed_total = active_request_count * context.max_kv_chunk_count + active_chunks_used = self.get_active_used(context) + active_chunks_needed_addl = active_chunks_needed_total - active_chunks_used + # pax("active_request_count, active_chunks_used", + # "active_chunks_needed_total, active_chunks_needed_addl") + return active_chunks_needed_addl + # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> @@ -139,17 +143,19 @@ def is_memory_available( (bool) Is memory available? """ # >>> - paused_avail = self.get_paused_avail(context) + # paused_avail = self.get_paused_avail(context) active_avail = self.get_active_avail(context) - if active_avail >= num_chunks and self.total_avail == 0: - pax("self, num_chunks, paused_avail, active_avail", { - "paused_used" : self.get_paused_used(context), - "active_used" : self.get_active_used(context), - "total_request_count" : context.total_request_count, - "paused_request_count" : context.paused_request_count, - }) + active_needed_addl = self.get_active_needed_addl(context) + # if active_avail >= num_chunks and self.total_avail == 0: + # pax("self, num_chunks", { + # "paused_used" : "%d / %d" % (self.get_paused_used(context), self.paused_count), + # "active_used" : "%d / %d" % (self.get_active_used(context), self.active_count), + # "total_request_count" : context.total_request_count, + # "paused_request_count" : context.paused_request_count, + # }, "active_addl") # <<< - return self.get_active_avail(context) >= num_chunks + # return self.get_active_avail(context) >= num_chunks + return active_avail >= num_chunks and active_needed_addl <= self.total_avail # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 051d317b56c..c5668370d0d 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1504,7 +1504,7 @@ def get_request_str(i): self.chunk_allocator.total_count - self.chunk_allocator.total_avail - 1, self.chunk_allocator.total_count - 1, self.chunk_allocator.get_paused_used(self), - self.chunk_allocator.total_count - self.chunk_allocator.active_count - 1, + self.chunk_allocator.paused_count, self.chunk_allocator.get_active_used(self), self.chunk_allocator.active_count, ) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 97305f45e9d..d142e47bba6 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1103,13 +1103,8 @@ def validate_args(args, defaults={}): assert not args.enable_cuda_graph, "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_chunk_size % 256 == 0, "chunk 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: @@ -1396,47 +1391,16 @@ 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', - # type=float, default=40., - # help='Total buffer size (GB) allocated for the chunked KV ' - # 'memory.') - # +++ group.add_argument('--inference-dynamic-batching-active-buffer-size-gb', type=float, default=40., - help='Total buffer size (GB) allocated for the chunked 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-chunk-size', type=int, default=256, help='KV cache chunk 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 blocked ' - # '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', - # 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`.') - # <<< group.add_argument('--inference-dynamic-batching-num-cuda-graphs', type=int, default=16, help='Maximum number of cuda graphs to capture, where the ' @@ -1449,15 +1413,6 @@ def _add_inference_args(parser): help='Track paused request ids by adding \'paused\' events ' 'to each request\'s event history. This has a very minor ' 'impact on latency.') - # >>> - # group.add_argument('--inference-dynamic-batching-unified-memory-level', - # type=int, default=0, 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. ' - # 'Eventually, additional levels will be included to ' - # 'control other tensors within the context.') - # <<< group.add_argument('--symmetric-ar-type', type=str, default=None, choices=['two_shot', "one_shot", "multimem_all_reduce", None], help='What type of symmetric all reduce to use. The default is none which is no use of symetric memory') From ee9a09a0549425c681d2c6509f859b4b8199a83e Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 6 Oct 2025 11:30:25 -0700 Subject: [PATCH 33/69] no adding requests if paused requests. --- .../contexts/dynamic_chunk_allocator.py | 99 +++++++++++-------- .../inference/contexts/dynamic_context.py | 88 ++--------------- .../core/inference/engines/dynamic_engine.py | 13 +-- 3 files changed, 71 insertions(+), 129 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_chunk_allocator.py b/megatron/core/inference/contexts/dynamic_chunk_allocator.py index 410e8a84e81..0de48cb5cf7 100644 --- a/megatron/core/inference/contexts/dynamic_chunk_allocator.py +++ b/megatron/core/inference/contexts/dynamic_chunk_allocator.py @@ -72,15 +72,15 @@ def get_active_used(self, context: "DynamicInferenceContext"): context.paused_request_count:context.total_request_count ].sum().item() - def get_paused_used(self, context: "DynamicInferenceContext"): - """Compute number of paused chunks used. + # def get_paused_used(self, context: "DynamicInferenceContext"): + # """Compute number of paused chunks used. - Args: - context (DynamicInferenceContext): Dynamic inference context. - """ - return context.request_kv_chunk_counts[ - :context.paused_request_count - ].sum().item() + # Args: + # context (DynamicInferenceContext): Dynamic inference context. + # """ + # return context.request_kv_chunk_counts[ + # :context.paused_request_count + # ].sum().item() def get_active_avail(self, context: "DynamicInferenceContext"): """Compute number of active chunks available. @@ -90,30 +90,29 @@ def get_active_avail(self, context: "DynamicInferenceContext"): """ return self.active_count - self.get_active_used(context) - def get_paused_avail(self, context: "DynamicInferenceContext"): - """Compute number of paused chunks available. + # def get_paused_avail(self, context: "DynamicInferenceContext"): + # """Compute number of paused chunks available. - Args: - context (DynamicInferenceContext): Dynamic inference context. - """ - return self.paused_count - self.get_paused_used(context) + # Args: + # context (DynamicInferenceContext): Dynamic inference context. + # """ + # return self.paused_count - self.get_paused_used(context) - def get_active_needed_addl(self, context: "DynamicInferenceContext"): - """Compute max number of additional chunks needed to complete currently - active requests. + # def get_active_needed_addl(self, context: "DynamicInferenceContext"): + # """Compute max number of additional chunks needed to complete currently + # active requests. - Args: - context (DynamicInferenceContext): Dynamic inference context. - """ - - active_request_count = context.total_request_count - context.paused_request_count - active_chunks_needed_total = active_request_count * context.max_kv_chunk_count - active_chunks_used = self.get_active_used(context) - active_chunks_needed_addl = active_chunks_needed_total - active_chunks_used - # pax("active_request_count, active_chunks_used", - # "active_chunks_needed_total, active_chunks_needed_addl") - return active_chunks_needed_addl + # Args: + # context (DynamicInferenceContext): Dynamic inference context. + # """ + # active_request_count = context.total_request_count - context.paused_request_count + # active_chunks_needed_total = active_request_count * context.max_kv_chunk_count + # active_chunks_used = self.get_active_used(context) + # active_chunks_needed_addl = active_chunks_needed_total - active_chunks_used + # # pax("active_request_count, active_chunks_used", + # # "active_chunks_needed_total, active_chunks_needed_addl") + # return active_chunks_needed_addl # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> @@ -128,6 +127,35 @@ def get_active_needed_addl(self, context: "DynamicInferenceContext"): # """ # return self.total_avail >= num_chunks # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ + # def is_memory_available( + # self, + # context: "DynamicInferenceContext", + # num_chunks: int, + # ) -> bool: + # """Check if memory chunks are available. + + # Args: + # context (DynamicInferenceContext): Dynamic inference context. + # num_chunks (int): Number of chunks to check. + + # Return: + # (bool) Is memory available? + # """ + # # >>> + # # paused_avail = self.get_paused_avail(context) + # active_avail = self.get_active_avail(context) + # active_needed_addl = self.get_active_needed_addl(context) + # # if active_avail >= num_chunks and self.total_avail == 0: + # # pax("self, num_chunks", { + # # "paused_used" : "%d / %d" % (self.get_paused_used(context), self.paused_count), + # # "active_used" : "%d / %d" % (self.get_active_used(context), self.active_count), + # # "total_request_count" : context.total_request_count, + # # "paused_request_count" : context.paused_request_count, + # # }, "active_addl") + # # <<< + # # return self.get_active_avail(context) >= num_chunks + # return active_avail >= num_chunks and active_needed_addl <= self.total_avail + # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ def is_memory_available( self, context: "DynamicInferenceContext", @@ -143,19 +171,10 @@ def is_memory_available( (bool) Is memory available? """ # >>> - # paused_avail = self.get_paused_avail(context) - active_avail = self.get_active_avail(context) - active_needed_addl = self.get_active_needed_addl(context) - # if active_avail >= num_chunks and self.total_avail == 0: - # pax("self, num_chunks", { - # "paused_used" : "%d / %d" % (self.get_paused_used(context), self.paused_count), - # "active_used" : "%d / %d" % (self.get_active_used(context), self.active_count), - # "total_request_count" : context.total_request_count, - # "paused_request_count" : context.paused_request_count, - # }, "active_addl") + # if context.paused_request_count != 0: + # return False # <<< - # return self.get_active_avail(context) >= num_chunks - return active_avail >= num_chunks and active_needed_addl <= self.total_avail + return self.get_active_avail(context) >= num_chunks # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index c5668370d0d..5cc09c21c3f 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -96,14 +96,8 @@ 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 n_bytes >= nquery: - # return "%d %s" % (n_bytes // nquery, suffix) - # +++ if round(n_bytes / nquery) >= 1: - # return "%d %s" % (round(n_bytes / nquery), suffix) return "%.3g %s" % (n_bytes / nquery, suffix) - # <<< raise Exception("something went wrong.") @@ -150,10 +144,6 @@ class DynamicInferenceContext(BaseInferenceContext): 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. - # >>> - # model (torch.nn.Module): Model used for inference. This is needed for - # computing `max_tokens` below. - # <<< chunk_size_tokens (int): Size of KV cache chunk size. tensor_model_parallel_size (Optional[int]): Tensor model parallel size. num_cuda_graphs (Optional[int]): Maximum number of cuda graphs to capture, @@ -222,45 +212,16 @@ def __init__( ) # Initialize chunk allocator. - # >>> - # active_buffer_size_bytes_0 = active_buffer_size_bytes - # <<< active_chunk_count_total = active_buffer_size_bytes // self.chunk_size_bytes self.chunk_allocator = ChunkAllocator(active_count=active_chunk_count_total) del active_chunk_count_total # no accidental use; use chunk_allocator - # >>> - # active_buffer_size_bytes = active_chunk_count_total * self.chunk_size_bytes active_buffer_size_bytes = self.chunk_allocator.active_count * self.chunk_size_bytes - # <<< - - # >>> - # active_buffer_size_bytes_1 = active_buffer_size_bytes - # pax("active_chunk_count_total", { - # "chunk_allocator" : self.chunk_allocator, - # }, "active_buffer_size_bytes_0, active_buffer_size_bytes_1") - # <<< # Set max_requests, max_tokens. - # >>> - # self.max_requests = active_chunk_count_total - # +++ - # self.max_requests = self.chunk_allocator.active_count - # +++ self.max_total_requests = self.chunk_allocator.total_count - 1 # -1 for dummy chunk self.max_active_requests = self.chunk_allocator.active_count - # <<< - # >>> - # pax({ - # "chunk_allocator" : self.chunk_allocator, - # "max_requests" : self.max_requests, - # }) - # <<< self.max_tokens = max_tokens - # >>> - # pax({"chunk_allocator": self.chunk_allocator}) - # <<< - # Initialize context state. self.params_dtype = params_dtype self.num_layers = num_layers @@ -408,46 +369,14 @@ def __init__( device=torch.cuda.current_device(), ) - # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> - # # Guaranteed active requests. - # # * See details in the class docstring above. `gtd_request_fraction` is - # # the fraction of chunks in the memory buffer that are reserved for - # # guaranteeing that some number of active requests can always proceed - # # with their generations. The number of chunks defined by - # # `buffer_guaranteed_fraction * chunk_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 chunks rather - # # than bytes is safer due to the non-linear impacts of a large - # # `chunk_size_tokens` or `max_kv_chunk_count`. When computing from - # # chunks, this space will always be less than `chunk_count_total`. When - # # computing from bytes, this space can unexpectedly be much larger than - # # `chunk_count_total`, resulting in stalled generations. - # gtd_chunk_count = int(buffer_guaranteed_fraction * chunk_count_total) - # gtd_chunk_count = min(gtd_chunk_count, chunk_count_total) - # self.gtd_request_count = max(1, gtd_chunk_count // self.max_kv_chunk_count) - # self.gtd_chunk_count = self.gtd_request_count * self.max_kv_chunk_count - # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - - # >>> - # # Store the dummy chunk idx reference for convenience - # self.dummy_chunk_idx = self.chunk_allocator.dummy_chunk_idx - # <<< - # Reset attention state. self.reset_attention_state() - # >>> # Print info. print("DynamicInferenceContext: allocated context with active buffer size %s (%d chunks)." % ( get_mem_size_str(active_buffer_size_bytes), self.chunk_allocator.active_count, )) - # <<< - - # >>> - # print("~~~\nexiting ..."); exit() - # <<< TOKEN_ROUNDER = 64 REQUEST_ROUNDER = 4 @@ -1006,6 +935,7 @@ def add_request( self.total_request_count - self.paused_request_count >= self.max_active_requests # or # self.get_active_request_count() >= self.chunk_allocator.active_count + or self.paused_request_count != 0 ): raise RequestOverflowError(request_id) # <<< @@ -1356,10 +1286,11 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T # >>> # if active_request_count + resume_request_count == 0: # if active_request_count == 0: - if resume_request_count >= 2: - pax("active_chunk_count_total, active_chunk_count_used, active_chunk_count_avail", - "paused_chunk_counts, paused_chunk_counts_cumsum", - "active_request_count, resume_request_count") + # if resume_request_count >= 2: + # if True: + # pax("active_chunk_count_avail", + # "paused_chunk_counts, paused_chunk_counts_cumsum", + # "active_request_count, resume_request_count") # <<< # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< @@ -1500,11 +1431,12 @@ def get_request_str(i): newly_resumed_request_ids.tolist(), self.request_ids[self.total_request_count:(self.total_request_count + finished_request_count)].tolist(), ) - alloc_str = "alloc: t %d/%d [ p %d/%d, a %d/%d ]" % ( + # alloc_str = "alloc: t %d/%d [ p %d/%d, a %d/%d ]" % ( + alloc_str = "alloc: t %d/%d [ a %d/%d ]" % ( self.chunk_allocator.total_count - self.chunk_allocator.total_avail - 1, self.chunk_allocator.total_count - 1, - self.chunk_allocator.get_paused_used(self), - self.chunk_allocator.paused_count, + # self.chunk_allocator.get_paused_used(self), + # self.chunk_allocator.paused_count, self.chunk_allocator.get_active_used(self), self.chunk_allocator.active_count, ) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 391aa53cb9b..b87416f0e4d 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -517,9 +517,6 @@ def post_process_requests( def schedule_waiting_requests(self): """Tries to schedule any requests in the waiting pool.""" - # >>> - # print(" > schedule_waiting_requests: %s." % self.waiting_request_ids) - # <<< for waiting_request_id in self.waiting_request_ids.copy(): waiting_request: DynamicInferenceRequest = self.requests[waiting_request_id] try: @@ -530,13 +527,7 @@ def schedule_waiting_requests(self): ) waiting_request.add_event_add() self.waiting_request_ids.popleft() - # >>> - # print(" > schedule_waiting_requests: %s ... (success)." % waiting_request_id) - # <<< except Exception as e: - # >>> - # print(" > schedule_waiting_requests: %s ... (%s)." % (waiting_request_id, e.__class__.__name__)) - # <<< break async def async_step( @@ -664,8 +655,8 @@ async def async_step( self.step_count += 1 # >>> - # if self.step_count >= 602: # 600, 300 - # raise Exception(f"step {self.step_count}.") + if self.step_count >= 1505: # 600, 305 + raise Exception(f"step {self.step_count}.") # <<< range_pop() return { From f52030453de056602ac8a005f89410f39923c1ed Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 7 Oct 2025 04:39:14 -0700 Subject: [PATCH 34/69] passed context to chunk allocator's constructor. --- .../inference/gpt/gpt_dynamic_inference.py | 13 +- .../contexts/dynamic_chunk_allocator.py | 177 +++--------------- .../inference/contexts/dynamic_context.py | 53 +----- .../core/inference/engines/dynamic_engine.py | 4 - 4 files changed, 40 insertions(+), 207 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index 837ee145b33..25c740d824c 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -66,8 +66,8 @@ import builtins builtins.pax = _pax -def print_mem(key): - print(f"+++++++++++ {key} ........... {get_mem_stats_str()}.") +# def print_mem(key): +# print(f"+++++++++++ {key} ........... {get_mem_stats_str()}.") # <<< @@ -135,6 +135,7 @@ def get_inference_context_active_buffer_size_bytes(model: torch.nn.Module) -> in return active_buffer_size_bytes +# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> # # >>> # def get_inference_context_max_tokens(model: torch.nn.Module) -> int: # # def get_inference_context_max_tokens(controller: torch.nn.Module) -> int: @@ -206,15 +207,13 @@ def get_inference_context_max_tokens(model: torch.nn.Module) -> int: correlation between `max_tokens` and request throughput. """ return 16384 +# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< def get_inference_context( requests: List[Request], sampling_params: SamplingParams, - # >>> - # buffer_size_bytes: int, model: torch.nn.Module, - # <<< calculate_max_sequence_length_from_requests: bool =True, ): """The inference context manages the KV cache and other inference state.""" @@ -424,7 +423,7 @@ def main(): ) # >>> - print_mem("initialize_megatron()") + # print_mem("initialize_megatron()") # <<< # Start Nsight profiler. @@ -473,7 +472,7 @@ def main(): ) # >>> - print_mem("engine") + # print_mem("engine") # exit() # <<< diff --git a/megatron/core/inference/contexts/dynamic_chunk_allocator.py b/megatron/core/inference/contexts/dynamic_chunk_allocator.py index 0de48cb5cf7..cbe88485f9c 100644 --- a/megatron/core/inference/contexts/dynamic_chunk_allocator.py +++ b/megatron/core/inference/contexts/dynamic_chunk_allocator.py @@ -15,31 +15,15 @@ class ChunkAllocator: - Releasing chunks back to the pool Args: + context (DynamicInferenceContext): Dynamic inference context. active_count (int): Total number of active chunks 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, active_count: int): - - # active_count -= 1 # subtract 1 for dummy_chunk_idx (see below) - # self.total_count = 2 * active_count - # self.active_count = active_count - # self.paused_count = active_count - # # >>> - # # self.active_avail = active_count - 1 # reserve last chunk for decode-only steps - # self.active_avail = active_count - # # <<< - # self.paused_avail = active_count - # self.dummy_chunk_idx = self.total_count - 1 - - # # Initialize chunk pool as a "stack" data structure - # self.chunk_bag = torch.arange( - # self.total_count, - # dtype=torch.int32, - # device=torch.cuda.current_device(), - # ) - def __init__(self, active_count: int): + def __init__(self, context: "DynamicInferenceContext", active_count: int): + + self.context = context active_count -= 1 # -1 for dummy_chunk_idx (see below) self.total_count = 2 * active_count + 1 # +1 for dummy_chunk_idx @@ -61,164 +45,53 @@ def __str__(self): f"; active {self.active_count}" ) - # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> - def get_active_used(self, context: "DynamicInferenceContext"): - """Compute number of active chunks used. - - Args: - context (DynamicInferenceContext): Dynamic inference context. - """ - return context.request_kv_chunk_counts[ - context.paused_request_count:context.total_request_count + def get_active_used(self): + """Compute number of active chunks used.""" + return self.context.request_kv_chunk_counts[ + self.context.paused_request_count:self.context.total_request_count ].sum().item() - # def get_paused_used(self, context: "DynamicInferenceContext"): - # """Compute number of paused chunks used. + def get_paused_used(self): + """Compute number of paused chunks used.""" + return self.context.request_kv_chunk_counts[ + :self.context.paused_request_count + ].sum().item() - # Args: - # context (DynamicInferenceContext): Dynamic inference context. - # """ - # return context.request_kv_chunk_counts[ - # :context.paused_request_count - # ].sum().item() + def get_active_avail(self): + """Compute number of active chunks available.""" + return self.active_count - self.get_active_used() - def get_active_avail(self, context: "DynamicInferenceContext"): - """Compute number of active chunks available. + def get_paused_avail(self): + """Compute number of paused chunks available.""" + return self.paused_count - self.get_paused_used() - Args: - context (DynamicInferenceContext): Dynamic inference context. - """ - return self.active_count - self.get_active_used(context) - - # def get_paused_avail(self, context: "DynamicInferenceContext"): - # """Compute number of paused chunks available. - - # Args: - # context (DynamicInferenceContext): Dynamic inference context. - # """ - # return self.paused_count - self.get_paused_used(context) - - # def get_active_needed_addl(self, context: "DynamicInferenceContext"): - # """Compute max number of additional chunks needed to complete currently - # active requests. - - # Args: - # context (DynamicInferenceContext): Dynamic inference context. - # """ - - # active_request_count = context.total_request_count - context.paused_request_count - # active_chunks_needed_total = active_request_count * context.max_kv_chunk_count - # active_chunks_used = self.get_active_used(context) - # active_chunks_needed_addl = active_chunks_needed_total - active_chunks_used - # # pax("active_request_count, active_chunks_used", - # # "active_chunks_needed_total, active_chunks_needed_addl") - # return active_chunks_needed_addl - # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - - # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> - # def is_memory_available(self, num_chunks: int) -> bool: - # """Check if memory chunks are available. - - # Args: - # num_chunks (int): Number of chunks to check. - - # Return: - # (bool) Is memory available? - # """ - # return self.total_avail >= num_chunks - # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - # def is_memory_available( - # self, - # context: "DynamicInferenceContext", - # num_chunks: int, - # ) -> bool: - # """Check if memory chunks are available. - - # Args: - # context (DynamicInferenceContext): Dynamic inference context. - # num_chunks (int): Number of chunks to check. - - # Return: - # (bool) Is memory available? - # """ - # # >>> - # # paused_avail = self.get_paused_avail(context) - # active_avail = self.get_active_avail(context) - # active_needed_addl = self.get_active_needed_addl(context) - # # if active_avail >= num_chunks and self.total_avail == 0: - # # pax("self, num_chunks", { - # # "paused_used" : "%d / %d" % (self.get_paused_used(context), self.paused_count), - # # "active_used" : "%d / %d" % (self.get_active_used(context), self.active_count), - # # "total_request_count" : context.total_request_count, - # # "paused_request_count" : context.paused_request_count, - # # }, "active_addl") - # # <<< - # # return self.get_active_avail(context) >= num_chunks - # return active_avail >= num_chunks and active_needed_addl <= self.total_avail - # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - def is_memory_available( - self, - context: "DynamicInferenceContext", - num_chunks: int, - ) -> bool: + def is_memory_available(self, num_chunks: int) -> bool: """Check if memory chunks are available. Args: - context (DynamicInferenceContext): Dynamic inference context. num_chunks (int): Number of chunks to check. Return: (bool) Is memory available? """ - # >>> - # if context.paused_request_count != 0: - # return False - # <<< - return self.get_active_avail(context) >= num_chunks - # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - - # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> - # def allocate_memory_chunks(self, num_chunks: int = 1) -> Optional[Tensor]: - # """Allocate memory chunks if available, else return None. - - # Args: - # num_chunks (int): Number of chunks to allocate. - - # Return: - # (Optional[Tensor]) Allocated chunk IDs. - # """ - # if self.is_memory_available(num_chunks): - # self.total_avail -= num_chunks - # return self.chunk_bag[self.total_avail : (self.total_avail + num_chunks)] - # else: - # return None - # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - def allocate_memory_chunks( - self, - context: "DynamicInferenceContext", - num_chunks: int, - ) -> Optional[Tensor]: + return self.get_active_avail() >= num_chunks + + def allocate_memory_chunks(self, num_chunks: int) -> Optional[Tensor]: """Allocate memory chunks if available, else return None. Args: - context (DynamicInferenceContext): Dynamic inference context. num_chunks (int): Number of chunks to allocate. Return: (Optional[Tensor]) Allocated chunk IDs. """ - if self.is_memory_available(context, num_chunks): + if self.is_memory_available(num_chunks): self.total_avail -= num_chunks chunk_ids = self.chunk_bag[self.total_avail : (self.total_avail + num_chunks)] - # >>> - # assert num_chunks == chunk_ids.numel() - if num_chunks != chunk_ids.numel(): - pax("self, num_chunks") - # <<< + assert num_chunks == chunk_ids.numel() return chunk_ids else: return None - # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< def release_memory_chunks(self, chunks: Tensor) -> None: """Release memory chunks. diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 5cc09c21c3f..448ef72af6f 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -213,7 +213,10 @@ def __init__( # Initialize chunk allocator. active_chunk_count_total = active_buffer_size_bytes // self.chunk_size_bytes - self.chunk_allocator = ChunkAllocator(active_count=active_chunk_count_total) + self.chunk_allocator = ChunkAllocator( + context=self, + active_count=active_chunk_count_total, + ) del active_chunk_count_total # no accidental use; use chunk_allocator active_buffer_size_bytes = self.chunk_allocator.active_count * self.chunk_size_bytes @@ -927,34 +930,15 @@ def add_request( # TODO : Should move this into some waiting queue if self.active_token_count + context_length > self.max_tokens: raise TokenOverflowError(request_id) - # >>> - # if self.total_request_count >= self.max_total_requests: - # raise RequestOverflowError(request_id) - # +++ if ( self.total_request_count - self.paused_request_count >= self.max_active_requests - # or - # self.get_active_request_count() >= self.chunk_allocator.active_count or self.paused_request_count != 0 ): raise RequestOverflowError(request_id) - # <<< # Preallocate chunks. num_chunks_needed = math.ceil(context_length / self.chunk_size_tokens) - # >>> - # new_chunk_ids = self.chunk_allocator.allocate_memory_chunks(num_chunks_needed) - # +++ - new_chunk_ids = self.chunk_allocator.allocate_memory_chunks(self, num_chunks_needed) - # <<< - # >>> - # if request_id == 4: - # pax({ - # "chunk_allocator" : self.chunk_allocator, - # "total_request_count" : self.total_request_count, - # "paused_request_count" : self.paused_request_count, - # }, "new_chunk_ids") - # <<< + new_chunk_ids = self.chunk_allocator.allocate_memory_chunks(num_chunks_needed) if new_chunk_ids is None: raise ChunkOverflowError(request_id) @@ -1263,7 +1247,7 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ resume_request_count = 0 if self.paused_request_count > 0: - active_chunk_count_avail = self.chunk_allocator.get_active_avail(self) + active_chunk_count_avail = self.chunk_allocator.get_active_avail() # pax("active_chunk_count_avail") paused_chunk_counts = self.request_kv_chunk_counts[:self.paused_request_count] @@ -1358,17 +1342,7 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T == 0 ), "The request_last_kv_chunk_offset should be 0 for the requests that just got resumed this step. " - # >>> - # chunk_ids = self.chunk_allocator.allocate_memory_chunks(resume_request_count) - # +++ - chunk_ids = self.chunk_allocator.allocate_memory_chunks(self, resume_request_count) - # <<< - # >>> - if chunk_ids is None: - pax({ - "chunk_allocator" : self.chunk_allocator, - }, "resume_request_count, chunk_ids") - # <<< + chunk_ids = self.chunk_allocator.allocate_memory_chunks(resume_request_count) row_idx = torch.arange( self.paused_request_count, self.paused_request_count + resume_request_count, @@ -1377,16 +1351,7 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T col_idx = self.request_kv_chunk_counts[ self.paused_request_count : (self.paused_request_count + resume_request_count) ] - # >>> - # self.request_to_kv_chunk_ids[row_idx, col_idx] = chunk_ids - # +++ - try: - self.request_to_kv_chunk_ids[row_idx, col_idx] = chunk_ids - except Exception as e: - pax({ - "request_to_kv_chunk_ids" : self.request_to_kv_chunk_ids, - }, "row_idx, col_idx, resume_request_count, chunk_ids") - # <<< + self.request_to_kv_chunk_ids[row_idx, col_idx] = chunk_ids self.request_kv_chunk_counts[ self.paused_request_count : (self.paused_request_count + resume_request_count) ] += 1 @@ -1437,7 +1402,7 @@ def get_request_str(i): self.chunk_allocator.total_count - 1, # self.chunk_allocator.get_paused_used(self), # self.chunk_allocator.paused_count, - self.chunk_allocator.get_active_used(self), + self.chunk_allocator.get_active_used(), self.chunk_allocator.active_count, ) print("++++++++++ %s" % " ..... ".join(( diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index b87416f0e4d..919a50075b2 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -654,10 +654,6 @@ async def async_step( print(output_str) self.step_count += 1 - # >>> - if self.step_count >= 1505: # 600, 305 - raise Exception(f"step {self.step_count}.") - # <<< range_pop() return { "active_requests": active_requests, From 12b986655f81a4c1b19ae28efeb3508d714e5189 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 7 Oct 2025 04:56:08 -0700 Subject: [PATCH 35/69] some clean up. --- .../inference/gpt/gpt_dynamic_inference.py | 1 + .../inference/contexts/dynamic_context.py | 95 +------------------ 2 files changed, 4 insertions(+), 92 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index 25c740d824c..331ccc87bdc 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -449,6 +449,7 @@ def main(): model = get_model() requests = build_requests(args, tokenizer) context = get_inference_context(requests, sampling_params, model) + pax("requests, context") controller = get_inference_controller(model, context) # Validate all context_length's <= max_tokens. diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 448ef72af6f..d7b230310b2 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -988,11 +988,6 @@ def add_request( self.total_request_count += 1 self.active_token_count += context_length - # >>> - # print(f"++++++++ chunk allocator ... {self.chunk_allocator}.") - # pax({"max_requests": self.max_requests}) - # <<< - def _move_book_keeping_tensors(self, src_idxs, dst_idxs, next_tokens): """ Swaps all the relevent booking tensors with src idxs to dst idxs @@ -1054,10 +1049,6 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T # finished_request_count are requests that have reached the termination criterion active_request_count = (active_requests_mask == 1).sum().item() finished_request_count = (active_requests_mask == 0).sum().item() - # >>> - # if finished_request_count > 0: - # raise Exception("huh?") - # <<< assert ( active_request_count + finished_request_count + self.paused_request_count == self.total_request_count @@ -1148,21 +1139,12 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T active_requests_requiring_new_chunk_count = ( (active_requests_requiring_new_chunk == 1).sum().item() ) - # >>> - # if active_requests_requiring_new_chunk_count > 0: - # pax("active_requests_requiring_new_chunk") - # <<< - # >>> [ new ] if active_requests_requiring_new_chunk_count > 0: newly_paused_request_ids = self.request_ids[ torch.nonzero(active_requests_requiring_new_chunk) + self.paused_request_count ] - # >>> - # pax("newly_paused_request_ids") - # <<< - # <<< # Swap unfinished active requests on the left side with paused requests on the right side # NOTE : We add paused request count because we concatenate @@ -1193,58 +1175,17 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T ) dst_idxs = torch.cat((active_request_ids_on_left, paused_requests_idxs_on_right)) src_idxs = torch.cat((paused_requests_idxs_on_right, active_request_ids_on_left)) - # >>> [ new ] - # newly_paused_request_ids = self.request_ids[dst_idxs] - # <<< self._move_book_keeping_tensors( src_idxs=src_idxs, dst_idxs=dst_idxs, next_tokens=next_tokens ) - # >>> [ orig ] - # newly_paused_request_ids = self.request_ids[dst_idxs] - # <<< - # >>> - # pax("active_request_ids_on_left, paused_requests_idxs_on_right") - # <<< self.paused_request_count += active_requests_requiring_new_chunk_count active_request_count -= active_requests_requiring_new_chunk_count - # >>> - # if active_requests_requiring_new_chunk_count > 0: - # pax({ - # "chunk_allocator" : self.chunk_allocator, - # "paused_request_count" : self.paused_request_count, - # }, "active_request_count") - # <<< - # 6. Now that we have the requests in following order [Paused, Active, Finished] # We determine how many requests we can resume and resume them # Assign released chunks to paused requests. # todo: @shanmugamr, un-pause requests using FIFO, rather than LIFO. - # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> - # num_non_gtd_chunks = max(0, self.chunk_allocator.chunk_count_avail - self.gtd_chunk_count) - # if num_non_gtd_chunks: - # # if we have non-gtd chunks, use them. Do not dip into the gtd-chunk pool - # resume_request_count = min(num_non_gtd_chunks, self.paused_request_count) - # else: - # # only dip into the gtd-chunk pool if we have run out of non-gtd-chunks and the active - # # request count has fallen below a certain threshold. - # resume_request_count = min( - # max(self.gtd_request_count - active_request_count, 0), self.paused_request_count - # ) - # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ - # if active_request_count < self.chunk_allocator.active_count: - # resume_request_count = min( - # # >>> - # # self.chunk_allocator.active_avail, - # # <<< - # self.chunk_allocator.total_avail, - # self.chunk_allocator.active_count - active_request_count, - # self.paused_request_count, - # ) - # else: - # resume_request_count = 0 - # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ resume_request_count = 0 if self.paused_request_count > 0: active_chunk_count_avail = self.chunk_allocator.get_active_avail() @@ -1254,43 +1195,14 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T paused_chunk_counts = paused_chunk_counts.flip(dims=[0]) paused_chunk_counts += 1 # +1 for newly added chunk paused_chunk_counts_cumsum = paused_chunk_counts.cumsum(dim=0) - # >>> - # paused_chunk_counts_cumsum += 1 # +1 for newly added chunk - # <<< resume_request_count = torch.nonzero( paused_chunk_counts_cumsum <= active_chunk_count_avail ).numel() - # for i in range(self.paused_request_count - 1, -1, -1): - # current_chunk_count - # pax("i") - # if resume_request_count != 1: - # pax("resume_request_count") - - # >>> - # if active_request_count + resume_request_count == 0: - # if active_request_count == 0: - # if resume_request_count >= 2: - # if True: - # pax("active_chunk_count_avail", - # "paused_chunk_counts, paused_chunk_counts_cumsum", - # "active_request_count, resume_request_count") - # <<< - # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - # >>> newly_resumed_request_ids = self.request_ids[(self.paused_request_count - resume_request_count):self.paused_request_count] # <<< - # >>> - # if resume_request_count != 0 and self.chunk_allocator.total_avail == 0: - if active_request_count + resume_request_count == 0: - pax({ - "chunk_allocator" : self.chunk_allocator, - "paused_request_count" : self.paused_request_count, - }, "active_request_count, resume_request_count") - # <<< - self.paused_request_count -= resume_request_count active_request_count += resume_request_count assert active_request_count > 0, "active_request_count == %d." % active_request_count @@ -1302,10 +1214,9 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T # requests are newly paused during this update. if newly_paused_request_ids is not None and resume_request_count > 0: newly_paused_request_ids = newly_paused_request_ids[:-resume_request_count] - # >>> - # if newly_paused_request_ids is not None and newly_paused_request_ids.numel() > 0: - # pax("newly_paused_request_ids") - # <<< + # >>> + # newly_resumed_request_ids = newly_resumed_request_ids[-resume_request_count:] + # <<< # 7. We make changes to the request book keeping tesnsors and setup the tokens for next iteration # All these active requests are in decode phase, so they need only 1 token per request From 185ffbf28712c6f39af1ca45f90d0414bb591447 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 7 Oct 2025 05:27:19 -0700 Subject: [PATCH 36/69] 4gb ran smoothly. --- .../inference/gpt/gpt_dynamic_inference.py | 4 +- .../inference/contexts/dynamic_context.py | 72 +++++++++---------- .../core/inference/engines/dynamic_engine.py | 7 +- 3 files changed, 45 insertions(+), 38 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index 331ccc87bdc..8bfa966f0cf 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -449,7 +449,9 @@ def main(): model = get_model() requests = build_requests(args, tokenizer) context = get_inference_context(requests, sampling_params, model) - pax("requests, context") + # >>> + # pax({"requests / len": len(requests)}) + # <<< controller = get_inference_controller(model, context) # Validate all context_length's <= max_tokens. diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index d7b230310b2..51852f7cb69 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1200,7 +1200,7 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T ).numel() # >>> - newly_resumed_request_ids = self.request_ids[(self.paused_request_count - resume_request_count):self.paused_request_count] + # newly_resumed_request_ids = self.request_ids[(self.paused_request_count - resume_request_count):self.paused_request_count] # <<< self.paused_request_count -= resume_request_count @@ -1286,41 +1286,41 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T ) # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> - def get_request_str(i): - r = self.request_ids[i].item() - q = self.request_query_lengths[i].item() - kv = self.request_kv_length_offsets[i].item() - # pax("request_id") - return "%d[%d:%d]" % (r, q, kv) - request_strs = [ - get_request_str(i) - for i in range(self.total_request_count + finished_request_count) - ] - # pax("request_strs") - request_str = "requests: %s -- %s -- %s" % ( - ", ".join(request_strs[:self.paused_request_count]), - ", ".join(request_strs[self.paused_request_count:self.total_request_count]), - ", ".join(request_strs[self.total_request_count:(self.total_request_count + finished_request_count)]), - ) - newly_str = "paused: %s, resumed: %s, finished: %s" % ( - "--" if newly_paused_request_ids is None else newly_paused_request_ids.tolist(), - newly_resumed_request_ids.tolist(), - self.request_ids[self.total_request_count:(self.total_request_count + finished_request_count)].tolist(), - ) - # alloc_str = "alloc: t %d/%d [ p %d/%d, a %d/%d ]" % ( - alloc_str = "alloc: t %d/%d [ a %d/%d ]" % ( - self.chunk_allocator.total_count - self.chunk_allocator.total_avail - 1, - self.chunk_allocator.total_count - 1, - # self.chunk_allocator.get_paused_used(self), - # self.chunk_allocator.paused_count, - self.chunk_allocator.get_active_used(), - self.chunk_allocator.active_count, - ) - print("++++++++++ %s" % " ..... ".join(( - request_str, - # newly_str, - alloc_str, - ))) + # def get_request_str(i): + # r = self.request_ids[i].item() + # q = self.request_query_lengths[i].item() + # kv = self.request_kv_length_offsets[i].item() + # # pax("request_id") + # return "%d[%d:%d]" % (r, q, kv) + # request_strs = [ + # get_request_str(i) + # for i in range(self.total_request_count + finished_request_count) + # ] + # # pax("request_strs") + # request_str = "requests: %s -- %s -- %s" % ( + # ", ".join(request_strs[:self.paused_request_count]), + # ", ".join(request_strs[self.paused_request_count:self.total_request_count]), + # ", ".join(request_strs[self.total_request_count:(self.total_request_count + finished_request_count)]), + # ) + # newly_str = "paused: %s, resumed: %s, finished: %s" % ( + # "--" if newly_paused_request_ids is None else newly_paused_request_ids.tolist(), + # newly_resumed_request_ids.tolist(), + # self.request_ids[self.total_request_count:(self.total_request_count + finished_request_count)].tolist(), + # ) + # # alloc_str = "alloc: t %d/%d [ p %d/%d, a %d/%d ]" % ( + # alloc_str = "alloc: t %d/%d [ a %d/%d ]" % ( + # self.chunk_allocator.total_count - self.chunk_allocator.total_avail - 1, + # self.chunk_allocator.total_count - 1, + # # self.chunk_allocator.get_paused_used(self), + # # self.chunk_allocator.paused_count, + # self.chunk_allocator.get_active_used(), + # self.chunk_allocator.active_count, + # ) + # print("++++++++++ %s" % " ..... ".join(( + # request_str, + # # newly_str, + # alloc_str, + # ))) # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< return newly_paused_request_ids diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 919a50075b2..3e8297815de 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -622,7 +622,7 @@ async def async_step( step_type = "decode" if is_decode_only else "non-decode" output_str = ( "* step %d | %s ... time: %.3f%s ... " - "reqs: %d [ active %d, paused %d, finished %d ] ... " + "reqs: %d [ active %d, paused %d, waiting %d, finished %d ] ... " "mem: tensors %d, alloc %.1f gb, res %.1f gb." % ( self.step_count, @@ -643,6 +643,7 @@ async def async_step( prev_total_request_count, prev_total_request_count - prev_paused_request_count, prev_paused_request_count, + len(self.waiting_request_ids), self.finished_request_count, mem["allocation.all.current"], mem["allocated_bytes.all.current"] / (1024**3), @@ -654,6 +655,10 @@ async def async_step( print(output_str) self.step_count += 1 + # >>> + # if self.step_count >= 100: # 600, 305 + # raise Exception(f"step {self.step_count}.") + # <<< range_pop() return { "active_requests": active_requests, From f45f3bfefefe7ff66cdf1094d6d8dead8df81e37 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Thu, 9 Oct 2025 06:02:48 -0700 Subject: [PATCH 37/69] fixed synthetic prompts to have expected length. --- .../inference/gpt/gpt_dynamic_inference.py | 32 ++++++++--- examples/inference/gpt/utils.py | 54 +++++++++++++++++-- .../contexts/dynamic_chunk_allocator.py | 1 + .../inference/contexts/dynamic_context.py | 10 +++- .../core/inference/engines/dynamic_engine.py | 21 +++++++- 5 files changed, 104 insertions(+), 14 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index 8bfa966f0cf..b1799bfcd99 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -137,6 +137,7 @@ def get_inference_context_active_buffer_size_bytes(model: torch.nn.Module) -> in # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> # # >>> +# @torch.inference_mode() # def get_inference_context_max_tokens(model: torch.nn.Module) -> int: # # def get_inference_context_max_tokens(controller: torch.nn.Module) -> int: # # <<< @@ -155,7 +156,10 @@ def get_inference_context_active_buffer_size_bytes(model: torch.nn.Module) -> in # num_tokens_min = num_tokens_step # num_tokens_max = 24000 # 1000000 # time_map = {} -# tbar = tqdm(range(num_tokens_min, num_tokens_max, num_tokens_step)) +# # >>> +# # tbar = tqdm(range(num_tokens_min, num_tokens_max, num_tokens_step)) +# tbar = tqdm([ 16384 ]) +# # <<< # # tbar = tqdm((3000,)) # # tbar = tqdm((3072,)) # # tbar = tqdm((1000,)) @@ -184,13 +188,20 @@ def get_inference_context_active_buffer_size_bytes(model: torch.nn.Module) -> in # # ) # t = time.time() # for _ in range(3): -# result = model( -# input_ids, -# position_ids, -# attention_mask=None, -# inference_context=None, -# runtime_gather_output=True, # Inference should always gather the logits -# ) +# mem0 = torch.cuda.memory_stats() +# with torch.inference_mode(): +# result = model( +# input_ids, +# position_ids, +# attention_mask=None, +# inference_context=None, +# runtime_gather_output=True, # Inference should always gather the logits +# ) +# mem1 = torch.cuda.memory_stats() +# pax("mem0, mem1", { +# "peak / 0" : mem0["allocated_bytes.all.peak"], +# "peak / 1" : mem1["allocated_bytes.all.peak"], +# }) # t = time.time() - t # time_map[num_tokens] = t # print("~~ x ~~") @@ -199,6 +210,7 @@ def get_inference_context_active_buffer_size_bytes(model: torch.nn.Module) -> in # [ print(y) for y in time_map.values() ] # raise Exception("matplotlib.") # return max_tokens +# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ def get_inference_context_max_tokens(model: torch.nn.Module) -> int: """Estimate `max_tokens` for the context. @@ -227,6 +239,10 @@ def get_inference_context( max_sequence_length = max_context_length + max_gen_length else: max_sequence_length = args.inference_max_seq_length + # >>> + pax("max_context_length, max_gen_length", + "calculate_max_sequence_length_from_requests, max_sequence_length") + # <<< # Active buffer size & max tokens. active_buffer_size_bytes = get_inference_context_active_buffer_size_bytes(model) diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py index c81553f5164..dc0efd59de0 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 @@ -235,11 +236,54 @@ def get_synthetic_requests(args: Namespace, tokenizer: Any) -> list[Request]: 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 ] + + # >>> + # pax({ + # "args.num_tokens_to_prompt" : args.num_tokens_to_prompt, + # }, "max_prompt_length, max_prompt_text, max_prompt_tokens", { + # "prompt_lengths" : "%d / %s" % (len(prompt_lengths), prompt_lengths), + # "prompt_tokens_list" : [ + # "%d[%d] / '%s' ... %s" % (prompt_lengths[i], len(tt), prompt_texts[i], tt) + # for i, tt in enumerate(prompt_tokens_list) + # ], + # }) + # <<< + # Init requests. + # >>> + # requests = [ + # Request("hi " * random.randint(*args.num_tokens_to_prompt), t, tokenizer) + # for t in time_offsets + # ] + # +++ + assert len(prompt_texts) == len(time_offsets) requests = [ - Request("hi " * random.randint(*args.num_tokens_to_prompt), t, tokenizer) - for t in time_offsets + Request(t, o, tokenizer) + for t, o in zip(prompt_texts, time_offsets) ] + # <<< + + # >>> + # pax("requests", { + # "requests / 0" : requests[0], + # "tokens" : [ "%d, '%s'" % (i, tokenizer.detokenize([i])) for i in requests[0].prompt_tokens ], + # }) + # <<< return requests @@ -346,7 +390,11 @@ def build_dynamic_engine_setup_prefix( # Buffer limits config buffer_limits_str = ( - f"bf {args.inference_dynamic_batching_active_buffer_size_gb:.0f} " + # >>> + # f"bf: {args.inference_dynamic_batching_active_buffer_size_gb:.0f} gb, " + f"bf: {get_mem_size_str(args.inference_dynamic_batching_active_buffer_size_gb*1024**3)}, " + # <<< + f"{context.chunk_allocator.active_count} chunks " f"[r {context.max_active_requests}, t {context.max_tokens}]" ) diff --git a/megatron/core/inference/contexts/dynamic_chunk_allocator.py b/megatron/core/inference/contexts/dynamic_chunk_allocator.py index cbe88485f9c..db7ec185f04 100644 --- a/megatron/core/inference/contexts/dynamic_chunk_allocator.py +++ b/megatron/core/inference/contexts/dynamic_chunk_allocator.py @@ -26,6 +26,7 @@ def __init__(self, context: "DynamicInferenceContext", active_count: int): self.context = context active_count -= 1 # -1 for dummy_chunk_idx (see below) + active_count = max(1, active_count) # need at least one chunk self.total_count = 2 * active_count + 1 # +1 for dummy_chunk_idx self.total_avail = self.total_count - 1 # -1 for dummy_chunk_idx self.active_count = active_count diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 51852f7cb69..7fd689f05d5 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -98,7 +98,7 @@ def get_mem_size_str(n_bytes: int) -> str: nquery = int(1024**exp) if round(n_bytes / nquery) >= 1: return "%.3g %s" % (n_bytes / nquery, suffix) - raise Exception("something went wrong.") + raise Exception(f"something went wrong, n_bytes={n_bytes}.") # pylint: disable=line-too-long @@ -217,7 +217,7 @@ def __init__( context=self, active_count=active_chunk_count_total, ) - del active_chunk_count_total # no accidental use; use chunk_allocator + del active_chunk_count_total # use self.chunk_allocator.active_count active_buffer_size_bytes = self.chunk_allocator.active_count * self.chunk_size_bytes # Set max_requests, max_tokens. @@ -300,6 +300,12 @@ def __init__( # Chunk ids. self.max_kv_chunk_count = math.ceil(self.max_sequence_length / self.chunk_size_tokens) + # >>> + pax({ + "max_sequence_length" : self.max_sequence_length, + "max_kv_chunk_count" : self.max_kv_chunk_count, + }) + # <<< self.request_to_kv_chunk_ids = torch.full( (self.max_total_requests, self.max_kv_chunk_count), -1, diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 3e8297815de..16c37496286 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -3,6 +3,7 @@ import asyncio import multiprocessing import os +import psutil import struct import time import warnings @@ -19,6 +20,7 @@ from megatron.core.inference.contexts.dynamic_context import ( ContextOverflowError, DynamicInferenceContext, + get_mem_size_str, WarmupEngineMode, ) from megatron.core.inference.data_parallel_inference_coordinator import ( @@ -56,6 +58,9 @@ HAVE_MSGPACK = False +process = psutil.Process(os.getpid()) + + def format_mem_bytes(mem_bytes): """Convert a byte count to a human-readable string in tb, gb, mb, kb, or bytes.""" for power, suffix in [(4, "tb"), (3, "gb"), (2, "mb"), (1, "kb"), (0, "bytes")]: @@ -620,10 +625,16 @@ async def async_step( context = self.context mem = torch.cuda.memory_stats() step_type = "decode" if is_decode_only else "non-decode" + # >>> + mem_info = process.memory_info() + # <<< output_str = ( "* step %d | %s ... time: %.3f%s ... " - "reqs: %d [ active %d, paused %d, waiting %d, finished %d ] ... " + "reqs: %d [ active %d/%d, paused %d/%d, waiting %d, finished %d ] ... " + # >>> "mem: tensors %d, alloc %.1f gb, res %.1f gb." + # "mem: alloc %.1f gb, res %.1f gb, cpu p %s/v %s." + # <<< % ( self.step_count, datetime.now().strftime("%H:%M:%S"), @@ -642,12 +653,20 @@ async def async_step( ), prev_total_request_count, prev_total_request_count - prev_paused_request_count, + context.chunk_allocator.active_count, prev_paused_request_count, + context.chunk_allocator.paused_count, len(self.waiting_request_ids), self.finished_request_count, + # >>> mem["allocation.all.current"], + # <<< mem["allocated_bytes.all.current"] / (1024**3), mem["reserved_bytes.all.current"] / (1024**3), + # >>> + # get_mem_size_str(mem_info.rss), + # get_mem_size_str(mem_info.vms), + # <<< ) ) if prev_is_decode_only: From eb7bb9d9423d73150694cfb09295ee65fea9a7be Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 13 Oct 2025 16:32:52 -0700 Subject: [PATCH 38/69] fixed resume_request_count. --- .../inference/gpt/gpt_dynamic_inference.py | 31 ++++++++----- .../inference/contexts/dynamic_context.py | 43 ++++++++++++++++--- .../core/inference/engines/dynamic_engine.py | 34 ++++++++++++--- 3 files changed, 84 insertions(+), 24 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index b1799bfcd99..0814d188eca 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -240,8 +240,8 @@ def get_inference_context( else: max_sequence_length = args.inference_max_seq_length # >>> - pax("max_context_length, max_gen_length", - "calculate_max_sequence_length_from_requests, max_sequence_length") + # pax("max_context_length, max_gen_length", + # "calculate_max_sequence_length_from_requests, max_sequence_length") # <<< # Active buffer size & max tokens. @@ -604,14 +604,25 @@ 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( + # 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/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 7fd689f05d5..f2623dc8086 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -301,10 +301,10 @@ def __init__( # Chunk ids. self.max_kv_chunk_count = math.ceil(self.max_sequence_length / self.chunk_size_tokens) # >>> - pax({ - "max_sequence_length" : self.max_sequence_length, - "max_kv_chunk_count" : self.max_kv_chunk_count, - }) + # pax({ + # "max_sequence_length" : self.max_sequence_length, + # "max_kv_chunk_count" : self.max_kv_chunk_count, + # }) # <<< self.request_to_kv_chunk_ids = torch.full( (self.max_total_requests, self.max_kv_chunk_count), @@ -1195,15 +1195,40 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T resume_request_count = 0 if self.paused_request_count > 0: active_chunk_count_avail = self.chunk_allocator.get_active_avail() + # >>> # pax("active_chunk_count_avail") + # <<< paused_chunk_counts = self.request_kv_chunk_counts[:self.paused_request_count] paused_chunk_counts = paused_chunk_counts.flip(dims=[0]) paused_chunk_counts += 1 # +1 for newly added chunk paused_chunk_counts_cumsum = paused_chunk_counts.cumsum(dim=0) - resume_request_count = torch.nonzero( - paused_chunk_counts_cumsum <= active_chunk_count_avail - ).numel() + # >>> + # resume_request_count = torch.nonzero( + # paused_chunk_counts_cumsum <= active_chunk_count_avail + # ).numel() + # +++ + resume_request_count = min( + torch.nonzero( + paused_chunk_counts_cumsum <= active_chunk_count_avail + ).numel(), + self.chunk_allocator.total_avail, + ) + # <<< + # >>> + # if resume_request_count > self.chunk_allocator.total_avail: + # pax({ + # "chunk_allocator" : self.chunk_allocator, + # "active_used" : "%d / %d" % ( + # self.chunk_allocator.get_active_used(), + # self.chunk_allocator.active_count, + # ), + # "paused_used" : "%d / %d" % ( + # self.chunk_allocator.get_paused_used(), + # self.chunk_allocator.paused_count, + # ), + # }, "resume_request_count") + # <<< # >>> # newly_resumed_request_ids = self.request_ids[(self.paused_request_count - resume_request_count):self.paused_request_count] @@ -1259,6 +1284,10 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T == 0 ), "The request_last_kv_chunk_offset should be 0 for the requests that just got resumed this step. " + # >>> + if resume_request_count > self.chunk_allocator.total_avail: + pax({"chunk_allocator": self.chunk_allocator}, "resume_request_count") + # <<< chunk_ids = self.chunk_allocator.allocate_memory_chunks(resume_request_count) row_idx = torch.arange( self.paused_request_count, diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 16c37496286..773cb6d056f 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -557,6 +557,15 @@ async def async_step( 2. Requests that ran in the last step and have now finished. 3. The step time in seconds. """ + # >>> + ec = os.environ["INFERENCE_EMPTY_CACHE"] + if ec == "1": + torch.cuda.empty_cache() + elif ec == "0": + pass + else: + raise Exception(f"specialize for ec '{ec}'.") + # <<< # Previous context state, for printing output below. prev_is_decode_only = self.context.is_decode_only() @@ -626,13 +635,20 @@ async def async_step( mem = torch.cuda.memory_stats() step_type = "decode" if is_decode_only else "non-decode" # >>> - mem_info = process.memory_info() + # mem_info = process.memory_info() # <<< output_str = ( "* step %d | %s ... time: %.3f%s ... " - "reqs: %d [ active %d/%d, paused %d/%d, waiting %d, finished %d ] ... " + # >>> + # "reqs: %d [ active %d/%d, paused %d/%d, waiting %d, finished %d ] ... " + "reqs: a %d/%d, p %d/%d, w %d, f %d ... " + # <<< + # >>> + "chunks: a %d/%d, p %d/%d ... " + # <<< # >>> "mem: tensors %d, alloc %.1f gb, res %.1f gb." + # "mem: tensors %d, alloc %f gb, res %f gb." # "mem: alloc %.1f gb, res %.1f gb, cpu p %s/v %s." # <<< % ( @@ -651,7 +667,9 @@ async def async_step( ), ) ), - prev_total_request_count, + # >>> + # prev_total_request_count, + # <<< prev_total_request_count - prev_paused_request_count, context.chunk_allocator.active_count, prev_paused_request_count, @@ -659,6 +677,12 @@ async def async_step( len(self.waiting_request_ids), self.finished_request_count, # >>> + context.chunk_allocator.get_active_used(), + context.chunk_allocator.active_count, + context.chunk_allocator.get_paused_used(), + context.chunk_allocator.paused_count, + # <<< + # >>> mem["allocation.all.current"], # <<< mem["allocated_bytes.all.current"] / (1024**3), @@ -674,10 +698,6 @@ async def async_step( print(output_str) self.step_count += 1 - # >>> - # if self.step_count >= 100: # 600, 305 - # raise Exception(f"step {self.step_count}.") - # <<< range_pop() return { "active_requests": active_requests, From 7a0934fffeeb005991cbcefb7bd23e8f9d031f9a Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Thu, 16 Oct 2025 08:47:31 -0700 Subject: [PATCH 39/69] clean up. --- .../inference/gpt/gpt_dynamic_inference.py | 113 ------------------ examples/inference/gpt/utils.py | 29 ----- .../inference/contexts/dynamic_context.py | 80 +------------ .../core/inference/engines/dynamic_engine.py | 32 ----- 4 files changed, 1 insertion(+), 253 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index 0814d188eca..b1d0a117c90 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -61,15 +61,6 @@ torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunState]) torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunDiagnostic]) -# >>> -from lutil import pax as _pax, get_mem_stats_str -import builtins -builtins.pax = _pax - -# def print_mem(key): -# print(f"+++++++++++ {key} ........... {get_mem_stats_str()}.") -# <<< - def add_dynamic_inference_args(parser: ArgumentParser) -> ArgumentParser: """Dynamic inference arguments.""" @@ -135,82 +126,6 @@ def get_inference_context_active_buffer_size_bytes(model: torch.nn.Module) -> in return active_buffer_size_bytes -# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> -# # >>> -# @torch.inference_mode() -# def get_inference_context_max_tokens(model: torch.nn.Module) -> int: -# # def get_inference_context_max_tokens(controller: torch.nn.Module) -> int: -# # <<< -# """Estimate `max_tokens` for the context. - -# This function performs a binary search over the number of input tokens and -# analyzes the latency to find the 'knee of the curve'. -# """ - -# import time - -# # args = get_args() -# # model = GPTInferenceWrapper(model, args, inference_context="not-actually-set") # None) - -# num_tokens_step = 1000 -# num_tokens_min = num_tokens_step -# num_tokens_max = 24000 # 1000000 -# time_map = {} -# # >>> -# # tbar = tqdm(range(num_tokens_min, num_tokens_max, num_tokens_step)) -# tbar = tqdm([ 16384 ]) -# # <<< -# # tbar = tqdm((3000,)) -# # tbar = tqdm((3072,)) -# # tbar = tqdm((1000,)) -# # num_tokens_list = list( -# for num_tokens in tbar: -# tbar.set_description(f"num_tokens {num_tokens}") -# input_ids = torch.full( -# (num_tokens,), -# 0, -# dtype=torch.long, -# device=torch.cuda.current_device(), -# ).unsqueeze(0) -# # position_ids = torch.arange( -# # num_tokens, -# position_ids = torch.zeros( -# (num_tokens,), -# dtype=torch.long, -# device=torch.cuda.current_device(), -# ).unsqueeze(0) -# # result = model.run_one_forward_step( -# # { -# # "tokens": input_ids, -# # "position_ids": position_ids, -# # "attention_mask": None, -# # } -# # ) -# t = time.time() -# for _ in range(3): -# mem0 = torch.cuda.memory_stats() -# with torch.inference_mode(): -# result = model( -# input_ids, -# position_ids, -# attention_mask=None, -# inference_context=None, -# runtime_gather_output=True, # Inference should always gather the logits -# ) -# mem1 = torch.cuda.memory_stats() -# pax("mem0, mem1", { -# "peak / 0" : mem0["allocated_bytes.all.peak"], -# "peak / 1" : mem1["allocated_bytes.all.peak"], -# }) -# t = time.time() - t -# time_map[num_tokens] = t -# print("~~ x ~~") -# [ print(x) for x in time_map.keys() ] -# print("~~ y ~~") -# [ print(y) for y in time_map.values() ] -# raise Exception("matplotlib.") -# return max_tokens -# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ def get_inference_context_max_tokens(model: torch.nn.Module) -> int: """Estimate `max_tokens` for the context. @@ -219,7 +134,6 @@ def get_inference_context_max_tokens(model: torch.nn.Module) -> int: correlation between `max_tokens` and request throughput. """ return 16384 -# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< def get_inference_context( @@ -239,10 +153,6 @@ def get_inference_context( max_sequence_length = max_context_length + max_gen_length else: max_sequence_length = args.inference_max_seq_length - # >>> - # pax("max_context_length, max_gen_length", - # "calculate_max_sequence_length_from_requests, max_sequence_length") - # <<< # Active buffer size & max tokens. active_buffer_size_bytes = get_inference_context_active_buffer_size_bytes(model) @@ -438,10 +348,6 @@ def main(): args_defaults={'no_load_rng': True, 'no_load_optim': True}, ) - # >>> - # print_mem("initialize_megatron()") - # <<< - # Start Nsight profiler. if os.environ.get("NSIGHT_PREFIX"): torch.cuda.cudart().cudaProfilerStart() @@ -465,9 +371,6 @@ def main(): model = get_model() requests = build_requests(args, tokenizer) context = get_inference_context(requests, sampling_params, model) - # >>> - # pax({"requests / len": len(requests)}) - # <<< controller = get_inference_controller(model, context) # Validate all context_length's <= max_tokens. @@ -490,11 +393,6 @@ def main(): track_paused_request_events=args.inference_dynamic_batching_track_paused_request_events, ) - # >>> - # print_mem("engine") - # exit() - # <<< - setup_prefix = build_dynamic_engine_setup_prefix(args, model, context, requests) print("~~~") print(setup_prefix) @@ -604,16 +502,6 @@ 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", @@ -622,7 +510,6 @@ def escape_str(s): f"steps: {engine.step_count:d}", f"capture {capture_str}", ))) - # <<< print("~~~") # Stop Nsight profiler. diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py index dc0efd59de0..971d99caa52 100644 --- a/examples/inference/gpt/utils.py +++ b/examples/inference/gpt/utils.py @@ -252,38 +252,12 @@ def get_synthetic_requests(args: Namespace, tokenizer: Any) -> list[Request]: prompt_tokens_list = [ max_prompt_tokens[:l] for l in prompt_lengths ] prompt_texts = [ tokenizer.detokenize(tt) for tt in prompt_tokens_list ] - # >>> - # pax({ - # "args.num_tokens_to_prompt" : args.num_tokens_to_prompt, - # }, "max_prompt_length, max_prompt_text, max_prompt_tokens", { - # "prompt_lengths" : "%d / %s" % (len(prompt_lengths), prompt_lengths), - # "prompt_tokens_list" : [ - # "%d[%d] / '%s' ... %s" % (prompt_lengths[i], len(tt), prompt_texts[i], tt) - # for i, tt in enumerate(prompt_tokens_list) - # ], - # }) - # <<< - # Init requests. - # >>> - # requests = [ - # Request("hi " * random.randint(*args.num_tokens_to_prompt), t, tokenizer) - # for t in time_offsets - # ] - # +++ assert len(prompt_texts) == len(time_offsets) requests = [ Request(t, o, tokenizer) for t, o in zip(prompt_texts, time_offsets) ] - # <<< - - # >>> - # pax("requests", { - # "requests / 0" : requests[0], - # "tokens" : [ "%d, '%s'" % (i, tokenizer.detokenize([i])) for i in requests[0].prompt_tokens ], - # }) - # <<< return requests @@ -390,10 +364,7 @@ def build_dynamic_engine_setup_prefix( # Buffer limits config buffer_limits_str = ( - # >>> - # f"bf: {args.inference_dynamic_batching_active_buffer_size_gb:.0f} gb, " f"bf: {get_mem_size_str(args.inference_dynamic_batching_active_buffer_size_gb*1024**3)}, " - # <<< f"{context.chunk_allocator.active_count} chunks " f"[r {context.max_active_requests}, t {context.max_tokens}]" ) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index f2623dc8086..2019a0aad6c 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -300,12 +300,6 @@ def __init__( # Chunk ids. self.max_kv_chunk_count = math.ceil(self.max_sequence_length / self.chunk_size_tokens) - # >>> - # pax({ - # "max_sequence_length" : self.max_sequence_length, - # "max_kv_chunk_count" : self.max_kv_chunk_count, - # }) - # <<< self.request_to_kv_chunk_ids = torch.full( (self.max_total_requests, self.max_kv_chunk_count), -1, @@ -1195,44 +1189,16 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T resume_request_count = 0 if self.paused_request_count > 0: active_chunk_count_avail = self.chunk_allocator.get_active_avail() - # >>> - # pax("active_chunk_count_avail") - # <<< - paused_chunk_counts = self.request_kv_chunk_counts[:self.paused_request_count] paused_chunk_counts = paused_chunk_counts.flip(dims=[0]) paused_chunk_counts += 1 # +1 for newly added chunk paused_chunk_counts_cumsum = paused_chunk_counts.cumsum(dim=0) - # >>> - # resume_request_count = torch.nonzero( - # paused_chunk_counts_cumsum <= active_chunk_count_avail - # ).numel() - # +++ resume_request_count = min( torch.nonzero( paused_chunk_counts_cumsum <= active_chunk_count_avail ).numel(), self.chunk_allocator.total_avail, ) - # <<< - # >>> - # if resume_request_count > self.chunk_allocator.total_avail: - # pax({ - # "chunk_allocator" : self.chunk_allocator, - # "active_used" : "%d / %d" % ( - # self.chunk_allocator.get_active_used(), - # self.chunk_allocator.active_count, - # ), - # "paused_used" : "%d / %d" % ( - # self.chunk_allocator.get_paused_used(), - # self.chunk_allocator.paused_count, - # ), - # }, "resume_request_count") - # <<< - - # >>> - # newly_resumed_request_ids = self.request_ids[(self.paused_request_count - resume_request_count):self.paused_request_count] - # <<< self.paused_request_count -= resume_request_count active_request_count += resume_request_count @@ -1245,9 +1211,6 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T # requests are newly paused during this update. if newly_paused_request_ids is not None and resume_request_count > 0: newly_paused_request_ids = newly_paused_request_ids[:-resume_request_count] - # >>> - # newly_resumed_request_ids = newly_resumed_request_ids[-resume_request_count:] - # <<< # 7. We make changes to the request book keeping tesnsors and setup the tokens for next iteration # All these active requests are in decode phase, so they need only 1 token per request @@ -1284,10 +1247,7 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T == 0 ), "The request_last_kv_chunk_offset should be 0 for the requests that just got resumed this step. " - # >>> - if resume_request_count > self.chunk_allocator.total_avail: - pax({"chunk_allocator": self.chunk_allocator}, "resume_request_count") - # <<< + assert resume_request_count <= self.chunk_allocator.total_avail chunk_ids = self.chunk_allocator.allocate_memory_chunks(resume_request_count) row_idx = torch.arange( self.paused_request_count, @@ -1320,44 +1280,6 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T self.request_last_kv_chunk_offset[self.paused_request_count : self.total_request_count] ) - # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> - # def get_request_str(i): - # r = self.request_ids[i].item() - # q = self.request_query_lengths[i].item() - # kv = self.request_kv_length_offsets[i].item() - # # pax("request_id") - # return "%d[%d:%d]" % (r, q, kv) - # request_strs = [ - # get_request_str(i) - # for i in range(self.total_request_count + finished_request_count) - # ] - # # pax("request_strs") - # request_str = "requests: %s -- %s -- %s" % ( - # ", ".join(request_strs[:self.paused_request_count]), - # ", ".join(request_strs[self.paused_request_count:self.total_request_count]), - # ", ".join(request_strs[self.total_request_count:(self.total_request_count + finished_request_count)]), - # ) - # newly_str = "paused: %s, resumed: %s, finished: %s" % ( - # "--" if newly_paused_request_ids is None else newly_paused_request_ids.tolist(), - # newly_resumed_request_ids.tolist(), - # self.request_ids[self.total_request_count:(self.total_request_count + finished_request_count)].tolist(), - # ) - # # alloc_str = "alloc: t %d/%d [ p %d/%d, a %d/%d ]" % ( - # alloc_str = "alloc: t %d/%d [ a %d/%d ]" % ( - # self.chunk_allocator.total_count - self.chunk_allocator.total_avail - 1, - # self.chunk_allocator.total_count - 1, - # # self.chunk_allocator.get_paused_used(self), - # # self.chunk_allocator.paused_count, - # self.chunk_allocator.get_active_used(), - # self.chunk_allocator.active_count, - # ) - # print("++++++++++ %s" % " ..... ".join(( - # request_str, - # # newly_str, - # alloc_str, - # ))) - # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - return newly_paused_request_ids def calculate_log_probs( diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 773cb6d056f..03e8f6e8967 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -557,15 +557,6 @@ async def async_step( 2. Requests that ran in the last step and have now finished. 3. The step time in seconds. """ - # >>> - ec = os.environ["INFERENCE_EMPTY_CACHE"] - if ec == "1": - torch.cuda.empty_cache() - elif ec == "0": - pass - else: - raise Exception(f"specialize for ec '{ec}'.") - # <<< # Previous context state, for printing output below. prev_is_decode_only = self.context.is_decode_only() @@ -634,23 +625,11 @@ async def async_step( context = self.context mem = torch.cuda.memory_stats() step_type = "decode" if is_decode_only else "non-decode" - # >>> - # mem_info = process.memory_info() - # <<< output_str = ( "* step %d | %s ... time: %.3f%s ... " - # >>> - # "reqs: %d [ active %d/%d, paused %d/%d, waiting %d, finished %d ] ... " "reqs: a %d/%d, p %d/%d, w %d, f %d ... " - # <<< - # >>> "chunks: a %d/%d, p %d/%d ... " - # <<< - # >>> "mem: tensors %d, alloc %.1f gb, res %.1f gb." - # "mem: tensors %d, alloc %f gb, res %f gb." - # "mem: alloc %.1f gb, res %.1f gb, cpu p %s/v %s." - # <<< % ( self.step_count, datetime.now().strftime("%H:%M:%S"), @@ -667,30 +646,19 @@ async def async_step( ), ) ), - # >>> - # prev_total_request_count, - # <<< prev_total_request_count - prev_paused_request_count, context.chunk_allocator.active_count, prev_paused_request_count, context.chunk_allocator.paused_count, len(self.waiting_request_ids), self.finished_request_count, - # >>> context.chunk_allocator.get_active_used(), context.chunk_allocator.active_count, context.chunk_allocator.get_paused_used(), context.chunk_allocator.paused_count, - # <<< - # >>> mem["allocation.all.current"], - # <<< mem["allocated_bytes.all.current"] / (1024**3), mem["reserved_bytes.all.current"] / (1024**3), - # >>> - # get_mem_size_str(mem_info.rss), - # get_mem_size_str(mem_info.vms), - # <<< ) ) if prev_is_decode_only: From f93e4de834d55511d7a368f2c171ca50284ce01d Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Thu, 16 Oct 2025 09:00:35 -0700 Subject: [PATCH 40/69] minor updates. --- examples/inference/gpt/gpt_dynamic_inference_12b.sh | 9 --------- examples/inference/gpt/gpt_dynamic_inference_357m.sh | 9 --------- 2 files changed, 18 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference_12b.sh b/examples/inference/gpt/gpt_dynamic_inference_12b.sh index 5ae70f22c89..6c0c6f4ed61 100644 --- a/examples/inference/gpt/gpt_dynamic_inference_12b.sh +++ b/examples/inference/gpt/gpt_dynamic_inference_12b.sh @@ -23,16 +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} -# --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} \ -# +++ : ${ACTIVE_BUFFER_SIZE_GB=50.} -# <<< # Cuda graphs. : ${ENABLE_CUDA_GRAPHS=1} diff --git a/examples/inference/gpt/gpt_dynamic_inference_357m.sh b/examples/inference/gpt/gpt_dynamic_inference_357m.sh index 0383405257d..31a245cb380 100644 --- a/examples/inference/gpt/gpt_dynamic_inference_357m.sh +++ b/examples/inference/gpt/gpt_dynamic_inference_357m.sh @@ -24,16 +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} -# --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} \ -# +++ : ${ACTIVE_BUFFER_SIZE_GB=50.} -# <<< # Cuda graphs. : ${ENABLE_CUDA_GRAPHS=1} From b4d32ef514d923083e96d3a06aa9e67bb1b4e46f Mon Sep 17 00:00:00 2001 From: Mcore Bot Date: Thu, 16 Oct 2025 16:17:54 +0000 Subject: [PATCH 41/69] chore: Format files --- .../contexts/dynamic_chunk_allocator.py | 30 +++++++------ .../inference/contexts/dynamic_context.py | 44 +++++++++---------- .../core/inference/engines/dynamic_engine.py | 4 +- 3 files changed, 40 insertions(+), 38 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_chunk_allocator.py b/megatron/core/inference/contexts/dynamic_chunk_allocator.py index db7ec185f04..5bf8a541f0a 100644 --- a/megatron/core/inference/contexts/dynamic_chunk_allocator.py +++ b/megatron/core/inference/contexts/dynamic_chunk_allocator.py @@ -25,19 +25,17 @@ def __init__(self, context: "DynamicInferenceContext", active_count: int): self.context = context - active_count -= 1 # -1 for dummy_chunk_idx (see below) - active_count = max(1, active_count) # need at least one chunk - self.total_count = 2 * active_count + 1 # +1 for dummy_chunk_idx - self.total_avail = self.total_count - 1 # -1 for dummy_chunk_idx + active_count -= 1 # -1 for dummy_chunk_idx (see below) + active_count = max(1, active_count) # need at least one chunk + self.total_count = 2 * active_count + 1 # +1 for dummy_chunk_idx + self.total_avail = self.total_count - 1 # -1 for dummy_chunk_idx self.active_count = active_count - self.paused_count = self.total_count - self.active_count - 1 # -1 for dummy_chunk_idx + self.paused_count = self.total_count - self.active_count - 1 # -1 for dummy_chunk_idx self.dummy_chunk_idx = self.total_count - 1 # Initialize chunk pool as a "stack" data structure self.chunk_bag = torch.arange( - self.total_count, - dtype=torch.int32, - device=torch.cuda.current_device(), + self.total_count, dtype=torch.int32, device=torch.cuda.current_device() ) def __str__(self): @@ -48,15 +46,19 @@ def __str__(self): def get_active_used(self): """Compute number of active chunks used.""" - return self.context.request_kv_chunk_counts[ - self.context.paused_request_count:self.context.total_request_count - ].sum().item() + return ( + self.context.request_kv_chunk_counts[ + self.context.paused_request_count : self.context.total_request_count + ] + .sum() + .item() + ) def get_paused_used(self): """Compute number of paused chunks used.""" - return self.context.request_kv_chunk_counts[ - :self.context.paused_request_count - ].sum().item() + return ( + self.context.request_kv_chunk_counts[: self.context.paused_request_count].sum().item() + ) def get_active_avail(self): """Compute number of active chunks available.""" diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 19a8054014e..b4dbe623585 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -235,15 +235,12 @@ def __init__( # Initialize chunk allocator. active_chunk_count_total = active_buffer_size_bytes // self.chunk_size_bytes - self.chunk_allocator = ChunkAllocator( - context=self, - active_count=active_chunk_count_total, - ) - del active_chunk_count_total # use self.chunk_allocator.active_count + self.chunk_allocator = ChunkAllocator(context=self, active_count=active_chunk_count_total) + del active_chunk_count_total # use self.chunk_allocator.active_count active_buffer_size_bytes = self.chunk_allocator.active_count * self.chunk_size_bytes # Set max_requests, max_tokens. - self.max_total_requests = self.chunk_allocator.total_count - 1 # -1 for dummy chunk + self.max_total_requests = self.chunk_allocator.total_count - 1 # -1 for dummy chunk self.max_active_requests = self.chunk_allocator.active_count self.max_tokens = max_tokens @@ -251,9 +248,9 @@ def __init__( self.params_dtype = params_dtype self.num_layers = num_layers self.max_sequence_length = max_sequence_length - assert has_unified_memory, ( - "CUDA unified memory must be available to use `DynamicInferenceContext`." - ) + assert ( + has_unified_memory + ), "CUDA unified memory must be available to use `DynamicInferenceContext`." self.unified_memory_mempool = create_unified_mempool() self.total_request_count = 0 @@ -350,7 +347,11 @@ def __init__( self.cuda_graph_token_counts = [self.max_active_requests] else: self.cuda_graph_token_counts = list( - range(self.cuda_graph_step_size, self.max_active_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_active_requests: self.cuda_graph_token_counts.append(self.max_active_requests) @@ -408,10 +409,10 @@ def __init__( self.use_flashinfer_fused_rope = use_flashinfer_fused_rope # Print info. - print("DynamicInferenceContext: allocated context with active buffer size %s (%d chunks)." % ( - get_mem_size_str(active_buffer_size_bytes), - self.chunk_allocator.active_count, - )) + print( + "DynamicInferenceContext: allocated context with active buffer size %s (%d chunks)." + % (get_mem_size_str(active_buffer_size_bytes), self.chunk_allocator.active_count) + ) TOKEN_ROUNDER = 64 REQUEST_ROUNDER = 4 @@ -757,7 +758,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_active_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}" @@ -1279,8 +1282,7 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T if active_requests_requiring_new_chunk_count > 0: newly_paused_request_ids = self.request_ids[ - torch.nonzero(active_requests_requiring_new_chunk) - + self.paused_request_count + torch.nonzero(active_requests_requiring_new_chunk) + self.paused_request_count ] # Swap unfinished active requests on the left side with paused requests on the right side @@ -1326,14 +1328,12 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T resume_request_count = 0 if self.paused_request_count > 0: active_chunk_count_avail = self.chunk_allocator.get_active_avail() - paused_chunk_counts = self.request_kv_chunk_counts[:self.paused_request_count] + paused_chunk_counts = self.request_kv_chunk_counts[: self.paused_request_count] paused_chunk_counts = paused_chunk_counts.flip(dims=[0]) - paused_chunk_counts += 1 # +1 for newly added chunk + paused_chunk_counts += 1 # +1 for newly added chunk paused_chunk_counts_cumsum = paused_chunk_counts.cumsum(dim=0) resume_request_count = min( - torch.nonzero( - paused_chunk_counts_cumsum <= active_chunk_count_avail - ).numel(), + torch.nonzero(paused_chunk_counts_cumsum <= active_chunk_count_avail).numel(), self.chunk_allocator.total_avail, ) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index ac3982bb026..3b1582a8514 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -4,7 +4,6 @@ import logging import multiprocessing import os -import psutil import struct import time import warnings @@ -13,6 +12,7 @@ from itertools import repeat from typing import Dict, List, Optional, Tuple, Union +import psutil import torch from torch import Tensor from torch.cuda.nvtx import range_pop, range_push @@ -20,10 +20,10 @@ from megatron.core import parallel_state from megatron.core.inference.contexts.dynamic_context import ( DynamicInferenceContext, - get_mem_size_str, MaxSequenceLengthOverflowError, TokenOverflowError, WarmupEngineMode, + get_mem_size_str, ) from megatron.core.inference.data_parallel_inference_coordinator import ( DataParallelInferenceCoordinator, From 55b14a030db80459e1ac96743a992f439a82156c Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Thu, 16 Oct 2025 10:39:37 -0700 Subject: [PATCH 42/69] fixed max_requests rename. --- .../inference/contexts/dynamic_context.py | 29 +++++++++---------- .../core/inference/engines/dynamic_engine.py | 6 ++-- 2 files changed, 16 insertions(+), 19 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 19a8054014e..80f91fe3cc5 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -161,9 +161,9 @@ class DynamicInferenceContext(BaseInferenceContext): chunk_size_tokens (int): Size of KV cache chunk size. tensor_model_parallel_size (Optional[int]): Tensor model parallel size. num_cuda_graphs (Optional[int]): Maximum number of cuda graphs to capture, - where the cuda graph batch sizes range from 1 to `max_requests` (as - computed below). Due to rounding, the actual number of cuda graphs may - not equal this argument. + 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 (bool): If True, only the last token logits are materialized in the context. use_cuda_graphs_for_non_decode_steps (bool): If True, use cuda graphs for non-decode @@ -242,7 +242,7 @@ def __init__( del active_chunk_count_total # use self.chunk_allocator.active_count active_buffer_size_bytes = self.chunk_allocator.active_count * self.chunk_size_bytes - # Set max_requests, max_tokens. + # Set max_total_requests, max_active_requests, max_tokens. self.max_total_requests = self.chunk_allocator.total_count - 1 # -1 for dummy chunk self.max_active_requests = self.chunk_allocator.active_count self.max_tokens = max_tokens @@ -717,7 +717,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. @@ -974,20 +974,21 @@ 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 ) chunks = math.ceil( (req.remaining_prompt_length + req.finished_chunk_token_count) / self.chunk_size_tokens ) - math.ceil(req.finished_chunk_token_count / self.chunk_size_tokens) - kv_cache_available = self.chunk_allocator.is_memory_available(chunks, safe=safe) + kv_cache_available = self.chunk_allocator.is_memory_available(chunks) 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: @@ -1027,9 +1028,7 @@ def add_request(self, req: DynamicInferenceRequest, chunk_length: Optional[int] num_chunks_needed = overall_required_chunks - already_allocated_chunks if num_chunks_needed > 0: - new_chunk_ids = self.chunk_allocator.allocate_memory_chunks( - num_chunks_needed, safe=not is_chunked_prefill - ) + new_chunk_ids = self.chunk_allocator.allocate_memory_chunks(num_chunks_needed) if new_chunk_ids is None or len(new_chunk_ids) != num_chunks_needed: raise ChunkOverflowError(req.request_id) @@ -1047,7 +1046,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: @@ -1145,7 +1144,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 diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index ac3982bb026..532bc2734f4 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -593,7 +593,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) @@ -636,9 +636,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: From 3bdd6611d8ada01d454e9320d7957bd30cbae845 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Wed, 29 Oct 2025 13:29:57 -0700 Subject: [PATCH 43/69] removed placeholder methods for computing buffer_size and max_tokens. --- .../inference/gpt/gpt_dynamic_inference.py | 31 +------------------ .../inference/contexts/dynamic_context.py | 7 +++-- .../core/inference/engines/dynamic_engine.py | 3 -- 3 files changed, 5 insertions(+), 36 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index 81ab22243af..9acf1a170b8 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -115,30 +115,6 @@ def get_model() -> MegatronModule: return model -def get_inference_context_active_buffer_size_bytes(model: torch.nn.Module) -> int: - """Estimate memory remaining for context's `memory_buffer`. - - The following implementation is a placeholder, using only - `args.inference_dynamic_batching_active_buffer_size_gb`. Eventually, this will - be computed by first computing memory used by the model, activations, and the - KV cache. - """ - args = get_args() - active_buffer_size_gb = args.inference_dynamic_batching_active_buffer_size_gb - active_buffer_size_bytes = int(active_buffer_size_gb * 1024**3) - return active_buffer_size_bytes - - -def get_inference_context_max_tokens(model: torch.nn.Module) -> int: - """Estimate `max_tokens` for the context. - - The following implementation is a placeholder, simply returning a hard-coded - value for `max_tokens`. Eventually, this will be computed by analyzing the - correlation between `max_tokens` and request throughput. - """ - return 16384 - - def get_inference_context( requests: List[Request], sampling_params: SamplingParams, @@ -157,10 +133,6 @@ def get_inference_context( else: max_sequence_length = args.inference_max_seq_length - # Active buffer size & max tokens. - active_buffer_size_bytes = get_inference_context_active_buffer_size_bytes(model) - max_tokens = get_inference_context_max_tokens(model) - # Inference context. context = DynamicInferenceContext( params_dtype=args.params_dtype, @@ -176,8 +148,7 @@ def get_inference_context( else None ), chunk_size_tokens=args.inference_dynamic_batching_chunk_size, - active_buffer_size_bytes=active_buffer_size_bytes, - max_tokens=max_tokens, + 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, diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index d4a490e1f73..638c8486321 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -152,7 +152,7 @@ class DynamicInferenceContext(BaseInferenceContext): num_attention_heads (int): Number of attention heads. max_sequence_length (int): Max possible sequence length (prompt + output) that will occur. - active_buffer_size_bytes (float): Buffer size reserved for active requests + 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. @@ -184,8 +184,8 @@ def __init__( kv_channels: int, num_attention_heads: int, max_sequence_length: int, - active_buffer_size_bytes: float, - max_tokens: int, + active_buffer_size_gb: float, + max_tokens: int = 16384, chunk_size_tokens: int = 256, tensor_model_parallel_size: Optional[int] = None, cache_mla_latent: bool = False, @@ -234,6 +234,7 @@ def __init__( ) # Initialize chunk allocator. + active_buffer_size_bytes = int(active_buffer_size_gb * 1024**3) active_chunk_count_total = active_buffer_size_bytes // self.chunk_size_bytes self.chunk_allocator = ChunkAllocator(context=self, active_count=active_chunk_count_total) del active_chunk_count_total # use self.chunk_allocator.active_count diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 2db5bb4db76..f4793478751 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -60,9 +60,6 @@ HAVE_MSGPACK = False -process = psutil.Process(os.getpid()) - - def format_mem_bytes(mem_bytes): """Convert a byte count to a human-readable string in tb, gb, mb, kb, or bytes.""" for power, suffix in [(4, "tb"), (3, "gb"), (2, "mb"), (1, "kb"), (0, "bytes")]: From 559daf36c3e7995197417d7a1294375e60e9c219 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Thu, 30 Oct 2025 09:05:25 -0700 Subject: [PATCH 44/69] re-added unified_memory_level control to dynamic context. --- .../inference/gpt/gpt_dynamic_inference.py | 6 +++ .../inference/contexts/dynamic_context.py | 38 ++++++++++++++++--- 2 files changed, 39 insertions(+), 5 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index 9acf1a170b8..2741e5ea42a 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -64,6 +64,12 @@ torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunState]) torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunDiagnostic]) +# >>> +from lutil import pax as _pax +import builtins +builtins.pax = _pax +# <<< + def add_dynamic_inference_args(parser: ArgumentParser) -> ArgumentParser: """Dynamic inference arguments.""" diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 638c8486321..9f237fb4647 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -249,10 +249,22 @@ def __init__( self.params_dtype = params_dtype self.num_layers = num_layers self.max_sequence_length = max_sequence_length - assert ( - has_unified_memory - ), "CUDA unified memory must be available to use `DynamicInferenceContext`." - self.unified_memory_mempool = create_unified_mempool() + # >>> + # assert ( + # has_unified_memory + # ), "CUDA unified memory must be available to use `DynamicInferenceContext`." + # self.unified_memory_mempool = create_unified_mempool() + # +++ + self.unified_memory_level = unified_memory_level + if unified_memory_level > 0: + if not has_unified_memory and torch.distributed.get_rank() == 0: + warnings.warn( + "Unified memory requested but not available; defaulting to GPU memory." + ) + self.unified_memory_level = 0 + else: + self.unified_memory_mempool = create_unified_mempool() + # <<< self.total_request_count = 0 self.active_token_count = 0 @@ -290,7 +302,16 @@ def __init__( self.token_to_local_position_within_kv_chunk = torch.empty_like(self.token_to_input_ids) # Memory buffer. - with torch.cuda.use_mem_pool(self.unified_memory_mempool): + # >>> + # with torch.cuda.use_mem_pool(self.unified_memory_mempool): + # +++ + ctx_manager = ( + torch.cuda.use_mem_pool(self.unified_memory_mempool) + if self.unified_memory_level > 0 + else nullcontext() + ) + with ctx_manager: + # <<< if cache_mla_latent: self.memory_buffer = torch.full( ( @@ -317,6 +338,13 @@ def __init__( dtype=self.params_dtype, device=torch.cuda.current_device(), ) + # >>> + pax("ctx_manager", { + "unified_memory_level" : self.unified_memory_level, + # "unified_memory_mempool" : self.unified_memory_mempool, + "memory buffer gb" : self.memory_buffer.numel() * self.memory_buffer.element_size() / 1024**3, + }) + # <<< # Chunk ids. self.max_kv_chunk_count = math.ceil(self.max_sequence_length / self.chunk_size_tokens) From be51f16055ade5fb97fd4de6490248c641bf9715 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Thu, 30 Oct 2025 09:08:35 -0700 Subject: [PATCH 45/69] clean up. --- .../inference/gpt/gpt_dynamic_inference.py | 6 ------ .../core/inference/contexts/dynamic_context.py | 18 ------------------ 2 files changed, 24 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index 2741e5ea42a..9acf1a170b8 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -64,12 +64,6 @@ torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunState]) torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunDiagnostic]) -# >>> -from lutil import pax as _pax -import builtins -builtins.pax = _pax -# <<< - def add_dynamic_inference_args(parser: ArgumentParser) -> ArgumentParser: """Dynamic inference arguments.""" diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 9f237fb4647..bfb5cc80e2d 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -249,12 +249,6 @@ def __init__( self.params_dtype = params_dtype self.num_layers = num_layers self.max_sequence_length = max_sequence_length - # >>> - # assert ( - # has_unified_memory - # ), "CUDA unified memory must be available to use `DynamicInferenceContext`." - # self.unified_memory_mempool = create_unified_mempool() - # +++ self.unified_memory_level = unified_memory_level if unified_memory_level > 0: if not has_unified_memory and torch.distributed.get_rank() == 0: @@ -264,7 +258,6 @@ def __init__( self.unified_memory_level = 0 else: self.unified_memory_mempool = create_unified_mempool() - # <<< self.total_request_count = 0 self.active_token_count = 0 @@ -302,16 +295,12 @@ def __init__( self.token_to_local_position_within_kv_chunk = torch.empty_like(self.token_to_input_ids) # Memory buffer. - # >>> - # with torch.cuda.use_mem_pool(self.unified_memory_mempool): - # +++ ctx_manager = ( torch.cuda.use_mem_pool(self.unified_memory_mempool) if self.unified_memory_level > 0 else nullcontext() ) with ctx_manager: - # <<< if cache_mla_latent: self.memory_buffer = torch.full( ( @@ -338,13 +327,6 @@ def __init__( dtype=self.params_dtype, device=torch.cuda.current_device(), ) - # >>> - pax("ctx_manager", { - "unified_memory_level" : self.unified_memory_level, - # "unified_memory_mempool" : self.unified_memory_mempool, - "memory buffer gb" : self.memory_buffer.numel() * self.memory_buffer.element_size() / 1024**3, - }) - # <<< # Chunk ids. self.max_kv_chunk_count = math.ceil(self.max_sequence_length / self.chunk_size_tokens) From 991c95ae6d6c0d011184977135f025cddec233e1 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Thu, 30 Oct 2025 09:16:57 -0700 Subject: [PATCH 46/69] added uvm level to print. --- examples/inference/gpt/utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py index c929bfb7d95..ac19f503ee6 100644 --- a/examples/inference/gpt/utils.py +++ b/examples/inference/gpt/utils.py @@ -352,6 +352,9 @@ def build_dynamic_engine_setup_prefix( else: cg_str = "--" + # Unified memory. + uvm_str = f"uvm {int(args.inference_dynamic_batching_unified_memory_level)}" + # Prompt description prompt_src_str = ( "cli" if args.prompts else @@ -380,6 +383,7 @@ def build_dynamic_engine_setup_prefix( get_model_size_str(model), "dynamic", cg_str, + uvm_str, request_str, buffer_limits_str, ] From 88e412fda3fde6fca3b0b06980ee7bcae3fba9ca Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Thu, 30 Oct 2025 09:28:47 -0700 Subject: [PATCH 47/69] format. --- megatron/core/inference/contexts/dynamic_context.py | 6 +++--- megatron/core/inference/engines/dynamic_engine.py | 2 -- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index bfb5cc80e2d..c3841a2db27 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. All rights reserved. +import logging import math import warnings from contextlib import nullcontext @@ -420,7 +421,7 @@ def __init__( self.use_flashinfer_fused_rope = use_flashinfer_fused_rope # Print info. - print( + logging.info( "DynamicInferenceContext: allocated context with active buffer size %s (%d chunks)." % (get_mem_size_str(active_buffer_size_bytes), self.chunk_allocator.active_count) ) @@ -993,8 +994,7 @@ 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.paused_request_count - < self.max_active_requests + 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 diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index f4793478751..641f47707e7 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -12,7 +12,6 @@ from itertools import repeat from typing import Dict, List, Optional, Tuple, Union -import psutil import torch from torch import Tensor from torch.cuda.nvtx import range_pop, range_push @@ -23,7 +22,6 @@ MaxSequenceLengthOverflowError, TokenOverflowError, WarmupEngineMode, - get_mem_size_str, ) from megatron.core.inference.data_parallel_inference_coordinator import ( DataParallelInferenceCoordinator, From d7c5eaec571952e5158289abf2f85243a3e916c1 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Thu, 30 Oct 2025 11:57:16 -0700 Subject: [PATCH 48/69] using correct sampling_params during build_requests(). --- .../inference/gpt/gpt_dynamic_inference.py | 16 ++++++++- examples/inference/gpt/utils.py | 12 +++++-- .../inference/contexts/dynamic_context.py | 9 +++-- .../core/inference/engines/dynamic_engine.py | 35 +++++++++++++++---- 4 files changed, 59 insertions(+), 13 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index 02f442af7f1..2f7c94bbeda 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -64,6 +64,12 @@ torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunState]) torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunDiagnostic]) +# >>> +from lutil import pax as _pax +import builtins +builtins.pax = _pax +# <<< + def add_dynamic_inference_args(parser: ArgumentParser) -> ArgumentParser: """Dynamic inference arguments.""" @@ -119,7 +125,7 @@ def get_model() -> MegatronModule: def get_inference_context( requests: List[Request], sampling_params: Optional[SamplingParams] = None, - calculate_max_sequence_length_from_requests: bool =True, + calculate_max_sequence_length_from_requests: bool = True ): """The inference context manages the KV cache and other inference state.""" @@ -254,6 +260,9 @@ def _add_request(): _request.prompt_text, _request.sampling_params, ) + # >>> + print("+++++++++++ added request %d / %d." % (num_requests_added, num_requests_total)) + # <<< _request.time_start = get_curr_time() _request.state = "started" num_requests_added += 1 @@ -292,6 +301,11 @@ def _add_request(): # Update requests. active_requests = result["active_requests"] finished_requests = result["finished_requests"] + # >>> + if finished_requests: + from lutil import pax + pax("finished_requests") + # <<< step_time = result["step_time"] if len(active_requests) > 0 or len(finished_requests) > 0: if is_decode_only: diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py index 61371f62f86..fcc9c8136dd 100644 --- a/examples/inference/gpt/utils.py +++ b/examples/inference/gpt/utils.py @@ -279,10 +279,18 @@ def get_synthetic_requests( # Init requests. assert len(prompt_texts) == len(time_offsets) requests = [ - Request(t, o, tokenizer) + Request(t, o, tokenizer, sampling_params=sampling_params) for t, o in zip(prompt_texts, time_offsets) ] + # >>> + # pax("max_prompt_length", { + # "old num_tokens_to_generate" : sampling_params.num_tokens_to_generate, + # "new num_tokens_to_generate" : requests[0].sampling_params.num_tokens_to_generate, + # "new num_tokens_to_generate'" : requests[-1].sampling_params.num_tokens_to_generate, + # "set(num_tokens_to_generate)" : "%d | %s" % (len(requests), set(r.sampling_params.num_tokens_to_generate for r in requests)), + # }) + # <<< return requests @@ -396,7 +404,7 @@ def build_dynamic_engine_setup_prefix( # Buffer limits config buffer_limits_str = ( f"bf: {get_mem_size_str(args.inference_dynamic_batching_active_buffer_size_gb*1024**3)}, " - f"{context.chunk_allocator.active_count} chunks " + f"{context.block_allocator.active_count} chunks " f"[r {context.max_active_requests}, t {context.max_tokens}]" ) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index aaf34d858c3..4f4b96c90aa 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -471,8 +471,8 @@ def __init__( # Print info. logging.info( - "DynamicInferenceContext: allocated context with active buffer size %s (%d chunks)." - % (get_mem_size_str(active_buffer_size_bytes), self.chunk_allocator.active_count) + "DynamicInferenceContext: allocated context with active buffer size %s (%d blocks)." + % (get_mem_size_str(active_buffer_size_bytes), self.block_allocator.active_count) ) TOKEN_ROUNDER = 64 @@ -1101,6 +1101,9 @@ def add_request(self, req: DynamicInferenceRequest, chunk_length: Optional[int] Return: None """ + # >>> + raise Exception("hi.") + # <<< if chunk_length is None: chunk_length = req.remaining_prompt_length @@ -1430,7 +1433,7 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T paused_block_counts_cumsum = paused_block_counts.cumsum(dim=0) resume_request_count = min( torch.nonzero(paused_chunk_counts_cumsum <= active_chunk_count_avail).numel(), - self.chunk_allocator.total_avail, + self.block_allocator.total_avail, ) self.paused_request_count -= resume_request_count diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index e9c2d7b07cc..7c375e7e0f3 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -404,13 +404,28 @@ def _add_request( len(request.prompt_tokens) + request.sampling_params.num_tokens_to_generate > self.context.max_sequence_length ): + # >>> + from lutil import pax + pax("request", { + "prompt" : len(request.prompt_tokens), + "gen" : request.sampling_params.num_tokens_to_generate, + "max_sequence_length" : self.context.max_sequence_length, + }) + # <<< request.status = Status.FAILED request.add_event_error_nontransient(MaxSequenceLengthOverflowError(request_id)) if len(request.prompt_tokens) > self.context.max_tokens and not self.enable_chunked_prefill: + # >>> + raise Exception("hi.") + # <<< request.status = Status.FAILED request.add_event_error_nontransient(TokenOverflowError(request_id)) + # >>> + # from lutil import pax + # pax("request") + # <<< if request.status != Status.FAILED: self.waiting_request_ids.append(request_id) @@ -646,6 +661,9 @@ def schedule_chunked_prefill(self): """ can_schedule = True while self.waiting_request_ids and can_schedule: + # >>> + # raise Exception("hi.") + # <<< can_schedule = False req = self.requests[self.waiting_request_ids[0]] @@ -662,6 +680,9 @@ def schedule_chunked_prefill(self): 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 + # >>> + raise Exception("hi.") + # <<< if request_can_be_added and kv_cache_available: if token_fully_can_be_added: self.context.chunked_prefill_request_id = -1 @@ -772,7 +793,7 @@ async def async_step( output_str = ( "* step %d | %s ... time: %.3f%s ... " "reqs: a %d/%d, p %d/%d, w %d, f %d ... " - "chunks: a %d/%d, p %d/%d ... " + "blocks: a %d/%d, p %d/%d ... " "mem: tensors %d, alloc %.1f gb, res %.1f gb." % ( self.step_count, @@ -791,15 +812,15 @@ async def async_step( ) ), prev_total_request_count - prev_paused_request_count, - context.chunk_allocator.active_count, + context.block_allocator.active_count, prev_paused_request_count, - context.chunk_allocator.paused_count, + context.block_allocator.paused_count, len(self.waiting_request_ids), self.finished_request_count, - context.chunk_allocator.get_active_used(), - context.chunk_allocator.active_count, - context.chunk_allocator.get_paused_used(), - context.chunk_allocator.paused_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), From bf1a342220fa556d5e21251689dc4699d68798c9 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Thu, 30 Oct 2025 12:07:32 -0700 Subject: [PATCH 49/69] running again. --- .../inference/gpt/gpt_dynamic_inference.py | 14 ------------- .../inference/contexts/dynamic_context.py | 11 ++++------ .../core/inference/engines/dynamic_engine.py | 21 ------------------- 3 files changed, 4 insertions(+), 42 deletions(-) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index 2f7c94bbeda..f8cd0e26c08 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -64,12 +64,6 @@ torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunState]) torch.serialization.add_safe_globals([megatron.core.rerun_state_machine.RerunDiagnostic]) -# >>> -from lutil import pax as _pax -import builtins -builtins.pax = _pax -# <<< - def add_dynamic_inference_args(parser: ArgumentParser) -> ArgumentParser: """Dynamic inference arguments.""" @@ -260,9 +254,6 @@ def _add_request(): _request.prompt_text, _request.sampling_params, ) - # >>> - print("+++++++++++ added request %d / %d." % (num_requests_added, num_requests_total)) - # <<< _request.time_start = get_curr_time() _request.state = "started" num_requests_added += 1 @@ -301,11 +292,6 @@ def _add_request(): # Update requests. active_requests = result["active_requests"] finished_requests = result["finished_requests"] - # >>> - if finished_requests: - from lutil import pax - pax("finished_requests") - # <<< step_time = result["step_time"] if len(active_requests) > 0 or len(finished_requests) > 0: if is_decode_only: diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 4f4b96c90aa..da41f166a9f 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1086,8 +1086,8 @@ def check_availability(self, req: DynamicInferenceRequest) -> (bool, bool, bool) self.active_token_count + req.remaining_prompt_length <= self.max_tokens ) blocks = math.ceil( - (req.remaining_prompt_length + req.finished_block_token_count) / self.block_size_tokens - ) - math.ceil(req.finished_block_token_count / self.block_size_tokens) + (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) return request_can_be_added, request_tokens_can_be_added, kv_cache_available @@ -1101,9 +1101,6 @@ def add_request(self, req: DynamicInferenceRequest, chunk_length: Optional[int] Return: None """ - # >>> - raise Exception("hi.") - # <<< if chunk_length is None: chunk_length = req.remaining_prompt_length @@ -1379,9 +1376,9 @@ 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_chunk_count > 0: + if active_requests_requiring_new_block_count > 0: newly_paused_request_ids = self.request_ids[ - torch.nonzero(active_requests_requiring_new_chunk) + self.paused_request_count + 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 diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 7c375e7e0f3..5619eeece21 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -404,28 +404,13 @@ def _add_request( len(request.prompt_tokens) + request.sampling_params.num_tokens_to_generate > self.context.max_sequence_length ): - # >>> - from lutil import pax - pax("request", { - "prompt" : len(request.prompt_tokens), - "gen" : request.sampling_params.num_tokens_to_generate, - "max_sequence_length" : self.context.max_sequence_length, - }) - # <<< request.status = Status.FAILED request.add_event_error_nontransient(MaxSequenceLengthOverflowError(request_id)) if len(request.prompt_tokens) > self.context.max_tokens and not self.enable_chunked_prefill: - # >>> - raise Exception("hi.") - # <<< request.status = Status.FAILED request.add_event_error_nontransient(TokenOverflowError(request_id)) - # >>> - # from lutil import pax - # pax("request") - # <<< if request.status != Status.FAILED: self.waiting_request_ids.append(request_id) @@ -661,9 +646,6 @@ def schedule_chunked_prefill(self): """ can_schedule = True while self.waiting_request_ids and can_schedule: - # >>> - # raise Exception("hi.") - # <<< can_schedule = False req = self.requests[self.waiting_request_ids[0]] @@ -680,9 +662,6 @@ def schedule_chunked_prefill(self): 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 - # >>> - raise Exception("hi.") - # <<< if request_can_be_added and kv_cache_available: if token_fully_can_be_added: self.context.chunked_prefill_request_id = -1 From 73b7041a011c353c715228335bb2d2bf7c986f4b Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Thu, 30 Oct 2025 12:08:30 -0700 Subject: [PATCH 50/69] minor update. --- examples/inference/gpt/utils.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/examples/inference/gpt/utils.py b/examples/inference/gpt/utils.py index fcc9c8136dd..73b5bd90fe5 100644 --- a/examples/inference/gpt/utils.py +++ b/examples/inference/gpt/utils.py @@ -283,14 +283,6 @@ def get_synthetic_requests( for t, o in zip(prompt_texts, time_offsets) ] - # >>> - # pax("max_prompt_length", { - # "old num_tokens_to_generate" : sampling_params.num_tokens_to_generate, - # "new num_tokens_to_generate" : requests[0].sampling_params.num_tokens_to_generate, - # "new num_tokens_to_generate'" : requests[-1].sampling_params.num_tokens_to_generate, - # "set(num_tokens_to_generate)" : "%d | %s" % (len(requests), set(r.sampling_params.num_tokens_to_generate for r in requests)), - # }) - # <<< return requests From f2548c61d63e1125d933c616d05a316685495bd5 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Fri, 31 Oct 2025 08:38:17 -0700 Subject: [PATCH 51/69] fixed test_dynamic_context.py. --- .../inference/contexts/dynamic_context.py | 2 +- .../contexts/test_dynamic_context.py | 155 +++++++++--------- 2 files changed, 79 insertions(+), 78 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index da41f166a9f..e19971284a9 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1429,7 +1429,7 @@ def update_requests(self, active_requests_mask: Tensor, new_tokens: Tensor) -> T paused_block_counts += 1 # +1 for newly added block paused_block_counts_cumsum = paused_block_counts.cumsum(dim=0) resume_request_count = min( - torch.nonzero(paused_chunk_counts_cumsum <= active_chunk_count_avail).numel(), + torch.nonzero(paused_block_counts_cumsum <= active_block_count_avail).numel(), self.block_allocator.total_avail, ) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 70c2bf6ed10..d0d5d096654 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -40,7 +40,7 @@ def _get_dynamic_context( kv_channels, num_attention_heads, max_sequence_length, - buffer_size_gb, + active_buffer_size_gb, block_size_tokens, buffer_guarenteed_fraction, buffer_overflow_factor, @@ -55,12 +55,8 @@ 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_guarenteed_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, use_flashinfer_fused_rope=None, # default to using flash-infer if available # this is for compatibility with the LTS environment ) @@ -80,7 +76,7 @@ def test_initialize_dynamic_context(self): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, + active_buffer_size_gb=0.03, buffer_guarenteed_fraction=0.1, block_size_tokens=128, max_requests_override=None, @@ -88,11 +84,12 @@ def test_initialize_dynamic_context(self): buffer_overflow_factor=None, ) - 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.block_allocator.paused_count == 490 + assert dynamic_context.max_total_requests == 980 + assert dynamic_context.max_active_requests == 490 + assert dynamic_context.max_tokens == 16384 # Check initializations to -1 assert torch.all(dynamic_context.request_ids == -1) @@ -107,8 +104,7 @@ def test_is_static_batching(self): num_attention_heads=8, max_sequence_length=512, num_cuda_graphs=None, - buffer_size_gb=1.0, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=1.0, block_size_tokens=128, ) assert not dynamic_context.is_static_batching() @@ -123,23 +119,17 @@ def test_is_memory_available(self): num_attention_heads=8, max_sequence_length=512, num_cuda_graphs=None, - buffer_size_gb=1.0, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=1.0, block_size_tokens=128, ) - 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 def test_request_overflow(self): self._setup_model_parallel_group(1, 1) @@ -151,12 +141,12 @@ def test_request_overflow(self): num_attention_heads=8, max_sequence_length=128, num_cuda_graphs=None, - buffer_size_gb=0.01, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=0.01, block_size_tokens=32, ) + 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, @@ -178,12 +168,9 @@ def test_token_overflow_error(self): num_attention_heads=8, max_sequence_length=512, num_cuda_graphs=None, - 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=20, # Setting a very low token limit ) with pytest.raises(TokenOverflowError): @@ -207,8 +194,7 @@ def test_reset(self): num_attention_heads=8, max_sequence_length=128, num_cuda_graphs=None, - buffer_size_gb=1.0, - buffer_guaranteed_fraction=0.1, + active_buffer_size_gb=1.0, block_size_tokens=128, ) @@ -231,7 +217,7 @@ def test_reset(self): 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.block_allocator.active_count = 1023 dynamic_context.memory_buffer.fill_(1) dynamic_context.request_to_kv_block_ids.fill_(1) @@ -258,8 +244,8 @@ def test_reset(self): 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) @@ -272,7 +258,7 @@ def test_allocate_and_release_memory_blocks(self): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, + active_buffer_size_gb=0.03, buffer_guarenteed_fraction=0.1, block_size_tokens=128, max_requests_override=None, @@ -282,18 +268,18 @@ def test_allocate_and_release_memory_blocks(self): assert dynamic_context.block_allocator.allocate_memory_blocks( 4 - ).cpu().detach().numpy().tolist() == [486, 487, 488, 489] - assert dynamic_context.block_allocator.block_count_avail == 486 + ).cpu().detach().numpy().tolist() == [976, 977, 978, 979] + assert dynamic_context.block_allocator.total_avail == 976 dynamic_context.block_allocator.release_memory_blocks( torch.tensor([488, 489], device='cuda') ) - assert dynamic_context.block_allocator.block_count_avail == 488 + assert dynamic_context.block_allocator.total_avail == 978 assert dynamic_context.block_allocator.allocate_memory_blocks(1).item() == 489 - assert dynamic_context.block_allocator.block_count_avail == 487 + assert dynamic_context.block_allocator.total_avail == 977 # 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 ) @@ -307,7 +293,7 @@ def test_add_request(self): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, + active_buffer_size_gb=0.03, buffer_guarenteed_fraction=0.1, block_size_tokens=128, max_requests_override=None, @@ -332,13 +318,13 @@ def test_add_request(self): assert dynamic_context.request_query_lengths[0] == context_length assert dynamic_context.request_kv_length_offsets[0] == 0 assert dynamic_context.request_to_kv_block_ids[0].cpu().detach().numpy().tolist() == [ - 488, - 489, + 978, + 979, -1, -1, ] assert dynamic_context.request_kv_block_counts[0] == 2 - assert dynamic_context.request_last_kv_block_id[0] == 489 + assert dynamic_context.request_last_kv_block_id[0] == 979 assert dynamic_context.request_last_kv_block_offset[0].item() == 15 assert torch.all( dynamic_context.token_to_pos_ids[0:context_length] @@ -356,13 +342,13 @@ def test_add_request(self): dynamic_context.token_to_block_idx[0:context_length][ 0 : dynamic_context.block_size_tokens ] - == 488 + == 978 ) assert torch.all( dynamic_context.token_to_block_idx[0:context_length][ dynamic_context.block_size_tokens : context_length ] - == 489 + == 979 ) assert torch.all( dynamic_context.token_to_local_position_within_kv_block[0:context_length] @@ -379,7 +365,7 @@ def test_update_request(self): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, + active_buffer_size_gb=0.03, buffer_guarenteed_fraction=0.1, block_size_tokens=128, max_requests_override=None, @@ -392,7 +378,7 @@ def test_update_request(self): 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 dynamic_context.update_requests( active_requests_mask=active_requests_mask, new_tokens=torch.tensor([0, 1, 2]) @@ -416,7 +402,7 @@ def test_update_request(self): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, + active_buffer_size_gb=0.03, buffer_guarenteed_fraction=0.1, block_size_tokens=128, max_requests_override=None, @@ -438,17 +424,17 @@ def test_update_request(self): ) 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 + 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 @@ -526,17 +512,17 @@ def test_update_request(self): 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], - ] + ] ) ) @@ -550,7 +536,7 @@ def test_release_memory_blocks_for_finished_requests(self): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, + active_buffer_size_gb=0.03, buffer_guarenteed_fraction=0.1, block_size_tokens=128, max_requests_override=None, @@ -560,12 +546,12 @@ def test_release_memory_blocks_for_finished_requests(self): # 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): @@ -589,7 +575,7 @@ def test_release_memory_blocks_for_finished_requests(self): 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 @pytest.mark.internal def test_finished_requests_with_multiple_blocks(self): @@ -601,7 +587,7 @@ def test_finished_requests_with_multiple_blocks(self): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, + active_buffer_size_gb=0.03, buffer_guarenteed_fraction=0.1, block_size_tokens=128, max_requests_override=None, @@ -611,12 +597,12 @@ def test_finished_requests_with_multiple_blocks(self): # 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 @@ -653,7 +639,7 @@ def test_finished_requests_with_multiple_blocks(self): 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 @pytest.mark.internal def test_calculate_and_store_log_probs(self): @@ -664,7 +650,7 @@ def test_calculate_and_store_log_probs(self): kv_channels=8, num_attention_heads=2, max_sequence_length=512, - buffer_size_gb=0.03, + active_buffer_size_gb=0.03, buffer_guarenteed_fraction=0.1, block_size_tokens=128, max_requests_override=None, @@ -886,8 +872,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): @@ -900,9 +886,7 @@ def init_contexts(*, unified_memory_level): kv_channels=8, num_attention_heads=2, 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, ) ) @@ -917,3 +901,20 @@ def init_contexts(*, unified_memory_level): # Unified memory test should succeed. init_contexts(unified_memory_level=1) +# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> +if __name__ == "__main__": + test = TestDynamicContext() + # test.test_initialize_dynamic_context() + # test.test_is_static_batching() + # test.test_is_memory_available() + # test.test_request_overflow() + # test.test_token_overflow_error() + # test.test_reset() + # test.test_allocate_and_release_memory_blocks() + # test.test_add_request() + # test.test_update_request() + # test.test_release_memory_blocks_for_finished_requests() + test.test_finished_requests_with_multiple_blocks() + test.test_calculate_and_store_log_probs() + test.test_unified_memory() +# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< From fa23ce8a2a4dc88016f2523204a1ad234fedf6fc Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Fri, 31 Oct 2025 10:42:08 -0700 Subject: [PATCH 52/69] default uvm level == 1; assert(max_tokens >= max_active_requests). --- .../inference/contexts/dynamic_context.py | 35 ++++- .../core/inference/engines/dynamic_engine.py | 7 + megatron/core/transformer/attention.py | 5 + megatron/training/arguments.py | 2 +- .../contexts/test_dynamic_context.py | 17 --- .../inference/engines/test_dynamic_engine.py | 135 +++++++++++++----- .../test_simple_text_generation_controller.py | 3 +- 7 files changed, 147 insertions(+), 57 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index e19971284a9..a5dd2102487 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -27,10 +27,14 @@ from .base_context import BaseInferenceContext from .dynamic_block_allocator import BlockAllocator -try: - from .fused_kv_append_kernel import triton_append_key_value_cache -except ImportError: - triton_append_key_value_cache = None +# >>> +# try: +# from .fused_kv_append_kernel import triton_append_key_value_cache +# except ImportError: +# triton_append_key_value_cache = None +# +++ +triton_append_key_value_cache = None +# <<< try: from packaging.version import Version as PkgVersion @@ -245,7 +249,7 @@ def __init__( materialize_only_last_token_logits: bool = True, 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, ): super().__init__(materialize_only_last_token_logits=materialize_only_last_token_logits) @@ -295,6 +299,12 @@ def __init__( self.max_active_requests = self.block_allocator.active_count self.max_tokens = 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 self.num_layers = num_layers @@ -617,6 +627,14 @@ def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) """ if triton_append_key_value_cache is not None and not self.cache_mla_latent: # currently does not support MLA latent cache + # >>> + pax("layer_number, key, value", { + "memory_buffer" : self.memory_buffer, + "padded_active_token_count" : self.padded_active_token_count, + "token_to_block_idx" : self.token_to_block_idx, + "token_to_local_position_within_kv_block" : self.token_to_local_position_within_kv_block, + }) + # <<< return triton_append_key_value_cache( layer_number=layer_number, key=key, @@ -1040,6 +1058,13 @@ def current_input_and_position_ids( (Tuple[Tensor, Tensor]) Flattened active input and position IDs. """ num_tokens = num_warmup_tokens or self.padded_active_token_count + # >>> + pax("num_warmup_tokens", { + "padded_active_token_count" : self.padded_active_token_count, + "inp" : self.token_to_input_ids[:num_tokens].unsqueeze(0), + "pos" : self.token_to_pos_ids[:num_tokens].unsqueeze(0), + }) + # <<< return ( self.token_to_input_ids[:num_tokens].unsqueeze(0), self.token_to_pos_ids[:num_tokens].unsqueeze(0), diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 5619eeece21..e32b4bf264b 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -211,6 +211,13 @@ def create_cuda_graphs(self, reset_context: bool = True): logging.info(f"{tbar_idx}/{len(context.cuda_graph_token_counts)}. {tbar_str}") # Forward pass -> logits. + # >>> + pax({ + "context" : self.context, + "block_allocator" : self.context.block_allocator, + }, "cuda_graph_token_count, input_ids, position_ids") + print("~~~~~~~~~~~~~~~~~~~~~~ cuda_graph_token_count : %d." % cuda_graph_token_count) + # <<< controller._dynamic_step_forward_logits(input_ids, position_ids) if reset_context: diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 659402a39b8..a4f5cf0b409 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -420,7 +420,9 @@ def _adjust_key_value_for_inference( rotary_pos_emb = (q_pos_emb, None) # key rotary emb has been applied # Append key/value data tensors to cache. + # >>> [ problematic ] inference_context.append_key_value_cache(self.layer_number, key, value) + # <<< _, max_seqlen_q = inference_context.cu_query_lengths() if getattr(self.config, "cache_mla_latents", None) and max_seqlen_q > 1: @@ -617,6 +619,9 @@ def flash_decode_and_prefill( if HAVE_FA3: output_total = flash_attn3_with_kvcache(**flash_attn_args) else: + # >>> + # pax("flash_attn_args") + # <<< output_total = flash_attn_with_kvcache(**flash_attn_args) return output_total diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 5081fee043a..8a74e744e54 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1460,7 +1460,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 d0d5d096654..6800685ecd4 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -901,20 +901,3 @@ def init_contexts(*, unified_memory_level): # Unified memory test should succeed. init_contexts(unified_memory_level=1) -# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> -if __name__ == "__main__": - test = TestDynamicContext() - # test.test_initialize_dynamic_context() - # test.test_is_static_batching() - # test.test_is_memory_available() - # test.test_request_overflow() - # test.test_token_overflow_error() - # test.test_reset() - # test.test_allocate_and_release_memory_blocks() - # test.test_add_request() - # test.test_update_request() - # test.test_release_memory_blocks_for_finished_requests() - test.test_finished_requests_with_multiple_blocks() - test.test_calculate_and_store_log_probs() - test.test_unified_memory() -# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 2227367469b..d83d94f67a1 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -39,6 +39,12 @@ from megatron.core.utils import is_fa_min_version from tests.unit_tests.test_utilities import Utils +# >>> +from lutil import pax as _pax +import builtins +builtins.pax = _pax +# <<< + def set_rounder(value): """Utility function to set the DynamicInferenceContext rounder.""" @@ -64,12 +70,11 @@ 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 @@ -103,13 +108,9 @@ def __post_init__(self): # 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 + if self.context_max_tokens is None: + self.context_max_tokens = self.num_requests * self.max_sequence_length @dataclass @@ -202,12 +203,9 @@ def _build_inference_context( num_attention_heads=transformer_config.num_query_groups, max_sequence_length=test_config.max_sequence_length, num_cuda_graphs=test_config.num_cuda_graphs, - 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, materialize_only_last_token_logits=test_config.materialize_only_last_token_logits, use_flashinfer_fused_rope=None, # default to using flash-infer if available @@ -410,13 +408,11 @@ def test_simple(self, num_cuda_graphs, cuda_graph_scope) -> None: # Run test. env = self._run_test( 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 # Validate generated tokens. @@ -448,8 +444,7 @@ def test_overflow_factor(self) -> None: # Run test. env = self._run_test( context_buffer_overflow_factor=0.1, - context_max_requests_override=None, - context_max_tokens_override=None, + context_max_tokens=None, ) # Validate max_requests, max_tokens. @@ -470,7 +465,7 @@ def test_request_overflow(self) -> 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=8, max_prompt_length=8, context_max_tokens=12 ) env = self._build_test_env(test_config) env.engine._add_request(env.requests[0]) @@ -489,7 +484,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]) @@ -507,8 +502,8 @@ def test_block_overflow(self) -> None: env = self._build_test_env(DynamicEngineTestConfig()) context = env.engine.context block_size_bytes = context.block_size_bytes - buffer_size_gb = (block_size_bytes + 1) / 1024**3 - test_config = DynamicEngineTestConfig(context_buffer_size_gb=buffer_size_gb) + active_buffer_size_gb = (block_size_bytes + 1) / 1024**3 + test_config = DynamicEngineTestConfig(context_active_buffer_size_gb=active_buffer_size_gb) env = self._build_test_env(test_config) env.engine._add_request(env.requests[0]) assert list(env.engine.waiting_request_ids) == [0] @@ -821,9 +816,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, ) @@ -852,21 +846,98 @@ def test_events(self): if __name__ == "__main__": test = TestDynamicInferenceEngine() - test.test_simple(4) + # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> + # 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.test_simple(None, "full") + test.test_simple(1, "full") + exit() + test.test_simple(4, "full") + test.test_simple(None, "full_iteration") + test.test_simple(1, "full_iteration") + test.test_simple(4, "full_iteration") 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_token_overflow_nontransient() 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_cuda_graph_token_counts() + test.test_cuda_graph_warmup(1, 8, WarmupEngineMode.DECODE) + test.test_cuda_graph_warmup(1, 8, WarmupEngineMode.NON_DECODE) + test.test_cuda_graph_warmup(2, 8, WarmupEngineMode.DECODE) + test.test_cuda_graph_warmup(2, 8, WarmupEngineMode.NON_DECODE) + test.test_cuda_graph_warmup(4, 8, WarmupEngineMode.DECODE) + test.test_cuda_graph_warmup(4, 8, WarmupEngineMode.NON_DECODE) + test.test_cuda_graph_warmup(8, 8, WarmupEngineMode.DECODE) + test.test_cuda_graph_warmup(8, 8, WarmupEngineMode.NON_DECODE) + test.test_cuda_graph_warmup(10, 16, WarmupEngineMode.DECODE) + test.test_cuda_graph_warmup(10, 16, WarmupEngineMode.NON_DECODE) + test.test_cuda_graph_warmup(12, 16, WarmupEngineMode.DECODE) + test.test_cuda_graph_warmup(12, 16, WarmupEngineMode.NON_DECODE) + test.test_cuda_graph_warmup(16, 16, WarmupEngineMode.DECODE) + test.test_cuda_graph_warmup(16, 16, WarmupEngineMode.NON_DECODE) + test.test_cuda_graph_warmup(20, 24, WarmupEngineMode.DECODE) + test.test_cuda_graph_warmup(20, 24, WarmupEngineMode.NON_DECODE) + test.test_cuda_graph_warmup(24, 24, WarmupEngineMode.DECODE) + test.test_cuda_graph_warmup(24, 24, WarmupEngineMode.NON_DECODE) + test.test_cuda_graph_warmup(28, 32, WarmupEngineMode.DECODE) + test.test_cuda_graph_warmup(28, 32, WarmupEngineMode.NON_DECODE) + test.test_cuda_graph_warmup(32, 32, WarmupEngineMode.DECODE) + test.test_cuda_graph_warmup(32, 32, WarmupEngineMode.NON_DECODE) test.test_generate_function() - asyncio.run(test.test_run_engine()) + test.test_run_engine() test.test_return_log_probs() - test.test_parallel_inference() - # test.test_events() # uncomment in megatron-core 0.16 + test.test_parallel_inference(1, 1, 1, False, False) + test.test_parallel_inference(1, 1, 1, False, True) + test.test_parallel_inference(1, 1, 1, True, False) + test.test_parallel_inference(1, 1, 1, True, True) + test.test_parallel_inference(1, 1, 2, False, False) + test.test_parallel_inference(1, 1, 2, False, True) + test.test_parallel_inference(1, 1, 2, True, False) + test.test_parallel_inference(1, 1, 2, True, True) + test.test_parallel_inference(1, 2, 1, False, False) + test.test_parallel_inference(1, 2, 1, False, True) + test.test_parallel_inference(1, 2, 1, True, False) + test.test_parallel_inference(1, 2, 1, True, True) + test.test_parallel_inference(1, 2, 2, False, False) + test.test_parallel_inference(1, 2, 2, False, True) + test.test_parallel_inference(1, 2, 2, True, False) + test.test_parallel_inference(1, 2, 2, True, True) + test.test_parallel_inference(2, 1, 1, False, False) + test.test_parallel_inference(2, 1, 1, False, True) + test.test_parallel_inference(2, 1, 1, True, False) + test.test_parallel_inference(2, 1, 1, True, True) + test.test_parallel_inference(2, 1, 2, False, False) + test.test_parallel_inference(2, 1, 2, False, True) + test.test_parallel_inference(2, 1, 2, True, False) + test.test_parallel_inference(2, 1, 2, True, True) + test.test_parallel_inference(2, 2, 1, False, False) + test.test_parallel_inference(2, 2, 1, False, True) + test.test_parallel_inference(2, 2, 1, True, False) + test.test_parallel_inference(2, 2, 1, True, True) + test.test_parallel_inference(2, 2, 2, False, False) + test.test_parallel_inference(2, 2, 2, False, True) + test.test_parallel_inference(2, 2, 2, True, False) + test.test_parallel_inference(2, 2, 2, True, True) + test.test_num_tokens_total() + test.test_events() + # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< test.teardown_method(None) print("~~~") print("success.") 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 f23a9782646..01ef8f5749f 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=1, 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 From c83b37312e2c19777e542207b6281288a7fc9224 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Fri, 31 Oct 2025 11:04:51 -0700 Subject: [PATCH 53/69] formalized DEFAULT_MAX_TOKENS. --- .../core/inference/contexts/dynamic_context.py | 17 +++++++++++------ megatron/core/transformer/attention.py | 5 ----- .../inference/engines/test_dynamic_engine.py | 9 --------- 3 files changed, 11 insertions(+), 20 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index a5dd2102487..195652dc490 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -211,7 +211,8 @@ class DynamicInferenceContext(BaseInferenceContext): 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. + primarily limited by prefill activation memory usage. (Defaults to + 16384). block_size_tokens (int): Size of KV cache block size. tensor_model_parallel_size (Optional[int]): Tensor model parallel size. num_cuda_graphs (Optional[int]): Maximum number of cuda graphs to capture, @@ -230,6 +231,10 @@ class DynamicInferenceContext(BaseInferenceContext): If None, defaults to using flash-infer if available. """ + DEFAULT_MAX_TOKENS = 16384 + TOKEN_ROUNDER = 64 + REQUEST_ROUNDER = 4 + def __init__( self, *, @@ -239,7 +244,7 @@ def __init__( num_attention_heads: int, max_sequence_length: int, active_buffer_size_gb: float, - max_tokens: int = 16384, + max_tokens: int = DEFAULT_MAX_TOKENS, block_size_tokens: int = 256, tensor_model_parallel_size: Optional[int] = None, cache_mla_latent: bool = False, @@ -297,7 +302,10 @@ def __init__( # 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 + self.max_tokens = max_tokens or self.DEFAULT_MAX_TOKENS + # >>> + # pax({"DEFAULT_MAX_TOKENS": self.DEFAULT_MAX_TOKENS}, "max_tokens", {"max_tokens'": self.max_tokens}) + # <<< assert self.max_tokens >= self.max_active_requests, ( f"max_tokens ({self.max_tokens}) must be >= " @@ -485,9 +493,6 @@ def __init__( % (get_mem_size_str(active_buffer_size_bytes), self.block_allocator.active_count) ) - TOKEN_ROUNDER = 64 - REQUEST_ROUNDER = 4 - @classmethod def round_up_tokens(cls, value, tp_size=None): """Round up to nearest multiple of `TOKEN_ROUNDER` (above) that is also divisible by tensor model parallel size.""" diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index a4f5cf0b409..659402a39b8 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -420,9 +420,7 @@ def _adjust_key_value_for_inference( rotary_pos_emb = (q_pos_emb, None) # key rotary emb has been applied # Append key/value data tensors to cache. - # >>> [ problematic ] inference_context.append_key_value_cache(self.layer_number, key, value) - # <<< _, max_seqlen_q = inference_context.cu_query_lengths() if getattr(self.config, "cache_mla_latents", None) and max_seqlen_q > 1: @@ -619,9 +617,6 @@ def flash_decode_and_prefill( if HAVE_FA3: output_total = flash_attn3_with_kvcache(**flash_attn_args) else: - # >>> - # pax("flash_attn_args") - # <<< output_total = flash_attn_with_kvcache(**flash_attn_args) return output_total diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index d83d94f67a1..0b7111b8603 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -72,8 +72,6 @@ class DynamicEngineTestConfig: 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_tokens: Optional[int] = None tensor_model_parallel_size: int = 1 pipeline_model_parallel_size: int = 1 @@ -105,13 +103,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 tokens. - if self.context_max_tokens is None: - self.context_max_tokens = self.num_requests * self.max_sequence_length - @dataclass class DynamicEngineTestEnv: From 0c387f56920b68180e270efb0e8a99feac01f003 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Fri, 31 Oct 2025 11:20:45 -0700 Subject: [PATCH 54/69] test simple runs with and without triton kernel. --- .../inference/contexts/dynamic_context.py | 32 +++++++++---------- .../core/inference/engines/dynamic_engine.py | 10 +++--- .../inference/engines/test_dynamic_engine.py | 4 +-- 3 files changed, 23 insertions(+), 23 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 195652dc490..167ddc9f060 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -28,12 +28,12 @@ from .dynamic_block_allocator import BlockAllocator # >>> -# try: -# from .fused_kv_append_kernel import triton_append_key_value_cache -# except ImportError: -# triton_append_key_value_cache = None +try: + from .fused_kv_append_kernel import triton_append_key_value_cache +except ImportError: + triton_append_key_value_cache = None # +++ -triton_append_key_value_cache = None +# triton_append_key_value_cache = None # <<< try: @@ -633,12 +633,12 @@ def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) if triton_append_key_value_cache is not None and not self.cache_mla_latent: # currently does not support MLA latent cache # >>> - pax("layer_number, key, value", { - "memory_buffer" : self.memory_buffer, - "padded_active_token_count" : self.padded_active_token_count, - "token_to_block_idx" : self.token_to_block_idx, - "token_to_local_position_within_kv_block" : self.token_to_local_position_within_kv_block, - }) + # pax("layer_number, key, value", { + # "memory_buffer" : self.memory_buffer, + # "padded_active_token_count" : self.padded_active_token_count, + # "token_to_block_idx" : self.token_to_block_idx, + # "token_to_local_position_within_kv_block" : self.token_to_local_position_within_kv_block, + # }) # <<< return triton_append_key_value_cache( layer_number=layer_number, @@ -1064,11 +1064,11 @@ def current_input_and_position_ids( """ num_tokens = num_warmup_tokens or self.padded_active_token_count # >>> - pax("num_warmup_tokens", { - "padded_active_token_count" : self.padded_active_token_count, - "inp" : self.token_to_input_ids[:num_tokens].unsqueeze(0), - "pos" : self.token_to_pos_ids[:num_tokens].unsqueeze(0), - }) + # pax("num_warmup_tokens", { + # "padded_active_token_count" : self.padded_active_token_count, + # "inp" : self.token_to_input_ids[:num_tokens].unsqueeze(0), + # "pos" : self.token_to_pos_ids[:num_tokens].unsqueeze(0), + # }) # <<< return ( self.token_to_input_ids[:num_tokens].unsqueeze(0), diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index e32b4bf264b..253368173ca 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -212,11 +212,11 @@ def create_cuda_graphs(self, reset_context: bool = True): # Forward pass -> logits. # >>> - pax({ - "context" : self.context, - "block_allocator" : self.context.block_allocator, - }, "cuda_graph_token_count, input_ids, position_ids") - print("~~~~~~~~~~~~~~~~~~~~~~ cuda_graph_token_count : %d." % cuda_graph_token_count) + # pax({ + # "context" : self.context, + # "block_allocator" : self.context.block_allocator, + # }, "cuda_graph_token_count, input_ids, position_ids") + # print("~~~~~~~~~~~~~~~~~~~~~~ cuda_graph_token_count : %d." % cuda_graph_token_count) # <<< controller._dynamic_step_forward_logits(input_ids, position_ids) diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 0b7111b8603..3b08d167f22 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -404,7 +404,7 @@ def test_simple(self, num_cuda_graphs, cuda_graph_scope) -> None: ) # Validate max_requests, max_tokens. - assert env.engine.context.max_tokens == 160 + assert env.engine.context.max_tokens == DynamicInferenceContext.DEFAULT_MAX_TOKENS # Validate generated tokens. expected_generated_tokens_list = [ @@ -856,7 +856,7 @@ def test_events(self): # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ # test.test_simple(None, "full") test.test_simple(1, "full") - exit() + print("~~~\nsuccess."); exit() test.test_simple(4, "full") test.test_simple(None, "full_iteration") test.test_simple(1, "full_iteration") From 717dfe93f0bdbb8026dd4da43bf3c460d7419be7 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Fri, 31 Oct 2025 11:23:36 -0700 Subject: [PATCH 55/69] test_dynamic_engine.py works. --- .../inference/contexts/dynamic_context.py | 22 ---- .../core/inference/engines/dynamic_engine.py | 7 -- .../inference/engines/test_dynamic_engine.py | 105 ------------------ 3 files changed, 134 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 167ddc9f060..c8bcacc3e5a 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -27,14 +27,10 @@ from .base_context import BaseInferenceContext from .dynamic_block_allocator import BlockAllocator -# >>> try: from .fused_kv_append_kernel import triton_append_key_value_cache except ImportError: triton_append_key_value_cache = None -# +++ -# triton_append_key_value_cache = None -# <<< try: from packaging.version import Version as PkgVersion @@ -303,9 +299,6 @@ def __init__( 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 - # >>> - # pax({"DEFAULT_MAX_TOKENS": self.DEFAULT_MAX_TOKENS}, "max_tokens", {"max_tokens'": self.max_tokens}) - # <<< assert self.max_tokens >= self.max_active_requests, ( f"max_tokens ({self.max_tokens}) must be >= " @@ -632,14 +625,6 @@ def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) """ if triton_append_key_value_cache is not None and not self.cache_mla_latent: # currently does not support MLA latent cache - # >>> - # pax("layer_number, key, value", { - # "memory_buffer" : self.memory_buffer, - # "padded_active_token_count" : self.padded_active_token_count, - # "token_to_block_idx" : self.token_to_block_idx, - # "token_to_local_position_within_kv_block" : self.token_to_local_position_within_kv_block, - # }) - # <<< return triton_append_key_value_cache( layer_number=layer_number, key=key, @@ -1063,13 +1048,6 @@ def current_input_and_position_ids( (Tuple[Tensor, Tensor]) Flattened active input and position IDs. """ num_tokens = num_warmup_tokens or self.padded_active_token_count - # >>> - # pax("num_warmup_tokens", { - # "padded_active_token_count" : self.padded_active_token_count, - # "inp" : self.token_to_input_ids[:num_tokens].unsqueeze(0), - # "pos" : self.token_to_pos_ids[:num_tokens].unsqueeze(0), - # }) - # <<< return ( self.token_to_input_ids[:num_tokens].unsqueeze(0), self.token_to_pos_ids[:num_tokens].unsqueeze(0), diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 253368173ca..5619eeece21 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -211,13 +211,6 @@ def create_cuda_graphs(self, reset_context: bool = True): logging.info(f"{tbar_idx}/{len(context.cuda_graph_token_counts)}. {tbar_str}") # Forward pass -> logits. - # >>> - # pax({ - # "context" : self.context, - # "block_allocator" : self.context.block_allocator, - # }, "cuda_graph_token_count, input_ids, position_ids") - # print("~~~~~~~~~~~~~~~~~~~~~~ cuda_graph_token_count : %d." % cuda_graph_token_count) - # <<< controller._dynamic_step_forward_logits(input_ids, position_ids) if reset_context: diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 3b08d167f22..5b8a86fb4aa 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -39,12 +39,6 @@ from megatron.core.utils import is_fa_min_version from tests.unit_tests.test_utilities import Utils -# >>> -from lutil import pax as _pax -import builtins -builtins.pax = _pax -# <<< - def set_rounder(value): """Utility function to set the DynamicInferenceContext rounder.""" @@ -833,102 +827,3 @@ def test_events(self): result_event_types = [[e.type.name for e in r.events] for r in env.requests] assert result_event_types == expected_event_types - - -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.test_simple(None, "full") - test.test_simple(1, "full") - print("~~~\nsuccess."); exit() - test.test_simple(4, "full") - test.test_simple(None, "full_iteration") - test.test_simple(1, "full_iteration") - test.test_simple(4, "full_iteration") - test.test_overflow_factor() - test.test_request_overflow() - test.test_token_overflow_transient() - test.test_token_overflow_nontransient() - test.test_block_overflow() - test.test_multi_add() - test.test_fixed_output_lengths() - test.test_cuda_graph_token_counts() - test.test_cuda_graph_warmup(1, 8, WarmupEngineMode.DECODE) - test.test_cuda_graph_warmup(1, 8, WarmupEngineMode.NON_DECODE) - test.test_cuda_graph_warmup(2, 8, WarmupEngineMode.DECODE) - test.test_cuda_graph_warmup(2, 8, WarmupEngineMode.NON_DECODE) - test.test_cuda_graph_warmup(4, 8, WarmupEngineMode.DECODE) - test.test_cuda_graph_warmup(4, 8, WarmupEngineMode.NON_DECODE) - test.test_cuda_graph_warmup(8, 8, WarmupEngineMode.DECODE) - test.test_cuda_graph_warmup(8, 8, WarmupEngineMode.NON_DECODE) - test.test_cuda_graph_warmup(10, 16, WarmupEngineMode.DECODE) - test.test_cuda_graph_warmup(10, 16, WarmupEngineMode.NON_DECODE) - test.test_cuda_graph_warmup(12, 16, WarmupEngineMode.DECODE) - test.test_cuda_graph_warmup(12, 16, WarmupEngineMode.NON_DECODE) - test.test_cuda_graph_warmup(16, 16, WarmupEngineMode.DECODE) - test.test_cuda_graph_warmup(16, 16, WarmupEngineMode.NON_DECODE) - test.test_cuda_graph_warmup(20, 24, WarmupEngineMode.DECODE) - test.test_cuda_graph_warmup(20, 24, WarmupEngineMode.NON_DECODE) - test.test_cuda_graph_warmup(24, 24, WarmupEngineMode.DECODE) - test.test_cuda_graph_warmup(24, 24, WarmupEngineMode.NON_DECODE) - test.test_cuda_graph_warmup(28, 32, WarmupEngineMode.DECODE) - test.test_cuda_graph_warmup(28, 32, WarmupEngineMode.NON_DECODE) - test.test_cuda_graph_warmup(32, 32, WarmupEngineMode.DECODE) - test.test_cuda_graph_warmup(32, 32, WarmupEngineMode.NON_DECODE) - test.test_generate_function() - test.test_run_engine() - test.test_return_log_probs() - test.test_parallel_inference(1, 1, 1, False, False) - test.test_parallel_inference(1, 1, 1, False, True) - test.test_parallel_inference(1, 1, 1, True, False) - test.test_parallel_inference(1, 1, 1, True, True) - test.test_parallel_inference(1, 1, 2, False, False) - test.test_parallel_inference(1, 1, 2, False, True) - test.test_parallel_inference(1, 1, 2, True, False) - test.test_parallel_inference(1, 1, 2, True, True) - test.test_parallel_inference(1, 2, 1, False, False) - test.test_parallel_inference(1, 2, 1, False, True) - test.test_parallel_inference(1, 2, 1, True, False) - test.test_parallel_inference(1, 2, 1, True, True) - test.test_parallel_inference(1, 2, 2, False, False) - test.test_parallel_inference(1, 2, 2, False, True) - test.test_parallel_inference(1, 2, 2, True, False) - test.test_parallel_inference(1, 2, 2, True, True) - test.test_parallel_inference(2, 1, 1, False, False) - test.test_parallel_inference(2, 1, 1, False, True) - test.test_parallel_inference(2, 1, 1, True, False) - test.test_parallel_inference(2, 1, 1, True, True) - test.test_parallel_inference(2, 1, 2, False, False) - test.test_parallel_inference(2, 1, 2, False, True) - test.test_parallel_inference(2, 1, 2, True, False) - test.test_parallel_inference(2, 1, 2, True, True) - test.test_parallel_inference(2, 2, 1, False, False) - test.test_parallel_inference(2, 2, 1, False, True) - test.test_parallel_inference(2, 2, 1, True, False) - test.test_parallel_inference(2, 2, 1, True, True) - test.test_parallel_inference(2, 2, 2, False, False) - test.test_parallel_inference(2, 2, 2, False, True) - test.test_parallel_inference(2, 2, 2, True, False) - test.test_parallel_inference(2, 2, 2, True, True) - test.test_num_tokens_total() - test.test_events() - # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - test.teardown_method(None) - print("~~~") - print("success.") From e4c6554e05b6d271f2613ed81c5bba6f547aa507 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Fri, 31 Oct 2025 11:38:44 -0700 Subject: [PATCH 56/69] format. --- tests/unit_tests/inference/contexts/test_dynamic_context.py | 6 ++---- tests/unit_tests/inference/engines/test_dynamic_engine.py | 5 +---- 2 files changed, 3 insertions(+), 8 deletions(-) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 6800685ecd4..83b43ce6e63 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -433,9 +433,7 @@ def test_update_request(self): ) dynamic_context.request_to_kv_block_ids[3][ 1 - ] = ( - dynamic_context.block_allocator.total_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 @@ -522,7 +520,7 @@ def test_update_request(self): [-1, -1, -1, -1], [-1, -1, -1, -1], [-1, -1, -1, -1], - ] + ] ) ) diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 5b8a86fb4aa..df64a65be6d 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -427,10 +427,7 @@ def test_simple(self, num_cuda_graphs, cuda_graph_scope) -> None: def test_overflow_factor(self) -> None: """Test overflow factor arg.""" # Run test. - env = self._run_test( - context_buffer_overflow_factor=0.1, - context_max_tokens=None, - ) + env = self._run_test(context_buffer_overflow_factor=0.1, context_max_tokens=None) # Validate max_requests, max_tokens. assert env.engine.context.max_requests == 420 From 769c1dbeb59f71955e32b93d34a6ab96fb6e117b Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 4 Nov 2025 02:41:19 -0800 Subject: [PATCH 57/69] fixed test_dynamic_engine.py. --- .../inference/engines/test_dynamic_engine.py | 50 +++++++------------ 1 file changed, 18 insertions(+), 32 deletions(-) diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index df64a65be6d..e3a391007ba 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -420,34 +420,13 @@ def test_simple(self, 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) -> None: - """Test overflow factor arg.""" - # Run test. - env = self._run_test(context_buffer_overflow_factor=0.1, context_max_tokens=None) - - # Validate max_requests, max_tokens. - assert env.engine.context.max_requests == 420 - assert env.engine.context.max_tokens == 420 - - @pytest.mark.internal - @pytest.mark.skipif( - not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" - ) - def test_request_overflow(self) -> None: - """Test request overflow.""" - self._run_test(context_max_requests_override=4) - @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) 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=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]) @@ -515,19 +494,22 @@ 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 ( @@ -573,7 +555,11 @@ 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 From c957a1d120ff9c2242dfe29a0767d0d8b45e6064 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 4 Nov 2025 02:50:14 -0800 Subject: [PATCH 58/69] fixed test_simple_text_generation_controller.py. --- .../test_simple_text_generation_controller.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 01ef8f5749f..46cb67efe80 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,7 +116,7 @@ def setup_model( kv_channels=transformer_config.kv_channels, num_attention_heads=transformer_config.num_attention_heads, max_sequence_length=2048, - active_buffer_size_gb=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 From da886a0efea4859b0ecbfdeac6b86474f32b5bad Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 4 Nov 2025 03:02:55 -0800 Subject: [PATCH 59/69] fixed test_dynamic_context.py. --- .../unit_tests/inference/contexts/test_dynamic_context.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 12cd508879d..604ca4a6c3f 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -168,7 +168,7 @@ def test_token_overflow_error(self): num_attention_heads=8, max_sequence_length=512, num_cuda_graphs=None, - active_buffer_size_gb=0.1, + active_buffer_size_gb=0.015, block_size_tokens=128, max_tokens=20, # Setting a very low token limit ) @@ -888,9 +888,9 @@ 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, active_buffer_size_gb=active_buffer_size_gb, unified_memory_level=unified_memory_level, From 4593de3851d7402bf525d0ba3ac3bcf713226389 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 4 Nov 2025 03:26:04 -0800 Subject: [PATCH 60/69] format. --- tests/unit_tests/inference/engines/test_dynamic_engine.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index e3a391007ba..cb99dac5fa2 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -507,8 +507,7 @@ def test_cuda_graph_token_counts(self) -> None: # Build cuda graphs (inside dynamic engine). env = self._build_test_env( DynamicEngineTestConfig( - context_active_buffer_size_gb=0.01, - num_cuda_graphs=num_cuda_graphs, + 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 @@ -556,9 +555,7 @@ def test_cuda_graph_warmup( # Initialize context. env = self._build_test_env( DynamicEngineTestConfig( - context_active_buffer_size_gb=0.0041, - num_cuda_graphs=8, - num_tokens_to_generate=1, + context_active_buffer_size_gb=0.0041, num_cuda_graphs=8, num_tokens_to_generate=1 ) ) From 9b6125d6667abde1357d3267a3b0b2d8bf84dbae Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Fri, 7 Nov 2025 09:49:16 -0800 Subject: [PATCH 61/69] minor fix. --- megatron/core/inference/contexts/dynamic_context.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index ed72618f929..e70323befc1 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -471,7 +471,7 @@ def __init__( 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_total_requests, block_size_tokens=self.block_size_tokens, @@ -479,7 +479,7 @@ def __init__( ) 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_total_requests, block_size_tokens=self.block_size_tokens, From feef5d5941dcf06884538c63dc552bd4c8e8694e Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Fri, 7 Nov 2025 12:09:00 -0800 Subject: [PATCH 62/69] fixed dynamic context instantiation from static engine. --- megatron/core/inference/contexts/dynamic_context.py | 7 ++----- megatron/core/inference/engines/static_engine.py | 2 +- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index e70323befc1..294c8d48892 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -532,7 +532,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, ): """ @@ -540,7 +540,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 @@ -552,10 +551,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, ) diff --git a/megatron/core/inference/engines/static_engine.py b/megatron/core/inference/engines/static_engine.py index d084528b8f2..cde6d44cbd9 100644 --- a/megatron/core/inference/engines/static_engine.py +++ b/megatron/core/inference/engines/static_engine.py @@ -98,7 +98,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 From ee078e533d7e0e472d8fd48d50cc46f728a80012 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Fri, 7 Nov 2025 13:00:51 -0800 Subject: [PATCH 63/69] fixed wandb tests. --- .../inference/contexts/dynamic_context.py | 12 ++++---- .../inference/test_wandb_logging.py | 30 ++++++++----------- 2 files changed, 17 insertions(+), 25 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 294c8d48892..daf2467b95d 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -1581,11 +1581,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. @@ -1607,7 +1607,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), @@ -1617,10 +1616,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/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, ) From 7cd2eddb1fd5fad800548d03863e2b2a1a91bc2f Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 10 Nov 2025 03:06:45 -0800 Subject: [PATCH 64/69] integrated mamba updates. --- .../inference/contexts/dynamic_context.py | 40 ++++++++----------- 1 file changed, 17 insertions(+), 23 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index c48f1011d20..55936852693 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -195,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 @@ -360,20 +360,6 @@ def __init__( ) assert self.block_size_bytes > 0 - # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> - # # 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 - # 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 - # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - mamba_states_memory_per_request = 0 if self.is_hybrid_model: mamba_states_memory_per_request += math.prod(mamba_conv_states_shape) @@ -381,6 +367,20 @@ def __init__( mamba_states_memory_per_request *= self.num_mamba_layers mamba_states_memory_per_request *= dtype_size_bytes + # 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}), " @@ -439,12 +439,6 @@ def __init__( 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) From 4db13d488bc9f17fc218afa2bf818a0d36fb05f0 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 10 Nov 2025 03:11:36 -0800 Subject: [PATCH 65/69] re-added max_tokens arg. --- examples/inference/gpt/gpt_dynamic_inference.py | 1 + megatron/training/arguments.py | 3 +++ 2 files changed, 4 insertions(+) diff --git a/examples/inference/gpt/gpt_dynamic_inference.py b/examples/inference/gpt/gpt_dynamic_inference.py index bd6eda2d819..fc2d17c1ebd 100644 --- a/examples/inference/gpt/gpt_dynamic_inference.py +++ b/examples/inference/gpt/gpt_dynamic_inference.py @@ -161,6 +161,7 @@ def get_inference_context( ), block_size_tokens=args.inference_dynamic_batching_block_size, 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, diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index c35f3bf0be4..cb2691c9d49 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1444,6 +1444,9 @@ def _add_inference_args(parser): type=int, default=256, help='KV cache block size. ' 'It should be a multiple of 256') + group.add_argument('--inference-dynamic-batching-max-tokens', + type=int, default=None, + 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 ' From 872acba4706456a5d5c12bceb6617cfc8c668691 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 10 Nov 2025 04:30:22 -0800 Subject: [PATCH 66/69] fixed test_dynamic_context.py. --- .../inference/contexts/dynamic_context.py | 14 +- .../contexts/test_dynamic_context.py | 144 +++++++----------- 2 files changed, 63 insertions(+), 95 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 55936852693..bc2eef21077 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -114,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 @@ -232,9 +232,9 @@ class DynamicInferenceContext(BaseInferenceContext): block_size_tokens (int): Size of KV cache block size. tensor_model_parallel_size (Optional[int]): Tensor model parallel size. num_cuda_graphs (Optional[int]): Maximum number of cuda graphs to capture, - where the cuda graph batch sizes range from 1 to `max_requests` (as - computed below). Due to rounding, the actual number of cuda graphs may - not equal this argument. + 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. @@ -549,16 +549,16 @@ def __init__( # Optional state tensors for hybrid models if self.is_hybrid_model: - self.mamba_metadata = MambaMetadata(max_requests=self.max_requests) + self.mamba_metadata = MambaMetadata(max_requests=self.max_total_requests) with ctx_manager: 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(), ) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 1e1d978fba7..644fa5a6f30 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -43,10 +43,7 @@ def _get_dynamic_context( max_sequence_length, 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, @@ -66,7 +63,7 @@ def _get_dynamic_context( use_cuda_graphs_for_non_decode_steps=not is_hybrid_model, active_buffer_size_gb=active_buffer_size_gb, block_size_tokens=block_size_tokens, - max_tokens=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), @@ -91,24 +88,24 @@ def test_initialize_dynamic_context(self, is_hybrid_model: bool): max_sequence_length=512, active_buffer_size_gb=0.03, block_size_tokens=128, - max_tokens_override=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 @@ -126,9 +123,7 @@ def test_is_static_batching(self): max_sequence_length=512, 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() @@ -144,9 +139,7 @@ def test_is_memory_available(self, is_hybrid_model): max_sequence_length=512, 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.active_count = 10 @@ -170,9 +163,7 @@ def test_request_overflow(self, is_hybrid_model: bool): max_sequence_length=128, 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, ) @@ -202,7 +193,7 @@ def test_token_overflow_error(self, is_hybrid_model: bool): max_sequence_length=512, active_buffer_size_gb=0.1, block_size_tokens=128, - max_tokens=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, ) @@ -211,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 ), @@ -231,9 +222,7 @@ def test_reset(self, is_hybrid_model: bool): max_sequence_length=128, 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, ) @@ -256,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.active_count = 1023 dynamic_context.memory_buffer.fill_(1) dynamic_context.request_to_kv_block_ids.fill_(1) if is_hybrid_model: @@ -307,16 +295,14 @@ def test_allocate_and_release_memory_blocks(self, is_hybrid_model): max_sequence_length=512, 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 ( @@ -327,16 +313,16 @@ 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( @@ -358,9 +344,7 @@ def test_add_request(self, is_hybrid_model: bool): max_sequence_length=512, 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 @@ -382,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( @@ -433,9 +417,7 @@ def test_update_request(self, is_hybrid_model: bool): max_sequence_length=512, 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, ) @@ -479,9 +461,7 @@ def test_update_request(self, is_hybrid_model: bool): max_sequence_length=512, 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, ) @@ -594,16 +574,16 @@ 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], - [-1, -1, -1, -1], - [-1, -1, -1, -1], - [-1, -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] ] ) ) @@ -612,16 +592,16 @@ 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], - [-1, -1, -1, -1], - [-1, -1, -1, -1], - [-1, -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] ] ) ) @@ -640,9 +620,7 @@ def test_release_memory_blocks_for_finished_requests(self, is_hybrid_model): max_sequence_length=512, 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, ) @@ -714,9 +692,7 @@ def test_finished_requests_with_multiple_blocks(self, is_hybrid_model): max_sequence_length=512, 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, ) @@ -787,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: @@ -805,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], ) @@ -866,9 +836,7 @@ def test_calculate_and_store_log_probs(self): max_sequence_length=512, 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 From 3ee18fe8ee90fa475359a75718b77b0a06de60ad Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 10 Nov 2025 05:22:12 -0800 Subject: [PATCH 67/69] second usage of 'with ctx_manager:' fails when using uvm. --- .../inference/contexts/dynamic_context.py | 22 +++++++++++++++++++ .../inference/engines/test_dynamic_engine.py | 6 +++++ 2 files changed, 28 insertions(+) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index bc2eef21077..43da67fde09 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -472,6 +472,20 @@ def __init__( dtype=self.params_dtype, device=torch.cuda.current_device(), ) + # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> + with ctx_manager: + # self.mamba_conv_states = torch.zeros( + # (self.num_mamba_layers, self.max_total_requests) + mamba_conv_states_shape, + # dtype=self.params_dtype, + # device=torch.cuda.current_device(), + # ) + # myt = torch.zeros( + # (10,), + # dtype=self.params_dtype, + # device=torch.cuda.current_device(), + # ) + print("hello, world.") + # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< # Block ids. self.max_kv_block_count = math.ceil(self.max_sequence_length / self.block_size_tokens) @@ -548,6 +562,14 @@ def __init__( ) # Optional state tensors for hybrid models + # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> + with ctx_manager: + self.mamba_conv_states = torch.zeros( + (self.num_mamba_layers, self.max_total_requests) + mamba_conv_states_shape, + dtype=self.params_dtype, + device=torch.cuda.current_device(), + ) + # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< if self.is_hybrid_model: self.mamba_metadata = MambaMetadata(max_requests=self.max_total_requests) diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index e60e3f6ddb8..134232a64bc 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -1040,3 +1040,9 @@ def test_chunked_prefill(self, model_provider: str): context_block_size_tokens=256, context_max_tokens_override=300, ) +# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> +if __name__ == "__main__": + test = TestDynamicInferenceEngine() + test.test_simple("mamba", None, "full") + print("~~~\nsuccess.") +# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< From e5a3276144ca5a3bba520de8d25c629a2dad8ce5 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 10 Nov 2025 05:45:19 -0800 Subject: [PATCH 68/69] using single 'with ctx_manager:' block. --- .../inference/contexts/dynamic_context.py | 55 +++++++------------ .../inference/engines/test_dynamic_engine.py | 8 +-- 2 files changed, 22 insertions(+), 41 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 43da67fde09..22a4b1b1423 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -440,12 +440,9 @@ def __init__( self.token_to_local_position_within_kv_block = torch.empty_like(self.token_to_input_ids) # 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( ( @@ -472,20 +469,6 @@ def __init__( dtype=self.params_dtype, device=torch.cuda.current_device(), ) - # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> - with ctx_manager: - # self.mamba_conv_states = torch.zeros( - # (self.num_mamba_layers, self.max_total_requests) + mamba_conv_states_shape, - # dtype=self.params_dtype, - # device=torch.cuda.current_device(), - # ) - # myt = torch.zeros( - # (10,), - # dtype=self.params_dtype, - # device=torch.cuda.current_device(), - # ) - print("hello, world.") - # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< # Block ids. self.max_kv_block_count = math.ceil(self.max_sequence_length / self.block_size_tokens) @@ -562,18 +545,11 @@ def __init__( ) # Optional state tensors for hybrid models - # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> - with ctx_manager: - self.mamba_conv_states = torch.zeros( - (self.num_mamba_layers, self.max_total_requests) + mamba_conv_states_shape, - dtype=self.params_dtype, - device=torch.cuda.current_device(), - ) - # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - if self.is_hybrid_model: - self.mamba_metadata = MambaMetadata(max_requests=self.max_total_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_total_requests) + mamba_conv_states_shape, dtype=self.params_dtype, @@ -585,8 +561,19 @@ def __init__( device=torch.cuda.current_device(), ) - else: - self.mamba_metadata = None + else: + self.mamba_metadata = None + + # 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 diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 134232a64bc..a168bd7c4b8 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -1038,11 +1038,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("mamba", None, "full") - print("~~~\nsuccess.") -# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< From f4f60a91e2088a460d6ed5c07b1e3b504d2ee8fa Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 10 Nov 2025 05:51:12 -0800 Subject: [PATCH 69/69] format. --- .../contexts/test_dynamic_context.py | 40 +++++++++---------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index 644fa5a6f30..00c0e2a9dab 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -574,16 +574,16 @@ def test_update_request(self, is_hybrid_model: bool): dynamic_context.request_to_kv_block_ids[0:10].cpu() == torch.tensor( [ - [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] + [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], ] ) ) @@ -592,16 +592,16 @@ def test_update_request(self, is_hybrid_model: bool): dynamic_context.request_to_kv_block_ids[0:10].cpu() == torch.tensor( [ - [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] + [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], ] ) )