megatron-core 0.16.0rc0.dev130734__tar.gz → 0.16.0rc0.dev130741__tar.gz

{megatron_core-0.16.0rc0.dev130734 → megatron_core-0.16.0rc0.dev130741}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev130734
+Version: 0.16.0rc0.dev130741
 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>
@@ -41,7 +41,7 @@ Requires-Dist: transformers; extra == "mlm"
 Provides-Extra: dev
 Requires-Dist: nvidia-modelopt[torch]; sys_platform != "darwin" and extra == "dev"
 Requires-Dist: transformer-engine[pytorch]<2.10.0,>=2.9.0a0; extra == "dev"
-Requires-Dist: nvidia-resiliency-ext<0.5.0,>=0.4.0a0; extra == "dev"
+Requires-Dist: nvidia-resiliency-ext; extra == "dev"
 Requires-Dist: tqdm; extra == "dev"
 Requires-Dist: einops~=0.8; extra == "dev"
 Requires-Dist: tensorstore!=0.1.46,!=0.1.72,~=0.1; extra == "dev"
@@ -59,13 +59,20 @@ Requires-Dist: wget; extra == "dev"
 Requires-Dist: onnxscript; extra == "dev"
 Provides-Extra: lts
 Requires-Dist: tqdm; extra == "lts"
-Requires-Dist: einops; extra == "lts"
-Requires-Dist: tensorstore!=0.1.46,!=0.1.72; extra == "lts"
-Requires-Dist: nvtx; extra == "lts"
-Requires-Dist: transformers; extra == "lts"
-Requires-Dist: zarr; extra == "lts"
+Requires-Dist: einops~=0.8; extra == "lts"
+Requires-Dist: tensorstore!=0.1.46,!=0.1.72,~=0.1; extra == "lts"
+Requires-Dist: nvtx~=0.2; extra == "lts"
+Requires-Dist: multi-storage-client~=0.27; extra == "lts"
+Requires-Dist: opentelemetry-api~=1.33.1; extra == "lts"
 Requires-Dist: setuptools<80.0.0; extra == "lts"
+Requires-Dist: mamba-ssm~=2.2; extra == "lts"
+Requires-Dist: causal-conv1d~=1.5; extra == "lts"
+Requires-Dist: nv-grouped-gemm~=1.1; extra == "lts"
+Requires-Dist: megatron-energon[av_decode]~=6.0; extra == "lts"
+Requires-Dist: av<16.0.0; extra == "lts"
+Requires-Dist: flashinfer-python; extra == "lts"
 Requires-Dist: wget; extra == "lts"
+Requires-Dist: onnxscript; extra == "lts"
 
 <div align="center">
 

{megatron_core-0.16.0rc0.dev130734 → megatron_core-0.16.0rc0.dev130741}/megatron/core/datasets/megatron_tokenizer.py

@@ -1,11 +1,14 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 import json
+import logging
 from abc import ABC, abstractmethod
 from collections import OrderedDict
 from typing import Any
 
 import numpy
 
+logger = logging.getLogger(__name__)
+
 
 class MegatronLegacyTokenizer(ABC):
     """Abstract class for tokenizer
@@ -20,6 +23,12 @@ class MegatronLegacyTokenizer(ABC):
     """
 
     def __init__(self, *tokenizer_paths: str, **tokenizer_options: Any):
+        # Deprecation warning
+        logger.warning(
+            "You’re using the legacy tokenizer system, which is deprecated "
+            "and will be removed in a future release. Please migrate to the new tokenizer system "
+            "(`megatron.core.tokenizers.MegatronTokenizer`)."
+        )
         self.unique_identifiers = OrderedDict()
         self.unique_identifiers["class"] = type(self).__name__
         self.unique_identifiers["tokenizer_path"] = list(tokenizer_paths)

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

@@ -0,0 +1,72 @@
+# 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.dev130741/megatron/core/inference/contexts/attention_context/mha_metadata.py

@@ -0,0 +1,220 @@
+# 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,
+        decode_only: bool = False,
+    ):
+        """
+        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]
+            decode_only: bool
+        """
+        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,
+        )
+
+        if decode_only:
+            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 = {
+            "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,
+        decode_only: bool = False,
+    ):
+        """
+        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]
+            decode_only: bool
+        """
+        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,
+            decode_only,
+        )
+
+    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,
+        decode_only: bool = False,
+    ):
+        """
+        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]
+            decode_only: bool
+        """
+        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,
+            decode_only,
+        )
+        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