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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev100285
+Version: 0.16.0rc0.dev101543
 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.dev101543}/megatron/core/inference/contexts/dynamic_context.py

@@ -16,17 +16,13 @@ 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
 from megatron.core.transformer import TransformerConfig
 from megatron.core.utils import divide as core_divide
 
-from .attention_context.mha_metadata import GraphedMHAMetadata, NonGraphedMHAMetadata
 from .base_context import BaseInferenceContext
 from .dynamic_block_allocator import BlockAllocator
 
@@ -326,20 +322,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
@@ -455,26 +447,30 @@ class DynamicInferenceContext(BaseInferenceContext):
             num_cuda_graphs is not None
         )
 
-        # Attention metadata initialization (tensors are now handled by MHAMetadata classes)
+        # `*_cudagraph_only` tensors are for use with cuda graphs to maintain
+        # consistent input shapes, which is required to use cuda graphs.
+        # During these steps, the `*_cudagraph_only`
+        # tensors are used, otherwise their same-name but un-suffixed
+        # corresponding tensors are used.
 
-        self.graph_attn_metadata = {}
-        self.non_graph_attn_metadata = {}
-        self.active_attn_metadata = None
-
-        self.graph_attn_metadata["mha_metadata"] = GraphedMHAMetadata(
-            block_count_total=block_count_total,
-            max_kv_block_count=self.max_kv_block_count,
-            max_requests=self.max_requests,
-            block_size_tokens=self.block_size_tokens,
-            max_seqlen=self.max_sequence_length,
+        self.query_seq_lengths_cudagraph_only = torch.full(
+            (self.max_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.kv_seq_lengths_cudagraph_only = torch.full(
+            (self.max_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.non_graph_attn_metadata["mha_metadata"] = NonGraphedMHAMetadata(
-            block_count_total=block_count_total,
-            max_kv_block_count=self.max_kv_block_count,
-            max_requests=self.max_requests,
-            block_size_tokens=self.block_size_tokens,
-            max_seqlen=self.max_sequence_length,
+        self.request_to_kv_block_ids_cudagraph_only = torch.full(
+            (self.max_requests, self.max_kv_block_count),
+            0,
+            dtype=torch.int,
+            device=torch.cuda.current_device(),
         )
 
         # Guaranteed active requests.
@@ -624,18 +620,11 @@ class DynamicInferenceContext(BaseInferenceContext):
 
     def cu_query_lengths(self) -> Tuple[Tensor, int]:
         """Cumulative query sequence lengths."""
-        return (
-            self.active_attn_metadata["mha_metadata"].state_data["cu_query_seq_lengths"],
-            self.active_attn_metadata["mha_metadata"].state_data["max_seqlen_q"],
-        )
+        return self.cu_query_seq_lengths, self.max_seqlen_q
 
-    def cu_kv_lengths(self) -> Tuple[Tensor, Tensor, int]:
+    def cu_kv_lengths(self) -> Tensor:
         """Cumulative key/value sequence lengths."""
-        return (
-            self.active_attn_metadata["mha_metadata"].state_data["cu_kv_seq_lengths"],
-            self.active_attn_metadata["mha_metadata"].state_data["kv_seq_lengths"],
-            self.active_attn_metadata["mha_metadata"].state_data["max_seqlen_k"],
-        )
+        return (self.cu_kv_seq_lengths, self.kv_seq_lengths, self.max_seqlen_k)
 
     def get_active_sequence_lengths(self) -> Tensor:
         """Total sequence length (query + key) for active requests."""
@@ -713,16 +702,12 @@ class DynamicInferenceContext(BaseInferenceContext):
             to blocks within the block-level memory buffer.
         """
         if self.cache_mla_latent:
-            return (
-                self.memory_buffer[layer_number - 1],
-                None,
-                self.active_attn_metadata["mha_metadata"].state_data["block_table"],
-            )
+            return (self.memory_buffer[layer_number - 1], None, self.block_table)
         else:
             return (
                 self.memory_buffer[0, layer_number - 1],
                 self.memory_buffer[1, layer_number - 1],
-                self.active_attn_metadata["mha_metadata"].state_data["block_table"],
+                self.block_table,
             )
 
     def apply_fused_qk_rotary_emb(
@@ -832,12 +817,17 @@ class DynamicInferenceContext(BaseInferenceContext):
 
     def reset_attention_state(self) -> None:
         """Reset state used within attention, after each step."""
-        # Attention metadata reset is now handled by MHAMetadata.reset()
-        for attn_metadata in self.non_graph_attn_metadata.values():
-            attn_metadata.reset()
-        for attn_metadata in self.graph_attn_metadata.values():
-            attn_metadata.reset()
-        self.active_attn_metadata = None
+        self.max_seqlen_q = None
+        self.max_seqlen_k = None
+        self.cu_query_seq_lengths = None
+        self.cu_query_seq_lengths_cudagraph_only.fill_(0)
+        self.query_seq_lengths_cudagraph_only.fill_(0)
+        self.cu_kv_seq_lengths = None
+        self.cu_kv_seq_lengths_cudagraph_only.fill_(0)
+        self.kv_seq_lengths = None
+        self.kv_seq_lengths_cudagraph_only.fill_(0)
+        self.request_to_kv_block_ids_cudagraph_only.fill_(0)
+        self.block_table = None
 
     def using_cuda_graph_this_step(self) -> bool:
         """Returns True if cuda graphs are being used for this step."""
@@ -937,30 +927,89 @@ class DynamicInferenceContext(BaseInferenceContext):
             self.active_token_count : self.padded_active_token_count
         ] = 0
 
-        real_req_batch_size = (
-            self.total_request_count - self.paused_request_count
-        )  # how many requests are indeed active
-        self.active_attn_metadata = (
-            self.graph_attn_metadata
-            if self.using_cuda_graph_this_step()
-            else self.non_graph_attn_metadata
-        )
-
         # Update cu_query_seq_lengths, max_seqlen_q.
-        active_slice = slice(self.paused_request_count, self.total_request_count)
-        query_lengths_view = self.request_query_lengths[active_slice]
-        request_kv_length_offsets_view = self.request_kv_length_offsets[active_slice]
-        request_to_kv_block_ids_view = self.request_to_kv_block_ids[active_slice]
-        self.active_attn_metadata["mha_metadata"].update(
-            request_query_lengths=query_lengths_view,
-            request_kv_length_offsets=request_kv_length_offsets_view,
-            request_to_kv_block_ids=request_to_kv_block_ids_view,
-            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()
+        query_lengths = self.request_query_lengths[
+            self.paused_request_count : self.total_request_count
+        ]
+        if self.is_decode_only() or self.using_cuda_graph_this_step():
+            self.query_seq_lengths_cudagraph_only[
+                0 : self.total_request_count - self.paused_request_count
+            ] = query_lengths
+            if self.is_decode_only():
+                self.cu_query_seq_lengths = None  # ensure no accidental use
+                self.max_seqlen_q = 1
+            else:
+                self.cu_query_seq_lengths_cudagraph_only[
+                    1 : self.padded_active_request_count + 1
+                ] = torch.cumsum(
+                    self.query_seq_lengths_cudagraph_only[: self.padded_active_request_count], dim=0
+                )
+
+                # The following will be passed to the FA kernel.
+                self.cu_query_seq_lengths = self.cu_query_seq_lengths_cudagraph_only[
+                    : (self.padded_active_request_count + 1)
+                ]
+                self.max_seqlen_q = self.padded_active_token_count
+        else:
+            cu_query_lengths = torch.cumsum(query_lengths, dim=0)
+            self.cu_query_seq_lengths = torch.full(
+                (self.total_request_count - self.paused_request_count + 1,),
+                0,
+                dtype=torch.int32,
+                device=torch.cuda.current_device(),
+            )
+            self.cu_query_seq_lengths[1:] = cu_query_lengths
+            self.max_seqlen_q = query_lengths.max().item()
+
+        kv_seq_lengths = self.request_kv_length_offsets + self.request_query_lengths
+        self.kv_seq_lengths = kv_seq_lengths[self.paused_request_count : self.total_request_count]
+        if self.is_decode_only() or self.using_cuda_graph_this_step():
+            # Re-assign `kv_seq_lengths` to be a view of the first
+            # `active_cuda_graph_request_count` tokens of `kv_seq_lengths_decode_only`,
+            # such that `kv_seq_lengths` has a static memory address and is therefore
+            # cuda graph compatible. This allows `kv_seq_lengths` to transition between,
+            # cuda graph sizes, which makes multi-batch-size cuda graphs possible.
+            self.kv_seq_lengths_cudagraph_only[
+                0 : self.total_request_count - self.paused_request_count
+            ] = self.kv_seq_lengths
+            self.kv_seq_lengths = self.kv_seq_lengths_cudagraph_only[
+                : self.padded_active_request_count
+            ]
+            self.max_seqlen_k = self.max_sequence_length
+            if self.is_decode_only():
+                self.cu_kv_seq_lengths = None  # ensure no accidental use
+            else:
+                cu_kv_lengths = torch.cumsum(self.kv_seq_lengths, dim=0)
+                # The following will be passed to the FA kernel.
+                self.cu_kv_seq_lengths_cudagraph_only[1 : cu_kv_lengths.size(0) + 1] = cu_kv_lengths
+                self.cu_kv_seq_lengths = self.cu_kv_seq_lengths_cudagraph_only[
+                    : (self.padded_active_request_count + 1)
+                ]
+        else:
+            self.cu_kv_seq_lengths = torch.full(
+                (self.total_request_count - self.paused_request_count + 1,),
+                0,
+                dtype=torch.int32,
+                device=torch.cuda.current_device(),
+            )
+            self.cu_kv_seq_lengths[1:] = torch.cumsum(self.kv_seq_lengths, dim=0)
+            self.max_seqlen_k = self.kv_seq_lengths.max().item()
+
+        # Update KV block IDs, block table.
+        request_to_kv_block_ids = self.request_to_kv_block_ids[
+            self.paused_request_count : self.total_request_count
+        ]
+        if self.is_decode_only() or self.using_cuda_graph_this_step():
+            self.request_to_kv_block_ids_cudagraph_only[
+                0 : self.total_request_count - self.paused_request_count
+            ] = request_to_kv_block_ids
+            self.block_table = self.request_to_kv_block_ids_cudagraph_only[
+                : self.padded_active_request_count
+            ]
+        else:
+            self.block_table = self.request_to_kv_block_ids[
+                self.paused_request_count : self.total_request_count
+            ]
 
     def reset(self) -> None:
         """Reset entire context.

{megatron_core-0.16.0rc0.dev100285 → megatron_core-0.16.0rc0.dev101543}/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.dev101543}/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.dev101543/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.dev101543}/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.dev101543}/megatron/core/package_info.py

@@ -4,7 +4,7 @@
 MAJOR = 0
 MINOR = 16
 PATCH = 0
-PRE_RELEASE = 'rc0.dev100285'
+PRE_RELEASE = 'rc0.dev101543'
 
 # 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.dev101543}/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.dev101543
 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.dev101543}/megatron_core.egg-info/SOURCES.txt

@@ -169,8 +169,6 @@ megatron/core/inference/contexts/dynamic_block_allocator.py
 megatron/core/inference/contexts/dynamic_context.py
 megatron/core/inference/contexts/fused_kv_append_kernel.py
 megatron/core/inference/contexts/static_context.py
-megatron/core/inference/contexts/attention_context/metadata_base.py
-megatron/core/inference/contexts/attention_context/mha_metadata.py
 megatron/core/inference/engines/__init__.py
 megatron/core/inference/engines/abstract_engine.py
 megatron/core/inference/engines/dynamic_engine.py

megatron_core-0.16.0rc0.dev100285/megatron/core/inference/contexts/attention_context/metadata_base.py

@@ -1,72 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-
-
-class MetadataBase:
-    """
-    Base class for attention metadata.
-    High-performance attention kernels often require input metadata in specific
-    formats—such as cumulative query lengths, cumulative key/value lengths,
-    and similar structures. Moreover, when using CUDA Graphs, these metadata
-    buffers must be statically allocated. This class serves as a unified container
-    that manages all such metadata in one place.
-    """
-
-    def __init__(self):
-        """
-        Initialize the metadata.
-        """
-        self.state_data = {}
-
-    def update(self, *args, **kwargs):
-        """
-        Construct the metadata from request states.
-        """
-        pass
-
-    def reset(self):
-        """
-        Reset the metadata.
-        """
-        pass
-
-    def tensor_copy_and_pad(
-        self,
-        tensor_buf,
-        unpadded_tensor,
-        real_batch_size,
-        padded_batch_size,
-        is_cumulative_tensor=False,
-        pad_value=0,
-    ):
-        """
-        Copy the unpadded tensor to the tensor_buf,
-        pad the tensor_buf with zero or the last value of the tensor,
-        depending on whether the tensor is cumulative.
-        Args:
-            tensor_buf: The destination tensor, at least padded_batch_size long.
-            unpadded_tensor: The tensor to copy, at least real_batch_size long.
-            real_batch_size: The real batch size.
-            padded_batch_size: Padded boundary of the tensor.
-            is_cumulative_tensor: Whether the tensor is cumulative.
-                If True, we pad the tensor_buf with the last value of the unpadded_tensor.
-            pad_value: The value to pad the tensor_buf with when the tensor is not cumulative.
-        """
-        assert real_batch_size <= padded_batch_size
-        assert tensor_buf.shape[0] >= padded_batch_size
-        assert unpadded_tensor.shape[0] >= real_batch_size
-        if is_cumulative_tensor:
-            if real_batch_size == 0:
-                value = pad_value
-            else:
-                value = unpadded_tensor[real_batch_size - 1]
-        else:
-            value = pad_value
-        tensor_buf[0:real_batch_size] = unpadded_tensor[:real_batch_size]
-        tensor_buf[real_batch_size:padded_batch_size] = value
-        return tensor_buf
-
-    def __str__(self):
-        """
-        Return a string representation of the metadata.
-        """
-        return "\n".join([f"{key}: {value}" for key, value in self.state_data.items()])