megatron-core 0.16.0rc0.dev115842__tar.gz → 0.16.0rc0.dev117916__tar.gz

{megatron_core-0.16.0rc0.dev115842 → megatron_core-0.16.0rc0.dev117916}/PKG-INFO

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

@@ -23,7 +23,6 @@ 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
 
@@ -448,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.
@@ -617,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."""
@@ -706,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(
@@ -825,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."""
@@ -930,29 +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,
-        )
-        # 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.dev115842 → megatron_core-0.16.0rc0.dev117916}/megatron/core/package_info.py

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

{megatron_core-0.16.0rc0.dev115842 → megatron_core-0.16.0rc0.dev117916}/megatron_core.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev115842
+Version: 0.16.0rc0.dev117916
 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.dev115842 → megatron_core-0.16.0rc0.dev117916}/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.dev115842/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()])

megatron_core-0.16.0rc0.dev115842/megatron/core/inference/contexts/attention_context/mha_metadata.py

@@ -1,210 +0,0 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-
-from typing import Optional
-
-import torch
-
-from .metadata_base import MetadataBase
-
-
-class MHAMetadata(MetadataBase):
-    """
-    Metadata for MHA layer using flash-attention.
-    """
-
-    def __init__(
-        self, block_count_total, max_kv_block_count, max_requests, block_size_tokens, max_seqlen
-    ):
-        super().__init__()
-        device = torch.cuda.current_device()
-        self.device = device
-        self.max_blocks = block_count_total
-        self.max_kv_blocks = max_kv_block_count
-        self.max_bs = max_requests
-        self.max_seqlen = max_seqlen
-        self._query_lengths_buf = torch.zeros(self.max_bs, dtype=torch.int32, device=device)
-        self._cu_query_seq_lengths_buf = torch.zeros(
-            self.max_bs + 1, dtype=torch.int32, device=device
-        )
-        self._cu_kv_seq_lengths_buf = torch.zeros(self.max_bs + 1, dtype=torch.int32, device=device)
-        self._kv_seq_lengths_buf = torch.zeros(self.max_bs, dtype=torch.int32, device=device)
-        self._block_table_buf = torch.zeros(
-            (self.max_bs, self.max_kv_blocks), dtype=torch.int32, device=device
-        )
-        self._max_seqlen_q = 0
-        self._max_seqlen_k = 0
-        self.state_data = {}
-
-    def update(
-        self,
-        request_query_lengths: torch.Tensor,
-        request_kv_length_offsets: torch.Tensor,
-        request_to_kv_block_ids: torch.Tensor,
-        padded_active_token_count: int,
-        real_batch_size: int,
-        padded_active_request_count: Optional[int] = None,
-    ):
-        """
-        Args:
-            request_query_lengths: (>real_batch_size,)
-            request_kv_length_offsets: (>real_batch_size,)
-            request_to_kv_block_ids: (>real_batch_size, max_kv_blocks)
-            padded_active_token_count: int
-            real_batch_size: int
-            padded_active_request_count: Optional[int]
-        """
-        if padded_active_request_count is None:
-            padded_active_request_count = real_batch_size
-
-        assert real_batch_size <= padded_active_request_count <= self.max_bs
-        assert request_query_lengths.shape[0] == real_batch_size
-        assert request_kv_length_offsets.shape[0] == real_batch_size
-        assert request_to_kv_block_ids.shape[0] == real_batch_size
-
-        self.tensor_copy_and_pad(
-            self._query_lengths_buf,
-            request_query_lengths,
-            real_batch_size,
-            padded_active_request_count,
-        )
-        self._cu_query_seq_lengths_buf[0] = 0
-        self.tensor_copy_and_pad(
-            self._cu_query_seq_lengths_buf[1:],
-            torch.cumsum(request_query_lengths, dim=0),
-            real_batch_size,
-            padded_active_request_count,
-            is_cumulative_tensor=True,
-        )
-        self.tensor_copy_and_pad(
-            self._kv_seq_lengths_buf,
-            request_kv_length_offsets + request_query_lengths,
-            real_batch_size,
-            padded_active_request_count,
-        )
-        self.tensor_copy_and_pad(
-            self._block_table_buf,
-            request_to_kv_block_ids,
-            real_batch_size,
-            padded_active_request_count,
-            pad_value=torch.tensor(self.max_kv_blocks, dtype=torch.int32, device=self.device).fill_(
-                -1
-            ),
-        )
-        self._cu_kv_seq_lengths_buf[0] = 0
-        self.tensor_copy_and_pad(
-            self._cu_kv_seq_lengths_buf[1:],
-            torch.cumsum(self._kv_seq_lengths_buf, dim=0),
-            real_batch_size,
-            padded_active_request_count,
-            is_cumulative_tensor=True,
-        )
-        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
-        self._max_seqlen_k = self.max_seqlen
-
-        self.state_data = {
-            "query_lengths": self._query_lengths_buf[:padded_active_request_count],
-            "cu_query_seq_lengths": self._cu_query_seq_lengths_buf[
-                : padded_active_request_count + 1
-            ],
-            "cu_kv_seq_lengths": self._cu_kv_seq_lengths_buf[: padded_active_request_count + 1],
-            "kv_seq_lengths": self._kv_seq_lengths_buf[:padded_active_request_count],
-            "block_table": self._block_table_buf[0:padded_active_request_count, :],
-            "max_seqlen_q": self._max_seqlen_q,
-            "max_seqlen_k": self._max_seqlen_k,
-        }
-
-    def reset(self):
-        """
-        Reset the metadata for the next batch.
-        """
-        self._query_lengths_buf.fill_(0)
-        self._cu_query_seq_lengths_buf.fill_(0)
-        self._cu_kv_seq_lengths_buf.fill_(0)
-        self._kv_seq_lengths_buf.fill_(0)
-        self._block_table_buf.fill_(0)
-        self._max_seqlen_q = 0
-        self._max_seqlen_k = 0
-
-
-class GraphedMHAMetadata(MHAMetadata):
-    """
-    Metadata for MHA layer using flash-attention with CUDA graphs.
-    """
-
-    def __init__(
-        self, block_count_total, max_kv_block_count, max_requests, block_size_tokens, max_seqlen
-    ):
-        super().__init__(
-            block_count_total, max_kv_block_count, max_requests, block_size_tokens, max_seqlen
-        )
-
-    def update(
-        self,
-        request_query_lengths: torch.Tensor,
-        request_kv_length_offsets: torch.Tensor,
-        request_to_kv_block_ids: torch.Tensor,
-        padded_active_token_count: int,
-        real_batch_size: int,
-        padded_active_request_count: Optional[int] = None,
-    ):
-        """
-        Args:
-            request_query_lengths: (>real_batch_size,)
-            request_kv_length_offsets: (>real_batch_size,)
-            request_to_kv_block_ids: (>real_batch_size, max_kv_blocks)
-            padded_active_token_count: int
-            real_batch_size: int
-            padded_active_request_count: Optional[int]
-        """
-        super().update(
-            request_query_lengths,
-            request_kv_length_offsets,
-            request_to_kv_block_ids,
-            padded_active_token_count,
-            real_batch_size,
-            padded_active_request_count,
-        )
-
-    def reset(self):
-        super().reset()
-
-
-class NonGraphedMHAMetadata(MHAMetadata):
-    """
-    Metadata for MHA layer using flash-attention without CUDA graphs.
-    """
-
-    def update(
-        self,
-        request_query_lengths: torch.Tensor,
-        request_kv_length_offsets: torch.Tensor,
-        request_to_kv_block_ids: torch.Tensor,
-        padded_active_token_count: int,
-        real_batch_size: int,
-        padded_active_request_count: Optional[int] = None,
-    ):
-        """
-        Args:
-            request_query_lengths: (>real_batch_size,)
-            request_kv_length_offsets: (>real_batch_size,)
-            request_to_kv_block_ids: (>real_batch_size, max_kv_blocks)
-            padded_active_token_count: int
-            real_batch_size: int
-            padded_active_request_count: Optional[int]
-        """
-        super().update(
-            request_query_lengths,
-            request_kv_length_offsets,
-            request_to_kv_block_ids,
-            padded_active_token_count,
-            real_batch_size,
-            padded_active_request_count,
-        )
-        if len(self.state_data["query_lengths"]) > 0:
-            self.state_data["max_seqlen_q"] = torch.max(self.state_data["query_lengths"]).item()
-            self.state_data["max_seqlen_k"] = torch.max(self.state_data["kv_seq_lengths"]).item()
-        else:
-            self.state_data["max_seqlen_q"] = 1
-            self.state_data["max_seqlen_k"] = 1