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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev101543
+Version: 0.16.0rc0.dev101856
 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">
 
@@ -168,7 +175,7 @@ pip install --no-build-isolation .[mlm,dev]
 
 ```
 Megatron-LM/
-├── megatron/                    
+├── megatron/
 │   ├── core/                    # Megatron Core (kernels, parallelism, building blocks)
 │   │   ├── models/              # Transformer models
 │   │   ├── transformer/         # Transformer building blocks
@@ -197,7 +204,7 @@ Megatron-LM/
 
 - **Training state-of-the-art foundation models** at scale with cutting-edge performance on latest NVIDIA hardware
 - **Research teams** exploring new architectures and training techniques
-- **Learning distributed training** concepts and best practices  
+- **Learning distributed training** concepts and best practices
 - **Quick experimentation** with proven model configurations
 
 **What you get:**
@@ -206,7 +213,7 @@ Megatron-LM/
 - End-to-end examples from data prep to evaluation
 - Research-focused tools and utilities
 
-### Megatron Core: Composable Library  
+### Megatron Core: Composable Library
 
 **Composable library** with GPU-optimized building blocks for custom training frameworks.
 
@@ -239,7 +246,7 @@ Megatron-LM/
 - **[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** - Training library with bidirectional Hugging Face ↔ Megatron checkpoint conversion, flexible training loops, and production-ready recipes
 - **[NeMo RL](https://github.com/NVIDIA-NeMo/RL)** - Scalable toolkit for efficient reinforcement learning with RLHF, DPO, and other post-training methods
 - **[NeMo Framework](https://docs.nvidia.com/nemo-framework/user-guide/latest/overview.html)** - Enterprise framework with cloud-native support and end-to-end examples
-- **[TensorRT Model Optimizer (ModelOpt)](https://github.com/NVIDIA/TensorRT-Model-Optimizer)** - Model optimization toolkit for quantization, pruning, and distillation
+- **[TensorRT Model Optimizer (ModelOpt)](https://github.com/NVIDIA/TensorRT-Model-Optimizer)** - Model optimization toolkit for quantization, pruning, distillation, speculative decoding, and more. Checkout end-to-end examples in [examples/post_training/modelopt](./examples/post_training/modelopt/).
 
 **Compatible with:** [Hugging Face Accelerate](https://github.com/huggingface/accelerate), [Colossal-AI](https://github.com/hpcaitech/ColossalAI), [DeepSpeed](https://github.com/microsoft/DeepSpeed)
 
@@ -326,7 +333,7 @@ Our codebase efficiently trains models from 2B to 462B parameters across thousan
 **Benchmark Configuration:**
 
 - **Vocabulary size**: 131,072 tokens
-- **Sequence length**: 4096 tokens  
+- **Sequence length**: 4096 tokens
 - **Model scaling**: Varied hidden size, attention heads, and layers to achieve target parameter counts
 - **Communication optimizations**: Fine-grained overlapping with DP (`--overlap-grad-reduce`, `--overlap-param-gather`), TP (`--tp-comm-overlap`), and PP (enabled by default)
 

{megatron_core-0.16.0rc0.dev101543 → megatron_core-0.16.0rc0.dev101856}/README.md

@@ -99,7 +99,7 @@ pip install --no-build-isolation .[mlm,dev]
 
 ```
 Megatron-LM/
-├── megatron/                    
+├── megatron/
 │   ├── core/                    # Megatron Core (kernels, parallelism, building blocks)
 │   │   ├── models/              # Transformer models
 │   │   ├── transformer/         # Transformer building blocks
@@ -128,7 +128,7 @@ Megatron-LM/
 
 - **Training state-of-the-art foundation models** at scale with cutting-edge performance on latest NVIDIA hardware
 - **Research teams** exploring new architectures and training techniques
-- **Learning distributed training** concepts and best practices  
+- **Learning distributed training** concepts and best practices
 - **Quick experimentation** with proven model configurations
 
 **What you get:**
@@ -137,7 +137,7 @@ Megatron-LM/
 - End-to-end examples from data prep to evaluation
 - Research-focused tools and utilities
 
-### Megatron Core: Composable Library  
+### Megatron Core: Composable Library
 
 **Composable library** with GPU-optimized building blocks for custom training frameworks.
 
@@ -170,7 +170,7 @@ Megatron-LM/
 - **[Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge)** - Training library with bidirectional Hugging Face ↔ Megatron checkpoint conversion, flexible training loops, and production-ready recipes
 - **[NeMo RL](https://github.com/NVIDIA-NeMo/RL)** - Scalable toolkit for efficient reinforcement learning with RLHF, DPO, and other post-training methods
 - **[NeMo Framework](https://docs.nvidia.com/nemo-framework/user-guide/latest/overview.html)** - Enterprise framework with cloud-native support and end-to-end examples
-- **[TensorRT Model Optimizer (ModelOpt)](https://github.com/NVIDIA/TensorRT-Model-Optimizer)** - Model optimization toolkit for quantization, pruning, and distillation
+- **[TensorRT Model Optimizer (ModelOpt)](https://github.com/NVIDIA/TensorRT-Model-Optimizer)** - Model optimization toolkit for quantization, pruning, distillation, speculative decoding, and more. Checkout end-to-end examples in [examples/post_training/modelopt](./examples/post_training/modelopt/).
 
 **Compatible with:** [Hugging Face Accelerate](https://github.com/huggingface/accelerate), [Colossal-AI](https://github.com/hpcaitech/ColossalAI), [DeepSpeed](https://github.com/microsoft/DeepSpeed)
 
@@ -257,7 +257,7 @@ Our codebase efficiently trains models from 2B to 462B parameters across thousan
 **Benchmark Configuration:**
 
 - **Vocabulary size**: 131,072 tokens
-- **Sequence length**: 4096 tokens  
+- **Sequence length**: 4096 tokens
 - **Model scaling**: Varied hidden size, attention heads, and layers to achieve target parameter counts
 - **Communication optimizations**: Fine-grained overlapping with DP (`--overlap-grad-reduce`, `--overlap-param-gather`), TP (`--tp-comm-overlap`), and PP (enabled by default)
 

{megatron_core-0.16.0rc0.dev101543 → megatron_core-0.16.0rc0.dev101856}/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.dev101543 → megatron_core-0.16.0rc0.dev101856}/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py

@@ -898,9 +898,10 @@ class MegatronFSDP(torch.nn.Module):
 
         # Register pre state_dict hook to ensure that the module parameters are
         # distributed before saving the state_dict.
-        self._state_dict_pre_hook = self.module.register_state_dict_pre_hook(
-            lambda *args, **kwargs: self._replace_param_with_distributed_if_needed()
-        )
+        for name, module in self.named_modules():
+            module.register_state_dict_pre_hook(
+                lambda *args, **kwargs: self._replace_param_with_distributed_if_needed()
+            )
 
     @contextmanager
     def no_sync(self):

{megatron_core-0.16.0rc0.dev101543 → megatron_core-0.16.0rc0.dev101856}/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py

@@ -2782,6 +2782,9 @@ class GradReducePipeline:
             outer_fsdp_group_grad_reduce (bool, optional): Whether to reduce gradients
                 across outer-DP groups. Defaults to False.
         """
+        # Sort parameters by their bucket IDs to ensure a deterministic processing order.
+        # Performing reduce-scatter operations out of order can lead to hangs.
+        params = sorted(list(params), key=lambda x: self.buffer.param_to_param_group[x])
         for param in params:
             bucket_id = self.buffer.param_to_param_group[param]
             param_group = self.buffer.parameter_groups[bucket_id]

{megatron_core-0.16.0rc0.dev101543 → megatron_core-0.16.0rc0.dev101856}/megatron/core/enums.py

@@ -20,15 +20,17 @@ class ModelType(enum.Enum):
 
 
 class Fp8Recipe(str, enum.Enum):
-    """FP8 recipe names: delayed, tensorwise, mxfp8, blockwise."""
+    """FP8 recipe names: delayed, tensorwise, mxfp8, blockwise, custom."""
 
     delayed = "delayed"
     tensorwise = "tensorwise"
     mxfp8 = "mxfp8"
     blockwise = "blockwise"
+    custom = "custom"
 
 
 class Fp4Recipe(str, enum.Enum):
-    """FP4 recipe names: nvfp4."""
+    """FP4 recipe names: nvfp4, custom."""
 
     nvfp4 = "nvfp4"
+    custom = "custom"

{megatron_core-0.16.0rc0.dev101543 → megatron_core-0.16.0rc0.dev101856}/megatron/core/extensions/transformer_engine_spec_provider.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import warnings
 from typing import Optional, Tuple

{megatron_core-0.16.0rc0.dev101543 → megatron_core-0.16.0rc0.dev101856}/megatron/core/fp4_utils.py

@@ -7,6 +7,7 @@ from contextlib import nullcontext
 import torch
 
 from megatron.core.enums import Fp4Recipe
+from megatron.core.fp8_utils import _get_custom_recipe
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import is_te_min_version
 
@@ -47,6 +48,28 @@ def is_nvfp4tensor(tensor: torch.Tensor) -> bool:
     return HAVE_TE_FP4_TENSOR_CLASS and isinstance(tensor, FP4_TENSOR_CLASS)
 
 
+def get_fp4_align_size(fp4_recipe: Fp4Recipe) -> int:
+    """
+    Get the alignment size required for FP4 GEMM.
+    FP4 GEMM requires Blackwell and later architectures.
+
+    The value 32 is a hardware requirement: TMA (Tensor Memory Accelerator) requires
+    a 16-byte aligned address for efficient memory access. Since FP4 uses 4 bits per value,
+    16 bytes (128 bits) corresponds to 32 FP4 values. Therefore, the alignment size for FP4
+    is 32. With this alignment, NVFP4 GEMM can be performed efficiently.
+
+    Note that since we are also random hadamard transform for NVFP4 training, we want
+    fused group nvfp4 quantize plus hadamard transform. Hadamard transform will leverage
+    tensor core instructions for better performance, while group quantize kernels also
+    prefer a more aligned size in token dimension M. Therefore, we apply align size 64
+    here for better performance in MOE.
+
+    Paper link: https://arxiv.org/pdf/2509.25149
+    """
+    # pylint: disable=unused-argument
+    return 64
+
+
 def dequantize_fp4_tensor(fp4_tensor: torch.Tensor) -> torch.Tensor:
     """Dequantize a fp4 tensor to a higher precision tensor."""
     if is_te_min_version("2.7.0.dev0"):
@@ -70,9 +93,11 @@ if HAVE_TE:
                         Transformer Engine. Please make sure you are using TE version 
                         >= 2.7.0.dev0."""
                     )
+            elif config.fp4_recipe == Fp4Recipe.custom:
+                fp4_recipe = _get_custom_recipe(config.fp4_quantizer_factory)
             else:
                 raise ValueError(
-                    "NVFP4BlockScaling is the only supported FP4 recipe. "
+                    "NVFP4BlockScaling and custom are the only supported FP4 recipes. "
                     "Please make sure you are using a compatible TE version >= 2.7.0.dev0."
                 )
         else:

{megatron_core-0.16.0rc0.dev101543 → megatron_core-0.16.0rc0.dev101856}/megatron/core/fp8_utils.py

@@ -2,14 +2,21 @@
 
 """Utility functions related to FP8 that are used throughout Megatron core"""
 
+import importlib
 import weakref
 from contextlib import nullcontext
 from functools import wraps
-from typing import List, Optional
+from typing import List, Optional, Union
 
 import torch
 
-from megatron.core.enums import Fp8Recipe
+from megatron.core.enums import Fp4Recipe, Fp8Recipe
+from megatron.core.tensor_parallel import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+    gather_from_sequence_parallel_region,
+    reduce_scatter_to_sequence_parallel_region,
+)
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import get_te_version, is_te_min_version
 
@@ -104,6 +111,53 @@ def dequantize_fp8_tensor(fp8_tensor: torch.Tensor) -> torch.Tensor:
         return fp8_tensor.from_float8()
 
 
+def _resolve_callable_from_python_import_path(dotted_path: str):
+    """Resolve a Python import path like 'pkg.mod.func' to a callable.
+
+    Raises ValueError with clear message on failure.
+    """
+    if not isinstance(dotted_path, str) or not dotted_path:
+        raise ValueError(
+            "fp8_quantizer_factory must be a non-empty string with format 'pkg.mod.func'."
+        )
+
+    parts = dotted_path.rsplit(".", 1)
+    if len(parts) == 1:
+        raise ValueError(f"Invalid fp8_quantizer_factory '{dotted_path}'. Expected 'pkg.mod.func'.")
+    module_path, attr = parts[0], parts[1]
+
+    try:
+        mod = importlib.import_module(module_path)
+    except Exception as exc:
+        raise ValueError(
+            f"Failed to import module '{module_path}' for fp8_quantizer_factory: {exc}"
+        ) from exc
+
+    fn = getattr(mod, attr, None)
+    if fn is None:
+        raise ValueError(
+            f"Attribute '{attr}' not found in module '{module_path}' for fp8_quantizer_factory."
+        )
+    if not callable(fn):
+        raise ValueError(
+            f"Resolved attribute '{module_path}.{attr}' is not callable for fp8_quantizer_factory."
+        )
+    return fn
+
+
+def _get_custom_recipe(quantizer_factory_python_path: str) -> Union[Fp8Recipe, Fp4Recipe]:
+    quantizer_factory = _resolve_callable_from_python_import_path(quantizer_factory_python_path)
+    try:
+        custom_recipe = transformer_engine.common.recipe.CustomRecipe(qfactory=quantizer_factory)
+    except AttributeError:
+        raise ValueError(
+            """CustomRecipe recipe is not available in this version of 
+            Transformer Engine. Please make sure you are using TE version 
+            >= 2.9.0.dev0."""
+        )
+    return custom_recipe
+
+
 def get_fp8_align_size(fp8_recipe: Fp8Recipe) -> int:
     """Get the alignment size required for fp8 GEMM."""
     if fp8_recipe == Fp8Recipe.mxfp8:
@@ -112,6 +166,27 @@ def get_fp8_align_size(fp8_recipe: Fp8Recipe) -> int:
         return 16
 
 
+def is_column_parallel_linear(module):
+    """Returns whether the given module is a ColumnParallelLinear layer."""
+    if HAVE_TE and (
+        isinstance(module, TEColumnParallelLinear)
+        or isinstance(module, TELayerNormColumnParallelLinear)
+    ):
+        return True
+    elif isinstance(module, ColumnParallelLinear):
+        return True
+    return False
+
+
+def is_row_parallel_linear(module):
+    """Returns whether the given module is a RowParallelLinear layer."""
+    if HAVE_TE and isinstance(module, TERowParallelLinear):
+        return True
+    elif isinstance(module, RowParallelLinear):
+        return True
+    return False
+
+
 """
 The code below abstracts the functionalities needed for implementing "--fp8-param-gather" into
 several functions. It provides different implementations for each function based on different
@@ -466,6 +541,8 @@ if HAVE_TE:
                 fp8_recipe = transformer_engine.common.recipe.MXFP8BlockScaling(
                     fp8_format=fp8_format
                 )
+            elif config.fp8_recipe == Fp8Recipe.custom:
+                fp8_recipe = _get_custom_recipe(config.fp8_quantizer_factory)
             else:
                 raise ValueError(
                     "Float8CurrentScaling, MXFP8BlockScaling, Float8BlockwiseScaling and "
@@ -587,6 +664,18 @@ if HAVE_TE:
             if not FP8GlobalStateManager.is_fp8_enabled():
                 return original_forward(input_tensor, *args, **kwargs)
 
+            # With sequence parallelism we need to all-gather before padding
+            # and reduce-scatter after unpadding
+            if is_sequence_parallel := getattr(module, "sequence_parallel", False):
+                if is_column_parallel_linear(module):
+                    input_tensor = gather_from_sequence_parallel_region(
+                        input_tensor, group=module.tp_group
+                    )
+
+                # Disable sequence parallelism on the module because we are handling the
+                # all-gather and reduce-scatter externally
+                module.sequence_parallel = False
+
             seq_len, batch_size, hidden_size = input_tensor.shape
             # Reshape to (S, B*H) to pad sequence dimension
             input_2d = input_tensor.reshape(seq_len, -1)
@@ -612,6 +701,16 @@ if HAVE_TE:
             unpadded_output_2d = _unpad_func(output_2d, [seq_len])
             unpadded_output = unpadded_output_2d.reshape(seq_len, batch_size, output_hidden_size)
 
+            if is_sequence_parallel:
+                # Reduce-scatter after unpadding
+                if is_row_parallel_linear(module):
+                    unpadded_output = reduce_scatter_to_sequence_parallel_region(
+                        unpadded_output, group=module.tp_group
+                    )
+
+                # Reset sequence parallelism flag on the module
+                module.sequence_parallel = True
+
             if other_outputs:
                 return (unpadded_output,) + other_outputs
             else:

{megatron_core-0.16.0rc0.dev101543 → megatron_core-0.16.0rc0.dev101856}/megatron/core/inference/async_stream.py

@@ -9,6 +9,7 @@ import asyncio
 from typing import Any, AsyncGenerator, Callable, Optional, Type, Union
 
 from megatron.core.inference.inference_request import InferenceRequest
+from megatron.core.utils import get_asyncio_loop
 
 STOP_ITERATION = Exception()
 
@@ -20,12 +21,17 @@ class AsyncStream:
     Adopted from https://github.com/vllm-project/vllm/blob/eb881ed006ca458b052905e33f0d16dbb428063a/vllm/v1/engine/async_stream.py # pylint: disable=line-too-long
     """
 
-    def __init__(self, request_id: int, cancel: Callable[[str], None]) -> None:
+    def __init__(
+        self,
+        request_id: int,
+        cancel: Callable[[str], None],
+        loop: Optional[asyncio.AbstractEventLoop] = None,
+    ) -> None:
         self._request_id = request_id
         self._cancel = cancel
         self._queue: asyncio.Queue = asyncio.Queue()
         self._finished = False
-        self._loop = asyncio.get_running_loop()
+        self._loop = get_asyncio_loop(loop)
 
     def put(self, item: Union[InferenceRequest, Exception]) -> None:
         """Adds a new value to the stream"""

{megatron_core-0.16.0rc0.dev101543 → megatron_core-0.16.0rc0.dev101856}/megatron/core/inference/communication_utils.py

@@ -71,8 +71,7 @@ def broadcast_from_last_pipeline_stage(
             tensor.shape
         ), f"Expected tensor of shape {size} but got {list(tensor.shape)}"
         assert dtype == tensor.dtype, f"Expected tensor of type {dtype} but got {tensor.dtype}"
-        _is_cuda(tensor)
-        assert tensor.is_contiguous()
+        _is_cuda_contiguous(tensor)
     else:
         tensor = torch.empty(size, dtype=dtype, device=torch.cuda.current_device())
 

megatron_core-0.16.0rc0.dev101856/megatron/core/inference/contexts/attention_context/mamba_metadata.py

@@ -0,0 +1,106 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+import torch
+
+
+class MambaMetadata:
+    """Manages the metadata tensors required for Mamba layers during inference."""
+
+    def __init__(self, max_requests: int):
+        """
+        Initializes the Mamba slot allocator.
+
+        Args:
+            max_requests (int): The maximum number of concurrent requests.
+        """
+        self.max_requests = max_requests
+
+        # Metadata for mapping requests to slots in the static Mamba state buffer
+        self.request_to_mamba_state_idx = torch.full(
+            (self.max_requests,), -1, dtype=torch.int32, device=torch.cuda.current_device()
+        )
+
+        # Separate mapping used only for CUDA graph compatibility
+        self.request_to_mamba_state_idx_cudagraph_only = torch.full(
+            (self.max_requests,), -1, dtype=torch.int32, device=torch.cuda.current_device()
+        )
+
+        # Allocator for Mamba state slots
+        self.mamba_state_free_slots = torch.arange(
+            self.max_requests, dtype=torch.int32, device=torch.cuda.current_device()
+        )
+        self.mamba_state_free_slot_count = self.max_requests
+
+    def reset(self) -> None:
+        """
+        Resets all Mamba states and frees all allocated slots.
+        """
+        self.request_to_mamba_state_idx.fill_(-1)
+        self.request_to_mamba_state_idx_cudagraph_only.fill_(-1)
+
+        # Re-initialize the free slot pool
+        self.mamba_state_free_slots = torch.arange(
+            self.max_requests, dtype=torch.int32, device=torch.cuda.current_device()
+        )
+        self.mamba_state_free_slot_count = self.max_requests
+
+    def reset_cudagraph_mapping(self) -> None:
+        """
+        Resets only the CUDA graph mapping tensor.
+        """
+        self.request_to_mamba_state_idx_cudagraph_only.fill_(-1)
+
+    def update_cudagraph_mapping(
+        self, active_mamba_indices: torch.Tensor, num_active_requests: int
+    ) -> None:
+        """
+        Updates the dedicated CUDA graph mapping tensor with the indices
+        of currently active requests.
+
+        Args:
+            active_mamba_indices (Tensor): Tensor containing the Mamba slot indices
+                                           for active requests.
+            num_active_requests (int): The number of active requests.
+        """
+        self.request_to_mamba_state_idx_cudagraph_only[0:num_active_requests] = active_mamba_indices
+
+    def allocate_slot(self) -> int:
+        """
+        Allocates a new slot for a request in the Mamba state buffers.
+
+        Returns:
+            int: The index of the allocated slot.
+            Returns None if no slots are available.
+        """
+        if self.mamba_state_free_slot_count == 0:
+            return None
+
+        # Get a free slot
+        self.mamba_state_free_slot_count -= 1
+        mamba_idx = self.mamba_state_free_slots[self.mamba_state_free_slot_count]
+
+        return mamba_idx
+
+    def free_slots(self, request_indices: torch.Tensor) -> None:
+        """
+        Frees the Mamba state slots associated with the given request indices.
+
+        Args:
+            request_indices (Tensor): A 1D tensor of request indices to free.
+        """
+        # Get the Mamba state indices for finished requests
+        mamba_indices_to_free = self.request_to_mamba_state_idx[request_indices]
+
+        # Filter out any invalid indices (e.g., -1)
+        mamba_indices_to_free = mamba_indices_to_free[mamba_indices_to_free != -1]
+        num_to_free = len(mamba_indices_to_free)
+
+        if num_to_free > 0:
+            # Add the freed indices back to the free slot pool
+            start_idx = self.mamba_state_free_slot_count
+            end_idx = start_idx + num_to_free
+            self.mamba_state_free_slots[start_idx:end_idx] = mamba_indices_to_free
+            self.mamba_state_free_slot_count = end_idx
+
+        # Invalidate the Mamba state index for the finished requests
+        self.request_to_mamba_state_idx[request_indices] = -1

megatron_core-0.16.0rc0.dev101856/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()])