megatron-core 0.16.0rc0.dev100285__tar.gz → 0.16.0rc0.dev101328__tar.gz

{megatron_core-0.16.0rc0.dev100285 → megatron_core-0.16.0rc0.dev101328}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev100285
+Version: 0.16.0rc0.dev101328
 Summary: Megatron Core - a library for efficient and scalable training of transformer based models
 Author-email: NVIDIA <nemo-toolkit@nvidia.com>
 Maintainer-email: NVIDIA <nemo-toolkit@nvidia.com>

{megatron_core-0.16.0rc0.dev100285 → megatron_core-0.16.0rc0.dev101328}/megatron/core/inference/contexts/attention_context/mha_metadata.py

@@ -43,7 +43,6 @@ class MHAMetadata(MetadataBase):
         padded_active_token_count: int,
         real_batch_size: int,
         padded_active_request_count: Optional[int] = None,
-        decode_only: bool = False,
     ):
         """
         Args:
@@ -53,7 +52,6 @@ class MHAMetadata(MetadataBase):
             padded_active_token_count: int
             real_batch_size: int
             padded_active_request_count: Optional[int]
-            decode_only: bool
         """
         if padded_active_request_count is None:
             padded_active_request_count = real_batch_size
@@ -100,11 +98,9 @@ class MHAMetadata(MetadataBase):
             padded_active_request_count,
             is_cumulative_tensor=True,
         )
-
-        if decode_only:
+        self._max_seqlen_q = padded_active_token_count
+        if torch.all(self._query_lengths_buf[:padded_active_request_count] <= 1):
             self._max_seqlen_q = 1
-        else:
-            self._max_seqlen_q = max(2, padded_active_token_count)
         self._max_seqlen_k = self.max_seqlen
 
         self.state_data = {
@@ -152,7 +148,6 @@ class GraphedMHAMetadata(MHAMetadata):
         padded_active_token_count: int,
         real_batch_size: int,
         padded_active_request_count: Optional[int] = None,
-        decode_only: bool = False,
     ):
         """
         Args:
@@ -162,7 +157,6 @@ class GraphedMHAMetadata(MHAMetadata):
             padded_active_token_count: int
             real_batch_size: int
             padded_active_request_count: Optional[int]
-            decode_only: bool
         """
         super().update(
             request_query_lengths,
@@ -171,7 +165,6 @@ class GraphedMHAMetadata(MHAMetadata):
             padded_active_token_count,
             real_batch_size,
             padded_active_request_count,
-            decode_only,
         )
 
     def reset(self):
@@ -191,7 +184,6 @@ class NonGraphedMHAMetadata(MHAMetadata):
         padded_active_token_count: int,
         real_batch_size: int,
         padded_active_request_count: Optional[int] = None,
-        decode_only: bool = False,
     ):
         """
         Args:
@@ -201,7 +193,6 @@ class NonGraphedMHAMetadata(MHAMetadata):
             padded_active_token_count: int
             real_batch_size: int
             padded_active_request_count: Optional[int]
-            decode_only: bool
         """
         super().update(
             request_query_lengths,
@@ -210,7 +201,6 @@ class NonGraphedMHAMetadata(MHAMetadata):
             padded_active_token_count,
             real_batch_size,
             padded_active_request_count,
-            decode_only,
         )
         if len(self.state_data["query_lengths"]) > 0:
             self.state_data["max_seqlen_q"] = torch.max(self.state_data["query_lengths"]).item()

{megatron_core-0.16.0rc0.dev100285 → megatron_core-0.16.0rc0.dev101328}/megatron/core/inference/contexts/dynamic_context.py

@@ -16,10 +16,7 @@ from megatron.core.inference.inference_request import DynamicInferenceRequest
 from megatron.core.inference.model_inference_wrappers.inference_wrapper_config import (
     InferenceWrapperConfig,
 )
-from megatron.core.inference.unified_memory import (
-    UnifiedMemoryUnsupportedError,
-    create_unified_mempool,
-)
+from megatron.core.inference.unified_memory import create_unified_mempool, has_unified_memory
 from megatron.core.inference.utils import tensor_swap
 from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb
 from megatron.core.package_info import __version__ as mcore_version
@@ -326,20 +323,16 @@ class DynamicInferenceContext(BaseInferenceContext):
         self.params_dtype = params_dtype
         self.num_layers = num_layers
         self.max_sequence_length = max_sequence_length
-
-        # Unified memory.
         self.unified_memory_level = unified_memory_level
         if unified_memory_level > 0:
-            try:
-                self.unified_memory_mempool = create_unified_mempool()
-            except UnifiedMemoryUnsupportedError:
-                if torch.distributed.get_rank() == 0:
-                    warnings.warn(
-                        "Unified memory requested but not available; defaulting to GPU memory."
-                    )
+            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()
 
-        # Request and token counts.
         self.total_request_count = 0
         self.active_token_count = 0
         self.paused_request_count = 0
@@ -958,7 +951,6 @@ class DynamicInferenceContext(BaseInferenceContext):
             padded_active_token_count=self.padded_active_token_count,
             real_batch_size=real_req_batch_size,
             padded_active_request_count=self.padded_active_request_count,
-            decode_only=self.is_decode_only(),
         )
         # All attention metadata calculations are now handled by MHAMetadata.update()
 

{megatron_core-0.16.0rc0.dev100285 → megatron_core-0.16.0rc0.dev101328}/megatron/core/inference/engines/dynamic_engine.py

@@ -165,17 +165,6 @@ class DynamicInferenceEngine(AbstractEngine):
         context = self.context
         controller = self.controller
 
-        config = controller.inference_wrapped_model.inference_wrapper_config
-        moe_pad_experts = config.moe_pad_experts_for_cuda_graph_inference
-
-        if moe_pad_experts and context.non_decode_cuda_graphs:
-            context.non_decode_cuda_graphs = False
-            if torch.distributed.get_rank() == 0:
-                warnings.warn(
-                    "MoE models do not support non-decode cuda graphs. "
-                    "Forcing non_decode_cuda_graphs to False."
-                )
-
         time_start = time.time()
         mem_stats_start = torch.cuda.memory_stats()
 
@@ -185,18 +174,15 @@ class DynamicInferenceEngine(AbstractEngine):
             context.cuda_graph_token_counts,
         )
         for warmup_engine_mode in [WarmupEngineMode.DECODE, WarmupEngineMode.NON_DECODE]:
-            # Check whether to skip non-decode graphs.
+            # Iterate cuda graph dims.
             if (
                 warmup_engine_mode == WarmupEngineMode.NON_DECODE
                 and not context.non_decode_cuda_graphs
             ):
                 continue
-
             tbar = enumerate(context.cuda_graph_token_counts)
             if HAVE_TQDM:
                 tbar = tqdm(tbar, total=len(context.cuda_graph_token_counts))
-
-            # Iterate cuda graph dims.
             for tbar_idx, cuda_graph_token_count in tbar:
                 if (
                     cuda_graph_token_count == 1

{megatron_core-0.16.0rc0.dev100285 → megatron_core-0.16.0rc0.dev101328}/megatron/core/inference/text_generation_controllers/text_generation_controller.py

@@ -508,8 +508,7 @@ class TextGenerationController:
             inference_wrapper_config.moe_pad_experts_for_cuda_graph_inference
         )
         if moe_pad_experts_for_cuda_graph_inference:
-            assert warmup_engine_mode is not WarmupEngineMode.NON_DECODE
-            if context.is_decode_only():
+            if context.is_decode_only() or warmup_engine_mode is not None:
                 capacity_factor = model_config.num_moe_experts / model_config.moe_router_topk
                 set_decode_expert_padding(unwrapped_model, True, capacity_factor=capacity_factor)
             else:

megatron_core-0.16.0rc0.dev101328/megatron/core/inference/unified_memory.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+import os
+import warnings
+from pathlib import Path
+
+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
+
+_mempool_c_src = r"""
+#include <cuda_runtime_api.h>
+#include <cstddef>
+
+#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.
+  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);
+}
+"""
+
+# Avoid linting errors.
+has_unified_memory = False
+_alloc = None
+
+# 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")
+        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()
+        _alloc = CUDAPluggableAllocator(_so_path, "managed_malloc", "managed_free").allocator()
+        has_unified_memory = True
+    except (RuntimeError, ImportError, OSError):
+        warnings.warn("Failed to create unified memory mempool.")
+
+
+def create_unified_mempool():
+    """Create a unified memory mempool using CUDA managed memory."""
+    assert has_unified_memory
+    return MemPool(allocator=_alloc)

{megatron_core-0.16.0rc0.dev100285 → megatron_core-0.16.0rc0.dev101328}/megatron/core/models/gpt/gpt_model.py

@@ -375,13 +375,14 @@ class GPTModel(LanguageModule):
                 )
                 or self.config.flash_decode
             )
+            and rotary_pos_cos is not None
             and inference_context.is_static_batching()
         ):
             current_batch_size = input_ids.shape[0]
             sequence_len_offset = torch.tensor(
                 [inference_context.sequence_len_offset] * current_batch_size,
                 dtype=torch.int32,
-                device=torch.cuda.current_device(),
+                device=rotary_pos_cos.device,  # Co-locate this with the rotary tensors
             )
         else:
             sequence_len_offset = None

{megatron_core-0.16.0rc0.dev100285 → megatron_core-0.16.0rc0.dev101328}/megatron/core/package_info.py

@@ -4,7 +4,7 @@
 MAJOR = 0
 MINOR = 16
 PATCH = 0
-PRE_RELEASE = 'rc0.dev100285'
+PRE_RELEASE = 'rc0.dev101328'
 
 # Use the following formatting: (major, minor, patch, pre-release)
 VERSION = (MAJOR, MINOR, PATCH, PRE_RELEASE)

{megatron_core-0.16.0rc0.dev100285 → megatron_core-0.16.0rc0.dev101328}/megatron_core.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev100285
+Version: 0.16.0rc0.dev101328
 Summary: Megatron Core - a library for efficient and scalable training of transformer based models
 Author-email: NVIDIA <nemo-toolkit@nvidia.com>
 Maintainer-email: NVIDIA <nemo-toolkit@nvidia.com>

megatron_core-0.16.0rc0.dev100285/megatron/core/inference/unified_memory.py

@@ -1,127 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-
-import os
-import warnings
-from enum import Enum, auto
-from pathlib import Path
-
-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
-
-
-class CompilationState(Enum):
-    """Enum to distinguish between unified memory (UVM) compilation states."""
-
-    UNATTEMPTED = auto()  # Compilation has not been attempted.
-    FAILURE = auto()  # Compilation attempted, but failed.
-    SUCCESS = auto()  # Compilation attempted, and succeeded.
-
-
-# Compilation vars.
-_compilation_state = CompilationState.UNATTEMPTED
-_alloc = None  # must remain global until process exit.
-_mod = None  # must remain global until process exit.
-
-
-class UnifiedMemoryUnsupportedError(Exception):
-    """Unified memory is not supported on this system."""
-
-    pass
-
-
-def compile_allocator():
-    """Attempt to compile UVM allocator."""
-
-    global _compilation_state, _alloc, _mod
-
-    if _compilation_state != CompilationState.UNATTEMPTED:
-        return
-
-    _mempool_c_src = r"""
-    #include <cuda_runtime_api.h>
-    #include <cstddef>
-
-    #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.
-      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);
-    }
-    """
-
-    # 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")
-            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()
-            _alloc = CUDAPluggableAllocator(_so_path, "managed_malloc", "managed_free").allocator()
-            _compilation_state = CompilationState.SUCCESS
-        except (RuntimeError, ImportError, OSError):
-            warnings.warn("Failed to create unified memory mempool.")
-            _compilation_state = CompilationState.FAILURE
-
-
-def create_unified_mempool() -> "MemPool":
-    """Create a unified memory mempool using CUDA managed memory.
-
-    Returns:
-        (MemPool) Unified memory mempool.
-    """
-
-    # Attempt to compile allocator.
-    compile_allocator()
-
-    # Return mempool.
-    if _compilation_state != CompilationState.SUCCESS:
-        raise UnifiedMemoryUnsupportedError()
-    else:
-        return MemPool(allocator=_alloc)