token ids
self.sample_index = self._build_sample_index(
self.config.sequence_length - 3, 2 if self.config.classification_head else 1
diff --git a/megatron/core/datasets/blended_dataset.py b/megatron/core/datasets/blended_dataset.py
index 802a9770506..9b642ee1ff3 100644
--- a/megatron/core/datasets/blended_dataset.py
+++ b/megatron/core/datasets/blended_dataset.py
@@ -150,7 +150,10 @@ def _build_indices(self) -> Tuple[numpy.ndarray, numpy.ndarray]:
else:
cache_hit = False
- if not path_to_cache or (not cache_hit and torch.distributed.get_rank() == 0):
+ if not path_to_cache or (
+ not cache_hit
+ and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0)
+ ):
log_single_rank(
logger, logging.INFO, f"Build and save the {type(self).__name__} indices"
)
diff --git a/megatron/core/datasets/readme.md b/megatron/core/datasets/readme.md
index 452bf24e4a2..58721b7471b 100644
--- a/megatron/core/datasets/readme.md
+++ b/megatron/core/datasets/readme.md
@@ -192,12 +192,24 @@ To query the `BlendedDataset` for the _k_-th sample we do the following
To save time during initialization, each index is built/cached sequentially on one process rank and subsequently loaded in parallel on other process ranks. The cached indices are unique to a hash generated in the `BlendedDataset.__init__` function.
+## Offline cache preparation
+
+For GPT-style training, the dataset caches described above can be prepared ahead of time with `tools/prepare_cache.py` instead of waiting for rank 0 to build them during training startup.
+
+The script reuses the normal dataset construction path used by `pretrain_gpt.py` and `pretrain_mamba.py`, including `GPTDataset`, `BlendedDataset`, and `BlendedMegatronDatasetBuilder`. It accepts the usual dataset arguments, supports blends and per-split dataset definitions, and requires `--data-cache-path` so the generated cache can later be reused by training.
+
+This is especially useful for large blends or many file prefixes, where building the document, sample, and shuffle indices can take several minutes and leave all GPUs idle while rank 0 performs CPU-only work.
+
+If the later training job does not specify `--global-batch-size` (which is needed to determine the dataset size and splits), you should specify `--prepare-cache-world-size` to explicitly set the world size used during cache preparation.
+
+`tools/prepare_cache.py` does not support `--mock-data`, `--sft`, `--fim-data`, or `--step-batch-size-schedule`.
+
## Fast DataLoader initialization
-Especially for large-scale runs, DataLoader initialization can take several minutes, since it involves opening and memory-mapping multiple files and can significantly stress the filesystem. To speed up this process, we have developed the following three optimizations, controlled by configuration flags":
+Especially for large-scale runs, DataLoader initialization can take several minutes, since it involves opening and memory-mapping multiple files and can significantly stress the filesystem. To speed up this process, we have developed the following three optimizations, controlled by configuration flags:
- `--dataloader-fast-cache-load`: This option assumes that the dataset cache already exists in the specified `--data-cache-path`. When enabled, it speeds up the creation process by removing synchronization points and file check assertions.
- `--dataloader-defer-npy-index-mmap`: This option also assumes that the dataset cache already exists in the specified `--data-cache-path`. When enabled, it defers the memory mapping of the dataset indexes (.npy files) until their first access. We recommend using this configuration together with `--num-workers` > 0 so that the DataLoader prefetches the next batches of data, thereby hiding the cost of index memory mapping.
- - `--per-dataset-sequences-path`: With this configuration, we specify the JSON file generated by the `tools/build_sequences_per_dataset.py` script. This script generates a single file containing the required metadata from all the specified file prefixes. This configuration is especially useful when dealing with hundreds to thousands of file prefixes, since it requires only a single `open` operation instead of one per file prefix.
\ No newline at end of file
+ - `--per-dataset-sequences-path`: With this configuration, we specify the JSON file generated by the `tools/build_sequences_per_dataset.py` script. This script generates a single file containing the required metadata from all the specified file prefixes. This configuration is especially useful when dealing with hundreds to thousands of file prefixes, since it requires only a single `open` operation instead of one per file prefix.
diff --git a/megatron/core/dist_checkpointing/exchange_utils.py b/megatron/core/dist_checkpointing/exchange_utils.py
index 79f906b237a..7c7863532f6 100644
--- a/megatron/core/dist_checkpointing/exchange_utils.py
+++ b/megatron/core/dist_checkpointing/exchange_utils.py
@@ -11,7 +11,6 @@
import numpy as np
import torch
-from ..utils import get_pg_rank, get_pg_size, log_single_rank
from .core import CheckpointingException
from .dict_utils import nested_values
from .mapping import ShardedStateDict, ShardedTensor, is_main_replica
@@ -197,6 +196,8 @@ def determine_main_replica_uniform_distribution(
parallelization. Returns None if the process_group is trivial (1 rank)
"""
+ from ..utils import get_pg_size
+
if parallelization_group is None:
parallelization_group = torch.distributed.group.WORLD
group_size = get_pg_size(group=parallelization_group)
@@ -285,6 +286,8 @@ def exchange_loaded_tensors_gather_rounds(
needed by this rank to load a given state dict. Includes
previously loaded tensors (from `loaded_tensors` input)
"""
+ from ..utils import get_pg_rank, get_pg_size
+
if parallelization_group is None:
parallelization_group = torch.distributed.group.WORLD
main_rank_for_shard, _, shard_to_metadata, all_ranks_for_shard = shard_distribution
@@ -398,6 +401,8 @@ def exchange_loaded_tensors_gather_object(
previously loaded tensors (from `loaded_tensors` input)
"""
+ from ..utils import log_single_rank
+
all_loaded_tensors_list = [None] * torch.distributed.get_world_size(group=parallelization_group)
torch.distributed.all_gather_object(
all_loaded_tensors_list, loaded_tensors, group=parallelization_group
@@ -431,6 +436,8 @@ def exchange_loaded_objects_gather_object(
Dict[_ShardId, Any]: dictionary mapping shard ids to objects needed by this rank to
load a given state dict.
"""
+ from ..utils import log_single_rank
+
all_loaded_objects_list = [None] * torch.distributed.get_world_size()
torch.distributed.all_gather_object(all_loaded_objects_list, loaded_objects, group=None)
all_loaded_objects_list = cast(List[Dict[_ShardId, Any]], all_loaded_objects_list)
diff --git a/megatron/core/dist_checkpointing/serialization.py b/megatron/core/dist_checkpointing/serialization.py
index a0426ec5f80..1d42a03c0c5 100644
--- a/megatron/core/dist_checkpointing/serialization.py
+++ b/megatron/core/dist_checkpointing/serialization.py
@@ -21,7 +21,6 @@
from .core import CheckpointingConfig, save_config
from .dict_utils import merge
from .mapping import (
- CheckpointingException,
CommonStateDict,
ShardedObject,
ShardedStateDict,
@@ -30,7 +29,6 @@
)
from .state_dict_utils import load_preprocess, save_preprocess
from .strategies.async_utils import AsyncRequest
-from .strategies.base import AsyncSaveShardedStrategy
from .strategies.common import load_common, save_common
from .strategies.torch import TorchDistLoadShardedStrategy, TorchDistSaveShardedStrategy
from .utils import extract_sharded_base, force_all_tensors_to_non_fp8
@@ -38,8 +36,10 @@
StrictHandling,
determine_global_metadata,
parse_strict_flag,
+ save_integrity_manifest,
validate_integrity_and_strict_load,
verify_checkpoint,
+ verify_integrity_manifest,
)
logger = logging.getLogger(__name__)
@@ -58,6 +58,7 @@ def load(
common_strategy: None = None,
validate_access_integrity: bool = True,
strict: Union[str, StrictHandling] = StrictHandling.ASSUME_OK_UNEXPECTED,
+ verify_integrity: bool = False,
) -> Union[StateDict, Tuple[StateDict, Set[str], Set[str]]]:
"""Loading entrypoint.
@@ -91,6 +92,10 @@ def load(
incur any performance overhead. Other recommended values
are: `False` (StrictHandling.LOG_UNEXPECTED) which logs only unexpected keys
or `StrictHandling.RETURN_ALL` which returns all mismatch keys.
+ verify_integrity (bool, optional): if True, re-hashes every checkpoint file
+ and compares against the SHA-256 manifest. Raises `CheckpointingException` on any
+ mismatch. Requires that the checkpoint was previously saved with
+ `verify_integrity=True`.
Returns:
StateDict or Tuple[StateDict, Set[str], Set[str]]: in most cases only
@@ -99,6 +104,8 @@ def load(
assert common_strategy is None
verify_checkpoint(checkpoint_dir)
+ if verify_integrity:
+ verify_integrity_manifest(checkpoint_dir)
if sharded_strategy is None:
sharded_strategy = TorchDistLoadShardedStrategy()
@@ -141,7 +148,12 @@ def load(
ckpt_sharded_metadata,
)
- async_strategy = getattr(common_state_dict.get("args"), "async_strategy", "nvrx")
+ ckpt_args = common_state_dict.get("args")
+ async_strategy = (
+ getattr(ckpt_args, "async_strategy", "mcore")
+ if getattr(ckpt_args, "async_save", False)
+ else "mcore"
+ )
loaded_state_dict = sharded_strategy.load(sharded_state_dict, checkpoint_dir, async_strategy)
merge(common_state_dict, loaded_state_dict)
@@ -297,6 +309,7 @@ def save(
] = None,
content_metadata: Optional[dict] = None,
async_strategy: Optional[str] = "nvrx",
+ verify_integrity: bool = False,
) -> Optional[AsyncRequest]:
"""Saving entrypoint.
@@ -342,6 +355,11 @@ def save(
modify the original state dict
content_metadata (dict, optional): metadata to identify the checkpoint content.
Useful for framework specific versioning.
+ verify_integrity (bool, optional): if True, compute SHA-256 hashes for every
+ file in the checkpoint directory after all data has been written. This manifest can
+ later be verified on load with `load(..., verify_integrity=True)`.
+ Adds I/O overhead proportional to the total checkpoint size (one extra
+ read pass over all files on rank 0).
Returns:
AsyncRequest (optional): if `async_sharded_save` is True, returns
@@ -388,17 +406,22 @@ def metadata_finalize_fn():
)
torch.distributed.barrier()
+ def integrity_finalize_fn():
+ if torch.distributed.get_rank() == 0:
+ save_integrity_manifest(checkpoint_dir)
+ torch.distributed.barrier()
+
if not async_sharded_save:
sharded_strategy.save(sharded_state_dict, checkpoint_dir)
metadata_finalize_fn()
+ if verify_integrity:
+ integrity_finalize_fn()
return None
- if not isinstance(sharded_strategy, AsyncSaveShardedStrategy):
- raise CheckpointingException(
- f'Cannot apply async_save to non-async strategy {sharded_strategy}'
- )
async_request = sharded_strategy.async_save(sharded_state_dict, checkpoint_dir, async_strategy)
async_request.finalize_fns.append(metadata_finalize_fn)
+ if verify_integrity:
+ async_request.finalize_fns.append(integrity_finalize_fn)
return async_request
diff --git a/megatron/core/dist_checkpointing/strategies/base.py b/megatron/core/dist_checkpointing/strategies/base.py
index eb20e145ffb..c438382ed14 100644
--- a/megatron/core/dist_checkpointing/strategies/base.py
+++ b/megatron/core/dist_checkpointing/strategies/base.py
@@ -2,18 +2,22 @@
""" Strategies base interfaces. """
+import logging
from abc import ABC, abstractmethod
-from collections import defaultdict
from enum import Enum
from pathlib import Path
-from typing import Any, DefaultDict, Union
+from typing import Union
-from ..mapping import CheckpointingException, ShardedStateDict
-from .async_utils import AsyncCallsQueue, AsyncRequest
+from ..mapping import ShardedStateDict
+from .async_utils import AsyncRequest
+from .torch import TorchDistLoadShardedStrategy, TorchDistSaveShardedStrategy
+
+logger = logging.getLogger(__name__)
class StrategyAction(Enum):
- """Specifies save vs load action."""
+ """Specifies save vs load and sharded vs common action.
+ To be removed in future releases."""
LOAD_COMMON = 'load_common'
LOAD_SHARDED = 'load_sharded'
@@ -21,53 +25,34 @@ class StrategyAction(Enum):
SAVE_SHARDED = 'save_sharded'
-default_strategies: DefaultDict[str, dict[tuple, Any]] = defaultdict(dict)
-
-async_calls = AsyncCallsQueue()
-
-
def get_default_strategy(action: StrategyAction, backend: str, version: int):
"""Retrieves a default strategy for a given action, backend and version."""
- error_hint: str = ""
- try:
- error_hint = ' Please use PyTorch version >=2.1'
- from .torch import register_default_torch_strategies
-
- register_default_torch_strategies()
- except ImportError as e:
- raise CheckpointingException(
- f'Cannot import a default strategy for: {(action.value, backend, version)}. '
- f'Error: {e}. Hint: {error_hint}'
- ) from e
- try:
- return default_strategies[action.value][(backend, version)]
- except KeyError as e:
- raise CheckpointingException(
- f'Cannot find a default strategy for: {(action.value, backend, version)}'
- ) from e
-
-
-def register_default_strategy(
- action: StrategyAction,
- backend: str,
- version: int,
- strategy: Union['SaveStrategyBase', 'LoadStrategyBase'],
-):
- """Adds a given strategy to the registry of default strategies.
-
- Args:
- action (StrategyAction): specifies save/load and sharded
- backend (str): backend that the strategy becomes a default for
- version (int): version that the strategy becomes a default for
- strategy (SaveStrategyBase, LoadStrategyBase): strategy to register
- """
- default_strategies[action.value][(backend, version)] = strategy
+
+ logger.warning(
+ 'megatron.core.dist_checkpointing.strategies.base.get_default_strategy'
+ ' is deprecated and will be removed in the future releases. Please use'
+ ' TorchDistLoadShardedStrategy() and TorchDistSaveShardedStrategy()'
+ ' to get the default load and save sharded strategies.'
+ )
+ if backend != 'torch_dist':
+ logger.warning(f'{backend} is not supported. `torch_dist` backend will be used.')
+ if action == StrategyAction.LOAD_SHARDED:
+ return TorchDistLoadShardedStrategy()
+ else:
+ assert action == StrategyAction.SAVE_SHARDED, f'{action} is not supported'
+ return TorchDistSaveShardedStrategy()
class LoadStrategyBase(ABC):
"""Base class for a load strategy. Requires implementing checks for compatibility with a
given checkpoint version."""
+ def __init__(self):
+ logger.warning(
+ "LoadStrategyBase & LoadShardedStrategy are deprecated "
+ "and will be removed in future releases."
+ )
+
@abstractmethod
def check_backend_compatibility(self, loaded_backend):
"""Verifies if this strategy is compatible with `loaded_backend`."""
@@ -89,6 +74,10 @@ class SaveStrategyBase(ABC):
version of the saved format."""
def __init__(self, backend: str, version: int):
+ logger.warning(
+ "SaveStrategyBase & SaveShardedStrategy are deprecated "
+ "and will be removed in future releases."
+ )
self.backend = backend
self.version = version
@@ -102,7 +91,7 @@ def __str__(self):
class LoadShardedStrategy(LoadStrategyBase):
- """Load strategy for sharded tensors"""
+ """Base class for load strategies to be removed in future releases."""
@abstractmethod
def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Union[str, Path]):
@@ -145,7 +134,7 @@ def remove_sharded_tensors(self, checkpoint_dir: Union[str, Path], key_prefix: s
class SaveShardedStrategy(SaveStrategyBase):
- """Save strategy for sharded tensors"""
+ """Base class for save strategies to be removed in future releases."""
@abstractmethod
def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Union[str, Path]):
@@ -154,7 +143,7 @@ def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Union[str,
class AsyncSaveShardedStrategy(SaveShardedStrategy):
- """Save strategy suitable for async save."""
+ """Save strategy suitable for async save. To be removed in future releases."""
@abstractmethod
def async_save(
@@ -174,6 +163,9 @@ def async_save(
def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Union[str, Path]):
"""Each async strategy can be trivially used as a sync strategy."""
+ logger.warning(
+ "AsyncSaveShardedStrategy is deprecated and will be removed in future releases."
+ )
async_request = self.async_save(sharded_state_dict, checkpoint_dir)
async_request.execute_sync()
del async_request
diff --git a/megatron/core/dist_checkpointing/strategies/common.py b/megatron/core/dist_checkpointing/strategies/common.py
index 0ae800e46f8..3fdab41b4b0 100644
--- a/megatron/core/dist_checkpointing/strategies/common.py
+++ b/megatron/core/dist_checkpointing/strategies/common.py
@@ -42,9 +42,9 @@ def load_common(checkpoint_dir: str):
try:
if MultiStorageClientFeature.is_enabled():
msc = MultiStorageClientFeature.import_package()
- return msc.torch.load(load_path, map_location='cpu', weights_only=False)
+ return msc.torch.load(load_path, map_location='cpu')
else:
- return torch.load(load_path, map_location='cpu', weights_only=False)
+ return torch.load(load_path, map_location='cpu')
except FileNotFoundError as e:
err_msg = f'Common file {load_path} does not exist'
if MultiStorageClientFeature.is_enabled():
diff --git a/megatron/core/dist_checkpointing/strategies/fully_parallel.py b/megatron/core/dist_checkpointing/strategies/fully_parallel.py
index a85efdaa10a..db3c8ee6cae 100644
--- a/megatron/core/dist_checkpointing/strategies/fully_parallel.py
+++ b/megatron/core/dist_checkpointing/strategies/fully_parallel.py
@@ -23,10 +23,9 @@
exchange_loaded_objects_gather_object,
)
from megatron.core.dist_checkpointing.mapping import ShardedStateDict, StateDict, is_main_replica
-from megatron.core.dist_checkpointing.strategies.base import (
- AsyncSaveShardedStrategy,
- LoadShardedStrategy,
- SaveShardedStrategy,
+from megatron.core.dist_checkpointing.strategies.torch import (
+ TorchDistLoadShardedStrategy,
+ TorchDistSaveShardedStrategy,
)
from megatron.core.dist_checkpointing.utils import (
_sharded_object_id,
@@ -38,14 +37,13 @@
determine_global_metadata,
validate_sharding_integrity,
)
-from megatron.core.utils import get_pg_rank, get_pg_size
logger = logging.getLogger(__name__)
T = TypeVar('T', ShardedObject, ShardedTensor)
-class FullyParallelSaveStrategyWrapper(AsyncSaveShardedStrategy):
+class FullyParallelSaveStrategyWrapper:
"""Wraps arbitrary strategy and distributes the save during `save`.
The save distribution happens without any *data* communication.
@@ -60,7 +58,7 @@ class FullyParallelSaveStrategyWrapper(AsyncSaveShardedStrategy):
described in `distribute_shards_to_ranks`.
Args:
- strategy (SaveShardedStrategy): base strategy to wrap
+ strategy (TorchDistSaveShardedStrategy): base strategy to wrap
parallelization_group (ProcessGroup, optional): process group to use for save
distribution. Note that this doesn't have to match exactly the
data distribution, but should cover the replication pattern
@@ -72,16 +70,20 @@ class FullyParallelSaveStrategyWrapper(AsyncSaveShardedStrategy):
def __init__(
self,
- strategy: SaveShardedStrategy,
+ strategy: TorchDistSaveShardedStrategy,
parallelization_group: Optional[torch.distributed.ProcessGroup] = None,
do_cache_distribution: bool = False,
+ backend: str = "torch_dist",
+ version: int = 1,
):
- super().__init__(strategy.backend, strategy.version)
+ """ """
self.base_strategy = strategy
if parallelization_group is None:
parallelization_group = torch.distributed.group.WORLD
self.parallelization_group = parallelization_group
self.do_cache_distribution = do_cache_distribution
+ self.backend = backend
+ self.version = version
self.cached_distribution: Optional[ShardDistribution] = None
@@ -92,10 +94,6 @@ def async_save(
async_strategy: str = "nvrx",
):
""" """
- if not isinstance(self.base_strategy, AsyncSaveShardedStrategy):
- raise CheckpointingException(
- f'Cannot apply async_save to non-async base strategy {self.base_strategy}'
- )
self.apply_saving_parallelization(sharded_state_dict)
return self.base_strategy.async_save(sharded_state_dict, checkpoint_dir, async_strategy)
@@ -140,19 +138,14 @@ def apply_saving_parallelization(self, sharded_state_dict: ShardedStateDict) ->
end = time()
logger.debug(f"parallel save sharding, time: {end - start}")
- @property
- def can_handle_sharded_objects(self):
- """ """
- return self.base_strategy.can_handle_sharded_objects
-
-class FullyParallelLoadStrategyWrapper(LoadShardedStrategy):
+class FullyParallelLoadStrategyWrapper:
"""Wraps arbitrary load strategy and distributes the load during `load`.
See `load` method docs for details.
Args:
- strategy (LoadShardedStrategy): base strategy to wrap
+ strategy (TorchDistLoadShardedStrategy): base strategy to wrap
parallelization_group (ProcessGroup, optional): process group to use for load
distribution. Note that this doesn't have to match exactly the
data distribution, but should cover the replication pattern
@@ -174,12 +167,11 @@ class FullyParallelLoadStrategyWrapper(LoadShardedStrategy):
def __init__(
self,
- strategy: LoadShardedStrategy,
+ strategy: TorchDistLoadShardedStrategy,
parallelization_group: Optional[torch.distributed.ProcessGroup] = None,
do_cache_distribution: bool = False,
exchange_algo: str = 'broadcast',
):
- super().__init__()
self.base_strategy = strategy
if parallelization_group is None:
parallelization_group = (
@@ -197,7 +189,7 @@ def load(
self,
sharded_state_dict: ShardedStateDict,
checkpoint_dir: Path,
- async_strategy: str = "nvrx",
+ async_strategy: str = "mcore",
) -> StateDict:
"""Distributes the load and calls underlying strategy only for parts of the state dict.
@@ -227,6 +219,7 @@ def load(
a state dict that would be loaded with the underlying strategy
without this wrapper.
"""
+ from megatron.core.utils import get_pg_size
loaded_state_dict = {}
@@ -403,11 +396,6 @@ def apply_loading_parallelization(
return precomputed_distribution
- @property
- def can_handle_sharded_objects(self):
- """ """
- return self.base_strategy.can_handle_sharded_objects
-
def load_tensors_metadata(self, checkpoint_dir: Path):
""" """
return self.base_strategy.load_tensors_metadata(checkpoint_dir)
@@ -416,14 +404,6 @@ def load_sharded_metadata(self, checkpoint_dir: Path):
""" """
return self.base_strategy.load_sharded_metadata(checkpoint_dir)
- def check_backend_compatibility(self, loaded_version):
- """ """
- return self.base_strategy.check_backend_compatibility(loaded_version)
-
- def check_version_compatibility(self, loaded_version):
- """ """
- return self.base_strategy.check_version_compatibility(loaded_version)
-
def distribute_main_replicas_with_precomputed_distribution(
sharded_state_dict: ShardedStateDict,
@@ -455,6 +435,8 @@ def distribute_main_replicas_with_precomputed_distribution(
rank1: A: 1, B: 0, C: 1
rank2: A: 1, B: 1, C: 0
"""
+ from megatron.core.utils import get_pg_rank, get_pg_size
+
if parallelization_group is None:
parallelization_group = torch.distributed.group.WORLD
if get_pg_size(group=parallelization_group) <= 1:
diff --git a/megatron/core/dist_checkpointing/strategies/nvrx.py b/megatron/core/dist_checkpointing/strategies/nvrx.py
new file mode 100644
index 00000000000..e4b1c1a9a86
--- /dev/null
+++ b/megatron/core/dist_checkpointing/strategies/nvrx.py
@@ -0,0 +1,87 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Helpers for interacting with the experimental nvidia-resiliency-ext API."""
+
+from importlib import import_module
+from typing import Any, Callable, Dict
+
+try:
+ from packaging.version import Version as PkgVersion
+
+ HAVE_PACKAGING = True
+except ImportError:
+ HAVE_PACKAGING = False
+
+NVRX_MIN_VERSION = "0.6.0"
+
+
+def has_nvrx_async_support() -> bool:
+ """Checks whether the NVRx async checkpointing symbols Megatron uses are importable."""
+ try:
+ core = import_module("nvidia_resiliency_ext.checkpointing.async_ckpt.core")
+ cached_metadata_reader = import_module(
+ "nvidia_resiliency_ext.checkpointing.async_ckpt.cached_metadata_filesystem_reader"
+ )
+ filesystem_async = import_module(
+ "nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async"
+ )
+ state_dict_saver = import_module(
+ "nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver"
+ )
+ except (ImportError, ModuleNotFoundError):
+ return False
+
+ required_symbols = (
+ getattr(core, "AsyncCallsQueue", None),
+ getattr(core, "AsyncRequest", None),
+ getattr(cached_metadata_reader, "CachedMetadataFileSystemReader", None),
+ getattr(filesystem_async, "FileSystemWriterAsync", None),
+ getattr(filesystem_async, "get_write_results_queue", None),
+ getattr(state_dict_saver, "CheckpointMetadataCache", None),
+ getattr(state_dict_saver, "save_state_dict_async_finalize", None),
+ getattr(state_dict_saver, "save_state_dict_async_plan", None),
+ )
+ assert (
+ is_nvrx_min_version()
+ ), f"Minimum required nvidia-resiliency-ext package version is {NVRX_MIN_VERSION}."
+
+ return all(symbol is not None for symbol in required_symbols) and hasattr(
+ filesystem_async, "_results_queue"
+ )
+
+
+def make_nvrx_async_request(
+ async_request_cls: type,
+ async_fn: Callable[..., Any],
+ async_fn_args: Any,
+ finalize_fns: list[Callable[..., Any]],
+ async_fn_kwargs: Dict[str, Any] | None = None,
+ preload_fn: Callable[..., Any] | None = None,
+):
+ """Builds an AsyncRequest using the expected NVRx API."""
+ return async_request_cls(
+ async_fn,
+ async_fn_args,
+ finalize_fns,
+ async_fn_kwargs=async_fn_kwargs or {},
+ preload_fn=preload_fn,
+ )
+
+
+def is_nvrx_min_version(version: str = NVRX_MIN_VERSION) -> bool:
+ """Check if minimum version of `NVRx` is installed."""
+ if not HAVE_PACKAGING:
+ raise ImportError(
+ "packaging is not installed. Please install it with `pip install packaging`."
+ )
+
+ try:
+ import nvidia_resiliency_ext as nvrx
+
+ HAVE_NVRX = True
+ except (ImportError, ModuleNotFoundError):
+ HAVE_NVRX = False
+
+ nvrx_version = str(nvrx.__version__) if HAVE_NVRX else "0.0.0"
+
+ return PkgVersion(nvrx_version) >= PkgVersion(version)
diff --git a/megatron/core/dist_checkpointing/strategies/torch.py b/megatron/core/dist_checkpointing/strategies/torch.py
index 000aa4d9265..31782acb851 100644
--- a/megatron/core/dist_checkpointing/strategies/torch.py
+++ b/megatron/core/dist_checkpointing/strategies/torch.py
@@ -1,17 +1,17 @@
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
""" Strategies using PyTorch distributed.checkpoint as an underlying format. """
+import inspect
import io
import os
import pickle
import warnings
-from abc import ABC
from collections import defaultdict
from contextlib import contextmanager
from itertools import product
from logging import getLogger
from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union, cast
import torch
from packaging.version import Version as PkgVersion
@@ -49,22 +49,19 @@
is_main_replica,
)
from .async_utils import AsyncRequest
-from .base import (
- AsyncSaveShardedStrategy,
- LoadShardedStrategy,
- StrategyAction,
- register_default_strategy,
-)
from .checkpointable import CheckpointableShardedTensor, LocalShardsContainer
+from .nvrx import has_nvrx_async_support, make_nvrx_async_request
-try:
+if TYPE_CHECKING:
from nvidia_resiliency_ext.checkpointing.async_ckpt.core import AsyncRequest as NVRxAsyncRequest
from nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver import (
CheckpointMetadataCache,
)
-except (ImportError, ModuleNotFoundError):
- CheckpointMetadataCache = ABC
- NVRxAsyncRequest = ABC
+else:
+ CheckpointMetadataCache = Any
+ NVRxAsyncRequest = Any
+
+HAVE_NVRX = has_nvrx_async_support()
try:
if not torch.cuda.is_available():
@@ -103,17 +100,8 @@ class MCoreSavePlan:
pass
-def register_default_torch_strategies():
- """Register default strategies related to PyT Distributed backend."""
- register_default_strategy(
- StrategyAction.LOAD_SHARDED, 'torch_dist', 1, TorchDistLoadShardedStrategy()
- )
- register_default_strategy(
- StrategyAction.SAVE_SHARDED, 'torch_dist', 1, TorchDistSaveShardedStrategy()
- )
-
-
logger = getLogger(__name__)
+_logged_mcore_async_deprecation = False
def flatten_state_dict(
@@ -596,7 +584,7 @@ def commit_tensor(self, read_item: ReadItem, tensor: torch.Tensor) -> None:
return super().commit_tensor(read_item, tensor)
-class TorchDistSaveShardedStrategy(AsyncSaveShardedStrategy):
+class TorchDistSaveShardedStrategy:
"""Async save strategy for the PyT Distributed format.
The idea is to translate MCore ShardedTensors into PyT ShardedTensors
@@ -612,6 +600,7 @@ def __init__(
thread_count: int = 1,
cached_metadata: bool = False,
separation_hint: Optional[str] = None,
+ cpu_shm_mode: bool = False,
):
"""Adds parameters specific to PyT Distributed format
Args:
@@ -626,8 +615,13 @@ def __init__(
gathering local metadata every checkpointing invocation
separation_hint(str, optional): If provided, all tensors whose keys have this
prefix will be saved to a separate file.
+ cpu_shm_mode (bool, optional): Copy GPU tensors to CPU shared-memory in the
+ training process before handing off to the async worker. Avoids CUDA IPC /
+ NVLink fabric handles in the worker subprocess. Only applies with nvrx async
+ strategy.
"""
- super().__init__(backend, version)
+ self.backend = backend
+ self.version = version
self.keep_only_main_replica = keep_only_main_replica
self.thread_count = thread_count
@@ -651,9 +645,16 @@ def __init__(
self.cached_global_metadata: Optional[Metadata] = None
self.separation_hint = separation_hint
+ self.cpu_shm_mode = cpu_shm_mode
self.validated_loaded_metadata_reuse = False
+ def save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path):
+ """Sync save always uses the built-in implementation."""
+ async_request = self.async_save(sharded_state_dict, checkpoint_dir, async_strategy="mcore")
+ async_request.execute_sync()
+ del async_request
+
def async_save(
self,
sharded_state_dict: ShardedStateDict,
@@ -668,11 +669,14 @@ def async_save(
Returns: None
"""
+ global _logged_mcore_async_deprecation
if async_strategy == "mcore":
- logger.warning(
- "MCore's async save is deprecated and will be removed in the future releases. "
- "Please, use NVRx async solution by setting `async_strategy` to `nvrx`."
- )
+ if not _logged_mcore_async_deprecation:
+ logger.warning(
+ "MCore's async save is deprecated and will be removed in the future releases. "
+ "Please, use NVRx async solution by setting `async_strategy` to `nvrx`."
+ )
+ _logged_mcore_async_deprecation = True
# Translate the state dict
(sharded_state_dict, flat_mapping, rename_mapping) = (
@@ -698,10 +702,24 @@ def async_save(
if async_strategy == "nvrx":
if self._metadata_cache is None:
self._metadata_cache = checkpointable_metadata_cache()
- if self.cached_global_metadata is not None:
+ if self.cached_global_metadata is not None and hasattr(
+ self._metadata_cache, "set_cached_global_metadata"
+ ):
self._metadata_cache.set_cached_global_metadata(self.cached_global_metadata)
# Define additional arguments
async_writer_kwargs["use_cached_data_structure"] = self.use_cached_ckpt_structure
+ if self.cpu_shm_mode:
+ if (
+ "use_cpu_shm_for_gpu_tensors"
+ in inspect.signature(async_writer.__init__).parameters
+ ):
+ async_writer_kwargs["use_cpu_shm_for_gpu_tensors"] = True
+ else:
+ raise AssertionError(
+ "Installed nvidia-resiliency-ext does not support "
+ "use_cpu_shm_for_gpu_tensors. Update nvidia-resiliency-ext "
+ "to enable cpu_shm_mode."
+ )
state_dict_saver_kwargs["enable_cache"] = self.use_cached_ckpt_structure
state_dict_saver_kwargs["metadata_cache"] = self._metadata_cache
else:
@@ -803,17 +821,24 @@ def _get_save_and_finalize_callbacks(
def finalize_fn():
save_state_dict_async_finalize(*save_state_dict_ret)
- return async_request(save_fn, save_args, [finalize_fn], preload_fn=preload_fn)
-
- def can_handle_sharded_objects(self):
- return True
+ return make_nvrx_async_request(
+ async_request, save_fn, save_args, [finalize_fn], preload_fn=preload_fn
+ )
def _get_filesystem_reader(
- checkpoint_dir: Union[str, Path], cache_metadata: bool = False, async_strategy: str = "nvrx"
+ checkpoint_dir: Union[str, Path], cache_metadata: bool = False, async_strategy: str = "mcore"
) -> FileSystemReader:
if MultiStorageClientFeature.is_enabled():
msc = MultiStorageClientFeature.import_package()
+ if cache_metadata:
+ warnings.warn(
+ "MSC is enabled: returning msc.torch.MultiStorageFileSystemReader instead of "
+ "CachedMetadataFileSystemReader. The cache_metadata=True request "
+ "(e.g. ckpt_assume_constant_structure=True) will be ignored and metadata "
+ "will be re-read on every load. Pass --enable-msc only when this is intended.",
+ stacklevel=2,
+ )
return msc.torch.MultiStorageFileSystemReader(checkpoint_dir, thread_count=2)
if cache_metadata:
@@ -823,19 +848,18 @@ def _get_filesystem_reader(
return FileSystemReader(checkpoint_dir)
-class TorchDistLoadShardedStrategy(LoadShardedStrategy):
+class TorchDistLoadShardedStrategy:
"""Basic load strategy for the PyT Distributed format."""
def __init__(self, cache_metadata: bool = False):
self.cached_global_metadata: Optional[Metadata] = None
self.cache_metadata = cache_metadata
- super().__init__()
def load(
self,
sharded_state_dict: ShardedStateDict,
checkpoint_dir: Path,
- async_strategy: str = "nvrx",
+ async_strategy: str = "mcore",
) -> StateDict:
"""Translates MCore ShardedTensors to PyT ShardedTensors & loads from PyT Distributed fmt.
@@ -867,7 +891,7 @@ def load(
fsr = _get_filesystem_reader(
checkpoint_dir, cache_metadata=self.cache_metadata, async_strategy=async_strategy
)
- checkpoint.load_state_dict(
+ checkpoint.load(
pyt_state_dict,
fsr,
planner=MCoreLoadPlanner(
@@ -876,6 +900,7 @@ def load(
flatten_state_dict=False,
flatten_sharded_tensors=False,
),
+ no_dist=True,
)
if self.cache_metadata:
@@ -1012,15 +1037,6 @@ def remove_sharded_tensors(self, checkpoint_dir: str, key_prefix: str):
else:
fs_writer.fs.rm_file(old_path)
- def can_handle_sharded_objects(self):
- return True
-
- def check_backend_compatibility(self, loaded_version):
- pass # TODO
-
- def check_version_compatibility(self, loaded_version):
- pass # TODO
-
def get_async_strategy(async_strategy: str = "nvrx", module: str = None) -> tuple:
"""Returns async strategy and related async imported modules"""
@@ -1059,9 +1075,8 @@ def get_async_strategy(async_strategy: str = "nvrx", module: str = None) -> tupl
async_strategy = "nvrx"
except (ImportError, ModuleNotFoundError):
raise ModuleNotFoundError(
- "nvidia-resiliency-ext package is not installed. "
- "Please, install nvidia-resiliency-ext package or set `async_strategy` to `mcore` "
- "to enable async save strategy."
+ "A compatible `nvidia-resiliency-ext` installation is required for "
+ '`async_strategy="nvrx"`. Please install it or set `async_strategy` to `mcore`.'
)
elif async_strategy == "mcore":
# do mcore async imports
diff --git a/megatron/core/dist_checkpointing/validation.py b/megatron/core/dist_checkpointing/validation.py
index 89ecba1a968..b0cbae618a7 100644
--- a/megatron/core/dist_checkpointing/validation.py
+++ b/megatron/core/dist_checkpointing/validation.py
@@ -1,10 +1,13 @@
# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+import hashlib
+import json
import logging
+import os
from collections import Counter, defaultdict
from enum import Enum
from pathlib import Path
-from typing import TYPE_CHECKING, List, Optional, Set, Tuple, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Union
import numpy as np
import torch
@@ -22,6 +25,7 @@
ShardedStateDict,
is_main_replica,
)
+from megatron.core.msc_utils import MultiStorageClientFeature
if TYPE_CHECKING:
from megatron.core.dist_checkpointing.serialization import CkptShardedMetadata
@@ -34,6 +38,10 @@
# list of lists of global saved/loaded ShardedBase objects (each element corresponds to global rank)
_GlobalMetadata = List[_LocalMetadata]
+INTEGRITY_FNAME = 'integrity.json'
+_HASH_ALGORITHM = 'sha256'
+_READ_CHUNK_SIZE = 1 << 20 # 1 MiB
+
class StrictHandling(Enum):
"""Determines handling of load mismatch (non-empty "unexpected" or "missing" keys).
@@ -199,7 +207,12 @@ def verify_checkpoint(checkpoint_dir: str):
Args:
checkpoint_dir (str): checkpoint directory
"""
- if not Path(checkpoint_dir).exists():
+ if MultiStorageClientFeature.is_enabled():
+ msc = MultiStorageClientFeature.import_package()
+ isdir = msc.os.path.isdir(str(checkpoint_dir), strict=False)
+ else:
+ isdir = os.path.isdir(checkpoint_dir)
+ if not isdir:
raise CheckpointingException(f'Checkpoint directory {checkpoint_dir} does not exist')
if not check_is_distributed_checkpoint(checkpoint_dir):
@@ -483,3 +496,148 @@ def determine_global_metadata(
global_metadata = [None] * torch.distributed.get_world_size()
torch.distributed.all_gather_object(global_metadata, local_metadata)
return local_metadata, global_metadata # type: ignore[return-value]
+
+
+def _compute_file_hash(file_path: str) -> str:
+ """Return the SHA-256 hex digest of `file_path`, read in streaming chunks.
+ Args:
+ file_path: absolute path to the file to hash.
+ Returns:
+ Lowercase hex-encoded SHA-256 digest string.
+ """
+ h = hashlib.sha256()
+ if MultiStorageClientFeature.is_enabled():
+ msc = MultiStorageClientFeature.import_package()
+ with msc.open(file_path, 'rb') as f:
+ for chunk in iter(lambda: f.read(_READ_CHUNK_SIZE), b''):
+ h.update(chunk)
+ else:
+ with open(file_path, 'rb') as f:
+ for chunk in iter(lambda: f.read(_READ_CHUNK_SIZE), b''):
+ h.update(chunk)
+ return h.hexdigest()
+
+
+def save_integrity_manifest(checkpoint_dir: str) -> None:
+ """Hash every file in `heckpoint_dir` and write an integrity manifest.
+ The manifest lists each filename (relative to `checkpoint_dir`)
+ together with its SHA-256 digest. The manifest file itself is excluded
+ from the listing.
+ Args:
+ checkpoint_dir: directory that contains the checkpoint files.
+ """
+ manifest: Dict[str, str] = {}
+
+ if MultiStorageClientFeature.is_enabled():
+ msc = MultiStorageClientFeature.import_package()
+ ckpt_path = msc.Path(checkpoint_dir)
+ for entry in sorted(ckpt_path.iterdir(), key=lambda p: str(p)):
+ if entry.name != INTEGRITY_FNAME:
+ manifest[entry.name] = _compute_file_hash(str(entry))
+ else:
+ ckpt_path = Path(checkpoint_dir)
+ for entry in sorted(ckpt_path.iterdir()):
+ if entry.is_file() and entry.name != INTEGRITY_FNAME:
+ manifest[entry.name] = _compute_file_hash(str(entry))
+
+ integrity_path = os.path.join(checkpoint_dir, INTEGRITY_FNAME)
+ payload = {'algorithm': _HASH_ALGORITHM, 'files': manifest}
+
+ if MultiStorageClientFeature.is_enabled():
+ msc = MultiStorageClientFeature.import_package()
+ with msc.open(integrity_path, 'w') as f:
+ json.dump(payload, f, indent=2)
+ else:
+ with open(integrity_path, 'w') as f:
+ json.dump(payload, f, indent=2)
+
+ logger.info("Saved integrity manifest with %d file(s) to %s", len(manifest), integrity_path)
+
+
+def _verify_integrity_manifest_impl(checkpoint_dir: str) -> None:
+ """Single-process implementation of integrity verification.
+ Reads ``integrity.json``, recomputes each file's hash, and raises
+ `megatron.core.dist_checkpointing.core.CheckpointingException`
+ on any mismatch or missing file.
+ Args:
+ checkpoint_dir: checkpoint directory to verify.
+ Raises:
+ CheckpointingException: if the manifest is absent, uses an unsupported
+ algorithm, or any file's hash does not match.
+ """
+ integrity_path = os.path.join(checkpoint_dir, INTEGRITY_FNAME)
+
+ if MultiStorageClientFeature.is_enabled():
+ msc = MultiStorageClientFeature.import_package()
+ if not msc.os.path.exists(integrity_path):
+ raise CheckpointingException(
+ f'Integrity manifest not found at {integrity_path}. '
+ 'The checkpoint must be saved with integrity verification enabled '
+ '(save_integrity=True) before it can be verified on load.'
+ )
+ with msc.open(integrity_path) as f:
+ manifest_data = json.load(f)
+ else:
+ if not os.path.exists(integrity_path):
+ raise CheckpointingException(
+ f'Integrity manifest not found at {integrity_path}. '
+ 'The checkpoint must be saved with integrity verification enabled '
+ '(save_integrity=True) before it can be verified on load.'
+ )
+ with open(integrity_path) as f:
+ manifest_data = json.load(f)
+
+ algorithm = manifest_data.get('algorithm', _HASH_ALGORITHM)
+ if algorithm != _HASH_ALGORITHM:
+ raise CheckpointingException(
+ f'Unsupported hash algorithm in integrity manifest: {algorithm!r}. '
+ f'Expected: {_HASH_ALGORITHM!r}.'
+ )
+
+ manifest: Dict[str, str] = manifest_data['files']
+ mismatches = []
+
+ for filename, expected_hash in manifest.items():
+ full_path = os.path.join(checkpoint_dir, filename)
+ try:
+ actual_hash = _compute_file_hash(full_path)
+ except (FileNotFoundError, OSError) as exc:
+ mismatches.append(f' {filename}: file missing or unreadable ({exc})')
+ continue
+ if actual_hash != expected_hash:
+ mismatches.append(
+ f' {filename}: hash mismatch '
+ f'(expected {expected_hash[:16]}..., got {actual_hash[:16]}...)'
+ )
+
+ if mismatches:
+ raise CheckpointingException(
+ f'Checkpoint integrity verification failed for {len(mismatches)} '
+ f'file(s) in {checkpoint_dir}:\n' + '\n'.join(mismatches)
+ )
+
+ logger.info("Checkpoint integrity verified: %d file(s) OK in %s", len(manifest), checkpoint_dir)
+
+
+def verify_integrity_manifest(checkpoint_dir: str) -> None:
+ """Verify checkpoint files against their recorded SHA-256 hashes.
+ Args:
+ checkpoint_dir: checkpoint directory to verify.
+ Raises:
+ CheckpointingException: if ``integrity.json`` is absent or any file's
+ hash no longer matches the stored value.
+ """
+ import torch
+
+ if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1:
+ error_payload = [None]
+ if torch.distributed.get_rank() == 0:
+ try:
+ _verify_integrity_manifest_impl(checkpoint_dir)
+ except CheckpointingException as exc:
+ error_payload = [str(exc)]
+ torch.distributed.broadcast_object_list(error_payload, src=0)
+ if error_payload[0] is not None:
+ raise CheckpointingException(error_payload[0])
+ else:
+ _verify_integrity_manifest_impl(checkpoint_dir)
diff --git a/megatron/core/distributed/README.md b/megatron/core/distributed/README.md
index c4a75284414..489e381f9e0 100644
--- a/megatron/core/distributed/README.md
+++ b/megatron/core/distributed/README.md
@@ -1,11 +1,27 @@
-## How to use pytorch FSDP2?
+# Distributed Data Parallelism
-Add these flag to enable Torch FSDP2.
+This module contains algorithms, data structures, and utilities used for different types of distributed data parallelism, such as DDP and FSDP.
+
+## Distributed Data Parallelism
+
+This is the default data parallelism used with all parallelism topologies in Megatron-LM.
+
+## Megatron-FSDP
+
+To use Megatron-FSDP in Megatron-LM, enable the following arguments:
+
+```
+--use-megatron-fsdp
+--ckpt-format fsdp_dtensor
+--init-model-with-meta-device
+```
+
+## FSDP2
+
+To use FSDP2 in Megatron-LM, enable the following arguments:
```
--use-torch-fsdp2
--no-gradient-accumulation-fusion
--ckpt-format torch_dist
```
-
-It is worth noting that CUDA_MAX_CONNECTIONS=1 should not be enabled to ensure that the communication of FSDP and the computation on the primary stream can be fully parallelized.
diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py
index 6855c32c15a..e313113a448 100644
--- a/megatron/core/distributed/distributed_data_parallel.py
+++ b/megatron/core/distributed/distributed_data_parallel.py
@@ -7,14 +7,14 @@
import torch
from ..config_logger import has_config_logger_enabled, log_config_to_disk
-from ..fp8_utils import is_float8tensor, post_all_gather_processing
+from ..optimizer.param_layout import FullParamLayout
from ..process_groups_config import ProcessGroupCollection
from ..transformer.cuda_graphs import is_graph_capturing
from ..transformer.transformer_config import TransformerConfig
from ..utils import log_single_rank
from .data_parallel_base import _BaseDataParallel
from .distributed_data_parallel_config import DistributedDataParallelConfig
-from .param_and_grad_buffer import _ParamAndGradBuffer, partition_buckets
+from .param_and_grad_buffer import _ParamAndGradBuffer, group_params_for_buffers, partition_buckets
logger = logging.getLogger(__name__)
@@ -35,6 +35,9 @@ class DistributedDataParallel(_BaseDataParallel):
use standard bucketing policy: assign parameters to smaller buckets and all-reduce
per bucket _if_ overlap_grad_reduce is True and pp_rank is 0.
pg_collection: Optional unified process group for distributed training.
+ full_param_layout: Optional FullParamLayout providing pre-computed layouts for all
+ dtype groups. When provided, each buffer uses the corresponding PerBufferParamLayout
+ instead of computing a default one.
"""
@@ -45,6 +48,7 @@ def __init__(
module: torch.nn.Module,
disable_bucketing: bool = False,
pg_collection: Optional[ProcessGroupCollection] = None,
+ full_param_layout: Optional[FullParamLayout] = None,
):
super().__init__(config=config, module=module)
if has_config_logger_enabled(config):
@@ -103,11 +107,10 @@ def __init__(
self.param_to_bucket_group = {}
- # Group parameters by their gradient type.
+ # Collect all trainable parameters.
param_to_name = {}
- dense_params = []
- expert_parallel_params = []
self.params_with_grad = []
+ all_params = []
for name, param in self.module.named_parameters():
if not param.requires_grad:
continue
@@ -118,142 +121,52 @@ def __init__(
param.grad_added_to_main_grad = False
param_to_name[param] = name
+ all_params.append(param)
+
+ # Group parameters by (param_dtype, grad_dtype, is_expert_parallel).
+ buffer_groups = group_params_for_buffers(all_params, self.ddp_config.grad_reduce_in_fp32)
+
+ # Auto-compute layouts when using distributed optimizer but no layout was provided.
+ # This maintains backward compatibility for callers that create DDP directly
+ # without pre-computing layouts (e.g., tests, external code).
+ if full_param_layout is None and self.ddp_config.use_distributed_optimizer:
+ log_single_rank(
+ logger,
+ logging.WARNING,
+ "DistributedDataParallel: full_param_layout not provided with "
+ "use_distributed_optimizer=True. Auto-computing layout inside DDP. "
+ "Callers should pre-compute layouts via "
+ "DistributedOptimizer.compute_full_param_layout() and pass them in.",
+ )
+ from ..optimizer.distrib_optimizer import DistributedOptimizer
+
+ full_param_layout = DistributedOptimizer.compute_full_param_layout(
+ all_params,
+ self.bucket_size,
+ self.intra_dp_cp_group.size(),
+ self.ddp_config,
+ expert_data_parallel_world_size=self.intra_expt_dp_group.size(),
+ )
- if getattr(param, 'allreduce', True):
- dense_params.append((param, name))
- else:
- expert_parallel_params.append((param, name))
-
- def _allocate_buffers_for_parameters(
- input_params, data_parallel_group, gradient_scaling_factor
- ):
- param_and_grad_dtype_to_params = {}
- param_and_grad_dtype_to_offsets = {}
- param_and_grad_dtype_to_indices = {}
-
- # Group parameters by their gradient type.
- for param, param_name in input_params:
- assert param.requires_grad
-
- param_dtype = param.dtype
- if is_float8tensor(param):
- # Currently TE's Float8Tensor is a wrapper of torch.Tensor. It has a "fake"
- # dtype (usually a higher precision dtype such as bfloat16), but its actual
- # data is stored in the form of a torch uint8 tensor within the Float8Tensor's
- # ".data" attribute. Therefore, when creating the param buffer for fp8 params,
- # it is necessary to use torch.uint8, not the "fake" dtype got from
- # "param.dtype".
- param_dtype = torch.uint8
- grad_dtype = torch.float if self.ddp_config.grad_reduce_in_fp32 else param.dtype
-
- params = param_and_grad_dtype_to_params.get((param_dtype, grad_dtype), [])
- params.append((param, param_name))
- param_and_grad_dtype_to_params[(param_dtype, grad_dtype)] = params
-
- # Get the index of each param among the params with same dtype, if a param is fp8,
- # use its "fake" high precision dtype to find which params have same dtype with it.
- # For example:
- # Case 1:
- # params = [p1(bf16), p2(bf16), p3(bf16), p4(bf16)]
- # param_and_grad_dtype_to_indices = {
- # (torch.bfloat16, torch.float32): [0, 1, 2, 3],
- # }
- # Case 2:
- # params = [p1(bf16), p2(fp8), p3(fp8), p4(bf16)]
- # param_and_grad_dtype_to_indices = {
- # (torch.bfloat16, torch.float32): [0, 3],
- # (torch.uint8, torch.float32): [1, 2],
- # }
- # We need these indices to load a non-native-fp8 checkpoint in native-fp8 mode.
- offset = param_and_grad_dtype_to_offsets.get((param.dtype, grad_dtype), 0)
- param_and_grad_dtype_to_offsets[(param.dtype, grad_dtype)] = offset + 1
- indices = param_and_grad_dtype_to_indices.get((param_dtype, grad_dtype), [])
- indices.append(offset)
- param_and_grad_dtype_to_indices[(param_dtype, grad_dtype)] = indices
-
- if not config.calculate_per_token_loss:
- target_gradient_scaling_factor = 1.0 / self.dp_cp_group.size()
- if self.ddp_config.average_in_collective:
- if self.ddp_config.num_distributed_optimizer_instances == 1:
- # Collective is averaging gradients in collective with data_parallel_group.
- assert (
- gradient_scaling_factor / data_parallel_group.size()
- == target_gradient_scaling_factor
- )
- else:
- # For non-expert parameters, gradient_scaling_factor is 1.
- # For expert parameters, gradient_scaling_factor is edp_size/dp_size.
- assert (gradient_scaling_factor == 1) or (
- gradient_scaling_factor
- == (self.expt_dp_group.size() / self.dp_cp_group.size())
- )
- else:
- assert gradient_scaling_factor == target_gradient_scaling_factor
-
- # Allocate the grad buffers and map the grads.
- buffers = []
- pg_collection = ProcessGroupCollection()
- pg_collection.tp = self.tp_group
- pg_collection.dp_cp = self.dp_cp_group
- for (param_dtype, grad_dtype), params in param_and_grad_dtype_to_params.items():
- buffers.append(
- _ParamAndGradBuffer(
- self.ddp_config,
- param_dtype,
- grad_dtype,
- params,
- data_parallel_group,
- self.bucket_size,
- param_to_name,
- gradient_scaling_factor,
- param_and_grad_dtype_to_indices[(param_dtype, grad_dtype)],
- self.ddp_config.nccl_ub,
- pg_collection,
- )
- )
-
- # In some scenarios, we want to put buckets from different buffers into a group so that
- # their communication can be aggregated. For example, when there are both fp8 buffers
- # and bf16 buffers in the model and vpp is enabled, each model chunk will have an fp8
- # bucket and a bf16 bucket, which doubles the number of communication kernels, and
- # because of the use of CUDA_DEVICE_MAX_CONNECTIONS=1, having multiple back-to-back
- # communications will prevent the overlap of the communication kernels with computation
- # kernels.
- # If bucketing is explicitly disabled, then put all buckets in a buffer into a single
- # bucket group.
- bucket_groups = partition_buckets(buffers, force_single_bucket_group=disable_bucketing)
-
- if self.ddp_config.num_distributed_optimizer_instances > 1:
+ # When a full_param_layout is provided, verify that the grouping is consistent
+ # with the layout (same buffer keys, same params per key, same param_indices).
+ if full_param_layout is not None:
+ assert set(buffer_groups.keys()) == set(full_param_layout.layouts.keys()), (
+ f"Buffer keys from param grouping {set(buffer_groups.keys())} do not match "
+ f"full_param_layout keys {set(full_param_layout.layouts.keys())}"
+ )
+ for buffer_key, (params, param_indices) in buffer_groups.items():
+ layout = full_param_layout.layouts[buffer_key]
+ assert set(params) == set(
+ layout.param_index_map.keys()
+ ), f"Params for {buffer_key} do not match between grouping and layout"
assert (
- self.ddp_config.use_distributed_optimizer
- ), 'Partial DistOpt cannot be used without DistOpt'
- communication_stream = torch.cuda.Stream(device=torch.cuda.current_device())
- for bucket_group in bucket_groups:
- bucket_group.inter_distributed_optimizer_instance_group = (
- self.inter_dist_opt_group
- )
- bucket_group.communication_stream = communication_stream
+ param_indices == layout.param_indices
+ ), f"param_indices for {buffer_key} do not match between grouping and layout"
- # Set `next_param_gather_bucket_group` for different bucket groups by iterating through
- # buckets in reverse order (since all-gathers happen in reverse order of buckets).
- # Note: overlap_param_gather covers both the distributed optimizer and the
- # layer-wise optimizer cases; the latter sets overlap_param_gather=True
- # without use_distributed_optimizer.
- if self.ddp_config.overlap_param_gather:
- num_bucket_groups = len(bucket_groups)
- for i in range(1, num_bucket_groups):
- bucket_groups[num_bucket_groups - i].next_param_gather_bucket_group = (
- bucket_groups[num_bucket_groups - i - 1]
- )
-
- # Create map from param to bucket group, used in pre_hook.
- for bucket_group in bucket_groups:
- for bucket in bucket_group.buckets:
- for param in bucket.params_list:
- self.param_to_bucket_group[param] = bucket_group
-
- return buffers, bucket_groups
+ self.full_param_layout = full_param_layout
+ # Compute gradient scaling factors.
if config.calculate_per_token_loss:
assert (
not self.ddp_config.average_in_collective
@@ -290,20 +203,132 @@ def _allocate_buffers_for_parameters(
gradient_scaling_factor = 1.0 / data_parallel_world_size
expert_gradient_scaling_factor = 1.0 / data_parallel_world_size
- # Allocate the param+grad buffers for dense params' grads.
- self.buffers, self.bucket_groups = _allocate_buffers_for_parameters(
- dense_params, self.intra_dp_cp_group, gradient_scaling_factor=gradient_scaling_factor
- )
+ # Allocate buffers for each group.
+ self.buffers = []
+ self.expert_parallel_buffers = []
+ pg_collection = ProcessGroupCollection(tp=self.tp_group, dp_cp=self.dp_cp_group)
+ for buffer_key, (params, param_indices) in buffer_groups.items():
+ if buffer_key.is_expert_parallel:
+ data_parallel_group = self.intra_expt_dp_group
+ scaling_factor = expert_gradient_scaling_factor
+ else:
+ data_parallel_group = self.intra_dp_cp_group
+ scaling_factor = gradient_scaling_factor
- # Allocate separate param+grad buffers for expert parallel params' grads.
- self.expert_parallel_buffers, self.expert_parallel_bucket_groups = (
- _allocate_buffers_for_parameters(
- expert_parallel_params,
- self.intra_expt_dp_group,
- gradient_scaling_factor=expert_gradient_scaling_factor,
+ if not config.calculate_per_token_loss:
+ target_gradient_scaling_factor = 1.0 / self.dp_cp_group.size()
+ if self.ddp_config.average_in_collective:
+ if self.ddp_config.num_distributed_optimizer_instances == 1:
+ # Collective is averaging gradients in collective with data_parallel_group.
+ assert (
+ scaling_factor / data_parallel_group.size()
+ == target_gradient_scaling_factor
+ )
+ else:
+ # For non-expert parameters, gradient_scaling_factor is 1.
+ # For expert parameters, gradient_scaling_factor is edp_size/dp_size.
+ assert (scaling_factor == 1) or (
+ scaling_factor == (self.expt_dp_group.size() / self.dp_cp_group.size())
+ )
+ else:
+ assert scaling_factor == target_gradient_scaling_factor
+
+ param_layout = (
+ full_param_layout.layouts.get(buffer_key) if full_param_layout is not None else None
)
+ params_with_names = [(p, param_to_name[p]) for p in params]
+ buffer = _ParamAndGradBuffer(
+ self.ddp_config,
+ buffer_key.param_dtype,
+ buffer_key.grad_dtype,
+ params_with_names,
+ data_parallel_group,
+ self.bucket_size,
+ param_to_name,
+ scaling_factor,
+ param_indices,
+ self.ddp_config.nccl_ub,
+ pg_collection,
+ param_layout=param_layout,
+ )
+ if buffer_key.is_expert_parallel:
+ self.expert_parallel_buffers.append(buffer)
+ else:
+ self.buffers.append(buffer)
+
+ # In some scenarios, we want to put buckets from different buffers into a group so that
+ # their communication can be aggregated. For example, when there are both fp8 buffers
+ # and bf16 buffers in the model and vpp is enabled, each model chunk will have an fp8
+ # bucket and a bf16 bucket, which doubles the number of communication kernels, and
+ # because of the use of CUDA_DEVICE_MAX_CONNECTIONS=1, having multiple back-to-back
+ # communications will prevent the overlap of the communication kernels with computation
+ # kernels.
+ # If bucketing is explicitly disabled, then put all buckets in a buffer into a single
+ # bucket group.
+ self.bucket_groups = partition_buckets(
+ self.buffers,
+ force_single_bucket_group=disable_bucketing,
+ reduce_scatter_with_fp32_accumulation=(
+ self.ddp_config.reduce_scatter_with_fp32_accumulation
+ ),
+ )
+ self.expert_parallel_bucket_groups = partition_buckets(
+ self.expert_parallel_buffers,
+ force_single_bucket_group=disable_bucketing,
+ reduce_scatter_with_fp32_accumulation=(
+ self.ddp_config.reduce_scatter_with_fp32_accumulation
+ ),
)
+ if self.ddp_config.num_distributed_optimizer_instances > 1:
+ assert (
+ self.ddp_config.use_distributed_optimizer
+ ), 'Partial DistOpt cannot be used without DistOpt'
+ for bucket_groups in [self.bucket_groups, self.expert_parallel_bucket_groups]:
+ communication_stream = torch.cuda.Stream(device=torch.cuda.current_device())
+ for bucket_group in bucket_groups:
+ bucket_group.inter_distributed_optimizer_instance_group = (
+ self.inter_dist_opt_group
+ )
+ bucket_group.communication_stream = communication_stream
+
+ # Set `next_param_gather_bucket_group` for different bucket groups by iterating through
+ # buckets in reverse order (since all-gathers happen in reverse order of buckets).
+ # Note: overlap_param_gather covers both the distributed optimizer and the
+ # layer-wise optimizer cases; the latter sets overlap_param_gather=True
+ # without use_distributed_optimizer.
+ if self.ddp_config.overlap_param_gather:
+ for bucket_groups in [self.bucket_groups, self.expert_parallel_bucket_groups]:
+ num_bucket_groups = len(bucket_groups)
+ for i in range(1, num_bucket_groups):
+ bucket_groups[num_bucket_groups - i].next_param_gather_bucket_group = (
+ bucket_groups[num_bucket_groups - i - 1]
+ )
+
+ # Set `previous_grad_reduce_bucket_group` so each bucket group can drain its predecessor's
+ # reduce-scatter at dispatch time. Only needed for reduce_scatter_with_fp32_accumulation,
+ # which holds an intermediate all-to-all output tensor pinned until .wait() runs; without
+ # this draining, all such tensors stay live until end-of-step. The fp32-accum path asserts
+ # num_distributed_optimizer_instances == 1 elsewhere, so we only link in that case.
+ # Grad-reduce dispatches happen in forward order of bucket_groups during backward (buckets
+ # closer to the output finish their gradients first), so bucket_groups[i]'s immediate
+ # predecessor in dispatch order is bucket_groups[i-1].
+ if (
+ self.ddp_config.overlap_grad_reduce
+ and self.ddp_config.reduce_scatter_with_fp32_accumulation
+ and self.ddp_config.num_distributed_optimizer_instances == 1
+ ):
+ for bucket_groups in [self.bucket_groups, self.expert_parallel_bucket_groups]:
+ for i in range(1, len(bucket_groups)):
+ bucket_groups[i].previous_grad_reduce_bucket_group = bucket_groups[i - 1]
+
+ # Create map from param to bucket group, used in pre_hook.
+ for bucket_groups in [self.bucket_groups, self.expert_parallel_bucket_groups]:
+ for bucket_group in bucket_groups:
+ for bucket in bucket_group.buckets:
+ for param in bucket.params_list:
+ self.param_to_bucket_group[param] = bucket_group
+
# Delete references to weight_tensor if they exist since we don't want two parameter copies
# if we re-mapped parameters (which happens when we use the distributed optimizer).
# This is a temporary workaround around a TE bug that is fixed with
@@ -464,6 +489,24 @@ def no_sync(self):
for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
bucket_group.is_last_microbatch = True
+ def _start_bucket_group_param_sync(
+ self, bucket_group: '_ParamAndGradBucketGroup', force_sync: bool
+ ) -> None:
+ """Dispatch one bucket group's param all-gather + run the FP8 / MXFP8 / FP4
+ post-all-gather work the synchronous path needs.
+
+ Factored out of :meth:`start_param_sync` so callers that own a subset
+ of bucket groups (e.g. a chained ``LayerWiseDistributedOptimizer`` +
+ ``DistributedOptimizer`` pair) can sync only their own buckets without
+ losing the post-processing that follows the collective.
+ """
+ bucket_group.start_param_sync(force_sync=force_sync)
+
+ if self.ddp_config.overlap_param_gather:
+ return
+
+ bucket_group._post_param_sync()
+
def start_param_sync(self, *unused, force_sync: bool = False, force_dispatch: bool = False):
"""
Initiates param sync (all-gather) communication operations for all model parameters.
@@ -484,43 +527,7 @@ def start_param_sync(self, *unused, force_sync: bool = False, force_dispatch: bo
return
for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
- bucket_group.start_param_sync(force_sync=force_sync)
-
- if not self.ddp_config.overlap_param_gather:
- # For MXFP8 params, we need to copy the all-gathered param data from the buffer to
- # the param.data, since param buffer is not mapped to model params for MXFP8 case.
- # The paramaters are cast from bf16 to MXFP8 during copy.
- # In the case of "overlap_param_gather=True", the param copy is done
- # in "finish_param_sync" stage after zeroing the shared gardient buffers.
- if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag:
- for bucket in bucket_group.buckets:
- is_bf16_weight_bucket = False
- for param in bucket.params:
- # Skip copying since bf16 weights in the mxfp8 model
- # are already mapped to param.data.
- if not is_float8tensor(param):
- is_bf16_weight_bucket = True
- break
- param_start, param_end = bucket.param_to_index[param]
- param_slice = bucket.param_data.view(-1)[param_start:param_end]
- param.data.copy_(param_slice.view(param.data.shape))
- if is_bf16_weight_bucket:
- continue
- # All-gathered params are not needed after being copied to param.data.
- # Zero out the param buffer (shared with grad buffer) for gradient
- # accumulation. We cannot zero out the entire grad buffer because one grad
- # buffer may correspond to multiple param buffers. If we zero out the entire
- # grad buffer, it would clear the data of those param buffers that have not
- # yet completed AG.
- bucket.param_data.zero_()
- else:
- fp8_params = []
- for bucket in bucket_group.buckets:
- for param in bucket.params:
- if is_float8tensor(param):
- fp8_params.append(param)
- if len(fp8_params) > 0:
- post_all_gather_processing(fp8_params)
+ self._start_bucket_group_param_sync(bucket_group, force_sync=force_sync)
def start_grad_sync(self, *unused):
"""
diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py
index ee592368b0c..b14b472d0ce 100644
--- a/megatron/core/distributed/distributed_data_parallel_config.py
+++ b/megatron/core/distributed/distributed_data_parallel_config.py
@@ -5,6 +5,8 @@
import torch
+from ..utils import is_torch_min_version
+
@dataclass
class DistributedDataParallelConfig:
@@ -48,6 +50,11 @@ class DistributedDataParallelConfig:
value of max(40000000, 1000000 * dp_size) parameters (larger DP sizes need larger
buckets to ensure collectives do not become latency-bound)."""
+ num_buckets: Optional[int] = None
+ """Number of buckets for data-parallel communication. Should only specify one of
+ `bucket_size` and `num_buckets`. If `num_buckets` is specified, `bucket_size`
+ will be determined at runtime."""
+
pad_buckets_for_high_nccl_busbw: bool = False
"""If true, make sure the bucket size is divisible by a large power of 2 (2^16) to
ensure NCCL collectives have high bus bandwidth at large DP counts, since NCCL
@@ -72,6 +79,10 @@ class DistributedDataParallelConfig:
"""If true, keep the compute param in fp8 (do not use any other intermediate dtype) and
perform the param all-gather in fp8."""
+ fp4_param_gather: bool = False
+ """If true, keep the compute param in fp4 (do not use any other intermediate dtype) and
+ perform the param all-gather in fp4."""
+
reuse_grad_buf_for_mxfp8_param_ag: bool = False
"""If true, reuse the grad buffer for param AG when using mxfp8 recipe. Should be
set to True only when fp8_recipe is mxfp8 and fp8_param_gather is True."""
@@ -143,7 +154,9 @@ class DistributedDataParallelConfig:
If True, use all-gather during the initial Megatron-FSDP parameter
synchronization step. This can increase overlap between the first
parameter all-gather and computation, helping to better hide the
- initial communication cost.
+ initial communication cost. Should be deactivated when using
+ full-iteration CG, or partial CG if AG/RS is launched beyond the
+ CG capture scope but is waited on during the capture scope.
"""
outer_dp_sharding_strategy: str = 'no_shard'
@@ -196,6 +209,34 @@ class DistributedDataParallelConfig:
No additional memory is allocated when `grad_comm_dtype == main_grads_dtype`.
"""
+ megatron_fsdp_use_decoupled_grad: bool = False
+ """If true, Megatron-FSDP's ParamAndGradBuffer uses the precision-aware optimizer
+ gradient path (e.g. `decoupled_grad` on optimizer parameters) instead of casting
+ main gradients to parameter dtype for `.grad`.
+ """
+
+ megatron_fsdp_cuda_graph_mode: bool = False
+ """If set to True, Megatron-FSDP will practice CUDA graph-safe operations, such as
+ not dereferencing `param.grad` after the optimizer step to preserve references for
+ CUDA graph replay. Can affect memory utilization in some cases, such as when the
+ gradient shard is not a view of the Megatron-FSDP sharded gradient buffer, so
+ FusedAdam(use_decoupled_grad=True) + megatron_fsdp_use_decoupled_grad=True or
+ setting megatron_fsdp_main_params_dtype == megatron_fsdp_main_grads_dtype is
+ recommended to avoid casting the gradient to the parameter precision and creating
+ a casted-copy of the gradient shard that cannot be dereferenced due to replay.
+ """
+
+ megatron_fsdp_enable_fine_grained_param_gather: bool = False
+ """If set to True, enables fine-grained parameter gathering for Megatron-FSDP.
+ This feature increases the overlap between parameter all-gather and forward computation,
+ at the cost of more frequent communication calls.
+ For MXFP8, this approach helps save memory during fine-grained activation
+ recomputation, because MXFP8 forward and backward passes use different
+ parameter representations (rowwise data for forward, colwise data for backward).
+ In this mode, only the rowwise parameters of modules involved in recomputation
+ will be unsharded.
+ """
+
def __post_init__(self):
import os
@@ -203,7 +244,7 @@ def __post_init__(self):
if self.reuse_grad_buf_for_mxfp8_param_ag:
assert self.fp8_param_gather, "Reuse grad buffer only when keeping params in MXFP8."
- if self.nccl_ub:
+ if self.nccl_ub and not is_torch_min_version("2.11.0a0"):
if 'expandable_segments:True' in os.getenv('PYTORCH_CUDA_ALLOC_CONF', '').split(','):
raise ValueError(
"PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True is currently not supported "
@@ -215,3 +256,7 @@ def __post_init__(self):
"Only need to explicitly specify param_name patterns for FP32 local accumulation "
"if .main_grads aren't already in FP32"
)
+
+ if self.num_buckets is not None:
+ assert self.bucket_size is None, "Cannot specify both num_buckets and bucket_size"
+ assert self.num_buckets > 0, "num_buckets must be greater than 0"
diff --git a/megatron/core/distributed/finalize_model_grads.py b/megatron/core/distributed/finalize_model_grads.py
index ca6bdd354ce..778d1b75412 100644
--- a/megatron/core/distributed/finalize_model_grads.py
+++ b/megatron/core/distributed/finalize_model_grads.py
@@ -1,7 +1,7 @@
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
from functools import partial
-from typing import Callable, List, Optional, Union
+from typing import Callable, Dict, List, Optional, Union
import torch
from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
@@ -275,6 +275,44 @@ def _allreduce_position_embedding_grads(
)
+def _allreduce_router_grads(model: List[torch.nn.Module], config: TransformerConfig):
+ """
+ All-reduce router grads.
+
+ Reduce grads across all the pp stages to ensure that parameters of the router stay in sync.
+ """
+
+ if parallel_state.get_pipeline_model_parallel_world_size() > 1:
+ grads_dict: Dict[str, List[torch.Tensor]] = {}
+ for model_chunk in model:
+ for name, param in get_attr_wrapped_model(model_chunk, 'named_parameters')():
+ if param.requires_grad and getattr(param, 'flextron_router_pp_sync', False):
+ grad = param.main_grad
+ if name in grads_dict:
+ # Add all the virtual PP rank's gradients to
+ # the first local virtual PP rank.
+ grads_dict[name][0].add_(grad)
+ # Append to the end for later update after cross-rank reduce.
+ grads_dict[name].append(grad)
+ else:
+ grads_dict[name] = [grad]
+
+ if grads_dict:
+ # All-reduce the gradient on the first VPP rank.
+ grads = [param_grad[0] for _, param_grad in grads_dict.items()]
+ coalesced = _flatten_dense_tensors(grads)
+ torch.distributed.all_reduce(
+ coalesced, group=parallel_state.get_pipeline_model_parallel_group()
+ )
+ for buf, synced in zip(grads, _unflatten_dense_tensors(coalesced, grads)):
+ buf.copy_(synced)
+
+ # Update the gradients on other VPP ranks.
+ for grads in grads_dict.values():
+ for grad in grads[1:]:
+ grad.copy_(grads[0])
+
+
def reset_model_temporary_tensors(config: TransformerConfig, model: List[torch.nn.Module]):
"""
Reset the temporary tensors of the model.
@@ -290,7 +328,11 @@ def reset_model_temporary_tensors(config: TransformerConfig, model: List[torch.n
module.reset_global_aux_loss_tracker()
-def _update_router_expert_bias(model: List[torch.nn.Module], config: TransformerConfig):
+def _update_router_expert_bias(
+ model: List[torch.nn.Module],
+ config: TransformerConfig,
+ tp_dp_cp_group: Optional[torch.distributed.ProcessGroup] = None,
+):
"""
Update the expert bias of the router for a global batch.
This requires all-reduce of local_tokens_per_expert across TPxCPxDP ranks
@@ -312,7 +354,10 @@ def _update_router_expert_bias(model: List[torch.nn.Module], config: Transformer
stacked_tokens_per_expert = torch.stack(tokens_per_expert_list, dim=0)
stacked_expert_bias = torch.stack(expert_bias_list, dim=0)
stacked_updated_expert_bias = get_updated_expert_bias(
- stacked_tokens_per_expert, stacked_expert_bias, config.moe_router_bias_update_rate
+ stacked_tokens_per_expert,
+ stacked_expert_bias,
+ config.moe_router_bias_update_rate,
+ tp_dp_cp_group=tp_dp_cp_group,
)
for expert_bias, updated_expert_bias in zip(expert_bias_list, stacked_updated_expert_bias):
@@ -410,6 +455,7 @@ def finalize_model_grads(
"""
config = get_model_config(model[0])
+ tp_dp_cp_group = None
if pg_collection is not None:
assert hasattr(pg_collection, 'tp')
assert hasattr(pg_collection, 'pp')
@@ -428,6 +474,11 @@ def finalize_model_grads(
"If you don't need pos_embd_group, you need to explicitly set it to None."
)
assert hasattr(pg_collection, 'dp_cp')
+ if config.moe_router_enable_expert_bias:
+ assert hasattr(pg_collection, 'tp_dp_cp') and pg_collection.tp_dp_cp is not None, (
+ "pg_collection must have tp_dp_cp when " "moe_router_enable_expert_bias is enabled."
+ )
+ tp_dp_cp_group = pg_collection.tp_dp_cp
tp_group = pg_collection.tp
pp_group = pg_collection.pp
embd_group = pg_collection.embd
@@ -457,6 +508,9 @@ def finalize_model_grads(
if config.timers is not None:
config.timers('conditional-embedder-grads-all-reduce').stop()
+ if getattr(config, 'flextron', False):
+ _allreduce_router_grads(model, config)
+
# All-reduce layer-norm grads (for sequence parallelism) and non-tensor parallel modules.
if config.timers is not None:
config.timers('non-tensor-parallel-grads-all-reduce', log_level=1).start(
@@ -478,7 +532,11 @@ def finalize_model_grads(
config.timers('embedding-grads-all-reduce').stop()
if config.moe_router_enable_expert_bias:
- _update_router_expert_bias(model, config)
+ if pg_collection is None:
+ tp_dp_cp_group = parallel_state.get_tensor_and_data_parallel_group(
+ with_context_parallel=True
+ )
+ _update_router_expert_bias(model, config, tp_dp_cp_group=tp_dp_cp_group)
reset_model_temporary_tensors(config, model)
@@ -495,7 +553,10 @@ def finalize_model_grads(
# all-reduce across DP ranks.
torch.distributed.all_reduce(num_tokens, group=dp_cp_group)
+
+ # Clamp to avoid div-by-zero without a host-side branch on a device tensor,
+ # which would otherwise cause a sync that is illegal during CUDA graph capture.
+ safe_num_tokens = torch.clamp(num_tokens, min=1)
+ scaling = 1.0 / safe_num_tokens
for model_chunk in model:
- if num_tokens > 0:
- scaling = 1.0 / num_tokens
- model_chunk.scale_gradients(scaling)
+ model_chunk.scale_gradients(scaling)
diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py
index 8993620c779..ea6b695988f 100644
--- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py
+++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py
@@ -14,7 +14,7 @@
import logging
import random
-from typing import List, Optional
+from typing import Dict, List, Optional
try:
import einops
@@ -26,6 +26,7 @@
import numpy as np
import torch
import torch.distributed as dist
+from torch import nn
try:
from torch.distributed import DeviceMesh
@@ -38,7 +39,6 @@
from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
from megatron.core.distributed.data_parallel_base import _BaseDataParallel
from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig
-from megatron.core.extensions.transformer_engine import TELinear
from megatron.core.process_groups_config import ProcessGroupCollection
from megatron.core.transformer.transformer_config import TransformerConfig
from megatron.core.transformer.transformer_layer import TransformerLayer
@@ -64,6 +64,32 @@ class FullyShardedDataParallel(_BaseDataParallel):
Fully Sharded Data Parallel (FSDP) wrapper for the Megatron model.
"""
+ # Module type registry (forked from Megatron-Bridge param_mapping utilities).
+ _MODULE_TYPE_REGISTRY: Dict[str, set] = {
+ "column": {
+ "ColumnParallelLinear",
+ "TEColumnParallelLinear",
+ "TELayerNormColumnParallelLinear",
+ "TEColumnParallelGroupedLinear",
+ "VocabParallelEmbedding",
+ "DotProductAttention", # for attention sink only
+ "TEDotProductAttention", # for attention sink only
+ },
+ "row": {"RowParallelLinear", "TERowParallelLinear", "TERowParallelGroupedLinear"},
+ "replicated": {
+ # Normalization layers
+ "TENorm",
+ "FusedLayerNorm",
+ "WrappedTorchNorm",
+ "LayerNorm",
+ "RMSNorm",
+ "L2Norm",
+ # Other non-parallel modules
+ "IdentityOp",
+ "TopKRouter",
+ },
+ }
+
def __init__(
self,
config: TransformerConfig,
@@ -80,6 +106,8 @@ def __init__(
if has_config_logger_enabled(config):
log_config_to_disk(config, locals(), prefix=type(self).__name__)
+ self.num_moe_experts = getattr(config, "num_moe_experts", None)
+
self.ddp_config = ddp_config
log_single_rank(
logger,
@@ -127,8 +155,28 @@ def __init__(
else:
self.fsdp_unit_modules = []
- self._fix_tensor_parallel_attributes(module)
+ self._annotate_tensor_parallelism(module)
+ if config.overlap_moe_expert_parallel_comm:
+ assert not ddp_config.fsdp_double_buffer, (
+ "1F1B overlap with FSDP does not support double buffer. "
+ "Please set fsdp_double_buffer=False in the ddp config."
+ )
+ assert config.cuda_graph_impl in ("none", "full_iteration"), (
+ "1F1B overlap with FSDP does not support per-layer CUDA graphs "
+ f"(cuda_graph_impl={config.cuda_graph_impl!r}). "
+ "Use cuda_graph_impl='full_iteration' or disable CUDA graphs "
+ "(cuda_graph_impl='none')."
+ )
+
+ if (
+ config.overlap_moe_expert_parallel_comm
+ and ddp_config.data_parallel_sharding_strategy == "optim_grads_params"
+ ):
+ assert self.fsdp_unit_modules == [TransformerLayer], (
+ "EP overlap with FSDP currently requires fsdp_unit_modules "
+ f"to be [TransformerLayer], got {self.fsdp_unit_modules}."
+ )
super().__init__(
config=config,
module=MegatronFSDP(
@@ -141,8 +189,21 @@ def __init__(
dist_index=self.megatron_fsdp_dist_index,
calculate_per_token_loss=config.calculate_per_token_loss,
init_model_with_meta_device=config.init_model_with_meta_device,
+ # EP overlap schedule calls sub-modules directly instead of
+ # TransformerLayer.forward(), so fine-grained hooks are needed
+ # to manage _training_state and all-gather each sub-module's
+ # parameters individually. This applies to all sharding
+ # strategies (not only optim_grads_params) because the hooks
+ # also maintain per-module training-state bookkeeping that the
+ # gradient-reduction pipeline relies on.
enable_fine_grained_param_gather_hook=(
- config.fp8_recipe == "mxfp8" and ddp_config.fp8_param_gather
+ (config.fp8_recipe == "mxfp8" and ddp_config.fp8_param_gather)
+ or config.overlap_moe_expert_parallel_comm
+ or self.ddp_config.megatron_fsdp_enable_fine_grained_param_gather
+ ),
+ enable_fine_grained_param_gather_backward_hook=(
+ config.overlap_moe_expert_parallel_comm
+ and ddp_config.data_parallel_sharding_strategy == "optim_grads_params"
),
),
)
@@ -154,6 +215,7 @@ def __init__(
self.scale_gradients = self.module.scale_gradients
self.zero_grad_buffer = self.module.zero_grad_buffer
self.broadcast_params = self.module.broadcast_params
+ self.synchronize_param_gather = self.module.synchronize_param_gather
self.module.state_dict_for_save_checkpoint = self.module.state_dict
self.state_dict_for_save_checkpoint = self.state_dict
self.module.config = config
@@ -182,43 +244,75 @@ def load_state_dict(self, state_dict, strict=True):
self.module.load_state_dict(custom_state_dict, strict=strict)
- def _fix_tensor_parallel_attributes(self, module):
- is_expert_param = lambda n, p: ".experts." in n
- is_router_param = lambda n, p: ".router.weight" in n
+ def _detect_parallelism_type(self, param_name: str, module: nn.Module) -> Optional[str]:
+ """
+ Infer tensor-parallelism type for a parameter under a given module
+ (forked from Megatron-Bridge).
+
+ Returns:
+ "column", "row", or "replicated" if a type can be inferred, else None.
+ """
+ module_type = type(module).__name__
+
+ # Handle fused modules like TELayerNormColumnParallelLinear
+ # These modules have both column-parallel weights (weight, bias)
+ # and replicated layer norm weights (layer_norm_weight, layer_norm_bias)
+ if module_type == "TELayerNormColumnParallelLinear":
+ # Check the actual parameter name to determine the correct parallelism type
+ if param_name.endswith("layer_norm_weight") or param_name.endswith("layer_norm_bias"):
+ return "replicated"
+ # All other parameters (weight, bias) are column-parallel
+ return "column"
+
+ # Check registry first
+ for parallelism, types in self._MODULE_TYPE_REGISTRY.items():
+ if module_type in types:
+ if parallelism == "row" and "bias" in param_name:
+ return "replicated"
+ return parallelism
+
+ # Fallback to inspecting module attributes
+ if hasattr(module, "tensor_model_parallel"):
+ if not module.tensor_model_parallel:
+ return "replicated"
+
+ # Check partition dimension
+ partition_dim = getattr(module, "partition_dim", None)
+ if partition_dim == 0:
+ return "column"
+ elif partition_dim == 1:
+ if "bias" in param_name:
+ return "replicated"
+ return "row"
+
+ # Fallback for normalization layers
+ if any(norm in module_type for norm in ["Norm", "Normalization"]):
+ return "replicated"
+
+ # Check parallel_mode for TELinear
+ if module_type == "TELinear":
+ if module.parallel_mode == "column":
+ return "column"
+ elif module.parallel_mode == "row":
+ if "bias" in param_name:
+ return "replicated"
+ return "row"
+ else:
+ return "replicated"
- if parallel_state.get_tensor_model_parallel_group():
- tp_size = parallel_state.get_tensor_model_parallel_group().size()
- else:
- tp_size = 1
+ return None
- if parallel_state.get_expert_tensor_parallel_group():
- expt_tp_size = parallel_state.get_expert_tensor_parallel_group().size()
- else:
- expt_tp_size = 1
-
- param_to_direct_module = {}
- for name, m in module.named_modules():
- for p in m.parameters(recurse=False):
- param_to_direct_module[p] = (name, m)
-
- for name, param in module.named_parameters():
- if is_expert_param(name, param) and expt_tp_size > 1:
- setattr(param, "_mcore_tp", True)
- if "linear_fc1.weight" in name:
- setattr(param, "_tp_partition_dim", 0)
- elif "linear_fc2.weight" in name:
- setattr(param, "_tp_partition_dim", 1)
-
- if not is_expert_param(name, param) and tp_size > 1:
- m_name, direct_module = param_to_direct_module[param]
- if isinstance(direct_module, (TELinear,)):
- parallel_mode = getattr(direct_module, "parallel_mode", None)
- if parallel_mode is None:
- setattr(param, "_mcore_tp", True)
- setattr(param, "_tp_duplicated", True)
- elif is_router_param(name, param):
- setattr(param, "_mcore_tp", True)
- setattr(param, "_tp_duplicated", True)
+ def _annotate_tensor_parallelism(self, root_module: nn.Module) -> None:
+ """Annotate parameters under root_module with inferred tensor-parallel metadata.
+
+ Each parameter that can be classified will get a `_tensor_parallel_mode` attribute
+ set to one of: "column", "row", or "replicated".
+ """
+ for submodule in root_module.modules():
+ for name, param in submodule.named_parameters(recurse=False):
+ detected_type = self._detect_parallelism_type(name, submodule)
+ if detected_type is not None:
+ setattr(param, "_tensor_parallel_mode", detected_type)
def _init_dist_index(self, pg_collection):
"""
@@ -283,8 +377,14 @@ def _init_dist_index(self, pg_collection):
single_rank_group = dist.new_group(ranks=[dist.get_rank()])
expt_tp_group = single_rank_group
+ # Extract AG groups from pg_collection for explicit passing
+ dp_cp_ag = getattr(pg_collection, 'dp_cp_ag', None) if pg_collection is not None else None
+ expt_dp_ag = (
+ getattr(pg_collection, 'expt_dp_ag', None) if pg_collection is not None else None
+ )
+
if enable_hsdp:
- if expt_dp_group is not None:
+ if self.num_moe_experts is not None:
expt_mesh = _get_hsdp_tp_mesh(
outer_fsdp_group, expt_dp_group, expt_tp_group, ep_size=ep_group.size()
)
@@ -311,9 +411,11 @@ def _init_dist_index(self, pg_collection):
hybrid_fsdp_group=hybrid_fsdp_group,
hybrid_fsdp_expt_group=hybrid_fsdp_expt_group,
expt_device_mesh=expt_device_mesh,
+ fsdp_group_ag=dp_cp_ag,
+ expt_fsdp_group_ag=expt_dp_ag,
)
else:
- if ep_group is not None:
+ if self.num_moe_experts is not None:
expt_mesh = _get_dp_tp_mesh(expt_dp_group, expt_tp_group, ep_size=ep_group.size())
expt_device_mesh = DeviceMesh.from_group(
[expt_dp_group, expt_tp_group],
@@ -335,6 +437,8 @@ def _init_dist_index(self, pg_collection):
dp_shard_dim="dp_cp",
tp_dim="tp",
expt_device_mesh=expt_device_mesh,
+ fsdp_group_ag=dp_cp_ag,
+ expt_fsdp_group_ag=expt_dp_ag,
)
self.tp_group = tp_group
diff --git a/megatron/core/distributed/fsdp/src/README.md b/megatron/core/distributed/fsdp/src/README.md
index dc984967e88..d3422d03abb 100644
--- a/megatron/core/distributed/fsdp/src/README.md
+++ b/megatron/core/distributed/fsdp/src/README.md
@@ -1,6 +1,6 @@
-# 🚀 Megatron-FSDP
+# Megatron-FSDP
@@ -12,38 +12,16 @@
## ✨ What is Megatron-FSDP?
-**Megatron-FSDP** is an NVIDIA-developed PyTorch extension that provides a high-performance implementation of Fully Sharded Data Parallelism (FSDP). It offers seamless cross-compatibility with major deep learning frameworks and parallelism libraries, making it easy to scale your PyTorch models across multiple GPUs and nodes.
+**Megatron-FSDP** is an NVIDIA-developed distributed parallelism library written in native PyTorch that provides a high-performance implementation of **Fully Sharded Data Parallelism (FSDP)**. It offers seamless cross-compatibility with various deep learning frameworks and parallelism libraries such as Megatron-Core, and is performance-optimized to support training and inference of extremely large PyTorch models at data-center scale on NVIDIA GPUs.
-Megatron-FSDP can provide up to 25% speed up and 23% memory savings compared to FSDP2.
+For comprehensive information about Megatron-FSDP, refer to: [Megatron-FSDP | Megatron-Core Developer Guide](https://docs.nvidia.com/megatron-core/developer-guide/latest/)
-### Compatibility
+### 🧩 Compatibility
-- **[PyTorch DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html)**
+- PyTorch **[DeviceMesh](https://docs.pytorch.org/docs/stable/distributed.html#devicemesh)**, **[DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html)**, and **[Distributed Checkpoint (DCP)](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html)**
- **[Megatron Core](https://github.com/NVIDIA/Megatron-LM)**
- **[TransformerEngine](https://github.com/NVIDIA/TransformerEngine)**
-
-## ✨ Features
-
-- **Easy Integration**: Simple `fully_shard` function for quick model parallelization
-- **High Performance**: Optimized for NVIDIA GPUs with efficient memory management
-- **Cross-Framework**: Works seamlessly with PyTorch, Huggingface Transformers, Megatron-LM, Megatron Bridge and TransformerEngine
-- **Scalable**: Supports both single-node multi-GPU and multi-node distributed training
-- **Flexible Configuration**: Configurable sharding strategies and process groups
-
-## ⚡ Optimizations
-
-- **Advanced Bucketing**: Data-type aware bucketing system to minimize the overhead of collective operations
-- **Buffer Management**: Zero copy communication is achieved by reorganizing the storage of parameters and main grad with `ParamAndGradBuffer` class
-- **Communication Overlapping**: Improved communication overlap of paramter all-gather and gradient reduce-scatter
-- **FP8 Mixed Precision with Transformer Engine**: Compatibility with Transformer Engine enables efficient FP8 mixed precision training
-- **Gradient accumulate fusion support with Transformer Engine**: Remove the explicit gradient copy to the communication buffer in backwards pass
-
-### Advanced Collective Communication
-- **SM Usage Reduction with SHARP**: FSDP's `All-Gather` (AG) and `Reduce-Scatter` (RS) collectives are designed to overlap with compute kernels. However, standard NCCL communication kernels can consume a significant number of GPU SMs (e.g., 16-32 SMs), "stealing" resources from compute (GEMM) kernels and reducing overall TFLOPS.
-- **In-Switch Processing**: We leverage **SHARP** (Scalable Hierarchical Aggregation and Reduction Protocol) to offload these collective operations. SHARP performs aggregation and reduction computations directly on the network switches (InfiniBand or NVLink Switch) instead of on the GPU SMs. This dramatically reduces the SM consumption for communication to **1-6 SM** freeing up GPU resources for compute. It also provides lower communication latency, especially in large, scaled-out workloads.
-- **Symmetric Optimizations for MNNVL**: We support **symmetric-based optimizations**, introduced in NCCL v2.27, which enable switch offloading for **Multi-Node NVLink (MNNVL)** systems such as GB200/GB300. This allows the same SM-saving benefits over the high-bandwidth NVLink fabric itself.
-- **Hierarchical Collectives**: When an FSDP sharding domain spans both NVLink and InfiniBand, the library utilizes **hierarchical SHARP collectives** (e.g., NVL-SHARP + IB-SHARP) to optimize the communication path across the entire system topology.
-
+- **[NVIDIA NeMo Framework Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo)**
## 📦 Installation
@@ -56,226 +34,57 @@ pip install megatron-fsdp
## 🚀 Quick Start
-### Basic Usage
-
-Transform your PyTorch model to use Fully Sharded Data Parallelism with just a few lines:
-
-```python
-import torch
-from megatron_fsdp import (
- fully_shard_model,
- fully_shard_optimizer,
-)
-
-"""
-Enable FSDP with Megatron-FSDP via the `fully_shard_*` API.
-"""
-# Shard your model.
-model = fully_shard_model(
- model,
- fsdp_unit_modules=[
- YourModelLayerClass,
- "import.path.to.model.class.YourModelLayerClass",
- ],
- ...
-)
-# Shard your optimizer.
-optimizer = fully_shard_optimizer(
- torch.optim.Adam(model.parameters(), lr=1e-3)
-)
-
-# Your model is now ready for distributed training!
-```
-
-### Comparison with FSDP-2
-
-`fully_shard` / `fully_shard_model` / `fully_shard_optimizer` are simple entrypoints into `MegatronFSDP`.
-
-- No need to call `fully_shard` on all the sub-modules, just pass your sub-module classes or import paths to `fully_shard`!
-- Seamlessly preserves the identity of your training loop with only a few lines of code and multiple options for initialization:
- - `fully_shard_*` is a two-line change when sharding the model and optimizer separately.
- - `fully_shard` is a one-line change for previously-initialized models and optimizers.
-
-Compare this with FSDP2:
-
-```python
-import torch
-from torch.distributed.fsdp import fully_shard
-
-# Your existing model and optimizer.
-model = YourModel()
-optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
-
-# Enable FSDP with FSDP2.
-for module in model.modules():
- # Sub-Modules to shard.
- if isinstance(module, YourModelLayerClass):
- fully_shard(module)
-fully_shard(model)
-
-# Your model is now ready for distributed training!
-```
-
-### `torch.compile` Compatibility
-
-Megatron-FSDP is compatible with `torch.compile`, but this feature is still experimental and may introduce performance regressions in some workloads.
-
-## 📖 Megatron-FSDP Comprehensive Walkthrough
-
-### Import `megatron_fsdp`.
-
```python
import torch
from megatron_fsdp import (
fully_shard_model,
fully_shard_optimizer,
- MixedPrecisionPolicy,
-)
-```
-
-### Set up a distributed environment using `DeviceMesh`.
-
-`DeviceMesh` simplifies the construction of complex arrangements of devices
-to support various parallelisms.
-
-```python
-from torch.distributed.device_mesh import DeviceMesh
-
-# Initialize DeviceMesh.
-device_mesh = torch.distributed.device_mesh.init_device_mesh(
- "cuda",
- mesh_shape=(dp_outer_size, dp_shard_size, cp_size, tp_size),
- mesh_dim_names=("dp_outer", "dp_shard", "cp", "tp"),
-)
-# Only relevant when using HSDP, where we also need the full DP group for data parallelism,
-# This sub-mesh can be provided to distributed samplers or dataloaders.
-device_mesh[("dp_outer", "dp_shard")]._flatten("dp")
-# Only required if using CP. Otherwise, just pass dp_shard to FSDP.
-device_mesh[("dp_shard", "cp")]._flatten("dp_shard_cp")
-# Only required if using HSDP. Otherwise, don't pass hybrid_fsdp_group.
-device_mesh[("dp_outer", "dp_shard", "cp")]._flatten("hsdp")
-hsdp_group = device_mesh["hsdp"].get_group()
-
-# Initialize DeviceMesh for expert parallel (EP) modules when using FSDP + EP.
-expert_device_mesh = torch.distributed.device_mesh.init_device_mesh(
- "cuda",
- mesh_shape=(dp_outer_size, expt_dp_shard_size, expt_tp_size),
- mesh_dim_names=("dp_outer", "dp_shard_cp", "tp"),
-)
-expert_device_mesh[("dp_outer", "dp_shard_cp")].flatten("hsdp")
-hsdp_expt_group = expert_device_mesh["hsdp"].get_group()
-```
-
-### Convert models into fully-sharded `MegatronFSDP` models with `fully_shard_model`.
-
-This wraps the model in a MegatronFSDP class that schedules the sharding
-lifecycle of the model parameters and gradients during training and inference.
-
-```python
-model = fully_shard_model(
- # PyTorch (Root) Module
- model,
- # Sharded Modules
- fsdp_unit_modules=[...],
- # Device Mesh
- device_mesh=device_mesh
- # Always required for FSDP or HSDP.
- dp_shard_dim="dp_shard_cp",
- # Set this required argument to use HSDP instead of FSDP. Otherwise, set this to None.
- dp_outer_dim="dp_outer",
- # Only required for TP-sensitive models (i.e. Megatron-LM / TransformerEngine)
- # or when using DTensor-based TP. Otherwise, set this to None.
- tp_dim="tp",
- # Only required when using HSDP. Otherwise, set this to None.
- hybrid_fsdp_group=hsdp_group,
- # Only required when using HSDP + EP. Otherwise, set this to None.
- hybrid_fsdp_expt_group=hsdp_expt_group,
- # Only required for FSDP + EP. Otherwise, set this to None.
- expt_device_mesh=expt_device_mesh,
- # FSDP Sharding Strategy: no_shard (0) / optim (1) / optim_grads (2) / optim_grads_params (3)
- zero_dp_strategy=3,
- outer_dp_sharding_strategy=1,
- # Initialize the model on devices in shards to avoid OOM. Requires device("meta")-init for model.
- init_model_with_meta_device=True,
- # Mixed-Precision Policy for controlling compute and communication precision in Megatron-FSDP.
- mixed_precision_policy=MixedPrecisionPolicy(),
- # Sync parameters and gradients each step. Allows for gradient transformations after backward pass,
- # and synchronizes parameters and gradients across HSDP groups, but deactivates compute-communication
- # overlap going into the subsequent training step.
- sync_model_each_microbatch=True,
- # Preprocess state dict for DCP checkpointing. Required for Torch Distributed Checkpoint.
- preproc_state_dict_for_dcp_ckpt=True,
-)
-```
-
-The original `torch.nn.Module` can be accessed at `MegatronFSDP.module`.
-
-### Initialize and fully-shard your optimizer on the `MegatronFSDP` model.
-
-Initialize your optimizer on the Megatron-FSDP model distributed `Parameter`(s).
-If your optimizer has already been initialized, either use the `fully_shard`
-entrypoint, or use `optimizer.add_param_group({"params": model.parameters()})`
-after resetting your optimizer state via `optimizer.param_groups.clear()`
-and `optimizer.state.clear()`.
-
-```python
-optimizer = torch.optim.Optimizer(model.parameters())
-```
-
-`fully_shard_optimizer` modifies your `optimizer.step()`, `optimizer.zero_grad()`,
-and distributed optimizer parameters to punctually trigger scheduled FSDP operations
-for Megatron-FSDP.
-
-```python
-fully_shard_optimizer(
- # PyTorch Optimizer
- optimizer,
- # Preprocess state dict for DCP checkpointing.
- # Required for Torch Distributed Checkpoint.
- preproc_state_dict_for_dcp_ckpt=True,
)
-```
-
-Extended arguments to `step()` and `zero_grad()` control these FSDP operations:
-```python
- optimizer.step(
- ...,
- # Sync all gradients before the optimizer step. Alternatively enabled using
- # `sync_model_each_microbatch=True` in MegatronFSDP.
- sync_grad_before_optimizer_step=True,
- # After `optimizer.step()`, install optimized weights into MegatronFSDP's buffers.
- install_optimized_model_weights=True,
- )
-
- optimizer.zero_grad(
- ...,
- # Also zero out MegatronFSDP's gradient accumulation buffers.
- zero_grad_buffer=True
- )
-```
-
-### `MegatronFSDP` Distributed Checkpointing
+# Initialize Torch Distributed.
+torch.distributed.init_process_group()
+torch.cuda.set_device(torch.distributed.get_rank())
-Distributed checkpoints can be saved and loaded using Torch DCP. Alternatively,
-you can load non-distributed checkpoints before fully-sharding your model with
-any existing checkpoint utility compatible with PyTorch Modules.
-
-```python
-# Save model and optimizer state.
-torch.distributed.checkpoint.save(
- {"model": model.state_dict(), "optimizer": optimizer.state_dict()},
- checkpoint_id=str(CKPT_DIR)
+# Fully-shard the model.
+model = torch.nn.Transformer()
+fsdp_model = fully_shard_model(
+ module=model,
+ fsdp_unit_modules=[
+ torch.nn.TransformerEncoder,
+ torch.nn.TransformerDecoder
+ ]
)
-# Load model and optimizer state.
-ckpt_state_dict = {"model": model.state_dict(), "optimizer": optimizer.state_dict()}
-torch.distributed.checkpoint.load(state_dict=ckpt_state_dict, checkpoint_id=str(CKPT_DIR))
-# `model.load_state_dict(strict=False)` is only necessary to ignore TE FP8 extra state
-# that is missing from the DCP checkpoint but present in TEBaseModule.
-# Megatron-FSDP does not support TE FP8 extra state checkpointing with DCP.
-model.load_state_dict(ckpt_state_dict["model"], strict=False)
-optimizer.load_state_dict(ckpt_state_dict["optimizer"])
+# Fully-shard the optimizer.
+toy_adam = torch.optim.AdamW(params=fsdp_model.parameters(), lr=0.01)
+optimizer = fully_shard_optimizer(optimizer=toy_adam)
+
+# Forward pass.
+inp = torch.randn(1, 512, 512).to("cuda")
+tgt = torch.randn(1, 512, 512).to("cuda")
+output = fsdp_model(inp, inp)
+
+# Backward pass.
+torch.nn.functional.mse_loss(output, tgt).backward()
+
+# Optimizer step.
+optimizer.step()
+optimizer.zero_grad()
+
+# Checkpoint the model and optimizer.
+torch.distributed.checkpoint.save({
+ "model": fsdp_model.state_dict(),
+ "optimizer": optimizer.state_dict(),
+}, checkpoint_id="ckpt/")
+
+# Load the saved checkpoint.
+ckpt = {
+ "model": fsdp_model.state_dict(),
+ "optimizer": optimizer.state_dict(),
+}
+torch.distributed.checkpoint.load(state_dict=ckpt, checkpoint_id="ckpt/")
+fsdp_model.load_state_dict(ckpt["model"], strict=False)
+optimizer.load_state_dict(ckpt["optimizer"])
```
## ⚙️ `fully_shard` / `MegatronFSDP` API - Advanced Features
@@ -305,17 +114,17 @@ Megatron-FSDP's `fully_shard_*` API has a comprehensive set of arguments for fin
- Defaults to `False`.
- Note that the `device` argument which installs your model on a specific device or rank will be deactivated when `init_model_with_meta_device=True`.
- `mixed_precision_policy` takes a `megatron_fsdp.MixedPrecisionPolicy` that configures mixed-precision compute and communication for Megatron-FSDP. Configuration options include:
- - `main_params_dtype` controls the data-type for parameters used in distributed optimization or quantization.
+ - `main_params_dtype` controls the data-type for parameters responsible for distributed checkpointing, distributed optimization, and quantization.
- Defaults to `torch.float32`.
- If set to `None`, the native model compute parameter data-type will be utilized.
- - Requires specification (cannot be `None`) when using `FP8` parameters with Megatron-FSDP.
+ - Requires specification (cannot be `None`) when using quantized parameters with Megatron-FSDP.
- `main_grads_dtype` controls the data-type for gradients used in distributed optimization.
- - Defaults to `None`, the model native gradient data-type will be utilized.
+ - Defaults to `None`, in which the model native gradient data-type will be utilized.
- While `torch.float32` (or higher) is recommended for accuracy at scale, as `main_grads_dtype` controls the data-type for gradient accumulation, `None` is more flexible and uses pre-determined parameter gradient logic in mixed-precision scenarios, such as `BF16` for `FP8`/`FP4` parameters quantized via TransformerEngine.
- - `grad_comm_dtype` controls the data-type for gradient communications (RS / AR) when reducing gradients. Lower precision `grad_comm_dtype` improves (communication) performance, but may increase memory utilization or sacrifice gradient precision in certain cases.
- - Defaults to `None`, the `main_grads_dtype` data-type will be utilized, and no additional memory is allocated when `grad_comm_dtype == main_grads_dtype`.
- - If using HSDP (either DP-Replicate or DP-Outer in `outer_dp_sharding_strategy`), `no_shard`, `optim`, or a `FixedPoolAllocator` (`fsdp_double_buffer`), allocating `dtype`-custom gradient communication buffers (per FSDP group) adds memory overhead of up to 10% or more, and users should consider the performance-memory trade-off when using this feature.
- - If using NCCL UBR v2.27+ (`nccl_ub=True`), gradient reduction may be performed in high-precision depending on the network domain (NVLink or IB), and can enable mixed-precision communication and accumulation, e.g. setting grad_comm_dtype to `BF16` can support `FP32` reduction even though we have `BF16` input and output communication buffers. Otherwise, gradients will be reduced in `grad_comm_dtype` (and accumulated in `main_grads_dtype`) as usual.
+ - `grad_comm_dtype` controls the data-type for gradient communications when reducing gradients. Lower precision `grad_comm_dtype` improves (communication) performance, but may increase memory utilization or sacrifice gradient precision in certain cases.
+ - Defaults to `None`, in which the `main_grads_dtype` data-type will be utilized. No additional memory is allocated when `grad_comm_dtype == main_grads_dtype`.
+ - If using HSDP (either DP-Replicate or DP-Outer in `outer_dp_sharding_strategy`), `no_shard`, or `optim`, allocating `dtype`-custom gradient communication buffers may increase per-unit memory overhead, so users should consider the performance-memory trade-off when using this feature.
+ - If using NCCL user buffer registration `v2.27+`, gradient reduction may be performed in high-precision depending on the network domain (NVLink or IB), and can enable mixed-precision communication and accumulation, e.g. setting grad_comm_dtype to `BF16` can support `FP32` reduction even though we have `BF16` input and output communication buffers. Otherwise, gradients will be reduced in `grad_comm_dtype` (and accumulated in `main_grads_dtype`) as usual.
- `overlap_grad_reduce` and `overlap_param_gather` will overlap gradient [`reduce-scatter`](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#reducescatter) and parameter [`all-gather`](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allgather) group communications with backward and forward compute with asynchronous calls and pre-fetching. (In the case of `no_shard`, parameters are not gathered but gradient [`all-reduce`](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html#allreduce) is overlapped.)
- Both default to `True`.
- `sync_model_each_microbatch` will trigger a `wait` (`MegatronFSDP.finish_grad_sync()`) on gradient reduction, parameter de-allocation, and optimizer parameter / gradient installation (in preparation for `optimizer.step()`) after every forward-backward pass. When using HSDP, parameters and gradients will be all-gathered and reduced respectively on the "outer" DP group each training step instead of each optimization cycle. This behavior is desirable for a transparent and user-friendly sharded training loop where post-backward transformations on the gradient and a clean compute / memory state are necessary within and between training iterations, but damages performance in situations where optimization is delayed (e.g. gradient accumulation) when the communications of the previous training iteration can be overlapped with the compute of the next training iteration. Will also override `is_last_microbatch` / `microbatch_count` logic in `MegatronFSDP`.
@@ -326,14 +135,29 @@ Megatron-FSDP's `fully_shard_*` API has a comprehensive set of arguments for fin
- Defaults to `False`.
- `keep_fp8_transpose_cache` will keep the fp8 transpose cache when using `MegatronFSDP`. This option will cause (number of parameter $\times$ 1 Byte) of memory overhead, but can skip the weight transpose operation in the backward propagation. This feature will not give any benefit from the Blackwell architecture.
- Defaults to `False`.
+- `use_decoupled_grad` installs the reduced gradient into a separate buffer: `Parameter.decoupled_grad`. This buffer is utilized by specific optimizers, such as TransformerEngine's `FusedAdam`, and can be used to temporarily store your gradient for custom `torch.nn.Optimizer`(s).
+ - Defaults to `False`.
+ - Required for `transformer_engine.pytorch.optimizers.FusedAdam`.
- `nccl_ub` will allocate and register the NCCL userbuffer for param and grad buffers. This option enables an SM-efficient NCCL algorithm that could improve the performance of overlapped computations. This flag will be much more effective when used together with SHARP if the FSDP communication includes both NVL and IB domains. Enabling this option will cause additional memory overhead due to the requirement to enable the `fsdp_double_buffer` option.
- **Only effective when using with Megatron-Core.**
- Defaults to `False`.
- By default we try to use NCCL window (symmetric) registration if it is available. If not it falls back to conventional local registration.
-- `fsdp_manual_registration` will manually register the FSDP communication buffers with the NCCL user buffer. For symmetric registration with large models, the registration itself can take a significant amount of time. This option minimizes the number of registration calls to reduce the registration time. However, with this option enabled, you need to manually call the `ParamAndGradBuffer.manual_buffer_registration()` function after the first iteration. This is already implemented in the Megatron-LM training loop. In other use cases, users are expected to call this function themselves.
+- `fsdp_manual_registration` will manually register the FSDP communication buffers with the NCCL user buffer. For symmetric registration with large models, the registration itself can take a significant amount of time. This option minimizes the number of registration calls to reduce the registration time. However, with this option enabled, you need to manually call the `ParamAndGradBuffer.manual_buffer_registration()` function after the first iteration. This is already implemented in the Megatron-LM training loop. In other use cases, users are expected to call this function themselves.
+ - This is an example of required modification in the training loop.
+ ```python
+ def train(...):
+ ...
+ # After the first iteration, user need to call the
+ # ParamAndGradBuffer.manual_buffer_registration() function in the training loop
+ if (iteration == start_iteration + 1):
+ if isinstance(model, megatron_FSDP) and model.ddp_config.fsdp_manual_registration:
+ param_and_grad_buffer = getattr(model, "param_and_grad_buffer", None)
+ if param_and_grad_buffer is not None:
+ param_and_grad_buffer.manual_buffer_registration()
+ ```
- **Only effective when using with Megatron-Core.**
- This option is only effective when `nccl_ub` is enabled.
- - Defaults to `False`.
+ - Defaults to `False`, but will be automatically enabled in Megatron-LM.
- `disable_symmetric_registration` will disable NCCL window (i.e. symmetric) registration when using `nccl_ub`.
- Defaults to `False`.
- `fsdp_double_buffer` will use persistently allocated double buffers for temporarily-defined memory needed in `MegatronFSDP` communications. Having persistent double buffers may increase peak VRAM utilization, but is required to register NCCL user buffers (`nccl_ub=True`) for `MegatronFSDP`. Currently, this is only supported for simple repetitive model structures such as GPT.
@@ -347,7 +171,7 @@ Megatron-FSDP natively supports mixed-precision activations and parameter shardi
- Within the [`transformer_engine.pytorch.autocast(recipe: transformer_engine.common.recipe.Recipe)`](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/pytorch.html#transformer_engine.pytorch.autocast) context, model activations are converted based on the recipe.
- Within the [`transformer_engine.pytorch.quantized_model_init(recipe: transformer_engine.common.recipe.Recipe)`](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/pytorch.html#transformer_engine.pytorch.quantized_model_init) context, TransformerEngine native modules (e.g. [`transformer_engine.pytorch.TransformerLayer`](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/pytorch.html#transformer_engine.pytorch.TransformerLayer)) have their parameters converted based on the recipe.
- - Requires FP8 model activations, i.e. `transformer_engine.pytorch.autocast`.
+ - Requires quantized model activations, i.e. `transformer_engine.pytorch.autocast`.
```python
# FP8 Recipe
@@ -382,4 +206,4 @@ with transformer_engine.pytorch.autocast(recipe=fp8_recipe):
mfsdp_model(x).sum().backward()
```
-ℹ️ `TransformerEngine` kernels have a fair bit of configuration constraints when using FP8-quantized parameters, such as using fused QKV parameters or defining activations and parameters with shapes compatible to FP8 CuBLAS kernels on supported hardware from NVIDIA. To properly initialize `TransformerLayer`, you can refer to the toy model used in our FP8 unit tests: `Megatron-LM/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py::TestMegatronFsdpFullyShard::test_fully_shard_te_quantized`.
\ No newline at end of file
+ℹ️ `TransformerEngine` kernels have various constraints related to quantized Tensors, such as using fused QKV parameters or defining activations and parameters with shapes compatible to CuBLAS kernels on supported hardware from NVIDIA. To properly initialize `TransformerLayer`, you can refer to the example model used in our unit tests: `Megatron-LM/tests/unit_tests/distributed/fsdp/test_mfsdp_fully_shard.py::TestMegatronFsdpFullyShard::test_fully_shard_te_quantized`.
\ No newline at end of file
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py
index 4ad5a8dddac..8947c8fe174 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py
@@ -5,6 +5,8 @@
import torch
+from .utils import is_torch_min_version
+
@dataclass
class DistributedDataParallelConfig:
@@ -86,7 +88,9 @@ class DistributedDataParallelConfig:
If True, use all-gather during the initial Megatron-FSDP parameter
synchronization step. This can increase overlap between the first
parameter all-gather and computation, helping to better hide the
- initial communication cost.
+ initial communication cost. Should be deactivated when using
+ full-iteration CG, or partial CG if AG/RS is launched beyond the
+ CG capture scope but is waited on during the capture scope.
"""
fsdp_db_use_persist_buf_on_alloc_fail: bool = False
@@ -145,11 +149,39 @@ class DistributedDataParallelConfig:
No additional memory is allocated when `grad_comm_dtype == main_grads_dtype`.
"""
+ megatron_fsdp_use_decoupled_grad: bool = False
+ """If true, Megatron-FSDP's ParamAndGradBuffer uses the precision-aware optimizer
+ gradient path (e.g. `decoupled_grad` on optimizer parameters) instead of casting
+ main gradients to parameter dtype for `.grad`.
+ """
+
+ megatron_fsdp_cuda_graph_mode: bool = False
+ """If set to True, Megatron-FSDP will practice CUDA graph-safe operations, such as
+ not dereferencing `param.grad` after the optimizer step to preserve references for
+ CUDA graph replay. Can affect memory utilization in some cases, such as when the
+ gradient shard is not a view of the Megatron-FSDP sharded gradient buffer, so
+ FusedAdam(use_decoupled_grad=True) + megatron_fsdp_use_decoupled_grad=True or
+ setting megatron_fsdp_main_params_dtype == megatron_fsdp_main_grads_dtype is
+ recommended to avoid casting the gradient to the parameter precision and creating
+ a casted-copy of the gradient shard that cannot be dereferenced due to replay.
+ """
+
+ megatron_fsdp_enable_fine_grained_param_gather: bool = False
+ """If set to True, enables fine-grained parameter gathering for Megatron-FSDP.
+ This feature increases the overlap between parameter all-gather and forward computation,
+ at the cost of more frequent communication calls.
+ For MXFP8, this approach helps save memory during fine-grained activation
+ recomputation, because MXFP8 forward and backward passes use different
+ parameter representations (rowwise data for forward, colwise data for backward).
+ In this mode, only the rowwise parameters of modules involved in recomputation
+ will be unsharded.
+ """
+
def __post_init__(self):
import os
"""Check the validity of the config."""
- if self.nccl_ub:
+ if self.nccl_ub and not is_torch_min_version("2.11.0a0"):
if 'expandable_segments:True' in os.getenv('PYTORCH_CUDA_ALLOC_CONF', '').split(','):
raise ValueError(
"PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True is currently not supported "
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py
index ee6a6013b2d..8b87899c234 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py
@@ -81,6 +81,8 @@ def fully_shard_model(
hybrid_fsdp_group: Optional[torch.distributed.ProcessGroup] = None,
hybrid_fsdp_expt_group: Optional[torch.distributed.ProcessGroup] = None,
expt_device_mesh: Optional[DeviceMesh] = None,
+ fsdp_group_ag: Optional[torch.distributed.ProcessGroup] = None,
+ expt_fsdp_group_ag: Optional[torch.distributed.ProcessGroup] = None,
fsdp_unit_modules: Optional[Sequence[Type[torch.nn.Module]] | Sequence[str]] = None,
zero_dp_strategy: str | int = 3,
outer_dp_sharding_strategy: str | int = 0,
@@ -101,6 +103,8 @@ def fully_shard_model(
fsdp_db_use_persist_buf_on_alloc_fail: bool = False,
disable_symmetric_registration: bool = False,
enable_fine_grained_param_gather: bool = False,
+ use_decoupled_grad: bool = False,
+ cuda_graph_mode: bool = False,
) -> torch.nn.Module:
"""
Fully-shard the model for Megatron-FSDP. This wraps the model in a MegatronFSDP
@@ -142,6 +146,17 @@ class that schedules the sharding lifecycle of the model parameters and gradient
Expert parallel device mesh object defining the topology for MoE distributed training.
Utilizes the mesh dimension names specified by the *_dim arguments.
+ fsdp_group_ag (Optional[torch.distributed.ProcessGroup]):
+ Independent all-gather process group for overlapping all-gather and reduce-scatter
+ operations. When provided, enables AG/RS overlap optimization for regular (non-expert)
+ parameters. Users should create this group with the same ranks as the dp-cp group.
+ Defaults to None.
+
+ expt_fsdp_group_ag (Optional[torch.distributed.ProcessGroup]):
+ Independent all-gather process group for expert parameters in MoE models. When provided,
+ enables AG/RS overlap optimization for expert parameters. Users should create this group
+ with the same ranks as the expert data parallel group. Defaults to None.
+
fsdp_unit_modules (Optional[Sequence[Type[torch.nn.Module]] | Sequence[str]]):
List of (sub-)module classes or (sub-)module class import paths that are "units",
which are torch.nn.Module(s) that are sharded and scheduled by Megatron-FSDP.
@@ -247,6 +262,21 @@ class that schedules the sharding lifecycle of the model parameters and gradient
unshards parameters per-Module instead of unsharding all sub-modules of an FSDP
unit module simultaneously. Defaults to False.
+ use_decoupled_grad (bool):
+ If true, reduced gradients are installed into `Parameter.decoupled_grad` instead
+ of `Parameter.grad`. Defaults to False.
+
+ cuda_graph_mode (bool):
+ If true, Megatron-FSDP will practice CUDA graph-safe operations, such as
+ not dereferencing `param.grad` after the optimizer step to preserve references
+ for CUDA graph replay. Can affect memory utilization in some cases, such as
+ when the gradient shard is not a view of the Megatron-FSDP sharded gradient
+ buffer, so `FusedAdam(use_decoupled_grad=True) + use_decoupled_grad=True` or
+ setting `megatron_fsdp_main_params_dtype == megatron_fsdp_main_grads_dtype`
+ is recommended to avoid casting the gradient to the parameter precision and
+ creating a casted-copy of the gradient shard that cannot be dereferenced due
+ to replay. Defaults to False.
+
Returns:
model (MegatronFSDP): The wrapped Megatron-FSDP model configured for FSDP.
"""
@@ -341,6 +371,8 @@ class that schedules the sharding lifecycle of the model parameters and gradient
fsdp_double_buffer=fsdp_double_buffer or nccl_ub,
fsdp_db_use_persist_buf_on_alloc_fail=fsdp_db_use_persist_buf_on_alloc_fail,
disable_symmetric_registration=disable_symmetric_registration,
+ megatron_fsdp_use_decoupled_grad=use_decoupled_grad,
+ megatron_fsdp_cuda_graph_mode=cuda_graph_mode,
)
# Create FSDPDistributedIndex.
@@ -362,6 +394,9 @@ class that schedules the sharding lifecycle of the model parameters and gradient
hsdp_outer_dp_shard=_outer_fsdp_sharding,
# Only required for Megatron-FSDP + EP.
expt_device_mesh=expt_device_mesh,
+ # AG groups for AG/RS overlap optimization.
+ fsdp_group_ag=fsdp_group_ag,
+ expt_fsdp_group_ag=expt_fsdp_group_ag,
)
# Wrap model in Megatron FSDP.
@@ -621,6 +656,8 @@ def fully_shard(
hybrid_fsdp_group: Optional[torch.distributed.ProcessGroup] = None,
hybrid_fsdp_expt_group: Optional[torch.distributed.ProcessGroup] = None,
expt_device_mesh: Optional[DeviceMesh] = None,
+ fsdp_group_ag: Optional[torch.distributed.ProcessGroup] = None,
+ expt_fsdp_group_ag: Optional[torch.distributed.ProcessGroup] = None,
fsdp_unit_modules: Optional[Sequence[Type[torch.nn.Module]] | Sequence[str]] = None,
zero_dp_strategy: str | int = 3,
outer_dp_sharding_strategy: str | int = 0,
@@ -641,6 +678,8 @@ def fully_shard(
fsdp_db_use_persist_buf_on_alloc_fail: bool = False,
disable_symmetric_registration: bool = False,
enable_fine_grained_param_gather: bool = False,
+ use_decoupled_grad: bool = False,
+ cuda_graph_mode: bool = False,
) -> tuple[MegatronFSDP, torch.optim.Optimizer]:
"""
Fully shard the model and the optimizer for Megatron-FSDP.
@@ -669,6 +708,8 @@ def fully_shard(
hybrid_fsdp_group=hybrid_fsdp_group,
hybrid_fsdp_expt_group=hybrid_fsdp_expt_group,
expt_device_mesh=expt_device_mesh,
+ fsdp_group_ag=fsdp_group_ag,
+ expt_fsdp_group_ag=expt_fsdp_group_ag,
fsdp_unit_modules=fsdp_unit_modules,
zero_dp_strategy=zero_dp_strategy,
outer_dp_sharding_strategy=outer_dp_sharding_strategy,
@@ -688,7 +729,8 @@ def fully_shard(
fsdp_double_buffer=fsdp_double_buffer,
fsdp_db_use_persist_buf_on_alloc_fail=fsdp_db_use_persist_buf_on_alloc_fail,
disable_symmetric_registration=disable_symmetric_registration,
- enable_fine_grained_param_gather=enable_fine_grained_param_gather,
+ use_decoupled_grad=use_decoupled_grad,
+ cuda_graph_mode=cuda_graph_mode,
)
# Extend optimizer methods to support Megatron-FSDP operations.
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
index f8640446814..6202601856e 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py
@@ -17,6 +17,7 @@
import logging
from contextlib import contextmanager
from enum import Enum, auto
+from functools import partial
from typing import Any, Dict, List, Optional, Tuple
import torch
@@ -46,11 +47,12 @@
try:
# Default to Megatron-LM FW.
- logger.info("Detected Megatron Core, using Megatron-FSDP with Megatron.")
from megatron.core.distributed.distributed_data_parallel_config import (
DistributedDataParallelConfig,
)
from megatron.core.utils import is_submodule
+
+ logger.info("Detected Megatron Core, using Megatron-FSDP with Megatron.")
except ImportError:
# Megatron-LM is not installed, use Megatron-FSDP as a standalone module.
logger.info("Megatron Core is not installed, Megatron-FSDP will run without Megatron Core.")
@@ -73,6 +75,34 @@ class TrainingState(Enum):
IDLE = auto()
+def setup_delayed_wgrad_acc_hook(module, grad_acc_func):
+ """Configure delayed wgrad gradient processing for MoE expert parameters.
+
+ When ``overlap_dispatch_backward_with_experts_wgrad`` is enabled on a TransformerLayer,
+ this function:
+ 1. Marks expert parameters so the normal post-accumulate-grad hook is skipped.
+ 2. Registers a callback on the MoE layer that invokes FSDP's gradient
+ reduce-scatter after the delayed wgrad computation completes.
+
+ Args:
+ module: The module being processed in the forward pre-hook. Only
+ ``TransformerLayer`` instances with the delayed wgrad config flag
+ enabled are affected; all other modules are no-ops.
+ process_post_backward_gradients_fn: The FSDP gradient processing function
+ (``_process_post_backward_gradients``) to be called after the delayed
+ wgrad computation finishes.
+ """
+ from functools import partial
+
+ need_backward_dw = getattr(module, "need_backward_dw", lambda: False)
+ if not need_backward_dw():
+ return
+
+ for param in module.parameters():
+ if getattr(param, 'skip_backward_post_hook', False):
+ param.post_wgrad_grad_acc_hook = partial(grad_acc_func, [param])
+
+
class MegatronFSDP(torch.nn.Module):
"""Fully Sharded Data Parallel training.
@@ -186,6 +216,7 @@ def __init__(
fsdp_db_use_persist_buf_on_alloc_fail: bool = False,
disable_symmetric_registration: bool = False,
enable_fine_grained_param_gather_hook: bool = False,
+ enable_fine_grained_param_gather_backward_hook: bool = False,
report_nan_in_param_grad: bool = False,
):
super().__init__()
@@ -238,6 +269,9 @@ def __init__(
self.calculate_per_token_loss = calculate_per_token_loss
self.init_model_with_meta_device = init_model_with_meta_device
self.enable_fine_grained_param_gather_hook = enable_fine_grained_param_gather_hook
+ self.enable_fine_grained_param_gather_backward_hook = (
+ enable_fine_grained_param_gather_backward_hook
+ )
self.report_nan_in_param_grad = report_nan_in_param_grad
# FSDPDistributedIndex stores the process groups and meshes used by Megatron-FSDP.
@@ -338,10 +372,16 @@ def _init_fsdp_param_and_grad_buffer(self):
else:
if self.ddp_config.average_in_collective:
gradient_scaling_factor = 1.0
- expert_gradient_scaling_factor = (
- self.dist_index.get_dp_group(is_expert_parallel=True).size()
- / self.dist_index.get_dp_group().size()
- )
+ expert_dp_group = self.dist_index.get_dp_group(is_expert_parallel=True)
+ if expert_dp_group is None:
+ # Dense model (no expert-parallel params): the expert scaling factor is
+ # never applied, but it is computed eagerly here. Fall back to 1.0 instead
+ # of dereferencing the missing expert data-parallel group.
+ expert_gradient_scaling_factor = 1.0
+ else:
+ expert_gradient_scaling_factor = (
+ expert_dp_group.size() / self.dist_index.get_dp_group().size()
+ )
else:
data_parallel_world_size = self.dist_index.get_dp_group().size()
gradient_scaling_factor = 1.0 / data_parallel_world_size
@@ -565,8 +605,10 @@ def _grad_acc(param):
return
# Sharded Gradient Buffer
- gbuf = group.hsdp_gbuf if group.hsdp_gbuf else group.main_grad_buffer
+ gbuf = group.hfsdp_helper_gbuf if group.hfsdp_helper_gbuf else group.main_grad_buffer
if gbuf.is_data_distributed:
+ # If TransformerEngine gradient accumulation is fused, then param.get_main_grad()
+ # already holds the wgrad and param.grad_added_to_main_grad=True.
if not param.grad_added_to_main_grad:
# Get `main_grad` will allocate bucket, check that the currently
# used main_grad buffer does not exceed the scope of two FSDP Unit
@@ -583,7 +625,6 @@ def _grad_acc(param):
param.main_grad.copy_(to_local_if_dtensor(param.grad))
del param.grad
else:
- # Prepare for fused wgrad accumulation.
param.main_grad.zero_()
# Unsharded Gradient Buffer
else:
@@ -621,9 +662,11 @@ def _post_backward_release_module(module, *unused):
# Release parameters for this module after backward.
release_module_parameters(module, bwd=True)
+ release_module_parameters(module, bwd=False)
# Transition this module back to the IDLE training state.
- module._training_state = TrainingState.IDLE
+ for sub_module in module.modules():
+ sub_module._training_state = TrainingState.IDLE
@torch.compiler.disable
def _process_post_backward_gradients(param_list):
@@ -662,6 +705,17 @@ def _process_post_backward_gradients(param_list):
"""
# Filter out shared parameters whose gradients are handled by the root hook.
param_list = [p for p in param_list if not getattr(p, "_is_shared", False)]
+
+ # Make sure for delayed wgrad params, the grad_acc_hooks are registered.
+ for p in param_list:
+ if getattr(p, 'skip_backward_post_hook', False):
+ assert hasattr(
+ p, 'post_wgrad_grad_acc_hook'
+ ), "Missing grad accumulation hook for delayed_wgrad_compute param."
+
+ if not param_list:
+ return
+
for param in param_list:
_grad_acc(param)
@@ -690,18 +744,7 @@ def _process_post_backward_gradients(param_list):
self._params_require_handle_grad.discard(param)
@torch.compiler.disable
- def _pre_forward_param_unshard(
- module: nn.Module,
- args: Optional[Tuple[Any, ...]] = None,
- kwargs: Optional[Dict[str, Any]] = None,
- ):
- # If args or kwargs are not passed, default to () and {}.
- # This matches PyTorch Module hook conventions:
- # torch.nn.Module._call_impl.inner()
- if args is None:
- args = ()
- if kwargs is None:
- kwargs = {}
+ def _pre_forward_param_unshard(module: nn.Module, *unused):
# Unshard the parameters before the forward pass.
input_training_state = module._training_state
fsdp_forward_prefetch = True
@@ -728,14 +771,14 @@ def _pre_forward_param_unshard(
prefetch=fsdp_forward_prefetch,
prefetch_order=PrefetchOrder.FORWARD_PASS_ORDER,
)
- return args, kwargs
+ return None
@torch.compiler.disable
def _register_post_backward_hook(
post_backward_hook: callable,
module: nn.Module,
- args: Optional[Tuple[Any, ...]] = None,
- kwargs: Optional[Dict[str, Any]] = None,
+ args: Tuple[Any, ...],
+ kwargs: Dict[str, Any],
):
"""
Register a post-backward hook for the given module by inserting an autograd
@@ -744,13 +787,6 @@ def _register_post_backward_hook(
since such operations can trigger an autograd error that
"the output is a view and is being modified in-place".
"""
- # If args or kwargs are not passed, default to () and {}.
- # This matches PyTorch Module hook conventions:
- # torch.nn.Module._call_impl.inner()
- if args is None:
- args = ()
- if kwargs is None:
- kwargs = {}
if not torch.is_grad_enabled():
# No gradients / backward pass, don't attach the post-backward hook.
return args, kwargs
@@ -840,16 +876,14 @@ def _pre_backward_param_unshard(module: nn.Module, *unused):
before the backward pass.
"""
# Set the module's training state to PRE_BACKWARD.
- module._training_state = TrainingState.PRE_BACKWARD
+ for sub_module in module.modules():
+ sub_module._training_state = TrainingState.PRE_BACKWARD
if isinstance(module, tuple(fsdp_unit_modules)):
param_list = list(module.parameters())
else:
param_list = list(module.parameters(recurse=False))
- if self.enable_fine_grained_param_gather_hook:
- param_list = list(module.parameters(recurse=False))
-
# All-gather / unshard the module parameters before the backward pass.
self.all_gather_and_wait_parameters_ready(
param_list, prefetch_order=PrefetchOrder.BACKWARD_PASS_ORDER, bwd=True
@@ -857,7 +891,7 @@ def _pre_backward_param_unshard(module: nn.Module, *unused):
self._root_pre_backward_hook_issued = False
- def _root_pre_backward(module: nn.Module, *unused):
+ def _root_pre_backward(module: nn.Module, *unused, skip_backward_hook: bool = False):
"""Marks the module's training state as PRE_BACKWARD before the
backprop, this function is registered on the root module.
@@ -871,11 +905,10 @@ def _root_pre_backward(module: nn.Module, *unused):
self._root_pre_backward_hook_issued = True
if self.data_parallel_sharding_strategy == "optim_grads_params":
- for module in root_module.modules():
- if isinstance(module, tuple(fsdp_unit_modules)):
- # Set PRE_BACKWARD state to skip resharding and forward pre-fetching
- # when performing activation recomputation / gradient checkpointing.
- module._training_state = TrainingState.PRE_BACKWARD
+ for sub_module in root_module.modules():
+ # Set PRE_BACKWARD state to skip resharding and forward pre-fetching
+ # when performing activation recomputation / gradient checkpointing.
+ sub_module._training_state = TrainingState.PRE_BACKWARD
# set all param buckets can be released
ag_pipeline = self.all_gather_pipeline
for bucket_id in range(ag_pipeline.num_buckets):
@@ -894,6 +927,8 @@ def _root_pre_backward(module: nn.Module, *unused):
param.grad_added_to_main_grad = False
# Queue the root post-backward hook to reduce leftover gradients after
# the backward pass.
+ if skip_backward_hook:
+ return
torch.autograd.Variable._execution_engine.queue_callback(_root_post_backward)
@torch.compiler.disable
@@ -981,10 +1016,23 @@ def _register_pre_backward_param_unshard_hook(module):
create_custom_backward_hook(module, _pre_backward_param_unshard)
)
+ # These hooks need to be exposed for manual management by 1F1B Overlapping
+ # and triggered by 1F1B Overlapped execution pipeline, except for
+ # `param_unshard` hook that needs to be installed at param level,
+ # such that non-overlapped params like embedding layer are also correctly
+ # unsharded.
+ self.post_forward_release_module = partial(_post_forward, input=None, output=None)
+ self.post_backward_release_module = _post_backward_release_module
+ self.pre_backward = partial(_root_pre_backward, module=None, skip_backward_hook=True)
+ self.post_backward = _root_post_backward
+
fsdp_modules = []
for name, module in root_module.named_modules():
+ # Set post backward hook for TE grouped gemm if enabled comm overlap
+ setup_delayed_wgrad_acc_hook(module, _process_post_backward_gradients)
if self.enable_fine_grained_param_gather_hook:
_register_pre_forward_param_unshard_hook(module)
+ if self.enable_fine_grained_param_gather_backward_hook:
_register_pre_backward_param_unshard_hook(module)
# Skip if the module is already registered in fsdp_modules.
@@ -1003,8 +1051,7 @@ def _register_pre_backward_param_unshard_hook(module):
module.register_forward_hook(_post_forward, prepend=False)
)
- if not self.enable_fine_grained_param_gather_hook:
- _register_pre_backward_param_unshard_hook(module)
+ _register_pre_backward_param_unshard_hook(module)
elif (
not self.ddp_config.keep_fp8_transpose_cache
and self.data_parallel_sharding_strategy == "optim_grads_params"
@@ -1036,9 +1083,16 @@ def _register_pre_backward_param_unshard_hook(module):
]
for param in grad_acc_param_list:
+ # Only register grad acc hook for parameters that require gradients.
+ if not param.requires_grad:
+ continue
self.grad_acc_hooks[f"grad_acc and reduce for {self.param_to_name[param]}"] = (
param.register_post_accumulate_grad_hook(
- lambda p: _process_post_backward_gradients([p])
+ lambda p: (
+ None
+ if getattr(p, 'skip_backward_post_hook', False)
+ else _process_post_backward_gradients([p])
+ )
)
)
@@ -1195,6 +1249,9 @@ def start_param_sync(self, *unused, force_sync: bool = False, force_dispatch: bo
"""
self._replace_param_with_raw_if_needed()
+ if self.data_parallel_sharding_strategy == "no_shard":
+ return
+
if not force_sync and self.ddp_config.overlap_param_gather:
# All-gather the first bucket before the forward pass.
if self.ddp_config.fsdp_all_gather_in_start_param_sync:
@@ -1239,7 +1296,7 @@ def synchronize_param_gather(self):
"""
Synchronize parameter all-gather operations for all model parameters.
"""
- self.all_gather_pipeline.reset()
+ self.all_gather_pipeline.reset(preserve_non_fsdp_units=True)
self._replace_param_with_distributed_if_needed()
def synchronize_gradient_reduce(self):
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/mixed_precision.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/mixed_precision.py
index 89c67f40d41..935508a57ab 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/mixed_precision.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/mixed_precision.py
@@ -276,9 +276,22 @@ def fp8_quantize(
fsdp_shard_model_params = [x[0] if x[1] is None else x for x in fsdp_shard_model_params]
if HAVE_TE_CAST_MASTER_WEIGHTS_TO_FP8:
- cast_master_weights_to_fp8(
- model_params, main_params, start_offsets, data_parallel_group, fsdp_shard_model_params
- )
+ args = [
+ model_params,
+ main_params,
+ start_offsets,
+ data_parallel_group,
+ fsdp_shard_model_params,
+ ]
+
+ # For newer TE versions (i.e., have post_all_gather_processing function), we keep the
+ # columnwise data and manually call post_all_gather_processing after all-gather, this
+ # makes fp8 params compatible with CUDA graph.
+ kwargs = {}
+ if HAVE_TE_POST_ALL_GATHER_PROCESSING:
+ kwargs["manual_post_all_gather_processing"] = True
+
+ cast_master_weights_to_fp8(*args, **kwargs)
else:
_fp8_quantize_fallback(
model_params, main_params, start_offsets, data_parallel_group, fsdp_shard_model_params
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/package_info.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/package_info.py
index a3a282a01c0..c3c7fc18ca5 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/package_info.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/package_info.py
@@ -2,7 +2,7 @@
MAJOR = 0
-MINOR = 3
+MINOR = 5
PATCH = 0
PRE_RELEASE = 'rc0'
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
index 684cd7a99eb..690ec263890 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py
@@ -51,9 +51,9 @@
FSDPDistributedIndex,
get_global_memory_buffer,
get_mcore_tensor_parallel_partition_dim,
- is_mcore_tensor_model_parallel,
is_mcore_tensor_parallel_duplicated,
log_single_rank,
+ using_tensor_parallel,
)
logger = logging.getLogger(__name__)
@@ -875,11 +875,6 @@ def __init__(
# NOTE: Specifying dp_rank is a tricky thing. Currently, only full-shard
# hybrid FSDP needs to do this to set dp rank that is different from the group rank.
if dp_rank is not None:
- logger.warning(
- f"[FSDP] DataParallelBuffer[{bucket_id}] initialized with dp_rank={dp_rank}, "
- f"native dp_rank={torch.distributed.get_rank(data_parallel_group)}, "
- f"global_rank={torch.distributed.get_rank()}"
- )
self.dp_rank = dp_rank
else:
self.dp_rank = torch.distributed.get_rank(data_parallel_group)
@@ -890,6 +885,7 @@ def __init__(
self.is_transpose_buffer = is_transpose_buffer
self.gradient_scaling_factor = gradient_scaling_factor
self.mem_alloc_context = mem_alloc_context if mem_alloc_context else nullcontext
+ self.chunk_size_factor = chunk_size_factor
# Setup the item index map, bucket index, and shard bucket index from
# the provided arguments, or build them if not provided.
@@ -1318,12 +1314,15 @@ class ParameterGroup:
Buffer used to store main model weights for data-parallel operations.
main_grad_buffer (Optional[DataParallelBuffer]):
Buffer used to store main gradients for data-parallel operations.
- hsdp_wbuf (Optional[DataParallelBuffer]):
- Buffer for weights used in Hybrid Sharded Data Parallel (HSDP).
- Exists only if full sharding (HFSDP) is enabled in HSDP.
- hsdp_gbuf (Optional[DataParallelBuffer]):
- Buffer for gradients used in HSDP.
- Exists only if full sharding (HFSDP) is enabled in HSDP.
+ hfsdp_helper_wbuf (Optional[DataParallelBuffer]):
+ Inner-DP helper buffer that owns persistent HFSDP parameter shards.
+ Created only when Hybrid FSDP (optimizer-state) full sharding is enabled.
+ hfsdp_helper_wtbuf (Optional[DataParallelBuffer]):
+ Inner-DP helper buffer that stores transpose weights for FP8/MXFP8.
+ Created only when Hybrid FSDP (optimizer-state) full sharding is enabled.
+ hfsdp_helper_gbuf (Optional[DataParallelBuffer]):
+ Inner-DP helper buffer that owns persistent HFSDP gradient shards.
+ Created only when Hybrid FSDP (optimizer-state) full sharding is enabled.
hsdp_comm_gbuf (Optional[DataParallelBuffer]):
Extra buffer to allocate buffers that enable custom gradient
communication data-types when using HSDP or HFSDP only.
@@ -1341,8 +1340,9 @@ class ParameterGroup:
transpose_weight_buffer: Optional[DataParallelBuffer] = None
main_weight_buffer: Optional[DataParallelBuffer] = None
main_grad_buffer: Optional[DataParallelBuffer] = None
- hsdp_wbuf: Optional[DataParallelBuffer] = None
- hsdp_gbuf: Optional[DataParallelBuffer] = None
+ hfsdp_helper_wbuf: Optional[DataParallelBuffer] = None
+ hfsdp_helper_wtbuf: Optional[DataParallelBuffer] = None
+ hfsdp_helper_gbuf: Optional[DataParallelBuffer] = None
hsdp_comm_gbuf: Optional[DataParallelBuffer] = None
@@ -1549,7 +1549,7 @@ def _does_param_require_new_bucket(param):
# Set aggregate buckets by FSDP units, i.e. buckets pertaining to the same
# FSDP unit module and are either expert or non-expert parameters should
# end up in the same bucket group for NCCL.
- # Non-FSDP unit parameters will be assigned to the identity bucket group.
+ # Non-FSDP unit module parameters will be assigned to the identity bucket group.
if bucket_group_by_fsdp_unit:
bucket_group_map = {}
@@ -1642,6 +1642,7 @@ def __init__(
)
self.ddp_config = ddp_config
+ self.use_decoupled_grad = ddp_config.megatron_fsdp_use_decoupled_grad
self.module = module
self.bucketing_policy = bucketing_policy
self.param_to_name = {p: name for name, p in self.module.named_parameters()}
@@ -1691,9 +1692,6 @@ def __init__(
if self.dist_index.get_fsdp_group(is_expert_parallel=True) is not None:
# Expert-DP group when using EP
self.ubr_groups.append(self.dist_index.get_fsdp_group(is_expert_parallel=True))
- if self.dist_index.get_outer_fsdp_group() is not None:
- # Outer/Inter-FSDP group when using hybrid FSDP
- self.ubr_groups.append(self.dist_index.get_outer_fsdp_group())
if (
self.dist_index.get_fsdp_group(
is_expert_parallel=False, independent_all_gather=True
@@ -1706,6 +1704,19 @@ def __init__(
is_expert_parallel=False, independent_all_gather=True
)
)
+ if (
+ self.dist_index.get_fsdp_group(is_expert_parallel=True, independent_all_gather=True)
+ is not None
+ ):
+ # Expert all-gather group used when overlapping all-gather and gradient reduction.
+ self.ubr_groups.append(
+ self.dist_index.get_fsdp_group(
+ is_expert_parallel=True, independent_all_gather=True
+ )
+ )
+ if self.dist_index.get_outer_fsdp_group() is not None:
+ # Outer/Inter-FSDP group when using hybrid FSDP (IB domain, registered last).
+ self.ubr_groups.append(self.dist_index.get_outer_fsdp_group())
log_single_rank(
logger,
@@ -1942,16 +1953,161 @@ def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params):
f"Invalid data_parallel_sharding_strategy: {data_parallel_sharding_strategy}"
)
- # Only create HSDP buffers if sharding on DP-Outer. Otherwise, no need to all-gather
- # parameters on DP-Outer, but still need to all-reduce gradients on DP-Outer.
- should_create_hfsdp_wbuf_and_gbuf = (
+ """
+ Hybrid FSDP (HFSDP) helper buffers for outer-DP optimizer-state sharding.
+
+ Design goal
+ ==========
+ This design extends Megatron-FSDP's Hybrid / Fully Sharded Data Parallelism
+ to support *outer* data-parallel (DP) optimizer-state sharding, without
+ complicating the per-rank model weight / grad views that are used by the
+ forward and backward passes.
+
+ Core idea
+ ==========
+ We introduce two(or three) persistent helper buffers:
+
+ - `hfsdp_helper_wbuf`: stores the *true* persistent parameter payload
+ in the inner-DP layout (Inner-DP param buffer).
+ - `hfsdp_helper_gbuf`: stores the *true* persistent gradient payload
+ in the inner-DP layout (Inner-DP grad buffer).
+ - `hfsdp_helper_wtbuf`: (optional) stores transpose weights in the inner-DP
+ layout for FP8 parameters that require transposition for efficient
+ mixed-precision matmuls.
+
+ These buffers own the real storage for parameters and gradients that
+ participate in HFSDP sharding. The existing model `weight` buffer and
+ `gradient` buffer are simplified to be pure "data-parallel buffers"
+ defined only over the DP dimension, and can alias (view into) the
+ helper buffers. In other words:
+
+ - Helper buffers: inner-DP-aware, persistent, sharded storage.
+ - Model buffers: outer-DP-oriented data-parallel views used for compute,
+ and to give the optimizer access to the relevant shards of the helper
+ buffers when needed.
+
+ By separating **storage** (helper buffers) from **compute views**
+ (model buffers), we can:
+
+ - Keep the model-side DP buffers conceptually simple (they do not need to
+ encode Inner-DP vs Outer-DP tiling).
+ - Implement fully sharded optimizer states over the DP mesh by sharding
+ optimizer states consistently with the model buffers and helper buffers.
+ - Control when and how data is synchronized between inner and outer DP
+ dimensions on each iteration.
+
+ Data flow per iteration
+ =======================
+ Compared to the usual Hybrid FSDP data flow (where the last micro-batch
+ backward issues a DP all-reduce on gradients), this design explicitly
+ uses parameter all-gather and gradient reduce-scatter between the DP and
+ inner-DP layouts:
+
+ 1. Parameters:
+ - Persistent parameter shards live in `hfsdp_helper_wbuf` in inner-DP
+ layout.
+ - At the beginning of each iteration (before the first micro-batch
+ forward), we all-gather DP-sharded parameters to form the inner-DP
+ parameter shards in `hfsdp_helper_wbuf`, following the standard
+ Hybrid FSDP pattern of making shards "bigger" for compute.
+ - The model weight buffer is set up as a DP-only view on top of
+ these inner-DP shards. It does not own persistent storage.
+
+ 2. Gradients:
+ - During backward for each micro-batch, gradients are accumulated into
+ the `hfsdp_helper_gbuf` in inner-DP layout.
+ - On the last micro-batch backward of the iteration, instead of a DP
+ all-reduce, we perform a reduce-scatter that maps the DP gradient
+ layout back into outer-DP gradient shards.
+ - Because the model gradient buffer is a view of `hfsdp_helper_gbuf`
+ in the current design, the reduced / scattered results effectively
+ update both the helper buffer and the model-gradient-buffer view.
+ - The optimizer then reads gradients from model-gradient-buffer view
+ (DP layout) to perform the update.
+
+ 3. Optimizer states:
+ - Optimizer states are constructed and kept sharded in the same DP /
+ inner-DP pattern as the helper buffers and their model-buffer views.
+ - Because parameters and gradients are stored persistently in helper
+ buffers and are sharded over DP, the optimizer only ever touches
+ fully sharded tensors. This enables *fully sharded* optimizer
+ states on the DP group (outer-DP sharding in HFSDP).
+ - After the optimizer step, updated parameter shards remain in
+ `hfsdp_helper_wbuf`. At the beginning of the next iteration, the
+ usual all-gather path re-exposes these updated outer-DP shards
+ through the model weight buffer.
+
+ Implementation details
+ ======================
+ - `hfsdp_helper_wbuf` / `hfsdp_helper_gbuf` (and optionally
+ `hfsdp_helper_wtbuf`) are allocated as the canonical storage for all
+ HFSDP-managed parameters / gradients. They encode the inner-DP
+ partitioning and are aligned with the DP device mesh used for HFSDP
+ optimizer sharding.
+
+ - The existing Megatron-FSDP weight and grad buffers are repurposed as
+ *outer-DP data-parallel buffers*. They:
+ - Have a shape / layout that only reflects the DP dimension.
+ - Implemented as views into the helper buffers to avoid extra copies.
+ - Serve as the interface tensors that the optimizer code reads from
+ and writes to.
+
+ - Synchronization between helper buffers and model buffers is explicit:
+ - "param sync" path:
+ - At initialization and at the beginning of each iteration, parameters
+ are exposed to the model buffer by all-gathering DP-sharded
+ parameters into inner-DP shards in `hfsdp_helper_wbuf` and then
+ viewing them through the model weight buffer.
+ - "grad sync" path:
+ - At the last micro-batch backward, we reduce-scatter gradients from
+ the inner-DP layout into model gradient buffer as DP shards. Because
+ the model gradient buffer is a view of `hfsdp_helper_gbuf`, the
+ reduced results are immediately visible through the model gradient buffer
+ as well.
+
+ - Outer-DP optimizer-state sharding:
+ - Optimizer state tensors are allocated with the same sharding pattern
+ as the model buffers along the DP dimension. Each rank only owns the
+ local shard of:
+ * its parameters (from `model weight buffer`),
+ * its gradients (from `model gradient buffer`),
+ * and the corresponding optimizer states.
+ - This is conceptually similar to ZeRO-style optimizer sharding, but
+ implemented on top of Megatron-FSDP's buffer / device-mesh abstractions,
+ using the helper buffers as the single source of truth for persistent
+ data.
+
+ Notes for maintainers
+ =====================
+ - When adding new parameters to HFSDP, register them with the helper
+ buffers first. The model weight / grad DP buffers should be treated as
+ *views* for compute, not as owners of persistent storage.
+
+ - When changing the DP mesh, inner/outer-DP dimension mapping, or the
+ sharding strategy, verify that:
+ - The helper buffer partitioning matches the intended HFSDP sharding
+ (inner-DP).
+ - The optimizer state partitioning is kept consistent with the model
+ weight and gradient buffer sharding.
+ - The synchronization paths correctly map between inner-DP layout
+ (helper buffers) and DP-only layout (model buffers), including the
+ parameter all-gather at the beginning of the iteration and the
+ gradient reduce-scatter on the last micro-batch backward.
+
+ - Any logic that assumes the model weight or gradient buffers own
+ persistent data should be updated to read/write from the helper
+ buffers instead. The model buffers are intentionally simplified to
+ keep the HFSDP optimizer sharding logic centralized in the helper
+ layer.
+ """
+ should_create_hfsdp_helper_buffers = (
self.dist_index.use_hybrid_fsdp
and self.ddp_config.outer_dp_sharding_strategy != "no_shard"
)
# DP-Outer sharding is only supported for fully-sharded DP-Shard.
# NOTE(@cspades): Important guard for HFSDP functionality!
if (
- should_create_hfsdp_wbuf_and_gbuf
+ should_create_hfsdp_helper_buffers
and self.ddp_config.data_parallel_sharding_strategy != "optim_grads_params"
):
raise NotImplementedError(
@@ -1978,6 +2134,7 @@ def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params):
name="fsdp_fp8_transpose_params",
fsdp_param_groups=self.parameter_groups,
size=UB_BUFFER_NUM,
+ fallback_to_persistent_buffer=self.ddp_config.fsdp_db_use_persist_buf_on_alloc_fail,
)
self.main_grad_alloc = FixedPoolAllocator(
name="fsdp_grads",
@@ -2018,13 +2175,13 @@ def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params):
# For all bucket groups (partitioned parameter groups)...
for group_id, group in enumerate(self.parameter_groups):
main_buf_extra_kwargs = {}
- if should_create_hfsdp_wbuf_and_gbuf:
+ if should_create_hfsdp_helper_buffers:
# DP-Outer + DP-Shard
main_buf_dp_group = self.dist_index.get_dp_group(
is_expert_parallel=group.is_expert_param
)
# DP-Shard
- hsdp_buf_dp_group = self.dist_index.get_fsdp_group(
+ inner_dp_group = self.dist_index.get_fsdp_group(
is_expert_parallel=group.is_expert_param
)
main_buf_extra_kwargs["dp_rank"] = self.dist_index.get_logical_hybrid_fsdp_rank(
@@ -2036,14 +2193,14 @@ def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params):
is_expert_parallel=group.is_expert_param
)
- # When --create-all-gather-group is enabled, use a separate process group for
- # all-gather operations (model_weight_buffer) to enable overlap with gradient reduction
- # operations (main_grad_buffer). This avoids head-of-line blocking between forward
- # all-gather and backward reduce-scatter on the same communicator.
+ # Use separate process group for all-gather operations (model_weight_buffer)
+ # to enable overlap with gradient reduction operations (main_grad_buffer).
+ # This avoids head-of-line blocking between forward all-gather and backward
+ # reduce-scatter on the same communicator.
model_wbuf_dp_group = main_buf_dp_group
- if not group.is_expert_param and not should_create_hfsdp_wbuf_and_gbuf:
+ if not should_create_hfsdp_helper_buffers:
ag_group = self.dist_index.get_fsdp_group(
- is_expert_parallel=False, independent_all_gather=True
+ is_expert_parallel=group.is_expert_param, independent_all_gather=True
)
if ag_group is not None:
model_wbuf_dp_group = ag_group
@@ -2174,71 +2331,33 @@ def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params):
buffer_size[group.main_grad_buffer.dtype] += group.main_grad_buffer.data_size
# Initialize the HSDP weight and grad buffers if hsdp full sharding is enabled.
- if should_create_hfsdp_wbuf_and_gbuf:
+ if should_create_hfsdp_helper_buffers:
# Initialize the HSDP weight buffer.
wbuf = group.model_weight_buffer
- group.hsdp_wbuf = DataParallelBuffer(
- self.ddp_config,
- group.params,
+ group.hfsdp_helper_wbuf = _create_hfsdp_helper_buffer(
+ group.model_weight_buffer,
+ inner_dp_group=inner_dp_group,
is_data_distributed=is_main_weight_buffer_distributed
- and hsdp_buf_dp_group.size() > 1,
- dtype=wbuf.dtype,
- device=wbuf.device,
- data_parallel_group=hsdp_buf_dp_group,
- is_transpose_buffer=False,
- temporary_bucket_allocator=self.weight_alloc,
- bucket_id=group_id,
- chunk_size_factor=group.chunk_size_factor,
- mem_alloc_context=self.mem_alloc_context,
- item_index_map=wbuf.item_index_map,
- bucket_index=wbuf.bucket_index,
- shard_bucket_index=_get_dp_buffer_shard_bucket_index(
- wbuf.bucket_index,
- is_data_distributed=is_main_weight_buffer_distributed
- and hsdp_buf_dp_group.size() > 1,
- data_parallel_world_size=hsdp_buf_dp_group.size(),
- data_parallel_rank=hsdp_buf_dp_group.rank(),
- ),
+ and inner_dp_group.size() > 1,
)
if group.transpose_weight_buffer is not None:
- # TODO(@kunlunl, @cspades): Create a hybrid-sharded transpose buffer
- # to map fully-sharded transpose weights to partially-sharded transpose
- # weights before and after fully-distributed optimization.
- raise NotImplementedError(
- "HFSDP (HSDP + fully-sharded optimizer state) doesn't "
- "support FP8 recipes that require a transpose buffer."
+ group.hfsdp_helper_wtbuf = _create_hfsdp_helper_buffer(
+ group.transpose_weight_buffer,
+ inner_dp_group=inner_dp_group,
+ is_data_distributed=is_main_weight_buffer_distributed
+ and inner_dp_group.size() > 1,
)
if should_create_grad_buffer_or_main_weight_buffer:
- # Initialize the HSDP grad buffer.
- gbuf = group.main_grad_buffer
- group.hsdp_gbuf = DataParallelBuffer(
- self.ddp_config,
- group.params,
+ group.hfsdp_helper_gbuf = _create_hfsdp_helper_buffer(
+ group.main_grad_buffer,
+ inner_dp_group=inner_dp_group,
is_data_distributed=is_grad_buffer_distributed
- and hsdp_buf_dp_group.size() > 1,
- dtype=gbuf.dtype,
- device=gbuf.device,
- data_parallel_group=hsdp_buf_dp_group,
- is_transpose_buffer=False,
- temporary_bucket_allocator=self.main_grad_alloc,
- gradient_scaling_factor=gradient_scaling_factor,
- bucket_id=group_id,
- chunk_size_factor=group.chunk_size_factor,
- mem_alloc_context=self.mem_alloc_context,
- item_index_map=gbuf.item_index_map,
- bucket_index=gbuf.bucket_index,
- shard_bucket_index=_get_dp_buffer_shard_bucket_index(
- gbuf.bucket_index,
- is_data_distributed=is_grad_buffer_distributed
- and hsdp_buf_dp_group.size() > 1,
- data_parallel_world_size=hsdp_buf_dp_group.size(),
- data_parallel_rank=hsdp_buf_dp_group.rank(),
- ),
+ and inner_dp_group.size() > 1,
)
buffer_size[group.main_grad_buffer.dtype] -= group.main_grad_buffer.data_size
- buffer_size[group.main_grad_buffer.dtype] += group.hsdp_gbuf.data_size
+ buffer_size[group.main_grad_buffer.dtype] += group.hfsdp_helper_gbuf.data_size
# Only create an extra grad comm buffer for HSDP.
if should_create_grad_buffer_or_main_weight_buffer and self.dist_index.use_hybrid_fsdp:
@@ -2252,7 +2371,7 @@ def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params):
is_expert_parallel=group.is_expert_param
)
hfsdp_kwargs = {}
- if should_create_hfsdp_wbuf_and_gbuf:
+ if should_create_hfsdp_helper_buffers:
hfsdp_kwargs["item_index_map"] = gbuf.item_index_map
hfsdp_kwargs["bucket_index"] = gbuf.bucket_index
hfsdp_kwargs["shard_bucket_index"] = _get_dp_buffer_shard_bucket_index(
@@ -2313,29 +2432,17 @@ def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params):
wbuf = group.model_weight_buffer
if wbuf:
with self.mem_alloc_context():
- if group.hsdp_wbuf:
- # When using HSDP, the hybrid-sharded buffer shards across the FSDP group,
- # while the main buffer shards across the larger / more granular DP group.
- # The main weight buffer data is a shard of the hybrid-sharded buffer data.
- # Because the hybrid buffer data is persistently allocated, the weight and
- # gradient memory footprint is similar to not sharding on DP-Outer, i.e.
- # replicating on DP-Outer. However, optimizer states based on main buffer
- # weights (self.dist_main_weight) and gradients (self.dist_main_grad) will
- # be sharded persistently upon initialization.
- hsdp_wbuf = group.hsdp_wbuf
- hsdp_wbuf.init_data(
- torch.empty(
- hsdp_wbuf.data_size, dtype=hsdp_wbuf.dtype, device=self.device
- )
+ if group.hfsdp_helper_wbuf:
+ _init_hfsdp_helper_and_dp_buffer_data(
+ group.hfsdp_helper_wbuf,
+ wbuf,
+ mem_alloc=lambda size, dtype: torch.empty(
+ size, dtype=dtype, device=self.device
+ ),
+ outer_dp_group=self.dist_index.get_outer_fsdp_group(
+ is_expert_parallel=group.is_expert_param
+ ),
)
- outer_fsdp_group = self.dist_index.get_outer_fsdp_group()
- wbuf_data = hsdp_wbuf.data[
- # Requires FSDP sharding for (DP-Shard, DP-Outer) to cover DP-Shard.
- wbuf.data_size
- * outer_fsdp_group.rank() : wbuf.data_size
- * (outer_fsdp_group.rank() + 1)
- ]
- wbuf.init_data(wbuf_data)
else:
# When not using HSDP, the main buffer shards across the FSDP group.
wbuf.init_data(
@@ -2346,10 +2453,16 @@ def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params):
tbuf = group.transpose_weight_buffer
if tbuf:
with self.mem_alloc_context():
- if group.hsdp_wbuf:
- raise NotImplementedError(
- "HFSDP (HSDP + fully-sharded optimizer state) doesn't "
- "support FP8 recipes that require a transpose buffer."
+ if group.hfsdp_helper_wbuf:
+ _init_hfsdp_helper_and_dp_buffer_data(
+ group.hfsdp_helper_wtbuf,
+ tbuf,
+ mem_alloc=lambda size, dtype: torch.empty(
+ size, dtype=dtype, device=self.device
+ ),
+ outer_dp_group=self.dist_index.get_outer_fsdp_group(
+ is_expert_parallel=group.is_expert_param
+ ),
)
else:
# Initialize the transpose buffer.
@@ -2552,34 +2665,23 @@ def _alloc(dtype, size):
continue
# Allocate the main grad buffer data, and attach it to the main grad buffer.
with self.mem_alloc_context():
- if group.hsdp_gbuf:
- # When using HSDP, the hybrid-sharded buffer shards across the FSDP group,
- # while the main buffer shards across the larger / more granular DP group.
- # The main weight buffer data is a shard of the hybrid-sharded buffer data.
- # Because the hybrid buffer data is persistently allocated, the weight and
- # gradient memory footprint is similar to not sharding on DP-Outer, i.e.
- # replicating on DP-Outer. However, optimizer states based on main buffer
- # weights (self.dist_main_weight) and gradients (self.dist_main_grad) will
- # be sharded persistently upon initialization.
- hsdp_gbuf = group.hsdp_gbuf
- hsdp_gbuf.init_data(_alloc(hsdp_gbuf.dtype, hsdp_gbuf.data_size))
- outer_fsdp_group = self.dist_index.get_outer_fsdp_group()
- gbuf_data = hsdp_gbuf.data[
- # Requires FSDP sharding for (DP-Shard, DP-Outer) to cover DP-Shard.
- gbuf.data_size
- * outer_fsdp_group.rank() : gbuf.data_size
- * (outer_fsdp_group.rank() + 1)
- ]
- gbuf.init_data(gbuf_data)
- hsdp_gbuf.data.zero_()
+ if group.hfsdp_helper_gbuf:
+ _init_hfsdp_helper_and_dp_buffer_data(
+ group.hfsdp_helper_gbuf,
+ gbuf,
+ mem_alloc=_alloc,
+ outer_dp_group=self.dist_index.get_outer_fsdp_group(
+ is_expert_parallel=group.is_expert_param
+ ),
+ )
+ group.hfsdp_helper_gbuf.data.zero_()
else:
# When not using HSDP, the main buffer shards across the FSDP group.
gbuf.init_data(_alloc(gbuf.dtype, gbuf.data_size))
gbuf.data.zero_()
- gbuf.data.zero_()
for item_id, p in enumerate(group.params):
# Attach the main grad buffer data and metadata to the parameter.
- p._gbuf = group.hsdp_gbuf if group.hsdp_gbuf else gbuf
+ p._gbuf = group.hfsdp_helper_gbuf if group.hfsdp_helper_gbuf else gbuf
p._item_id = item_id
def main_grad_getter(p):
@@ -2628,10 +2730,17 @@ def _reset_parameters(self, old_params, new_params):
new_param.requires_grad_(old_param.requires_grad)
- for tp_attr in ["_mcore_tp", "_tp_partition_dim", "_tp_duplicated"]:
+ for tp_attr in ["_tensor_parallel_mode"]:
if getattr(old_param, tp_attr, None) is not None:
setattr(new_param, tp_attr, getattr(old_param, tp_attr))
+ # For FSDP with delayed_wgrad_compute, `skip_backward_post_hook` needs
+ # to be reset on new param for correct grad accumulation of wgrad computation.
+ setattr(
+ new_param,
+ 'skip_backward_post_hook',
+ getattr(old_param, 'skip_backward_post_hook', False),
+ )
for item_id, p in enumerate(self.params):
if p in param_map:
new_p = param_map[p]
@@ -2648,8 +2757,9 @@ def _reset_parameters(self, old_params, new_params):
group.transpose_weight_buffer,
group.main_weight_buffer,
group.main_grad_buffer,
- group.hsdp_wbuf,
- group.hsdp_gbuf,
+ group.hfsdp_helper_wbuf,
+ group.hfsdp_helper_wtbuf,
+ group.hfsdp_helper_gbuf,
]:
if buf is None:
continue
@@ -2668,19 +2778,36 @@ def zero_grad(self):
"""
Zero out the underlying grad_buffer and reset all buckets in preparation
for the next iteration of training.
- """
- for name, param in self.optimizer_named_parameters:
- param.grad = None
- if hasattr(param, "decoupled_grad"):
- param.decoupled_grad = None
- if name in self.dist_main_grad:
- self.dist_main_grad[name]._local_tensor = None
+ Gradient shards are dereferenced to free memory. However, dereferencing is
+ not compatible with (FWD-BWD / full-iteration) CUDA graph-ability, because
+ we need to preserve this reference to the sharded gradient generated during
+ CUDA graph replay (`setattr` in `update_main_grads` not executed during
+ CUDA graph replay, as it is not a CUDA kernel).
+
+ If the gradient is decoupled (precision-aware) or is equivalent to the
+ distributed optimizer parameter precision, the gradient shard is a view of
+ the Megatron-FSDP sharded gradient buffer. If not, then not dereferencing
+ this gradient shard will increase memory utilization as this gradient is a
+ persistent casted-copy of the accumulated gradient.
+ """
+ if not self.ddp_config.megatron_fsdp_cuda_graph_mode:
+ # Dereference the sharded gradient to reclaim memory
+ # unless a full-iteration CUDA graph is utilized.
+ for name, param in self.optimizer_named_parameters:
+ param.grad = None
+ if hasattr(param, "decoupled_grad"):
+ param.decoupled_grad = None
+ if name in self.dist_main_grad:
+ self.dist_main_grad[name]._local_tensor = None
+
+ # Zero the Megatron-FSDP sharded gradient buffer. If param.grad or param.decoupled_grad
+ # is a view of this buffer, they will be zero'd as well.
for group in self.parameter_groups:
if group.main_grad_buffer:
group.main_grad_buffer.data.zero_()
- if group.hsdp_gbuf:
- group.hsdp_gbuf.data.zero_()
+ if group.hfsdp_helper_gbuf:
+ group.hfsdp_helper_gbuf.data.zero_()
def _init_distributed_params(self):
"""
@@ -2781,9 +2908,7 @@ def set_param_attribute():
"partition_stride",
"is_embedding_or_output_parameter",
"is_embedding_parameter",
- "_mcore_tp",
- "_tp_duplicated",
- "_tp_partition_dim",
+ "_tensor_parallel_mode",
]:
if hasattr(orig_param, attr_name):
setattr(param, attr_name, getattr(orig_param, attr_name))
@@ -2816,11 +2941,6 @@ def update_main_grads(self):
from the main gradient buffer. If the model parameters are sharded,
we only need to update the gradient shard associated with the model
parameter shard, as both are sharded symmetrically.
-
- Checks if high-precision main weights are utilized for optimization.
- Otherwise, falls back to low-precision model weights, and further
- falls back to the original module parameters not managed by cFSDP
- in the case of no sharding / cFSDP OFF.
"""
for name, param in self.optimizer_named_parameters:
orig_param = param.orig_param
@@ -2841,10 +2961,11 @@ def update_main_grads(self):
optimizer_grad = group.main_grad_buffer.get_item(
item_id, only_shard=sharded_optimizer_state
)
- if group.main_weight_buffer is not None:
- if not getattr(self, "use_precision_aware_optimizer", False):
- # Convert the gradient to the main weight buffer dtype.
- optimizer_grad = optimizer_grad.to(param.dtype)
+ if group.main_weight_buffer is not None and not self.use_decoupled_grad:
+ # Convert the gradient to the main weight data-type for optimization.
+ # Not needed for decoupled gradients, because the precision-aware
+ # optimizer can apply gradients to parameters of different precision!
+ optimizer_grad = optimizer_grad.to(param.dtype)
if name not in self.dist_main_grad:
# Register the gradient as a distributed tensor.
@@ -2867,13 +2988,13 @@ def update_main_grads(self):
if optimizer_grad.numel() == 0:
grad = None
- # The presence of main_grad_buffer but no main_weight_buffer may imply
- # that a precision-aware optimizer is used.
- if getattr(self, "use_precision_aware_optimizer", False):
+ # If use_decoupled_grad (i.e. for precision-aware optimizers like TE FusedAdam),
+ # install the gradient into param.decoupled_grad.
+ if self.use_decoupled_grad:
setattr(param, "decoupled_grad", grad)
else:
# Attach the gradient to the optimizer parameter.
- setattr(param, "grad", grad.to(param.dtype) if grad is not None else None)
+ setattr(param, "grad", grad)
@property
def num_buckets(self):
@@ -3051,25 +3172,6 @@ def _batch_quantize_blockwise_fp8_params(
)
_fp8_quantize_params(dense_param_quantize_kwargs, expert_param_quantize_kwargs)
- @torch.no_grad()
- def copy_model_weights_to_main_weights(self):
- """Copy the model weights to the main weights."""
- for group in self.parameter_groups:
- mbuf = group.main_weight_buffer
- if mbuf is None:
- continue
- wbuf = group.model_weight_buffer
- if mbuf.is_data_distributed:
- copyin_data = wbuf.get_shard_from_local_buffer()
- else:
- copyin_data = wbuf.data
- assert mbuf.data.numel() == copyin_data.numel(), (
- f"Master weight buffer size {mbuf.data.numel()} does not match "
- f"model weight buffer size {copyin_data.numel()}"
- )
- # TODO(mxfp8): Make sure it's not a fp8 buf?
- mbuf.data.copy_(copyin_data.data)
-
def all_gather_parameters(self, async_op: bool = True):
"""All gather the parameters.
Args:
@@ -3157,7 +3259,7 @@ def all_reduce_gradients(self, async_op: bool = False):
all_reduce_ops = []
for g in self.parameter_groups:
gbuf = g.main_grad_buffer
- if gbuf is not None:
+ if gbuf is None:
continue
scaling_factor = gbuf.gradient_scaling_factor
if self.ddp_config.check_for_nan_in_grad:
@@ -3173,19 +3275,133 @@ def all_reduce_gradients(self, async_op: bool = False):
op.wait()
+def _create_hfsdp_helper_buffer(
+ dp_buffer: DataParallelBuffer,
+ inner_dp_group: torch.distributed.ProcessGroup,
+ is_data_distributed: bool,
+) -> DataParallelBuffer:
+ """
+ Create a Hybrid-FSDP helper DataParallelBuffer on the inner-DP group.
+
+ This helper buffer mirrors the metadata of the original fully
+ `dp_buffer` (bucket config, params, allocator, etc.), but binds it to
+ the `inner_dp_group` and computes a per-rank `shard_bucket_index`
+ appropriate for that group. The resulting buffer is used as the
+ HFSDP helper buffer that owns the persistent inner-DP shard of the
+ global bucket, while still sharing the same logical bucket indexing
+ (`bucket_index`) with the fully DP buffer.
+
+ Parameters
+ ==========
+ dp_buffer : DataParallelBuffer
+ The existing fully data-parallel buffer whose configuration
+ and bucket layout should be mirrored.
+ inner_dp_group : torch.distributed.ProcessGroup
+ The process group representing the inner-DP (HFSDP) data-parallel
+ group for this helper buffer.
+ is_data_distributed : bool
+ Whether the underlying data in this helper buffer is sharded
+ across ranks in `inner_dp_group`.
+
+ Returns
+ =======
+ DataParallelBuffer
+ A new DataParallelBuffer configured as the HFSDP helper buffer
+ for the given `inner_dp_group`, sharing the same bucket index
+ as `dp_buffer` but with an inner-DP `shard_bucket_index`.
+ """
+ helper_buffer = DataParallelBuffer(
+ dp_buffer.ddp_config,
+ dp_buffer.params,
+ is_data_distributed=is_data_distributed,
+ dtype=dp_buffer.dtype,
+ device=dp_buffer.device,
+ data_parallel_group=inner_dp_group,
+ is_transpose_buffer=dp_buffer.is_transpose_buffer,
+ temporary_bucket_allocator=dp_buffer.temporary_bucket_allocator,
+ bucket_id=dp_buffer.bucket_id,
+ chunk_size_factor=dp_buffer.chunk_size_factor,
+ mem_alloc_context=dp_buffer.mem_alloc_context,
+ item_index_map=dp_buffer.item_index_map,
+ bucket_index=dp_buffer.bucket_index,
+ # HFSDP helper buffer shares the same global bucket layout as the
+ # fully DP buffer, but computes its own shard_bucket_index because
+ # data is distributed across ranks in the inner-DP group.
+ shard_bucket_index=_get_dp_buffer_shard_bucket_index(
+ bucket_index=dp_buffer.bucket_index,
+ is_data_distributed=is_data_distributed,
+ data_parallel_world_size=inner_dp_group.size(),
+ data_parallel_rank=inner_dp_group.rank(),
+ ),
+ )
+
+ return helper_buffer
+
+
+def _init_hfsdp_helper_and_dp_buffer_data(
+ hfsdp_helper_buffer: DataParallelBuffer,
+ dp_buffer: DataParallelBuffer,
+ mem_alloc: Callable[[torch.dtype, int], torch.Tensor],
+ outer_dp_group: torch.distributed.ProcessGroup,
+) -> None:
+ """
+ Initialize storage for the HFSDP helper buffer and its corresponding
+ fully-DP DataParallelBuffer view.
+
+ The helper buffer is allocated as a single contiguous tensor that
+ stores all DP shards for the given bucket. Each rank in the outer-DP
+ group then takes its local slice of this storage and exposes it
+ through `dp_buffer`, so the fully-DP buffer becomes a view into the helper
+ buffer rather than owning separate storage.
+
+ Parameters
+ ==========
+ hfsdp_helper_buffer : DataParallelBuffer
+ The HFSDP helper buffer that owns the full inner-/outer-DP bucket
+ storage.
+ dp_buffer : DataParallelBuffer
+ The fully-DP DataParallelBuffer that should view its local shard
+ from `hfsdp_helper_buffer`.
+ mem_alloc : Callable[[torch.dtype, int], torch.Tensor]
+ Allocation function used to create the backing tensor for the
+ helper buffer (dtype, numel).
+ outer_dp_group : torch.distributed.ProcessGroup
+ Process group for the outer data-parallel dimension. Its rank and
+ world size determine which slice of the helper buffer this rank
+ sees through `dp_buffer`.
+ """
+ # Allocate contiguous storage for all outer-DP shards in the helper buffer.
+ hfsdp_helper_buffer.init_data(
+ mem_alloc(dtype=hfsdp_helper_buffer.dtype, size=hfsdp_helper_buffer.data_size)
+ )
+
+ rank = outer_dp_group.rank()
+ shard_size = dp_buffer.data_size
+ start = shard_size * rank
+ end = shard_size * (rank + 1)
+
+ # Each outer-DP rank takes a disjoint slice of the helper buffer as its
+ # local DP buffer view. This keeps `dp_buffer` as a view into the
+ # helper-owned storage.
+ dp_buffer_data = hfsdp_helper_buffer.data[start:end]
+ dp_buffer.init_data(dp_buffer_data)
+
+
class BucketStatus(Enum):
"""
An enumeration of possible statuses for a data-parallel communication bucket.
Attributes:
EMPTY (int): The bucket is empty and not in use.
+ PRESERVED (int): The bucket storage is retained but not ready for use.
COMMUNICATING (int): The bucket is currently being used for communication.
READY_TO_USE (int): The bucket is filled with data and ready for use.
"""
EMPTY = 1
- COMMUNICATING = 2
- READY_TO_USE = 3
+ PRESERVED = 2
+ COMMUNICATING = 3
+ READY_TO_USE = 4
class GradReducePipeline:
@@ -3342,9 +3558,11 @@ def _enforce_double_buffer_limit(self, add_buckets):
for _, _, bucket_id in reversed(self.grad_reduce_queue):
fsdp_unit_id = param_groups[bucket_id].fsdp_unit_id
double_buf_units.add(fsdp_unit_id)
- if len(double_buf_units) > 2:
+ if len(double_buf_units) > 1:
keep_n -= 1
- self.wait_for_previous_grad_reduce(keep_n)
+
+ with torch.cuda.stream(self.rs_stream):
+ self.wait_for_previous_grad_reduce(keep_n)
def get_ready_bucket_group_for_reduction(self, bucket_id: int) -> Optional[List[int]]:
"""Checks if all buckets in the bucket group containing the given bucket_id
@@ -3376,7 +3594,7 @@ def get_fsdp_buffer(self, bucket_id: int) -> DataParallelBuffer:
"""Get the FSDP buffer for the given bucket ID."""
param_group = self.buffer.parameter_groups[bucket_id]
if self.buffer.ddp_config.outer_dp_sharding_strategy != "no_shard":
- return param_group.hsdp_gbuf
+ return param_group.hfsdp_helper_gbuf
return param_group.main_grad_buffer
def _bucket_group_gradient_reduce(
@@ -3705,8 +3923,14 @@ def num_buckets(self):
"""Return the number of buckets."""
return self.buffer.num_buckets
- def reset(self):
- """Reset the pipeline state."""
+ def reset(self, preserve_non_fsdp_units: bool = True):
+ """Reset the pipeline state.
+
+ Non-FSDP-unit buckets are preserved by default because their params may
+ be read across module boundaries. Setting preserve_non_fsdp_units=False
+ releases all bucket storage and is intended only for debugging when the
+ model will not be reused.
+ """
if len(self.param_gather_event_map) > 0:
warnings.warn(
(
@@ -3718,12 +3942,26 @@ def reset(self):
while len(self.param_gather_event_map) > 0:
(bucket_id, bwd) = next(iter(self.param_gather_event_map))
self.wait_bucket_ready(bucket_id, bwd)
+
for bucket_id in range(self.num_buckets):
+ is_unit_bucket = self.buffer.parameter_groups[bucket_id].fsdp_unit_id is not None
for bwd in [False, True]:
- self.bucket_can_be_released[self.get_bucket_key(bucket_id, bwd)] = True
+ bucket_key = self.get_bucket_key(bucket_id, bwd)
+ # If preserve_non_fsdp_units is set, then do not release buckets
+ # associated with FSDP non-units. Instead, mark the bucket as PRESERVED
+ # (not NEW) so a later all-gather refreshes preserved non-unit bucket
+ # storage in place.
+ if preserve_non_fsdp_units and not is_unit_bucket:
+ self.bucket_status[bucket_key] = BucketStatus.PRESERVED
+ else:
+ self.bucket_can_be_released[bucket_key] = True
self.recycle_unused_buckets()
- assert all([status is BucketStatus.EMPTY for status in self.bucket_status.values()]), (
+ expected_statuses = (BucketStatus.EMPTY,)
+ if preserve_non_fsdp_units:
+ expected_statuses += (BucketStatus.PRESERVED,)
+
+ assert all(status in expected_statuses for status in self.bucket_status.values()), (
f"There are still working buckets, it is not safe to reset. "
f"bucket_status: {self.bucket_status}."
)
@@ -3851,11 +4089,13 @@ def need_skip_prefetch(bucket_id):
ag_buckets = list(sorted(set(ag_buckets)))
bucket_id = next_bucket_id(ag_buckets)
- # Only all-gather on buckets that have not been allocated yet.
+ # Only all-gather on buckets that have not been allocated yet or whose
+ # persistent storage was preserved but is not ready for use.
ag_buckets = [
bucket_id
for bucket_id in ag_buckets
- if self.bucket_status[self.get_bucket_key(bucket_id, bwd)] == BucketStatus.EMPTY
+ if self.bucket_status[self.get_bucket_key(bucket_id, bwd)]
+ in (BucketStatus.EMPTY, BucketStatus.PRESERVED)
]
if len(ag_buckets) == 0:
return
@@ -3875,19 +4115,22 @@ def need_skip_prefetch(bucket_id):
self.ag_stream if self.ag_stream is not None else torch.cuda.current_stream()
)
if outer_fsdp_group_param_gather:
- # TODO(@kunlunl): Support MXFP8 with HFSDP. Requires an HFSDP transpose buffer.
self.outer_fsdp_group_param_gather_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(self.outer_fsdp_group_param_gather_stream):
- outer_fsdp_group = self.buffer.dist_index.get_outer_fsdp_group()
+ is_expert_parallel = parameter_groups[buckets[0]].is_expert_param
+ outer_fsdp_group = self.buffer.dist_index.get_outer_fsdp_group(
+ is_expert_parallel=is_expert_parallel
+ )
with _coalescing_manager(outer_fsdp_group, async_ops=False):
for bucket_id in buckets:
- # All-gather the (DP-Outer, DP-Shard) weight shards from the DP-backed
- # main weight buffer into the (DP-Shard)-backed hybrid weight buffer.
- wbuf = self.buffer.parameter_groups[bucket_id].model_weight_buffer
- hsdp_wbuf = self.buffer.parameter_groups[bucket_id].hsdp_wbuf
+ inner_dp_wbuf = self.get_fsdp_buffer(bucket_id, bwd=bwd)
+ shard_size = inner_dp_wbuf.data_size // outer_fsdp_group.size()
+ rank = outer_fsdp_group.rank()
torch.distributed.all_gather_into_tensor(
- output_tensor=hsdp_wbuf.data,
- input_tensor=wbuf.data,
+ output_tensor=inner_dp_wbuf.data,
+ input_tensor=inner_dp_wbuf.data[
+ rank * shard_size : (rank + 1) * shard_size
+ ],
group=outer_fsdp_group,
)
# Wait for the DP-Outer group all-gather to finish.
@@ -3926,12 +4169,13 @@ def wait_bucket_ready(self, bucket_id, bwd, empty_ok=False):
if self.bucket_status[bucket_key] == BucketStatus.READY_TO_USE:
# Already ready to use.
return
- if self.bucket_status[bucket_key] == BucketStatus.EMPTY:
+ if self.bucket_status[bucket_key] in (BucketStatus.EMPTY, BucketStatus.PRESERVED):
if empty_ok:
return
- # Bucket shouldn't be empty, this implies that the bucket
- # was not allocated or NCCL operations are not complete.
- raise ValueError(f"Bucket {bucket_id} is empty.")
+ # Bucket should not be empty or merely preserved here; this implies that
+ # the bucket was not allocated, was not made ready for use, or NCCL
+ # operations are not complete.
+ raise ValueError(f"Bucket {bucket_id} is {self.bucket_status[bucket_key].name}.")
# Wait for asynchronous / overlapped NCCL operations to complete.
param_gather_event, mark_bucket_ready_to_use = self.param_gather_event_map.pop(bucket_key)
@@ -4008,9 +4252,9 @@ def get_fsdp_buffer(self, bucket_id: int, bwd=False) -> DataParallelBuffer:
param_group = self.buffer.parameter_groups[bucket_id]
if self.buffer.ddp_config.outer_dp_sharding_strategy != "no_shard":
if bwd and param_group.transpose_weight_buffer is not None:
- raise RuntimeError("Transpose buffer is not supported for HSDP")
+ return param_group.hfsdp_helper_wtbuf
else:
- return param_group.hsdp_wbuf
+ return param_group.hfsdp_helper_wbuf
if bwd and param_group.transpose_weight_buffer is not None:
return param_group.transpose_weight_buffer
else:
@@ -4022,7 +4266,10 @@ def async_bucket_gather(self, bucket_id, bwd) -> None:
bucket_key = self.get_bucket_key(bucket_id, bwd)
self.bucket_can_be_released[bucket_key] = False
- if self.bucket_status[bucket_key] != BucketStatus.EMPTY:
+ if self.bucket_status[bucket_key] in (
+ BucketStatus.COMMUNICATING,
+ BucketStatus.READY_TO_USE,
+ ):
return
self.bucket_status[bucket_key] = BucketStatus.COMMUNICATING
@@ -4425,43 +4672,49 @@ def make_fsdp_dtensor(
orig_param = param
# Handle tensor model parallel specific logic
- if is_mcore_tensor_model_parallel(param):
+ if not isinstance(param, DTensor) and using_tensor_parallel(
+ dist_index, is_expert_parallel=is_expert_param
+ ):
# Ensure parameter is not already a DTensor
assert not isinstance(param, DTensor), (
"[Megatron-FSDP] Parameter is already a DTensor, yet tensor_model_parallel " "is True."
)
-
+ # Verify a DeviceMesh TP dimension exists.
+ assert dist_index.tp_dim is not None, (
+ "[Megatron-FSDP] TP dimension is missing from DeviceMesh / FSDPDistributedIndex! "
+ "Required for Megatron-Core or TransformerEngine modules that use TP. "
+ "If TP=1, a trivial TP dimension of size 1 should be provided."
+ )
tp_mesh = dist_index.get_submesh(dist_index.tp_dim, is_expert_parallel=is_expert_param)
global_shape = list(param.shape)
- if tp_mesh.mesh.numel() > 1:
- if is_mcore_tensor_parallel_duplicated(param):
- placements = [Replicate()]
- if force_sync_tp_duplicated_param:
- if local_tensor.numel() > 0:
- torch.distributed.broadcast(
- local_tensor, group=tp_mesh.get_group(), group_src=0
- )
- elif run_check:
- # TODO: Implement consistency check for duplicated TP parameters
- pass
- else:
- tp_dim = get_mcore_tensor_parallel_partition_dim(param)
- assert tp_dim is not None, (
- "[Megatron-FSDP] Parameter is not tensor model parallel, "
- "yet tensor_model_parallel is True."
- )
- placements = [Shard(tp_dim)]
- global_shape[tp_dim] *= tp_mesh.mesh.numel()
-
- # Construct TP-sharded DTensor using Megatron-style placement
- param = DTensor.from_local(
- local_tensor=local_tensor,
- device_mesh=tp_mesh,
- placements=placements,
- run_check=run_check,
- shape=tuple(global_shape),
- stride=torch.empty(global_shape).stride(),
+ if is_mcore_tensor_parallel_duplicated(param):
+ placements = [Replicate()]
+ if force_sync_tp_duplicated_param:
+ if local_tensor.numel() > 0:
+ torch.distributed.broadcast(
+ local_tensor, group=tp_mesh.get_group(), group_src=0
+ )
+ elif run_check:
+ # TODO: Implement consistency check for duplicated TP parameters
+ pass
+ else:
+ tp_dim = get_mcore_tensor_parallel_partition_dim(param)
+ assert tp_dim is not None, (
+ "[Megatron-FSDP] Parameter is not tensor model parallel, "
+ "yet tensor_model_parallel is True."
)
+ placements = [Shard(tp_dim)]
+ global_shape[tp_dim] *= tp_mesh.mesh.numel()
+
+ # Construct TP-sharded DTensor using Megatron-style placement
+ param = DTensor.from_local(
+ local_tensor=local_tensor,
+ device_mesh=tp_mesh,
+ placements=placements,
+ run_check=run_check,
+ shape=tuple(global_shape),
+ stride=torch.empty(global_shape).stride(),
+ )
# Get FSDP-configured mesh and placements from provided param
device_mesh, placements = _get_fsdp_tensor_spec(
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py
index f18a21df6c1..cffdf09362f 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Iterable, List, Optional, Union
+from typing import Iterable, List, Union
import torch
import torch.distributed as dist
@@ -25,8 +25,6 @@
from torch.distributed.checkpoint.planner import TensorWriteData, WriteItem, WriteItemType
from torch.distributed.tensor.placement_types import Replicate, Shard, _StridedShard
-from .utils import get_mesh_names
-
def gather_and_compute_chunk_metadata(dtensor: DTensor) -> ChunkStorageMetadata:
"""
@@ -255,148 +253,145 @@ def preprocess_state_dict_for_uneven_dtensor(state_dict: dict) -> dict:
return state_dict
-def gather_uneven_dtensor_to_full_tensor(
- dtensor: DTensor, target_device: Optional[torch.device] = None
-) -> DTensor:
+def uneven_dtensor_to_full_tensor(dtensor: DTensor) -> torch.Tensor:
"""
- Gather an unevenly sharded DTensor distributed across multiple ranks,
- reconstructing the full (unsharded) tensor on each rank.
+ Gather a DTensor with potentially uneven sharding across ranks into a full tensor.
- This function handles uneven chunk sizes and offsets by collecting
- chunk metadata from all ranks, performing all-gather operations,
- and assembling the full tensor accordingly. The returned tensor
- is fully replicated across the given device mesh.
+ This function handles DTensors with uneven shards (where different ranks may have
+ different-sized chunks) by gathering chunk metadata and local tensors across all
+ ranks, then reconstructing the complete tensor.
Args:
- dtensor (DTensor): Distributed tensor with uneven sharding across ranks.
- target_device (Optional[torch.device]): If specified, move the resulting
- full tensor to this device. Otherwise, use the original device.
+ dtensor (DTensor): The distributed tensor to gather. Must have chunk metadata
+ available (either pre-existing or will be computed).
Returns:
- DTensor: Fully replicated DTensor representing the reconstructed full tensor.
+ torch.Tensor: The fully reconstructed tensor with shape matching the original
+ DTensor's global shape.
+
+ Raises:
+ TypeError: If input is not a DTensor.
+ ValueError: If chunk metadata is malformed (expected exactly one chunk per rank).
+ AssertionError: If an unexpected placement type is encountered after processing
+ Shard placements.
+
+ Note:
+ - This function performs collective operations (all_gather_object, all_gather)
+ across the device mesh, requiring synchronization across ranks.
+ - Works with Shard and _StridedShard placements, and expects Replicate placements
+ for non-sharded dimensions.
+ - The function modifies the DTensor in-place by adding chunk metadata if missing.
+
+ Example:
+ >>> mesh = DeviceMesh("cuda", [0, 1, 2, 3])
+ >>> # Create a DTensor with uneven sharding
+ >>> dtensor = DTensor(..., placements=[Shard(0)])
+ >>> full_tensor = gather_uneven_dtensor_to_full_tensor(dtensor)
+ >>> assert full_tensor.shape == dtensor.shape
"""
+ # Validate input type
if not isinstance(dtensor, DTensor):
- raise TypeError("Input must be a DTensor.")
-
- device_mesh = dtensor.device_mesh
- if not device_mesh.mesh_dim_names:
- process_group = device_mesh.get_group()
- else:
- # Check if the fully-flattened mesh exists first.
- full_flattened_mesh_dim_name = "_".join(device_mesh.mesh_dim_names)
- if full_flattened_mesh_dim_name in get_mesh_names(device_mesh):
- # Retrieve the existing flattened DeviceMesh ProcessGroup.
- try:
- # Two Cases: Name is a root dimension, or using the old DeviceMesh
- # API which allows us to get flattened dimensions.
- process_group = device_mesh[full_flattened_mesh_dim_name].get_group()
- except:
- # Name is a flattened dimension that cannot be retrieved from the
- # DeviceMesh.__getitem__, so fall-back to new DeviceMesh API.
- process_group = (
- device_mesh._get_root_mesh()
- ._flatten_mapping[full_flattened_mesh_dim_name]
- .get_group()
- )
- else:
- # Create the _-separated flattened DeviceMesh ProcessGroup.
- process_group = device_mesh._flatten().get_group()
+ raise TypeError(f"Input must be a DTensor, got {type(dtensor).__name__}.")
- # Collect chunk metadata for uneven shards (update if missing)
+ # Ensure chunk metadata is available for uneven shards
if not hasattr(dtensor._local_tensor, "__create_chunk_list__"):
update_uneven_dtensor_chunk_metadata(dtensor)
+ # Retrieve and validate chunk metadata
chunk_metadata_list = dtensor.__create_chunk_list__()
if len(chunk_metadata_list) != 1:
- raise ValueError(f"Expected exactly one chunk metadata, got {len(chunk_metadata_list)}.")
-
+ raise ValueError(
+ f"Expected exactly one chunk metadata per rank, got {len(chunk_metadata_list)}."
+ )
local_chunk_metadata = chunk_metadata_list[0]
- world_size = process_group.size()
-
- # Prepare local chunk info dictionary
- local_chunk_info = {
- "shape": list(dtensor.to_local().shape),
- "offset": getattr(local_chunk_metadata, "offsets", [0] * len(dtensor.shape)),
- "rank": process_group.rank(),
- }
-
- # Gather chunk info from all ranks
- all_chunk_info = [None] * world_size
- dist.all_gather_object(all_chunk_info, local_chunk_info, group=process_group)
-
- # Delegate to helper function
- return _assemble_full_tensor_from_uneven_chunks(
- dtensor, all_chunk_info, process_group, target_device
- )
+ # Prepare local chunk information for gathering
+ local_chunks_info = [
+ {
+ "shape": dtensor.to_local().shape,
+ "offset": getattr(local_chunk_metadata, "offsets", [0] * len(dtensor.shape)),
+ }
+ ]
+ local_buffer = dtensor.to_local().contiguous().view(-1)
+
+ # Iterate through device mesh dimensions and gather across sharded dimensions
+ for mesh_dim, placement in enumerate(dtensor.placements):
+ if isinstance(placement, (Shard, _StridedShard)):
+ # Get the process group for this mesh dimension
+ shard_group = dtensor.device_mesh.get_group(mesh_dim)
+
+ # Gather chunk metadata from all ranks in this dimension
+ group_chunks_info = [None] * shard_group.size()
+ dist.all_gather_object(group_chunks_info, local_chunks_info, group=shard_group)
+
+ # Prepare buffers for gathering tensors from all ranks
+ group_tensors = [
+ torch.empty(
+ sum(chunk["shape"].numel() for chunk in chunks_info),
+ dtype=dtensor.dtype,
+ device=dtensor.device,
+ )
+ for chunks_info in group_chunks_info
+ ]
-def _assemble_full_tensor_from_uneven_chunks(
- dtensor: DTensor,
- all_chunk_info: List[dict],
- process_group: torch.distributed.ProcessGroup,
- target_device: Optional[torch.device],
-) -> DTensor:
- """
- Assemble the full tensor from unevenly sized chunks gathered from all ranks.
-
- Args:
- dtensor (DTensor): The original distributed tensor.
- all_chunk_info (List[Dict]): List of shard info dicts from all ranks,
- including shapes and offsets.
- process_group: Process group for collective communication.
- target_device: Optional device to move the final full tensor onto.
-
- Returns:
- DTensor: Fully replicated tensor constructed by placing chunks at
- the appropriate offsets.
- """
- local_tensor = dtensor.to_local()
+ # Gather actual tensor data from all ranks
+ dist.all_gather(group_tensors, local_buffer, group=shard_group)
- # Check if the DTensor has any shard placements
- have_shard_placement = any(
- isinstance(placement, Shard) or isinstance(placement, _StridedShard)
- for placement in dtensor.placements
- )
+ # Flatten the gathered metadata and concatenate tensors
+ local_chunks_info = [item for sublist in group_chunks_info for item in sublist]
+ local_buffer = torch.cat(group_tensors)
+ elif not isinstance(placement, Replicate):
+ raise ValueError(
+ f"Unexpected placement {placement} at mesh dimension {mesh_dim}. "
+ f"Expected Shard, _StridedShard, or Replicate."
+ )
- if not have_shard_placement:
- # No sharding (replicated tensor), just clone and move if needed
- full_tensor = local_tensor.clone()
- if target_device:
- full_tensor = full_tensor.to(target_device)
- else:
- # Prepare empty buffers to receive tensors from each rank
- gathered_tensors = [
- torch.empty(rank_info["shape"], dtype=local_tensor.dtype, device=local_tensor.device)
- for rank_info in all_chunk_info
- ]
+ # Split the gathered buffer back into individual chunks
+ all_local_chunks = []
+ buffer_offset = 0
+ for chunk_info in local_chunks_info:
+ chunk_shape = chunk_info["shape"]
+ chunk_numel = chunk_shape.numel()
+ chunk_tensor = local_buffer[buffer_offset : buffer_offset + chunk_numel].view(chunk_shape)
+ all_local_chunks.append(chunk_tensor)
+ buffer_offset += chunk_numel
- # Gather local tensors from all ranks
- dist.all_gather(gathered_tensors, local_tensor, group=process_group)
+ # Reconstruct the full tensor by placing chunks at their correct offsets
+ full_tensor = torch.zeros(dtensor.shape, dtype=dtensor.dtype, device=dtensor.device)
+ for chunk_info, local_chunk in zip(local_chunks_info, all_local_chunks):
+ offset = chunk_info["offset"]
+ slices = tuple(slice(o, o + s) for o, s in zip(offset, local_chunk.shape))
+ full_tensor[slices] = local_chunk
- # Allocate full tensor buffer
- full_tensor = torch.empty(
- dtensor.shape, dtype=local_tensor.dtype, device=local_tensor.device
- )
+ return full_tensor
- # Copy each gathered shard into the full tensor at its offset
- for rank_info, local_shard in zip(all_chunk_info, gathered_tensors):
- offset = rank_info["offset"]
- slices = tuple(slice(o, o + s) for o, s in zip(offset, local_shard.shape))
- full_tensor[slices] = local_shard
- # Optionally move to target device
- if target_device is not None:
- full_tensor = full_tensor.to(target_device)
+def redistribute_uneven_dtensor_to_replicated(dtensor: DTensor) -> DTensor:
+ """
+ Redistribute an unevenly sharded DTensor to a fully replicated DTensor.
- # Free memory of gathered shards as they are copied
- del gathered_tensors
+ This function first gathers the unevenly sharded DTensor into a full tensor
+ and then redistributes it as a replicated DTensor across all ranks.
- # Wrap into a replicated DTensor and return
- return DTensor.from_local(
+ Args:
+ dtensor (DTensor): The unevenly sharded DTensor to redistribute.
+ Returns:
+ DTensor: A replicated DTensor with the same data as the input DTensor.
+ """
+ full_tensor = uneven_dtensor_to_full_tensor(dtensor)
+ replicated_dtensor = DTensor.from_local(
full_tensor,
placements=[Replicate()] * len(dtensor.placements),
device_mesh=dtensor.device_mesh,
)
+ return replicated_dtensor
+
+
+def gather_uneven_dtensor_to_full_tensor(dtensor: DTensor) -> DTensor:
+ """
+ Deprecated: use `redistribute_uneven_dtensor_to_replicated` instead.
+ """
+ return redistribute_uneven_dtensor_to_replicated(dtensor)
def _intersection(s1, s2):
@@ -476,7 +471,7 @@ def split_dtensor(
new_dtensor = DTensor.from_local(
sliced_tensor,
- shape=out_shape,
+ shape=tuple(out_shape),
stride=sliced_tensor.stride(),
placements=dtensor.placements,
device_mesh=dtensor.device_mesh,
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py
index b961a449d3e..f771c17c17d 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py
@@ -21,13 +21,6 @@
from importlib.metadata import version
from typing import Callable, Optional, Sequence, Union
-try:
- import megatron.core.parallel_state as parallel_state
-
- HAVE_MEGATRON_CORE = True
-except (ImportError, ModuleNotFoundError):
- HAVE_MEGATRON_CORE = False
-
try:
import einops
@@ -53,6 +46,13 @@
HAVE_TE = False
+try:
+ _torch_version = PkgVersion(torch.__version__)
+except Exception:
+ # This is a WAR for building docs, where torch is not actually imported
+ _torch_version = PkgVersion("0.0.0")
+
+
_MODEL_PARALLEL_RNG_TRACKER_NAME = "model-parallel-rng"
@@ -85,6 +85,13 @@ def is_te_min_version(vers, check_equality=True):
return te_version > PkgVersion(vers)
+def is_torch_min_version(version, check_equality=True):
+ """Check if minimum version of `torch` is installed."""
+ if check_equality:
+ return _torch_version >= PkgVersion(version)
+ return _torch_version > PkgVersion(version)
+
+
def is_submodule(module, parent_module, strict=True):
"""
Check if a module is a submodule of another module.
@@ -98,6 +105,23 @@ def is_submodule(module, parent_module, strict=True):
return False
+def find_megatron_fsdp(model):
+ """Walk the model wrapper chain to find a MegatronFSDP instance, if any."""
+ # Lazy import to avoid a circular import: megatron_fsdp.py transitively imports
+ # this module during its own initialization, so a top-level import of
+ # MegatronFSDP here would fail with a partially-initialized module error.
+ try:
+ from megatron.core.distributed.fsdp.src.megatron_fsdp.megatron_fsdp import MegatronFSDP
+ except (ImportError, ModuleNotFoundError):
+ return None
+ m = model
+ while m is not None:
+ if isinstance(m, MegatronFSDP):
+ return m
+ m = getattr(m, 'module', None)
+ return None
+
+
def get_mesh_names(
device_mesh: Optional[DeviceMesh] = None, only_submesh_dims: bool = False
) -> list[str]:
@@ -481,6 +505,8 @@ def __init__(
hybrid_fsdp_expt_group: Optional[torch.distributed.ProcessGroup] = None,
hsdp_outer_dp_shard: bool = False,
expt_device_mesh: Optional[DeviceMesh] = None,
+ fsdp_group_ag: Optional[torch.distributed.ProcessGroup] = None,
+ expt_fsdp_group_ag: Optional[torch.distributed.ProcessGroup] = None,
):
"""
Args:
@@ -502,6 +528,13 @@ def __init__(
just sharding across dp_shard ranks and replicating across dp_outer ranks.
expt_device_mesh (Optional[DeviceMesh]): The expert parallel device mesh
to use for the DistributedIndex.
+ fsdp_group_ag (Optional[torch.distributed.ProcessGroup]): Independent all-gather
+ process group for overlapping all-gather and reduce-scatter operations.
+ When provided, enables AG/RS overlap optimization for regular (non-expert)
+ parameters.
+ expt_fsdp_group_ag (Optional[torch.distributed.ProcessGroup]): Independent all-gather
+ process group for expert parameters in MoE models. When provided, enables AG/RS
+ overlap optimization for expert parameters.
"""
# Device mesh arguments.
self.device_mesh = device_mesh
@@ -514,10 +547,6 @@ def __init__(
self.hsdp_outer_dp_shard = hsdp_outer_dp_shard
self.expt_device_mesh = expt_device_mesh
- # Handling the situation where M-Core MoE EP=1
- if self.expt_device_mesh is None:
- self.expt_device_mesh = device_mesh
-
# Hybrid FSDP Process Groups
# Retrieve the FSDP process group from the DeviceMesh.
self.fsdp_group = (
@@ -525,13 +554,9 @@ def __init__(
if contains_submesh(self.device_mesh, self.dp_shard_dim)
else None
)
- # AG group comes from parallel_state, not the mesh
- # the purpose of this independent group is to overlap all-gather and gradient reduction.
- self.fsdp_group_ag = None
- if HAVE_MEGATRON_CORE and parallel_state.has_separate_all_gather_group():
- self.fsdp_group_ag = parallel_state.get_data_parallel_group(
- with_context_parallel=True, independent_all_gather=True
- )
+ # AG groups: supplied via ProcessGroupCollection (Megatron-FSDP entrypoint).
+ self.fsdp_group_ag = fsdp_group_ag
+ self.expt_fsdp_group_ag = expt_fsdp_group_ag
# Retrieve the outer-FSDP process group from the DeviceMesh.
self.outer_fsdp_group = (
self.device_mesh[self.dp_outer_dim].get_group()
@@ -630,7 +655,8 @@ def get_submesh(
"""
Retrieve an Megatron-FSDP-registered submesh by name(s).
"""
- if isinstance(mesh_dim_names, str):
+ if isinstance(mesh_dim_names, str) or mesh_dim_names is None:
+ # Create tuple from singleton dim or None.
mesh_dim_names = (mesh_dim_names,)
# Construct submesh identifier: (*mesh_dim_names, is_expert_parallel)
@@ -640,30 +666,22 @@ def get_submesh(
device_submesh = self.mesh_library.get(submesh_identifier, None)
if device_submesh is None:
+ device_mesh = self.expt_device_mesh if is_expert_parallel else self.device_mesh
# Warn about not specifying tp_dim for layers or frameworks that depend on this.
- if self.tp_dim is None and not is_expert_parallel:
+ if self.tp_dim is None:
logger.warning(
- "[FSDPDistributedIndex] Note: For TransformerEngine, or "
- "other machine learning frameworks like Megatron that assume "
+ "[FSDPDistributedIndex] For TransformerEngine, or other "
+ "machine learning frameworks like Megatron that assume "
"TP=1, you must specify tp_dim to use Megatron-FSDP. "
- "Create a trivial TP dimension by setting the TP dimension size "
- "to 1 in the DeviceMesh.\n"
- f"DeviceMesh: {self.device_mesh}"
- )
- elif self.tp_dim is None and is_expert_parallel:
- logger.warning(
- "[FSDPDistributedIndex] Note: For TransformerEngine, or "
- "other machine learning frameworks like Megatron that assume "
- "ETP=1, you must specify tp_dim to use Megatron-FSDP. "
- "Create a trivial ETP dimension by setting the ETP dimension size "
- "to 1 in the DeviceMesh.\n"
- f"DeviceMesh: {self.expt_device_mesh}"
+ "Create a trivial TP dimension by setting the TP dimension "
+ "size to 1 in the DeviceMesh.\n"
+ f"{'Expert ' if is_expert_parallel else ''}DeviceMesh: {device_mesh}"
)
-
raise ValueError(
f"[FSDPDistributedIndex][get_submesh] No submesh with "
f"mesh_dim_names={mesh_dim_names}, is_expert_parallel={is_expert_parallel} "
- f"has been registered with Megatron-FSDP."
+ f"has been registered with Megatron-FSDP.\n"
+ f"{'Expert ' if is_expert_parallel else ''}DeviceMesh: {device_mesh}"
)
return device_submesh
@@ -683,6 +701,8 @@ def get_fsdp_group(
) -> ProcessGroup:
"""Get the FSDP process group."""
if is_expert_parallel:
+ if independent_all_gather:
+ return self.expt_fsdp_group_ag
return self.expt_fsdp_group
if independent_all_gather:
return self.fsdp_group_ag
@@ -814,23 +834,35 @@ def is_mcore_tensor_model_parallel(param: torch.Tensor) -> bool:
"""
Check if the given parameter is Megatron-Core tensor model parallel.
"""
- return getattr(param, "_mcore_tp", False) or getattr(param, "tensor_model_parallel", False)
+ return get_mcore_tensor_parallel_partition_dim(param) is not None
def is_mcore_tensor_parallel_duplicated(param: torch.Tensor) -> bool:
"""
Check if the given parameter is Megatron-Core tensor model parallel and duplicated.
"""
- return getattr(param, "_tp_duplicated", False)
+ return get_mcore_tensor_parallel_partition_dim(param) is None
def get_mcore_tensor_parallel_partition_dim(param: torch.Tensor) -> Optional[int]:
"""
Get the partition dimension for a Megatron-Core tensor model parallel parameter.
"""
- if is_mcore_tensor_model_parallel(param):
- if hasattr(param, "_tp_partition_dim"):
- return param._tp_partition_dim
- else:
- return param.partition_dim
+ if hasattr(param, "_tensor_parallel_mode"):
+ if param._tensor_parallel_mode == "column":
+ return 0
+ elif param._tensor_parallel_mode == "row":
+ return 1
+ if getattr(param, "tensor_model_parallel", False):
+ partition_dim = getattr(param, "partition_dim", None)
+ if partition_dim is not None and partition_dim >= 0:
+ return int(partition_dim)
return None
+
+
+def using_tensor_parallel(dist_index, is_expert_parallel: bool = False) -> bool:
+ """
+ Check if tensor parallelism is being used based on the distributed index.
+ """
+ tp_mesh = dist_index.get_submesh(dist_index.tp_dim, is_expert_parallel=is_expert_parallel)
+ return tp_mesh.mesh.numel() > 1
diff --git a/megatron/core/distributed/fsdp/src/pyproject.toml b/megatron/core/distributed/fsdp/src/pyproject.toml
index 783030cc809..2845a14ab62 100644
--- a/megatron/core/distributed/fsdp/src/pyproject.toml
+++ b/megatron/core/distributed/fsdp/src/pyproject.toml
@@ -1,7 +1,7 @@
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
[build-system]
-requires = ["setuptools<80.0.0", "pybind11"]
+requires = ["setuptools>=80", "pybind11"]
build-backend = "setuptools.build_meta"
[tool.setuptools]
diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py
index f21fa0ef0d8..bbb6c14705e 100644
--- a/megatron/core/distributed/param_and_grad_buffer.py
+++ b/megatron/core/distributed/param_and_grad_buffer.py
@@ -20,12 +20,14 @@
from megatron.core.rerun_state_machine import get_rerun_state_machine
from megatron.core.utils import log_single_rank
+from ..fp4_utils import get_nvfp4_rowwise_packed_shape, is_nvfp4tensor
from ..fp8_utils import (
is_float8tensor,
is_mxfp8tensor,
modify_underlying_storage,
post_all_gather_processing,
)
+from ..optimizer.param_layout import pad_bucket_end, pad_param_start
from ..utils import is_torch_min_version, log_on_each_pipeline_stage
from .distributed_data_parallel_config import DistributedDataParallelConfig
from .reduce_scatter_with_fp32_accumulation import reduce_scatter_with_fp32_accumulation
@@ -122,7 +124,6 @@ def __init__(
self.layerwise_params_list = None
self.layerwise_param_flat_sizes = None
self.layerwise_gather_list = None
- self._layerwise_src_buffer = None
def set_layerwise_params_list(self, layerwise_params_list: List[List[torch.nn.Parameter]]):
"""Set per-rank parameter lists for layer-wise async all-gather.
@@ -199,6 +200,12 @@ def __init__(
self.params.add(param)
self.next_param_gather_bucket_group = None
+ # Set in DistributedDataParallel.__init__ when reduce_scatter_with_fp32_accumulation is on:
+ # points to the bucket group whose grad-reduce was dispatched immediately before mine in
+ # the backward pass. start_grad_sync drains this predecessor before dispatching its own
+ # collective, so the predecessor's intermediate all-to-all buffer is freed before the new
+ # one is allocated.
+ self.previous_grad_reduce_bucket_group = None
if self.ddp_config.num_distributed_optimizer_instances > 1:
self.inter_distributed_optimizer_instance_group = None
@@ -207,6 +214,16 @@ def __init__(
not self.ddp_config.reduce_scatter_with_fp32_accumulation
), "RS w/ FP32 accumulation not supported with num_distributed_optimizer_instances > 1"
+ reduction_collective = (
+ "reduce-scatter" if self.ddp_config.use_distributed_optimizer else "all-reduce"
+ )
+ log_single_rank(
+ logger,
+ logging.INFO,
+ f"Using {reduction_collective} for gradient reductions because "
+ f"{self.ddp_config.use_distributed_optimizer=}",
+ )
+
global dist_reduce_scatter_func
if self.ddp_config.reduce_scatter_with_fp32_accumulation:
dist_reduce_scatter_func = reduce_scatter_with_fp32_accumulation
@@ -235,6 +252,10 @@ def __init__(
self.param_gather_handle = None
self.param_gather_dispatched = False
self.grad_reduce_handle = None
+ # Per-iteration flag: True once finish_grad_sync has run this step. Lets a successor
+ # bucket group early-drain its predecessor without the end-of-step finalize loop
+ # double-waiting. Reset by `reset()`.
+ self.grad_reduce_finished = False
# Each time a local shard is created from bucket.param_data or bucket.grad_data, it
# introduces some CPU overheads. We use these two lists to cache the created local
@@ -255,6 +276,39 @@ def reset(self):
self.is_first_batch = False
self.per_param_grad_ready_counts = {}
self.is_last_microbatch = True
+ self.grad_reduce_finished = False
+
+ def _post_param_sync(self):
+ """Run post-processing after param all-gather completes."""
+ if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag:
+ for bucket in self.buckets:
+ is_bf16_weight_bucket = False
+ for param in bucket.params:
+ # Skip copying since bf16 weights in the mxfp8 model
+ # are already mapped to param.data.
+ if not is_float8tensor(param):
+ is_bf16_weight_bucket = True
+ break
+ param_start, param_end = bucket.param_to_index[param]
+ param_slice = bucket.param_data.view(-1)[param_start:param_end]
+ param.data.copy_(param_slice.view(param.data.shape))
+ if is_bf16_weight_bucket:
+ continue
+ # All-gathered params are not needed after being copied to param.data.
+ # Zero out the param buffer (shared with grad buffer) for gradient accumulation.
+ # We cannot zero out the entire grad buffer because one grad buffer may
+ # correspond to multiple param buffers. If we zero out the entire grad buffer,
+ # it would clear the data of those param buffers that have not yet completed AG.
+ bucket.param_data.zero_()
+ return
+
+ quantized_params = []
+ for bucket in self.buckets:
+ for param in bucket.params:
+ if is_float8tensor(param) or is_nvfp4tensor(param):
+ quantized_params.append(param)
+ if len(quantized_params) > 0:
+ post_all_gather_processing(quantized_params)
def check_grads(self, check_for_nan_or_inf, check_for_large):
"""
@@ -314,6 +368,7 @@ def start_param_sync(self, force_sync: bool = False):
if self.param_gather_handle is not None:
self.param_gather_handle.wait()
self.param_gather_handle = None
+ self._post_param_sync()
return
else:
assert self.param_gather_handle is None
@@ -321,8 +376,9 @@ def start_param_sync(self, force_sync: bool = False):
async_op = self.ddp_config.overlap_param_gather and not force_sync
if not self.ddp_config.use_distributed_optimizer:
- # Layer-wise optimizer path: use all_gather for variable-size
- # param gather.
+ # Legacy layer-wise optimizer path: use all_gather for variable-size
+ # param gather. Once all layerwise call sites set
+ # ddp_config.use_distributed_optimizer=True, this branch can be removed.
#
# Each rank may own a different number of params per bucket, so
# layerwise_param_flat_sizes can vary across ranks. PyTorch's NCCL
@@ -330,6 +386,12 @@ def start_param_sync(self, force_sync: bool = False):
# (falling back to grouped send/recv internally when sizes differ),
# so no manual padding is needed.
dp_size = self.intra_distributed_optimizer_instance_size
+ if dp_size == 1:
+ # Single-rank group (e.g., expt_dp_size == 1): no all-gather needed.
+ if force_sync and self.ddp_config.overlap_param_gather:
+ self._post_param_sync()
+ self.param_gather_dispatched = True
+ return
local_rank = self.intra_distributed_optimizer_instance_rank
group = self.intra_distributed_optimizer_instance_group
layerwise_work_handles = []
@@ -339,44 +401,38 @@ def start_param_sync(self, force_sync: bool = False):
param_dtype = bucket.params_list[0].dtype
if max(bucket.layerwise_param_flat_sizes) == 0:
- # All ranks have empty params for this bucket — skip.
bucket.layerwise_gather_list = None
continue
- # Flatten local params. Detach from the autograd graph because
- # start_param_sync can be called during the forward pass (where
- # autograd is active) and all_gather will write into gather_list
- # entries in-place.
local_size = bucket.layerwise_param_flat_sizes[local_rank]
+ total_gather_size = sum(bucket.layerwise_param_flat_sizes)
+
+ # Reuse grad_data as the all_gather receive buffer; it is idle
+ # during forward and grad_dtype.element_size >= param_dtype.
+ reuse_buf = bucket.grad_data.view(param_dtype)
+ assert reuse_buf.numel() >= total_gather_size
+
+ # Partition reuse_buf into contiguous per-rank receive slices.
+ gather_list = []
+ offset = 0
+ for i in range(dp_size):
+ size = bucket.layerwise_param_flat_sizes[i]
+ gather_list.append(reuse_buf[offset : offset + size])
+ offset += size
+ local_slot_view = gather_list[local_rank]
+
+ # Flatten local params and copy into the local rank's slot.
+ # Detach from autograd since start_param_sync may be called
+ # during the forward pass where autograd is active.
if local_size > 0:
flat_local_params = _flatten_dense_tensors(
bucket.layerwise_params_list[local_rank]
).detach()
- else:
- flat_local_params = torch.empty(
- 0, device=bucket.grad_data.device, dtype=param_dtype
- )
- # Keep flat_local_params alive until the async operation completes.
- bucket._layerwise_src_buffer = flat_local_params
-
- # Allocate per-rank receive buffers with actual sizes (no padding).
- # Reuse flat_local_params for local_rank's slot to avoid an extra allocation.
- gather_list = []
- for i in range(dp_size):
- if i == local_rank:
- gather_list.append(flat_local_params)
- else:
- gather_list.append(
- torch.empty(
- bucket.layerwise_param_flat_sizes[i],
- device=flat_local_params.device,
- dtype=flat_local_params.dtype,
- )
- )
+ local_slot_view.copy_(flat_local_params)
bucket.layerwise_gather_list = gather_list
work = torch.distributed.all_gather(
- gather_list, flat_local_params, group=group, async_op=async_op
+ gather_list, local_slot_view, group=group, async_op=async_op
)
if async_op and work is not None:
layerwise_work_handles.append(work)
@@ -397,7 +453,11 @@ def start_param_sync(self, force_sync: bool = False):
for updated_p, model_p in zip(updated_params, params):
model_p.data.copy_(updated_p)
bucket.layerwise_gather_list = None
- bucket._layerwise_src_buffer = None
+ # Zero out grad_data since it was reused as the all-gather
+ # receive buffer. Without this, accumulation into main_grad
+ # (a view into grad_data) would start from the result of the
+ # latest parameter all-gather instead of zero.
+ bucket.grad_data.zero_()
self.param_gather_handle = None
else:
# Standard distributed optimizer path: use _coalescing_manager.
@@ -427,6 +487,8 @@ def start_param_sync(self, force_sync: bool = False):
# (async_op=False) is used, `cm` is not None. Manually set to None for
# consistency with prior code.
self.param_gather_handle = None
+ if force_sync and self.ddp_config.overlap_param_gather:
+ self._post_param_sync()
self.param_gather_dispatched = True
def finish_param_sync(self, skip_next_bucket_dispatch: bool = False):
@@ -466,30 +528,7 @@ def finish_param_sync(self, skip_next_bucket_dispatch: bool = False):
else:
self.next_param_gather_bucket_group.start_param_sync()
- # For the mxfp8_param with "reuse_grad_buf_for_mxfp8_param_ag=True",
- # we need to copy the param_data from the shared_param/grad_buffer to param.data
- # after the param all-gather.
- if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag:
- for bucket in self.buckets:
- is_bf16_weight_bucket = False
- for param in bucket.params:
- # Skip copying since bf16 weights in the mxfp8 model
- # are already mapped to param.data.
- if not is_float8tensor(param):
- is_bf16_weight_bucket = True
- break
- param_start, param_end = bucket.param_to_index[param]
- param_slice = bucket.param_data.view(-1)[param_start:param_end]
- param.data.copy_(param_slice.view(param.data.shape))
- if is_bf16_weight_bucket:
- continue
- # All-gathered params are not needed after being copied to param.data.
- # Zero out the param buffer (shared with grad buffer) for gradient accumulation.
- # We cannot zero out the entire grad buffer because one grad buffer may
- # correspond to multiple param buffers. If we zero out the entire grad buffer,
- # it would clear the data of those param buffers that have not yet completed AG.
- bucket.param_data.zero_()
- elif not self.ddp_config.use_distributed_optimizer:
+ if not self.ddp_config.use_distributed_optimizer:
for bucket in self.buckets:
if bucket.layerwise_gather_list is None:
continue
@@ -507,15 +546,12 @@ def finish_param_sync(self, skip_next_bucket_dispatch: bool = False):
for updated_p, model_p in zip(updated_params, params):
model_p.data.copy_(updated_p)
bucket.layerwise_gather_list = None
- bucket._layerwise_src_buffer = None
- else:
- fp8_params = []
- for bucket in self.buckets:
- for param in bucket.params:
- if is_float8tensor(param):
- fp8_params.append(param)
- if len(fp8_params) > 0:
- post_all_gather_processing(fp8_params)
+ # Zero out grad_data since it was reused as the all-gather
+ # receive buffer. Without this, accumulation into main_grad
+ # (a view into grad_data) would start from the result of the
+ # latest parameter all-gather instead of zero.
+ bucket.grad_data.zero_()
+ self._post_param_sync()
def start_grad_sync(self, force_all_reduce: Optional[bool] = False):
"""
@@ -531,6 +567,23 @@ def start_grad_sync(self, force_all_reduce: Optional[bool] = False):
# already been dispatched.
return
+ # Drain the predecessor bucket group's reduce-scatter before allocating ours. Only
+ # linked under reduce_scatter_with_fp32_accumulation, which holds an intermediate
+ # all-to-all output tensor pinned until .wait() runs. We only drain when the
+ # predecessor has actually been dispatched this iteration (grad_reduce_handle set):
+ # backward param ordering does not always match bucket linkage order (e.g. NVFP4
+ # bucket layouts), so the predecessor may not have fired yet when we arrive here.
+ # In that case the predecessor will dispatch and drain on its own once its params
+ # become ready. The end-of-step finalize loop still catches any bucket that
+ # neither a successor nor itself drained.
+ if (
+ self.previous_grad_reduce_bucket_group is not None
+ and self.previous_grad_reduce_bucket_group.grad_reduce_handle is not None
+ ):
+ self.previous_grad_reduce_bucket_group.finish_grad_sync(
+ force_all_reduce=force_all_reduce
+ )
+
assert (
self.grad_reduce_handle is None
), "Should not have multiple communication calls outstanding at once"
@@ -676,6 +729,14 @@ def finish_grad_sync(self, force_all_reduce: Optional[bool] = False):
When ddp_config.overlap_grad_reduce is set to True, waits for asynchronous
communication call to complete. When ddp_config.overlap_grad_reduce is set to False,
makes synchronous call.
+
+ When ddp_config.overlap_grad_reduce is set to True, this method is idempotent
+ within an iteration: a second call is a no-op. This lets a successor bucket
+ group early-drain its predecessor at dispatch time (see
+ `previous_grad_reduce_bucket_group`) while still allowing the end-of-step
+ finalize loop to call this on every bucket without double-waiting. The
+ non-overlap path preserves its original per-call dispatch+wait behaviour
+ because it has no predecessor draining.
"""
self.param_gather_dispatched = False
# If overlap_grad_reduce is False, start (and finish) synchronous communication call here.
@@ -683,6 +744,8 @@ def finish_grad_sync(self, force_all_reduce: Optional[bool] = False):
self.start_grad_sync(force_all_reduce=force_all_reduce)
self._copy_back_extra_main_grads()
return
+ if self.grad_reduce_finished:
+ return
# If first batch, start asynchronous communication here. register_grad_ready() launches
# asynchronous communication only once self.golden_per_param_grad_ready_counts is
# populated at the end of this first batch.
@@ -693,6 +756,7 @@ def finish_grad_sync(self, force_all_reduce: Optional[bool] = False):
if self.ddp_config.num_distributed_optimizer_instances > 1:
torch.cuda.current_stream().wait_stream(self.communication_stream)
self._copy_back_extra_main_grads()
+ self.grad_reduce_finished = True
return
assert self.grad_reduce_handle is not None, (
f"Communication call has not been issued for this bucket "
@@ -702,6 +766,7 @@ def finish_grad_sync(self, force_all_reduce: Optional[bool] = False):
self.grad_reduce_handle.wait()
self.grad_reduce_handle = None
self._copy_back_extra_main_grads()
+ self.grad_reduce_finished = True
def free_overlap_buffers(self):
"""Free GPU buffers used by overlap param gather.
@@ -716,7 +781,6 @@ def free_overlap_buffers(self):
self.param_gather_handle = None
for bucket in self.buckets:
bucket.layerwise_gather_list = None
- bucket._layerwise_src_buffer = None
def _copy_back_extra_main_grads(self):
"""
@@ -758,6 +822,121 @@ def register_grad_ready(
self.start_grad_sync(force_all_reduce=force_all_reduce)
+def group_params_for_buffers(
+ params: List[torch.nn.Parameter], grad_reduce_in_fp32: bool
+) -> Dict['BufferKey', Tuple[List[torch.nn.Parameter], List[int]]]:
+ """Group parameters by buffer identity for buffer allocation.
+
+ Each distinct buffer is identified by a BufferKey with three dimensions:
+ - param_dtype: storage dtype (torch.uint8 for FP8/NVFP4 parameters, else param.dtype).
+ - grad_dtype: gradient reduction dtype (torch.float if grad_reduce_in_fp32, else param.dtype).
+ - is_expert_parallel: whether the parameter is expert-parallel (param.allreduce == False),
+ which requires a separate buffer with a different data-parallel group.
+
+ The param_indices track each parameter's position among same-dtype params (using
+ the "fake" high-precision dtype for FP8/NVFP4 params), needed for loading non-native-fp8
+ checkpoints in native-fp8 mode.
+
+ Args:
+ params: List of parameters to group.
+ grad_reduce_in_fp32: Whether gradients are reduced in FP32.
+
+ Returns:
+ Dict mapping BufferKey to (params_list, param_indices).
+ """
+ from ..optimizer.param_layout import BufferKey
+
+ key_to_params = {}
+ dtype_to_offsets = {}
+ key_to_indices = {}
+
+ for param in params:
+ assert param.requires_grad
+
+ param_dtype = param.dtype
+ if is_float8tensor(param) or is_nvfp4tensor(param):
+ param_dtype = torch.uint8
+ grad_dtype = torch.float if grad_reduce_in_fp32 else param.dtype
+ is_expert_parallel = not getattr(param, 'allreduce', True)
+ is_managed_by_layer_wise_optimizer = getattr(
+ param, 'is_managed_by_layer_wise_optimizer', False
+ )
+
+ key = BufferKey(
+ param_dtype, grad_dtype, is_expert_parallel, is_managed_by_layer_wise_optimizer
+ )
+ param_list = key_to_params.get(key, [])
+ param_list.append(param)
+ key_to_params[key] = param_list
+
+ # Use param.dtype (not param_dtype) so FP8/NVFP4 params share offsets with their
+ # logical high-precision dtype, needed for checkpoint compatibility.
+ offset_key = BufferKey(
+ param.dtype, grad_dtype, is_expert_parallel, is_managed_by_layer_wise_optimizer
+ )
+ offset = dtype_to_offsets.get(offset_key, 0)
+ dtype_to_offsets[offset_key] = offset + 1
+ indices = key_to_indices.get(key, [])
+ indices.append(offset)
+ key_to_indices[key] = indices
+
+ result = {}
+ for key, param_list in key_to_params.items():
+ result[key] = (param_list, key_to_indices[key])
+ return result
+
+
+def _compute_default_per_buffer_param_layout(
+ params: List[torch.nn.Parameter], bucket_size: Optional[int]
+) -> 'PerBufferParamLayout':
+ """Compute parameter layout for the non-distributed-optimizer case.
+
+ No padding is applied. Parameters are iterated in reverse order (backprop order)
+ and grouped into buckets of approximately `bucket_size` elements.
+
+ Args:
+ params: List of parameters to lay out.
+ bucket_size: Approximate number of elements per bucket, or None for a single bucket.
+
+ Returns:
+ PerBufferParamLayout with the computed mapping.
+ """
+ from ..optimizer.param_layout import PerBufferParamLayout
+
+ param_index_map = {}
+ bucket_indices = []
+ per_bucket_numel_unpadded = []
+
+ param_start_index = 0
+ bucket_start_index = 0
+ bucket_params = set()
+ bucket_id = 0
+
+ for param in params[::-1]:
+ this_numel = param.data.nelement()
+ param_end_index = param_start_index + this_numel
+ param_index_map[param] = (param_start_index, param_end_index, bucket_id)
+ bucket_params.add(param)
+
+ if bucket_size is not None and (param_end_index - bucket_start_index) >= bucket_size:
+ per_bucket_numel_unpadded.append(param_end_index - bucket_start_index)
+ bucket_indices.append((bucket_start_index, param_end_index))
+ bucket_start_index = param_end_index
+ bucket_params = set()
+ bucket_id += 1
+ param_start_index = param_end_index
+
+ if len(bucket_params) > 0:
+ per_bucket_numel_unpadded.append(param_end_index - bucket_start_index)
+ bucket_indices.append((bucket_start_index, param_end_index))
+
+ return PerBufferParamLayout(
+ param_index_map=param_index_map,
+ bucket_indices=bucket_indices,
+ per_bucket_numel_unpadded=per_bucket_numel_unpadded,
+ )
+
+
class _ParamAndGradBuffer:
"""
Groups parameters and gradients into a contiguous buffer, and then breaks the buffer into
@@ -793,6 +972,7 @@ def __init__(
param_indices: List[int],
nccl_ub: bool,
pg_collection: Optional[ProcessGroupCollection] = None,
+ param_layout: Optional['PerBufferParamLayout'] = None,
):
if pg_collection is None:
@@ -827,127 +1007,63 @@ def __init__(
# Data structures to store underlying buckets and relevant indexing data.
self.buckets = []
self.param_to_bucket = {} # Param -> bucket mapping.
- self.param_index_map = {} # Param -> location in buffer mapping (used in dist. optimizer).
- def _pad(number_to_be_padded: int, divisor: int) -> int:
- return int(math.ceil(number_to_be_padded / divisor) * divisor)
-
- def _pad_end_of_bucket_if_needed(bucket_end_index: int) -> int:
- """
- Pads end index of bucket if using distributed optimizer (to ensure uniform sharding).
- """
- if self.ddp_config.use_distributed_optimizer:
- # Workaround for TE bug causing cuBLAS to pick an incompatible algorithm.
- # This also helps cuBLAS pick more efficient algorithms for GEMMs.
- # We now ensure that all buckets start at a memory address that is 256-byte
- # aligned (128 values since params and grads use >= 16-bit precision).
- if self.ddp_config.pad_buckets_for_high_nccl_busbw:
- # Make sure the bucket size is divisible by a large power of 2 (2^16) to
- # ensure NCCL collectives have high bus bandwidth at large DP counts,
- # since NCCL message size (which for ring algorithms is bucket_size /
- # dp_size) apparently needs to be divisible by a power of 2 for high busbw.
- bucket_size_divisor = math.lcm(self.data_parallel_world_size, 128, 2**16)
- else:
- bucket_size_divisor = math.lcm(self.data_parallel_world_size, 128)
- return _pad(bucket_end_index, bucket_size_divisor)
- return bucket_end_index
-
- def _pad_start_of_param_if_needed(param_start_index: int) -> int:
- """
- Pads start index of param if using distributed optimizer (to ensure "good" alignment).
- """
- if self.ddp_config.use_distributed_optimizer:
- # Ensure that params start at 128-byte aligned addresses (64 values
- # since params are >= 16-bit precision).
- return _pad(param_start_index, 64)
- return param_start_index
-
- # First, figure out how many elements should be in the underlying buffer storage.
- # Note that if we need to split the buffer into smaller buckets, each of these
- # might need to be padded as well (if using the distributed optimizer).
- param_start_index = 0
- bucket_start_index = param_start_index
- bucket_params = set()
- self.bucket_indices = []
- per_bucket_numel_unpadded = []
- bucket_id = 0
-
- def _update_bucket_metadata(param_end_index: int) -> int:
- """
- Record metadata for the bucket starting at bucket_start_index and ending with the
- passed-in param_end_index. Returns the bucket's end_index.
- """
- nonlocal bucket_start_index, bucket_params, bucket_id
- per_bucket_numel_unpadded.append(param_end_index - bucket_start_index)
- bucket_end_index = _pad_end_of_bucket_if_needed(param_end_index)
-
- # Record metadata of new bucket.
- self.bucket_indices.append((bucket_start_index, bucket_end_index))
- bucket_start_index = bucket_end_index
-
- # Prepare for next bucket.
- bucket_params = set()
- bucket_id += 1
-
- # Return the potentially padded bucket_end_index.
- return bucket_end_index
-
- def _does_param_require_new_bucket(param):
- """
- Split shared embedding parameters into separate bucket if using distributed
- optimizer that makes use of reduce-scatters instead of all-reduces.
- This ensures that the first and last pipeline stage partition optimizer state
- for the shared embedding parameters the same way across DP replicas, allowing
- the DP reduce-scatter to be before the embedding all-reduce.
- """
- return (
- getattr(param, "shared_embedding", False)
- and self.ddp_config.use_distributed_optimizer
- )
-
- for param, _ in params_with_names[::-1]:
- # Iterate through parameters in reverse order to roughly follow backprop order.
-
- this_numel = param.data.nelement()
- param_start_index = _pad_start_of_param_if_needed(param_start_index)
-
- # Create bucket with collected parameters if current param needs its own bucket.
- if _does_param_require_new_bucket(param) and len(bucket_params) > 0:
- # Ensure this param accounts for the new padding introduced at end of
- # previous bucket.
- param_start_index = _update_bucket_metadata(param_start_index)
-
- param_end_index = param_start_index + this_numel
- self.param_index_map[param] = (param_start_index, param_end_index, bucket_id)
- bucket_params.add(param)
-
- # If we have enough elements already or the current param is part of the shared
- # embedding layer and needs a separate bucket, form a new bucket.
- if (
- bucket_size is not None and (param_end_index - bucket_start_index) >= bucket_size
- ) or _does_param_require_new_bucket(param):
- bucket_end_index = _update_bucket_metadata(param_end_index)
- param_start_index = bucket_end_index
- else:
- param_start_index = param_end_index
-
- # Add remaining params to a new bucket.
- if len(bucket_params) > 0:
- bucket_end_index = _update_bucket_metadata(param_end_index)
+ # Use the provided layout if given, otherwise compute the default (no-padding) layout.
+ if param_layout is None:
+ param_layout = _compute_default_per_buffer_param_layout(self.params, bucket_size)
+ self.param_index_map = param_layout.param_index_map
+ self.bucket_indices = param_layout.bucket_indices
+ per_bucket_numel_unpadded = param_layout.per_bucket_numel_unpadded
+
+ # Check if this buffer contains NVFP4 params.
+ #
+ # NVFP4 uses a dual-buffer layout: the param buffer stores packed bytes (half the
+ # logical numel) while the grad buffer uses the full numel. This is because NVFP4
+ # packs two FP4 values into a single uint8 byte for storage/communication, but
+ # gradients are computed and reduced in BF16 at full element count.
+ #
+ # Logical view: [v0, v1, v2, v3, ...] numel = N
+ #
+ # Param buffer (uint8): [byte0, byte1, ...] numel = N // 2
+ # ^^^^^ packs v0+v1
+ #
+ # Grad buffer: [g0, g1, g2, g3, ...] numel = N
+ #
+ # We therefore maintain two index maps:
+ # - param_index_map: offsets using full numel (from pre-computed layout).
+ # - nvfp4_packed_param_index_map: offsets into the packed param buffer (numel // 2).
+ #
+ # The packed index map is derived from param_index_map by iterating through
+ # the already-computed layout and halving numel for NVFP4 tensors.
+ #
+ self.has_nvfp4_params = any(is_nvfp4tensor(p) for p in self.params)
+ self.nvfp4_packed_param_index_map = None
+ self.nvfp4_packed_bucket_indices = None
+ if self.has_nvfp4_params:
+ self._compute_nvfp4_packed_layout(params_with_names)
# Next, create underlying storage for buffer (with numel elements that includes
# padding as necessary).
- self.numel = bucket_end_index
+ self.numel = self.bucket_indices[-1][1]
self.numel_unpadded = sum(per_bucket_numel_unpadded)
+ if self.has_nvfp4_params:
+ self.nvfp4_packed_numel = self.nvfp4_packed_bucket_indices[-1][1]
+ # nvfp4_packed_numel_unpadded is already set by _compute_nvfp4_packed_layout.
+
assert self.numel_unpadded <= self.numel
+ if self.has_nvfp4_params:
+ assert self.nvfp4_packed_numel_unpadded <= self.nvfp4_packed_numel
if self.ddp_config.use_distributed_optimizer:
assert self.numel % self.data_parallel_world_size == 0
+ if self.has_nvfp4_params:
+ assert self.nvfp4_packed_numel % self.data_parallel_world_size == 0
else:
assert self.numel == self.numel_unpadded
self.param_data = None
self.grad_data = None
self.extra_main_grads = []
+ self.nccl_mem_pool = None
if self.nccl_ub:
# If nccl_ub is True, use nccl_allocator to allocate memory for param_data/grad_data.
@@ -955,6 +1071,7 @@ def _does_param_require_new_bucket(param):
pool = nccl_allocator.create_nccl_mem_pool(
symmetric=not self.ddp_config.disable_symmetric_registration
)
+ self.nccl_mem_pool = pool
mem_alloc_context = functools.partial(
nccl_allocator.nccl_mem,
pool,
@@ -995,8 +1112,9 @@ def _does_param_require_new_bucket(param):
else:
# Only re-map param tensors if using distributed optimizer.
if self.ddp_config.use_distributed_optimizer:
+ numel = self.nvfp4_packed_numel if self.has_nvfp4_params else self.numel
self.param_data = torch.zeros(
- self.numel,
+ numel,
dtype=self.param_dtype,
device=torch.cuda.current_device(),
requires_grad=False,
@@ -1013,22 +1131,84 @@ def _does_param_require_new_bucket(param):
self.param_data_cpu = None
# Finally, map param.data and param.main_grad fields to buffers.
+ def _create_bucket(bucket_id, bucket_params, bucket_params_with_extra_main_grads):
+ """
+ Look up precomputed bucket indices and create a new bucket.
+
+ Args:
+ bucket_id: ID of the bucket to create.
+ bucket_params: List of parameters in this bucket.
+ bucket_params_with_extra_main_grads: List of parameters with
+ extra FP32 main_grads.
+
+ Returns:
+ A new _ParamAndGradBucket instance.
+ """
+ bucket_start_index, bucket_end_index = self.bucket_indices[bucket_id]
+ if self.has_nvfp4_params:
+ nvfp4_packed_start_index, nvfp4_packed_end_index = self.nvfp4_packed_bucket_indices[
+ bucket_id
+ ]
+ else:
+ nvfp4_packed_start_index, nvfp4_packed_end_index = None, None
+ return self._new_bucket(
+ bucket_params=bucket_params,
+ start_index=bucket_start_index,
+ end_index=bucket_end_index,
+ numel_unpadded=per_bucket_numel_unpadded[bucket_id],
+ bucket_id=bucket_id,
+ nvfp4_packed_start_index=nvfp4_packed_start_index,
+ nvfp4_packed_end_index=nvfp4_packed_end_index,
+ bucket_params_with_extra_main_grads=bucket_params_with_extra_main_grads,
+ )
+
bucket_params = []
bucket_params_with_extra_main_grads = []
- bucket_start_index = 0
cur_bucket_id = 0
for param, param_name in params_with_names[::-1]:
+ # Get parameter indices computed in previous loop.
param_start_index, param_end_index, bucket_id = self.param_index_map[param]
+ nvfp4_packed_param_start_index = None
+ if self.has_nvfp4_params:
+ nvfp4_packed_param_start_index, _, _ = self.nvfp4_packed_param_index_map[param]
# For MXFP8 param:
# we only need to map bf16 weights (layernorm, embedding, etc) to the buffer.
if not self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag or not is_mxfp8tensor(param):
if self.param_data is not None:
- new_param_data = self._get(
- param.data.shape, param_start_index, buffer_type=BufferType.PARAM
- )
- if is_float8tensor(param):
+ if is_nvfp4tensor(param):
+ # Remap the NVFP4 tensor's internal rowwise uint8 storage so it
+ # points into the contiguous DDP param buffer. This enables the
+ # all-gather to communicate packed NVFP4 bytes directly.
+ from ..fp4_utils import modify_nvfp4_rowwise_storage
+
+ packed_shape = get_nvfp4_rowwise_packed_shape(param.data.shape)
+ rowwise_bytes_view = self._get(
+ packed_shape,
+ nvfp4_packed_param_start_index,
+ buffer_type=BufferType.PARAM,
+ )
+ modify_nvfp4_rowwise_storage(param, rowwise_bytes_view)
+ elif is_float8tensor(param):
+ new_param_data = self._get(
+ param.data.shape,
+ (
+ nvfp4_packed_param_start_index
+ if self.has_nvfp4_params
+ else param_start_index
+ ),
+ buffer_type=BufferType.PARAM,
+ )
modify_underlying_storage(param, new_param_data)
else:
+ new_param_data = self._get(
+ param.data.shape,
+ (
+ nvfp4_packed_param_start_index
+ if self.has_nvfp4_params
+ else param_start_index
+ ),
+ buffer_type=BufferType.PARAM,
+ )
old_param_data = param.data
param.data = new_param_data
assert old_param_data._base is None
@@ -1036,10 +1216,10 @@ def _does_param_require_new_bucket(param):
param.data.detach().copy_(old_param_data)
del old_param_data
+ # Grad buffer always uses full-numel offsets from param_index_map.
param.main_grad = self._get(
param.data.shape, param_start_index, buffer_type=BufferType.GRAD
)
-
# Create FP32 copy of .main_grads if necessary.
promote_main_grads_to_higher_precision = False
for param_name_pattern in ddp_config.param_name_patterns_for_fp32_local_accumulation:
@@ -1064,18 +1244,11 @@ def _does_param_require_new_bucket(param):
self.extra_main_grads.append(param.main_grad)
if bucket_id != cur_bucket_id:
- bucket_end_index = _pad_end_of_bucket_if_needed(param_start_index)
self.buckets.append(
- self._new_bucket(
- bucket_params=bucket_params,
- start_index=bucket_start_index,
- end_index=bucket_end_index,
- numel_unpadded=per_bucket_numel_unpadded[cur_bucket_id],
- bucket_id=cur_bucket_id,
- bucket_params_with_extra_main_grads=bucket_params_with_extra_main_grads,
+ _create_bucket(
+ cur_bucket_id, bucket_params, bucket_params_with_extra_main_grads
)
)
- bucket_start_index = bucket_end_index
bucket_params = []
bucket_params_with_extra_main_grads = []
assert cur_bucket_id + 1 == len(self.buckets)
@@ -1096,18 +1269,9 @@ def _does_param_require_new_bucket(param):
torch.cuda.synchronize()
# Add remaining params to a new bucket.
if len(bucket_params) > 0:
- bucket_end_index = _pad_end_of_bucket_if_needed(param_end_index)
self.buckets.append(
- self._new_bucket(
- bucket_params=bucket_params,
- start_index=bucket_start_index,
- end_index=bucket_end_index,
- numel_unpadded=per_bucket_numel_unpadded[cur_bucket_id],
- bucket_id=cur_bucket_id,
- bucket_params_with_extra_main_grads=bucket_params_with_extra_main_grads,
- )
+ _create_bucket(cur_bucket_id, bucket_params, bucket_params_with_extra_main_grads)
)
-
# Log buckets for all PP stages.
log_strs = []
log_strs.append(
@@ -1132,6 +1296,93 @@ def _does_param_require_new_bucket(param):
dp_cp_group=self.dp_cp_group,
)
+ def _compute_nvfp4_packed_layout(self, params_with_names):
+ """Derive packed NVFP4 index map and bucket indices from the primary layout.
+
+ The primary layout (self.param_index_map, self.bucket_indices) uses full numel
+ for all params. NVFP4 tensors pack two FP4 values into one byte, so the param
+ buffer needs a separate "packed" index map where NVFP4 params occupy half the
+ space. Non-NVFP4 params keep their full numel in the packed space.
+
+ The same padding rules used by the primary layout are applied here:
+ - 64-element alignment at the start of each param.
+ - Bucket-end padding for DP-divisibility (when using distributed optimizer).
+
+ Sets:
+ self.nvfp4_packed_param_index_map: param -> (start, end, bucket_id)
+ self.nvfp4_packed_bucket_indices: list of (start, end) per bucket
+ self.nvfp4_packed_numel_unpadded: total unpadded elements across all buckets
+ """
+
+ def _pad_start_of_param(param_start_index: int) -> int:
+ if self.ddp_config.use_distributed_optimizer:
+ return pad_param_start(param_start_index)
+ return param_start_index
+
+ def _pad_end_of_bucket(bucket_end_index: int) -> int:
+ if self.ddp_config.use_distributed_optimizer:
+ return pad_bucket_end(
+ bucket_end_index,
+ self.data_parallel_world_size,
+ self.ddp_config.pad_buckets_for_high_nccl_busbw,
+ )
+ return bucket_end_index
+
+ self.nvfp4_packed_param_index_map = {}
+ self.nvfp4_packed_bucket_indices = []
+ nvfp4_packed_per_bucket_numel_unpadded = []
+
+ packed_param_start = 0
+ packed_bucket_start = 0
+ cur_bucket_id = 0
+
+ for param, _ in params_with_names[::-1]:
+ _, _, bucket_id = self.param_index_map[param]
+ param_numel = param.data.nelement()
+
+ packed_param_start = _pad_start_of_param(packed_param_start)
+
+ # Finalize previous bucket if we've moved to a new one.
+ if bucket_id != cur_bucket_id:
+ # Record unpadded numel, then pad the bucket end.
+ nvfp4_packed_per_bucket_numel_unpadded.append(
+ packed_param_start - packed_bucket_start
+ )
+ packed_bucket_end = _pad_end_of_bucket(packed_param_start)
+ self.nvfp4_packed_bucket_indices.append((packed_bucket_start, packed_bucket_end))
+ packed_bucket_start = packed_bucket_end
+ packed_param_start = packed_bucket_start
+ cur_bucket_id = bucket_id
+
+ # NVFP4 tensors use half the numel in the packed param buffer.
+ if is_nvfp4tensor(param):
+ assert (
+ param_numel % 2 == 0
+ ), f"NVFP4 requires even numel for packing, got {param_numel}"
+ packed_numel = param_numel // 2
+ else:
+ packed_numel = param_numel
+
+ packed_param_end = packed_param_start + packed_numel
+ self.nvfp4_packed_param_index_map[param] = (
+ packed_param_start,
+ packed_param_end,
+ bucket_id,
+ )
+ packed_param_start = packed_param_end
+
+ # Finalize last bucket.
+ if packed_param_start > packed_bucket_start:
+ nvfp4_packed_per_bucket_numel_unpadded.append(packed_param_start - packed_bucket_start)
+ packed_bucket_end = _pad_end_of_bucket(packed_param_start)
+ self.nvfp4_packed_bucket_indices.append((packed_bucket_start, packed_bucket_end))
+
+ assert len(self.nvfp4_packed_bucket_indices) == len(self.bucket_indices), (
+ f"Packed bucket count ({len(self.nvfp4_packed_bucket_indices)}) != "
+ f"primary bucket count ({len(self.bucket_indices)})"
+ )
+ self.nvfp4_packed_numel_unpadded = sum(nvfp4_packed_per_bucket_numel_unpadded)
+
def scale_gradients(self, scaling_factor: float) -> None:
"""Scale the gradient data by `scaling_factor`."""
self.grad_data *= scaling_factor
@@ -1144,11 +1395,13 @@ def _get(self, shape: torch.Size, start_index: int, buffer_type: BufferType) ->
`start_index`.
"""
end_index = start_index + shape.numel()
- assert end_index <= self.numel, "Requested tensor is out of buffer range"
if buffer_type == BufferType.PARAM:
+ numel = self.nvfp4_packed_numel if self.has_nvfp4_params else self.numel
+ assert end_index <= numel, "Requested tensor is out of param buffer range"
assert self.param_data is not None
buffer_tensor = self.param_data[start_index:end_index]
elif buffer_type == BufferType.GRAD:
+ assert end_index <= self.numel, "Requested tensor is out of grad buffer range"
buffer_tensor = self.grad_data[start_index:end_index]
else:
raise Exception("Illegal buffer type provided to GradBuffer._get() function")
@@ -1163,24 +1416,46 @@ def _new_bucket(
numel_unpadded: int,
bucket_id: int,
bucket_params_with_extra_main_grads: List[torch.Tensor],
+ nvfp4_packed_start_index: int = None,
+ nvfp4_packed_end_index: int = None,
) -> _ParamAndGradBucket:
"""
Helper function that creates a new bucket. Also updates param->bucket mapping.
+
+ For NVFP4 buffers, nvfp4_packed_start_index and nvfp4_packed_end_index
+ are provided separately because the param buffer uses packed numel while
+ the grad buffer uses full numel.
"""
# Assert that indices are correctly padded (if needed), and that bucket
# position is same as originally computed.
+
if self.ddp_config.use_distributed_optimizer:
assert start_index % self.data_parallel_world_size == 0
assert end_index % self.data_parallel_world_size == 0
assert (start_index, end_index) == self.bucket_indices[bucket_id]
+ if nvfp4_packed_start_index is not None:
+ assert (
+ nvfp4_packed_start_index,
+ nvfp4_packed_end_index,
+ ) == self.nvfp4_packed_bucket_indices[bucket_id]
# Get appropriate view into global _ParamAndGradBuffer.
+ # For NVFP4, param buffer uses packed offsets; otherwise same as start/end.
bucketed_param_data = None
if self.param_data is not None:
- bucketed_param_data = self._get(
- torch.Size([end_index - start_index]), start_index, buffer_type=BufferType.PARAM
- )
+ if nvfp4_packed_start_index is not None:
+ assert nvfp4_packed_end_index is not None
+ bucketed_param_data = self._get(
+ torch.Size([nvfp4_packed_end_index - nvfp4_packed_start_index]),
+ nvfp4_packed_start_index,
+ buffer_type=BufferType.PARAM,
+ )
+ else:
+ bucketed_param_data = self._get(
+ torch.Size([end_index - start_index]), start_index, buffer_type=BufferType.PARAM
+ )
+ # Grad buffer always uses full-numel offsets.
bucketed_grad_data = self._get(
torch.Size([end_index - start_index]), start_index, buffer_type=BufferType.GRAD
)
@@ -1243,7 +1518,9 @@ def reload_from_cpu(self, move_params: bool = True, move_grads: bool = True):
def partition_buckets(
- buffers: List[_ParamAndGradBuffer], force_single_bucket_group: bool = False
+ buffers: List[_ParamAndGradBuffer],
+ force_single_bucket_group: bool = False,
+ reduce_scatter_with_fp32_accumulation: bool = False,
) -> List[_ParamAndGradBucketGroup]:
"""
Automatically regroup the buckets of input buffers and return a list of bucket groups.
@@ -1283,12 +1560,16 @@ def partition_buckets(
if len(buffers) == 0:
return []
- dtype_to_buffer_map = {}
+ # At most one fp8 (uint8) buffer is allowed; Cases 2 and 3 below branch on
+ # whether one is present. Non-uint8 dtypes can legitimately appear in
+ # multiple buffers (e.g. LayerWise-managed bf16 weights + Adam-managed bf16
+ # biases share the bf16 ``param_dtype`` but live in separate buffers), so
+ # the uniqueness check is restricted to uint8.
+ fp8_buffer = None
for buffer in buffers:
- dtype = buffer.param_dtype
- # Make sure that the param_dtype of any two buffers is different.
- assert dtype not in dtype_to_buffer_map
- dtype_to_buffer_map[dtype] = buffer
+ if buffer.param_dtype == torch.uint8:
+ assert fp8_buffer is None
+ fp8_buffer = buffer
# Case 1: Put all buckets into a single bucket group if force_single_bucket_group is True.
if force_single_bucket_group:
@@ -1307,7 +1588,7 @@ def partition_buckets(
)
return [bucket_group]
- if torch.uint8 not in dtype_to_buffer_map:
+ if fp8_buffer is None:
# Case 2: When there is no fp8 buffer in the input buffers, let each bucket group have
# only one bucket.
bucket_groups = []
@@ -1331,11 +1612,36 @@ def partition_buckets(
non_fp8_buckets.append(bucket)
bucket_groups = []
- fp8_buffer = dtype_to_buffer_map[torch.uint8]
for bucket in fp8_buffer.buckets:
if len(bucket_groups) == len(fp8_buffer.buckets) - 1:
- # The last bucket group.
- group_buckets = [bucket] + non_fp8_buckets
+ # reduce_scatter_with_fp32_accumulation requires exactly one bucket
+ # per group (see assert in _ParamAndGradBucketGroup.reduce_scatter).
+ # Without this flag the non-FP8 buckets would be merged into the last
+ # FP8 group, violating that constraint. So we split them out into
+ # their own individual groups instead.
+ if reduce_scatter_with_fp32_accumulation:
+ bucket_groups.append(
+ _ParamAndGradBucketGroup(
+ [bucket],
+ buffer.ddp_config,
+ buffer.data_parallel_group,
+ buffer.data_parallel_world_size,
+ )
+ )
+ if non_fp8_buckets:
+ for non_fp8_bucket in non_fp8_buckets:
+ bucket_groups.append(
+ _ParamAndGradBucketGroup(
+ [non_fp8_bucket],
+ buffer.ddp_config,
+ buffer.data_parallel_group,
+ buffer.data_parallel_world_size,
+ )
+ )
+
+ continue # Skip the default bucket group creation below
+ else:
+ group_buckets = [bucket] + non_fp8_buckets
else:
# The first N-1 bucket groups.
group_buckets = [bucket]
diff --git a/megatron/core/energy_monitor.py b/megatron/core/energy_monitor.py
index 4334cfe3873..c6f14ee269a 100644
--- a/megatron/core/energy_monitor.py
+++ b/megatron/core/energy_monitor.py
@@ -59,7 +59,11 @@ def resume(self) -> None:
def _get_energy(self) -> int:
"""Get current energy consumption from NVML."""
try:
- return nvmlDeviceGetTotalEnergyConsumption(self._handle)
+ # Passing None to nvmlDeviceGetTotalEnergyConsumption can cause a core
+ # dump, so short circuit if self._handle is None.
+ if self._handle is not None:
+ return nvmlDeviceGetTotalEnergyConsumption(self._handle)
+ return self._last_energy
except NVMLError:
return self._last_energy # return *something* if it errors
diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
index 28d2f8894e3..85214a1fd57 100755
--- a/megatron/core/extensions/transformer_engine.py
+++ b/megatron/core/extensions/transformer_engine.py
@@ -1,5 +1,7 @@
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+from __future__ import annotations
+import copy
import dataclasses
import enum
import inspect
@@ -34,10 +36,12 @@
)
from megatron.core.process_groups_config import ProcessGroupCollection
from megatron.core.quantization.quant_config import QuantizationConfig
+from megatron.core.quantization.utils import get_quant_config_or_none
from megatron.core.tensor_parallel.layers import (
_initialize_affine_weight_cpu,
set_tensor_model_parallel_attributes,
)
+from megatron.core.tensor_parallel.mappings import gather_from_tensor_model_parallel_region
from megatron.core.tensor_parallel.random import (
get_cuda_rng_tracker,
get_data_parallel_rng_tracker_name,
@@ -45,7 +49,7 @@
)
from megatron.core.tensor_parallel.utils import divide
from megatron.core.transformer.enums import AttnMaskType
-from megatron.core.transformer.mlp import MLP
+from megatron.core.transformer.mlp import MLP, MLPSubmodules
from megatron.core.transformer.torch_norm import LayerNormInterface
from megatron.core.transformer.transformer_config import TransformerConfig
from megatron.core.transformer.utils import (
@@ -65,7 +69,7 @@
try:
import transformer_engine as te
- from transformer_engine.pytorch.fp8 import FP8GlobalStateManager, fp8_autocast
+ from transformer_engine.pytorch.fp8 import FP8GlobalStateManager, fp8_autocast, fp8_model_init
HAVE_TE = True
except ImportError:
@@ -125,6 +129,14 @@ class TEQuantizationRecipe:
If an amax reduction is applicable, such as in per-tensor quantization recipe,
whether to reduce only along TP groups.
"""
+ fp8_param: bool = False
+ """
+ If cast the initialized parameters to fp8 precision and all-gather weights in FP8.
+ """
+ fp4_param: bool = False
+ """
+ If cast the initialized parameters to fp4 precision and all-gather weights in FP4.
+ """
@classmethod
def parse_from_config(cls, quant_config: Dict[Any, Any]) -> "TEQuantizationRecipe":
@@ -207,6 +219,61 @@ def parse_from_config(quant_config: QuantizationConfig) -> "TEQuantizationParams
raise NotImplementedError(f"Unhandled configuration type {config_type}")
+def _get_fp8_model_init_for_quant_recipe(qrecipe: TEQuantizationRecipe):
+ if qrecipe.fp8_quantization_recipe is None and qrecipe.fp4_quantization_recipe is None:
+ enabled = False
+ quant_recipe = None
+ elif qrecipe.fp8_quantization_recipe is not None:
+ enabled = qrecipe.fp8_param
+ if qrecipe.fp8_format == "e4m3":
+ fp8_format = te.common.recipe.Format.E4M3
+ elif qrecipe.fp8_format == "hybrid":
+ fp8_format = te.common.recipe.Format.HYBRID
+ else:
+ raise ValueError(f"Unhandled fp8_format {qrecipe.fp8_format}")
+
+ if qrecipe.fp8_quantization_recipe == Fp8Recipe.custom:
+ from megatron.core.fp8_utils import _get_custom_recipe
+
+ assert qrecipe.custom_recipe_factory is not None
+ quant_recipe = _get_custom_recipe(qrecipe.custom_recipe_factory)
+ elif qrecipe.fp8_quantization_recipe == Fp8Recipe.tensorwise:
+ quant_recipe = te.common.recipe.Float8CurrentScaling(fp8_format=fp8_format)
+ elif qrecipe.fp8_quantization_recipe == Fp8Recipe.blockwise:
+ quant_recipe = te.common.recipe.Float8BlockScaling(fp8_format=fp8_format)
+ elif qrecipe.fp8_quantization_recipe == Fp8Recipe.mxfp8:
+ quant_recipe = te.common.recipe.MXFP8BlockScaling(fp8_format=fp8_format)
+ else:
+ raise ValueError(f"Unhandled fp8 recipe: {qrecipe.fp8_quantization_recipe}")
+ else:
+ # Fp4 configured.
+ enabled = qrecipe.fp4_param
+ if qrecipe.fp4_quantization_recipe == Fp4Recipe.custom:
+ from megatron.core.fp8_utils import _get_custom_recipe
+
+ assert qrecipe.custom_recipe_factory is not None
+ quant_recipe = _get_custom_recipe(qrecipe.custom_recipe_factory)
+ elif qrecipe.fp4_quantization_recipe == Fp4Recipe.nvfp4:
+ quant_recipe = te.common.recipe.NVFP4BlockScaling()
+ else:
+ raise ValueError(f"Unhandled fp4 recipe: {qrecipe.fp4_quantization_recipe}")
+
+ return fp8_model_init(
+ enabled=enabled,
+ recipe=quant_recipe,
+ preserve_high_precision_init_val=torch.is_grad_enabled(),
+ )
+
+
+def _get_fp8_model_init_for_quant_params(qparams: TEQuantizationParams | None, training: bool):
+ if qparams is None:
+ return nullcontext()
+ elif not training and qparams.evaluation_recipe is not None:
+ return _get_fp8_model_init_for_quant_recipe(qparams.evaluation_recipe)
+ else:
+ return _get_fp8_model_init_for_quant_recipe(qparams.training_recipe)
+
+
def _get_fp8_autocast_for_quant_recipe(qrecipe: TEQuantizationRecipe):
if FP8GlobalStateManager.is_fp8_enabled():
if not qrecipe.override_quantized_autocast:
@@ -706,7 +773,7 @@ def __init__(
output_size: int,
*,
parallel_mode: Optional[str],
- config: ModelParallelConfig,
+ config: TransformerConfig,
init_method: Callable,
bias: bool,
skip_bias_add: bool,
@@ -715,7 +782,12 @@ def __init__(
is_expert: bool = False,
symmetric_ar_type: Optional[str] = None,
tp_group: Optional[torch.distributed.ProcessGroup] = None,
+ name: str | None = None,
):
+ """
+ Args:
+ name (str | None): module instance name passed top-down from its paranet module
+ """
if not HAVE_TE:
raise ImportError(
"Transformer Engine is not installed. "
@@ -883,24 +955,31 @@ def __init__(
UserWarning,
)
- super().__init__(
- in_features=input_size,
- out_features=output_size,
- sequence_parallel=self.config.sequence_parallel,
- fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
- # Pass None if not initialized for backward compatibility with the ckpt converter.
- tp_group=tp_group_for_te if torch.distributed.is_initialized() else None,
- tp_size=tp_size,
- get_rng_state_tracker=(
- get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
- ),
- init_method=condition_init_method(config, init_method),
- bias=bias,
- return_bias=self.te_return_bias,
- parallel_mode=te_parallel_mode,
- **extra_kwargs,
- )
self.te_quant_params: Optional[TEQuantizationParams] = None
+ quant_config = get_quant_config_or_none(name, config.quant_recipe)
+ self.finish_init(quant_config)
+ init_quant_context = _get_fp8_model_init_for_quant_params(
+ self.te_quant_params, torch.is_grad_enabled()
+ )
+
+ with init_quant_context:
+ super().__init__(
+ in_features=input_size,
+ out_features=output_size,
+ sequence_parallel=self.config.sequence_parallel,
+ fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
+ # Pass None if not initialized for backward compatibility with the ckpt converter.
+ tp_group=tp_group_for_te if torch.distributed.is_initialized() else None,
+ tp_size=tp_size,
+ get_rng_state_tracker=(
+ get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
+ ),
+ init_method=condition_init_method(config, init_method),
+ bias=bias,
+ return_bias=self.te_return_bias,
+ parallel_mode=te_parallel_mode,
+ **extra_kwargs,
+ )
for param in self.parameters():
if is_expert:
@@ -932,7 +1011,7 @@ def will_execute_quantized(self, is_context_quantized: bool) -> bool:
self.te_quant_params, self.training, is_context_quantized
)
- def forward(self, x):
+ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]:
"""Forward."""
_is_first_microbatch = (
None if self.disable_parameter_transpose_cache else self.is_first_microbatch
@@ -993,7 +1072,12 @@ def __init__(
tp_comm_buffer_name: Optional[str] = None,
tp_group: Optional[torch.distributed.ProcessGroup] = None,
stride: int = 1,
+ name: str | None = None,
):
+ """
+ Args:
+ name (str | None): module instance name passed top-down from its paranet module
+ """
if not HAVE_TE:
raise ImportError(
"Transformer Engine is not installed. "
@@ -1111,30 +1195,37 @@ def __init__(
self.stride = stride
- super().__init__(
- in_features=input_size,
- out_features=output_size,
- eps=self.config.layernorm_epsilon,
- sequence_parallel=self.config.sequence_parallel,
- fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
- tp_group=tp_group if torch.distributed.is_initialized() else None,
- tp_size=self.config.tensor_model_parallel_size,
- get_rng_state_tracker=(
- get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
- ),
- init_method=(
- condition_init_method(config, init_method)
- if not config.use_cpu_initialization
- else lambda w: None
- ),
- bias=bias,
- return_bias=self.te_return_bias,
- parallel_mode="column",
- return_layernorm_output=False,
- zero_centered_gamma=self.config.layernorm_zero_centered_gamma,
- **extra_kwargs,
- )
self.te_quant_params: Optional[TEQuantizationParams] = None
+ quant_config = get_quant_config_or_none(name, config.quant_recipe)
+ self.finish_init(quant_config)
+ init_quant_context = _get_fp8_model_init_for_quant_params(
+ self.te_quant_params, torch.is_grad_enabled()
+ )
+
+ with init_quant_context:
+ super().__init__(
+ in_features=input_size,
+ out_features=output_size,
+ eps=self.config.layernorm_epsilon,
+ sequence_parallel=self.config.sequence_parallel,
+ fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
+ tp_group=tp_group if torch.distributed.is_initialized() else None,
+ tp_size=self.config.tensor_model_parallel_size,
+ get_rng_state_tracker=(
+ get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
+ ),
+ init_method=(
+ condition_init_method(config, init_method)
+ if not config.use_cpu_initialization
+ else lambda w: None
+ ),
+ bias=bias,
+ return_bias=self.te_return_bias,
+ parallel_mode="column",
+ return_layernorm_output=False,
+ zero_centered_gamma=self.config.layernorm_zero_centered_gamma,
+ **extra_kwargs,
+ )
# Set proper partition_stride
setattr(self.weight, 'partition_stride', stride)
@@ -1235,7 +1326,7 @@ def __init__(
input_size: int,
output_size: int,
*,
- config: ModelParallelConfig,
+ config: TransformerConfig,
init_method: Callable,
gather_output: bool,
bias: bool,
@@ -1245,7 +1336,12 @@ def __init__(
tp_comm_buffer_name: Optional[str] = None,
tp_group: Optional[torch.distributed.ProcessGroup] = None,
stride: int = 1,
+ name: str | None = None,
):
+ """
+ Args:
+ name (str | None): module instance name passed top-down from its paranet module
+ """
if not HAVE_TE:
raise ImportError(
"Transformer Engine is not installed. "
@@ -1277,6 +1373,7 @@ def __init__(
tp_comm_buffer_name=tp_comm_buffer_name,
symmetric_ar_type=config.symmetric_ar_type,
tp_group=tp_group,
+ name=name,
)
# Set proper partition_stride
@@ -1336,6 +1433,135 @@ def backward_dw(self):
super().backward_dw()
+class TELMHeadColumnParallelLinear(TEColumnParallelLinear):
+ """Wrapper for ``TEColumnParallelLinear`` used as the LM-head output projection under MXFP8.
+
+ Drop-in replacement for the ``tensor_parallel.ColumnParallelLinear`` LM head:
+ ``delay_wgrad_compute`` is forced off to mirror its no-op ``backward_dw``,
+ and ``get/set_extra_state`` match the bf16 LM head's state-dict shim. The
+ LM-head kwargs ``keep_master_weight_for_test``, ``skip_weight_param_allocation``,
+ ``defer_embedding_wgrad_compute`` buffers, and ``disable_grad_reduce`` are
+ accepted to preserve the ``ColumnParallelLinear`` signature but currently
+ raise when set non-default — TE will not support them natively, so they
+ would have to be implemented in this subclass, which has not been done yet.
+
+ Active only when ``fp8_output_proj=True`` with ``fp8_recipe='mxfp8'``.
+ """
+
+ def __init__(
+ self,
+ input_size,
+ output_size,
+ *,
+ config,
+ init_method,
+ bias=True,
+ gather_output=False,
+ stride=1,
+ keep_master_weight_for_test=False,
+ skip_bias_add=False,
+ skip_weight_param_allocation: bool = False,
+ embedding_activation_buffer=None,
+ grad_output_buffer=None,
+ is_expert: bool = False,
+ tp_comm_buffer_name: Optional[str] = None,
+ disable_grad_reduce: bool = False,
+ tp_group: Optional[torch.distributed.ProcessGroup] = None,
+ ):
+ from megatron.core.fp8_utils import is_mxfp8_output_proj_active
+
+ if not is_mxfp8_output_proj_active(config):
+ raise RuntimeError(
+ "TELMHeadColumnParallelLinear is only valid when fp8_output_proj=True, "
+ "fp8=True, and fp8_recipe='mxfp8'."
+ )
+ if keep_master_weight_for_test:
+ raise ValueError("TE output projection does not support keep_master_weight_for_test.")
+ if skip_weight_param_allocation:
+ raise ValueError("TE output projection does not support skip_weight_param_allocation.")
+ if embedding_activation_buffer is not None or grad_output_buffer is not None:
+ raise ValueError(
+ "TE MXFP8 output projection does not support defer_embedding_wgrad_compute."
+ )
+ if disable_grad_reduce:
+ raise ValueError("TE output projection does not support disable_grad_reduce.")
+
+ te_config = copy.copy(config)
+ # Match ColumnParallelLinear.backward_dw's no-op so the LM head keeps
+ # the same wgrad-timing behavior it had before this subclass existed.
+ te_config.delay_wgrad_compute = False
+
+ super().__init__(
+ input_size=input_size,
+ output_size=output_size,
+ config=te_config,
+ init_method=init_method,
+ gather_output=False,
+ bias=bias,
+ skip_bias_add=skip_bias_add,
+ is_expert=is_expert,
+ skip_weight_param_allocation=skip_weight_param_allocation,
+ tp_comm_buffer_name=tp_comm_buffer_name,
+ tp_group=tp_group,
+ stride=stride,
+ )
+
+ self.input_size = input_size
+ self.output_size = output_size
+ self.output_size_per_partition = divide(output_size, self.tp_size)
+ self.gather_output = gather_output
+ self.skip_bias_add = skip_bias_add
+ self.embedding_activation_buffer = None
+ self.grad_output_buffer = None
+ self.disable_grad_reduce = False
+ self.tp_group = self._tp_group
+
+ self._register_load_state_dict_pre_hook(
+ lambda state_dict, prefix, *args, **kwargs: state_dict.setdefault(
+ f"{prefix}_extra_state"
+ )
+ )
+
+ def get_extra_state(self):
+ """Return None to match ``ColumnParallelLinear``'s no-extra-state shim.
+
+ Keeps the LM-head state dict compatible across the bf16 / MXFP8 swap.
+ """
+ return None
+
+ def set_extra_state(self, state):
+ """No-op to match ``ColumnParallelLinear.set_extra_state`` (ignored)."""
+ return
+
+ def forward(
+ self,
+ input_: torch.Tensor,
+ weight: Optional[torch.Tensor] = None,
+ runtime_gather_output: Optional[bool] = None,
+ ):
+ """Run TE MXFP8 output projection. Returns ``(output, bias)``."""
+ from megatron.core.fp8_utils import get_fp8_context
+
+ if weight is not None and weight is not self.weight:
+ raise RuntimeError("TE MXFP8 output projection does not support runtime weight.")
+
+ with get_fp8_context(self.config):
+ torch.cuda.nvtx.range_push("mxfp8_output_proj_telinear")
+ try:
+ output_parallel, output_bias = super().forward(input_)
+ finally:
+ torch.cuda.nvtx.range_pop()
+
+ gather_output = self.gather_output
+ if runtime_gather_output is not None:
+ gather_output = runtime_gather_output
+ if gather_output:
+ output = gather_from_tensor_model_parallel_region(output_parallel, group=self.tp_group)
+ else:
+ output = output_parallel
+ return output, output_bias
+
+
class TERowParallelLinear(TELinear):
"""Wrapper for the Transformer-Engine's `Linear` layer
but specialized similar to megatron's `RowParallelLinear` layer."""
@@ -1345,7 +1571,7 @@ def __init__(
input_size: int,
output_size: int,
*,
- config: ModelParallelConfig,
+ config: TransformerConfig,
init_method: Callable,
bias: bool,
input_is_parallel: bool,
@@ -1353,7 +1579,12 @@ def __init__(
is_expert: bool,
tp_comm_buffer_name: Optional[str] = None,
tp_group: Optional[torch.distributed.ProcessGroup] = None,
+ name: str | None = None,
):
+ """
+ Args:
+ name (str | None): module instance name passed top-down from its paranet module
+ """
if not HAVE_TE:
raise ImportError(
"Transformer Engine is not installed. "
@@ -1385,6 +1616,7 @@ def __init__(
tp_comm_buffer_name=tp_comm_buffer_name,
symmetric_ar_type=config.symmetric_ar_type,
tp_group=tp_group,
+ name=name,
)
if config.use_cpu_initialization:
world_size = get_pg_size(tp_group)
@@ -1786,7 +2018,12 @@ def __init__(
is_expert: bool = False,
tp_comm_buffer_name: Optional[str] = None,
pg_collection: Optional[ProcessGroupCollection] = None,
+ name: str | None = None,
):
+ """
+ Args:
+ name (str | None): module instance name passed top-down from its paranet module
+ """
self.config = config
# TE returns a zero length Tensor when bias=False and
@@ -1800,9 +2037,13 @@ def __init__(
extra_kwargs = _get_extra_te_kwargs(config)
- if self.config.delay_wgrad_compute:
+ self.delay_wgrad_compute = (
+ self.config.delay_wgrad_compute
+ or self.config.overlap_dispatch_backward_with_experts_wgrad
+ )
+ if self.delay_wgrad_compute:
if is_te_min_version("2.3.0"):
- extra_kwargs["delay_wgrad_compute"] = self.config.delay_wgrad_compute
+ extra_kwargs["delay_wgrad_compute"] = True
else:
raise RuntimeError(
"Only TE with version >=2.3.0 supports delay_wgrad_compute now."
@@ -1843,24 +2084,40 @@ def __init__(
tp_size = 1
tp_group_for_te = None
- super().__init__(
- num_gemms=num_gemms,
- in_features=input_size,
- out_features=output_size,
- sequence_parallel=self.config.sequence_parallel,
- fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
- tp_group=tp_group_for_te if torch.distributed.is_initialized() else None,
- tp_size=tp_size,
- get_rng_state_tracker=(
- get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
- ),
- init_method=condition_init_method(config, init_method),
- bias=bias,
- return_bias=self.te_return_bias,
- parallel_mode=parallel_mode,
- **extra_kwargs,
- )
+ if is_te_min_version("2.14.0"):
+ extra_kwargs["single_grouped_weight"] = getattr(
+ config, "moe_single_grouped_weight", False
+ )
+ extra_kwargs["single_grouped_bias"] = getattr(
+ config, "moe_single_grouped_bias", False
+ )
+
self.te_quant_params: Optional[TEQuantizationParams] = None
+ quant_config = get_quant_config_or_none(name, config.quant_recipe)
+ self.finish_init(quant_config)
+ init_quant_context = _get_fp8_model_init_for_quant_params(
+ self.te_quant_params, torch.is_grad_enabled()
+ )
+
+ with init_quant_context:
+ super().__init__(
+ num_gemms=num_gemms,
+ in_features=input_size,
+ out_features=output_size,
+ sequence_parallel=self.config.sequence_parallel,
+ fuse_wgrad_accumulation=self.config.gradient_accumulation_fusion,
+ tp_group=tp_group_for_te if torch.distributed.is_initialized() else None,
+ tp_size=tp_size,
+ get_rng_state_tracker=(
+ get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() else None
+ ),
+ init_method=condition_init_method(config, init_method),
+ bias=bias,
+ return_bias=self.te_return_bias,
+ parallel_mode=parallel_mode,
+ **extra_kwargs,
+ )
+
for param in self.parameters():
setattr(param, "allreduce", not (is_expert and self.expert_parallel))
@@ -1879,6 +2136,10 @@ def __init__(
setattr(weight, "partition_dim", part_dim)
setattr(weight, "partition_stride", 1)
+ self._register_load_state_dict_pre_hook(
+ type(self)._normalize_grouped_parameter_keys, with_module=True
+ )
+
def merge_extra_states(
self,
state_dict,
@@ -1969,6 +2230,76 @@ def merge_extra_states(
self._register_load_state_dict_pre_hook(merge_extra_states, with_module=True)
+ def _normalize_grouped_parameter_keys(
+ self,
+ state_dict,
+ prefix,
+ local_metadata,
+ strict,
+ missing_keys,
+ unexpected_keys,
+ error_msgs,
+ ):
+ """Make grouped checkpoint keys compatible across parameter layouts.
+
+ Registered as a load_state_dict pre-hook to bridge checkpoints saved
+ in one layout (single grouped tensor vs per-GEMM indexed tensors)
+ and a model expecting the other.
+ """
+
+ def maybe_remap_param(param_name: str, single_grouped: bool) -> None:
+ grouped_key = f"{prefix}{param_name}"
+ indexed_keys = [
+ f"{prefix}{param_name}{gemm_idx}" for gemm_idx in range(self.num_gemms)
+ ]
+ has_grouped_key = grouped_key in state_dict
+ has_any_indexed_key = any(key in state_dict for key in indexed_keys)
+ has_all_indexed_keys = all(key in state_dict for key in indexed_keys)
+
+ if single_grouped:
+ if has_grouped_key or not has_all_indexed_keys:
+ return
+ state_dict[grouped_key] = torch.stack(
+ [state_dict.pop(key) for key in indexed_keys], dim=0
+ )
+ else:
+ if has_any_indexed_key or not has_grouped_key:
+ return
+ split_tensors = self._split_grouped_checkpoint_tensor(
+ state_dict.pop(grouped_key), grouped_key
+ )
+ for gemm_idx, tensor in enumerate(split_tensors):
+ state_dict[f"{prefix}{param_name}{gemm_idx}"] = tensor
+
+ maybe_remap_param("weight", getattr(self, "single_grouped_weight", False))
+ if self.use_bias:
+ maybe_remap_param("bias", getattr(self, "single_grouped_bias", False))
+
+ def _split_grouped_checkpoint_tensor(
+ self, tensor: torch.Tensor, checkpoint_key: str
+ ) -> list[torch.Tensor]:
+ """Split grouped checkpoint tensor into one tensor per GEMM."""
+ if hasattr(tensor, "split_into_quantized_tensors") and callable(
+ tensor.split_into_quantized_tensors
+ ):
+ grouped_tensors = getattr(tensor, "quantized_tensors", None)
+ if grouped_tensors is None:
+ grouped_tensors = tensor.split_into_quantized_tensors()
+ if len(grouped_tensors) != self.num_gemms:
+ raise RuntimeError(
+ f"Grouped checkpoint tensor {checkpoint_key} has {len(grouped_tensors)} "
+ f"groups, expected {self.num_gemms}."
+ )
+ return list(grouped_tensors)
+ if tensor.ndim > 0 and tensor.shape[0] == self.num_gemms:
+ return list(tensor.unbind(dim=0))
+ if tensor.ndim > 0 and tensor.shape[0] % self.num_gemms == 0:
+ return list(torch.chunk(tensor, self.num_gemms, dim=0))
+ raise RuntimeError(
+ f"Cannot split checkpoint tensor {checkpoint_key} with shape {tuple(tensor.shape)} "
+ f"into {self.num_gemms} GEMM shards."
+ )
+
def finish_init(self, quantization_config: QuantizationConfig):
"""Post-init of quantization override"""
if quantization_config is None:
@@ -2029,11 +2360,13 @@ def _encode_extra_state(self, state):
return state_serialized
def _decode_extra_state(self, state):
+ from megatron.core.safe_globals import SafeUnpickler
+
if isinstance(state, torch.Tensor):
# No FP8 is indicated by an empty tensor we don't need to unpickle.
if state.numel() == 0:
return
- return pickle.loads(state.detach().cpu().numpy().tobytes())
+ return SafeUnpickler(io.BytesIO(state.detach().cpu().numpy().tobytes())).load()
elif isinstance(state, io.BytesIO):
state.seek(0)
return torch.load(state, map_location="cuda", weights_only=False)
@@ -2090,6 +2423,21 @@ def _sharded_state_dict_grouped(
singleton_local_shards = (metadata or {}).get('singleton_local_shards', False)
sharded_state_dict = {}
full_state_dict = self.state_dict(prefix="", keep_vars=True)
+ grouped_split_cache = {}
+
+ def get_gemm_tensor(param_name: str, gemm_idx: int) -> torch.Tensor:
+ indexed_name = f"{param_name}{gemm_idx}"
+ if indexed_name in full_state_dict:
+ return full_state_dict[indexed_name]
+ if param_name not in full_state_dict:
+ raise KeyError(indexed_name)
+ if param_name not in grouped_split_cache:
+ grouped_split_cache[param_name] = self._split_grouped_checkpoint_tensor(
+ full_state_dict[param_name], param_name
+ )
+ grouped_splits = grouped_split_cache[param_name]
+ return grouped_splits[gemm_idx]
+
num_global_experts = get_pg_size(self._pg_collection.ep) * self.num_gemms
local_expert_indices_offset = get_pg_rank(self._pg_collection.ep) * self.num_gemms
ep_axis = len(sharded_offsets)
@@ -2097,11 +2445,11 @@ def _sharded_state_dict_grouped(
for gemm_idx in range(self.num_gemms):
global_expert_idx = local_expert_indices_offset + gemm_idx
state_dict = {
- f"{gemm_idx}.weight": full_state_dict[f"weight{gemm_idx}"],
+ f"{gemm_idx}.weight": get_gemm_tensor("weight", gemm_idx),
f"{gemm_idx}._extra_state": extra_states[gemm_idx],
}
if self.use_bias:
- state_dict[f"{gemm_idx}.bias"] = full_state_dict[f"bias{gemm_idx}"]
+ state_dict[f"{gemm_idx}.bias"] = get_gemm_tensor("bias", gemm_idx)
if singleton_local_shards:
expert_prefix = f"{global_expert_idx}.{prefix}"
new_sharded_offsets = sharded_offsets
@@ -2149,7 +2497,7 @@ def backward_dw(self):
Compute weight gradients during the backward pass
if delay_wgrad_compute is enabled.
"""
- if self.config.delay_wgrad_compute:
+ if self.delay_wgrad_compute:
super().backward_dw()
class TEColumnParallelGroupedLinear(TEGroupedLinear):
@@ -2171,7 +2519,12 @@ def __init__(
is_expert: bool,
tp_comm_buffer_name: Optional[str] = None,
pg_collection: Optional[ProcessGroupCollection] = None,
+ name: str | None = None,
):
+ """
+ Args:
+ name (str | None): module instance name passed top-down from its paranet module
+ """
super().__init__(
num_gemms=num_gemms,
input_size=input_size,
@@ -2184,6 +2537,7 @@ def __init__(
is_expert=is_expert,
tp_comm_buffer_name=tp_comm_buffer_name,
pg_collection=pg_collection,
+ name=name,
)
def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
@@ -2217,7 +2571,12 @@ def __init__(
is_expert: bool,
tp_comm_buffer_name: Optional[str] = None,
pg_collection: Optional[ProcessGroupCollection] = None,
+ name: str | None = None,
):
+ """
+ Args:
+ name (str | None): module instance name passed top-down from its paranet module
+ """
super().__init__(
num_gemms=num_gemms,
input_size=input_size,
@@ -2230,6 +2589,7 @@ def __init__(
is_expert=is_expert,
tp_comm_buffer_name=tp_comm_buffer_name,
pg_collection=pg_collection,
+ name=name,
)
def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
@@ -2537,8 +2897,215 @@ def forward(self, hidden_states: torch.Tensor, **kwargs) -> Tuple[Tensor, Option
return out, bias
+ @classmethod
+ def as_mlp_submodule(
+ cls,
+ submodules: MLPSubmodules,
+ config: TransformerConfig,
+ pg_collection: ProcessGroupCollection,
+ is_mtp_layer: bool,
+ is_expert: bool = False,
+ input_size: int | None = None,
+ ffn_hidden_size: int | None = None,
+ name: str | None = None,
+ ) -> MLP:
+ """Helper function to build an MLP as a TransformerLayer's mlp submodule."""
+ del is_mtp_layer
+ assert hasattr(
+ pg_collection, 'tp'
+ ), 'TP process group is required for TEFusedMLP in TransformerLayer'
+ return cls(
+ config=config,
+ submodules=submodules,
+ tp_group=pg_collection.tp,
+ is_expert=is_expert,
+ input_size=input_size,
+ ffn_hidden_size=ffn_hidden_size,
+ name=name,
+ )
+
+ class TEFusedMLPWithGroupedLinear(TEFusedMLP):
+ """Dense MLP using GroupedLinear(num_groups=1) to trigger
+ ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 fusion on SM100+ with MXFP8 recipe.
+
+ Subclass of TEFusedMLP -> does not modify TEFusedMLP or TEGroupedMLP.
+ The fused kernel fires automatically via the TE op fuser when it detects
+ the GroupedLinear -> ScaledSwiGLU -> GroupedLinear pattern with MXFP8 recipe.
+ """
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self._norm_seq: Optional[Tuple[te.pytorch.ops.Sequential]] = None
+ if not is_te_min_version("2.14.0"):
+ raise RuntimeError(
+ f"{self.__class__.__name__} requires Transformer Engine >= 2.14.0 "
+ "(needs pytorch.ops.GroupedLinear and pytorch.ops.ScaledSwiGLU)"
+ )
+ if self.config.add_bias_linear:
+ raise ValueError(
+ f"{self.__class__.__name__} does not support add_bias_linear=True; "
+ "the CuTeGEMM fused kernel requires bias-free linear layers."
+ )
+ if self.config.activation_func != F.silu or not self.config.gated_linear_unit:
+ raise ValueError(
+ f"{self.__class__.__name__} requires SwiGLU activation "
+ "(activation_func=F.silu, gated_linear_unit=True) "
+ "for the CuTeGEMM fused kernel, but got "
+ f"activation_func={self.config.activation_func}, "
+ f"gated_linear_unit={self.config.gated_linear_unit}."
+ )
+
+ def _make_fused_impl(self) -> te.pytorch.ops.Sequential:
+ """Construct fused module with GroupedLinear(num_groups=1) + ScaledSwiGLU."""
+
+ tp_world_size = get_tensor_model_parallel_world_size()
+ if tp_world_size > 1:
+ return super()._make_fused_impl()
+
+ fused_impl = te.pytorch.ops.Sequential()
+
+ # RNG state
+ rng_state_tracker_function = None
+ if get_cuda_rng_tracker().is_initialized():
+ rng_state_tracker_function = get_cuda_rng_tracker
+
+ # Check submodule types (same as TEFusedMLP)
+ if not isinstance(self.linear_fc1, te.pytorch.LayerNormLinear):
+ raise ValueError(
+ f"{self.__class__.__name__} expects FC1 to be "
+ "Transformer Engine LayerNormLinear, but found "
+ f"{self.linear_fc1.__class__.__name__}."
+ )
+ if not isinstance(self.linear_fc2, te.pytorch.Linear):
+ raise ValueError(
+ f"{self.__class__.__name__} expects FC2 to be "
+ "Transformer Engine Linear, but found "
+ f"{self.linear_fc2.__class__.__name__}."
+ )
+
+ # Norm op (same as TEFusedMLP)
+ norm_type = self.linear_fc1.normalization
+ norm_shape = self.linear_fc1.weight.size(1)
+ kwargs = {
+ "eps": self.linear_fc1.eps,
+ "device": "meta",
+ "dtype": self.linear_fc1.layer_norm_weight.dtype,
+ "zero_centered_gamma": self.linear_fc1.zero_centered_gamma,
+ }
+ op = None
+ if norm_type == "LayerNorm":
+ op = te.pytorch.ops.LayerNorm(norm_shape, **kwargs)
+ op.weight = self.linear_fc1.layer_norm_weight
+ op.bias = self.linear_fc1.layer_norm_bias
+ elif norm_type == "RMSNorm":
+ op = te.pytorch.ops.RMSNorm(norm_shape, **kwargs)
+ op.weight = self.linear_fc1.layer_norm_weight
+ else:
+ raise ValueError(f"Unsupported normalization ({norm_type})")
+ # Store norm in a separate Sequential applied OUTSIDE the MXFP8 autocast
+ # in forward(). Running norm inside MXFP8 context corrupts the saved rstd
+ # used in RMSNorm backward, causing gradient amplification up to 10^6.
+ # Wrapped in tuple to avoid nn.Module submodule registration (which would
+ # duplicate the shared norm weight in state_dict/parameters).
+ norm_seq = te.pytorch.ops.Sequential()
+ norm_seq.append(op)
+ self._norm_seq = (norm_seq,)
+
+ # GLU interleave size must match ScaledSwiGLU and the CuTe kernel.
+ _GLU_INTERLEAVE_SIZE = 32
+
+ # FC1: GroupedLinear(num_groups=1) instead of BasicLinear
+ weight = self.linear_fc1.weight
+ op = te.pytorch.ops.GroupedLinear(
+ num_groups=1,
+ in_features=weight.size(1),
+ out_features=weight.size(0) * tp_world_size,
+ device="meta",
+ dtype=weight.dtype,
+ bias=False,
+ rng_state_tracker_function=rng_state_tracker_function,
+ accumulate_into_main_grad=self.linear_fc1.fuse_wgrad_accumulation,
+ )
+ op.weight0 = weight
+ op._glu_interleave_size = _GLU_INTERLEAVE_SIZE # signals fuser_forward to interleave
+ fused_impl.append(op)
+
+ # ScaledSwiGLU with glu_interleave_size=32
+ # Required by ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8
+ fused_impl.append(te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=32))
+
+ # FC2: GroupedLinear(num_groups=1) instead of BasicLinear
+ weight = self.linear_fc2.weight
+ op = te.pytorch.ops.GroupedLinear(
+ num_groups=1,
+ in_features=weight.size(1),
+ out_features=weight.size(0),
+ device="meta",
+ dtype=weight.dtype,
+ bias=False,
+ rng_state_tracker_function=rng_state_tracker_function,
+ accumulate_into_main_grad=self.linear_fc2.fuse_wgrad_accumulation,
+ )
+ op.weight0 = weight
+ # FC2 has no SwiGLU — MXFP8 quantization done on-the-fly in fuser_forward.
+ # No _mxfp8_weight0 pre-computation to avoid ~28 GB persistent FP8 tensors.
+ fused_impl.append(op)
+
+ self._register_hooks_on_fused_impl(fused_impl)
+ return fused_impl
+
+ def forward(self, hidden_states: torch.Tensor, **kwargs) -> Tuple[Tensor, Optional[Tensor]]:
+ """Forward pass using GroupedLinear(num_groups=1) + ScaledSwiGLU."""
+
+ if get_tensor_model_parallel_world_size() > 1:
+ return super().forward(hidden_states, **kwargs)
+
+ orig_shape = hidden_states.shape
+ hidden_size = hidden_states.size(-1)
+ hidden_states_2d = hidden_states.view(-1, hidden_size)
+ total_tokens = hidden_states_2d.size(0)
+
+ tokens_per_expert = torch.full(
+ (1,), total_tokens, dtype=torch.long, device=hidden_states.device
+ )
+ scales = torch.ones(
+ total_tokens, device=hidden_states.device, dtype=hidden_states.dtype
+ )
+
+ # Build fused impl and cache recipe lazily on first forward pass.
+ # Both are created once and reused — avoids object creation every call.
+ if not hasattr(self, '_recipe'):
+ if os.getenv("FP4_RECIPE", "") == "nvfp4":
+ self._recipe = te.common.recipe.NVFP4BlockScaling()
+ else:
+ self._recipe = te.common.recipe.MXFP8BlockScaling()
+ recipe = self._recipe
+
+ if self._fused_impl is None:
+ with te.pytorch.quantized_model_init(enabled=True, recipe=recipe):
+ self._fused_impl = (self._make_fused_impl(),)
+
+ # Apply norm in BF16 OUTSIDE the MXFP8 autocast to preserve the rstd
+ # tensor used by RMSNorm backward (running it inside causes up to 10^6
+ # gradient amplification, and causes convergence issues).
+ normed = self._norm_seq[0](hidden_states_2d)
+
+ with te.pytorch.autocast(enabled=True, recipe=recipe):
+ out = self._fused_impl[0](normed, tokens_per_expert, scales, tokens_per_expert)
+
+ out = out.view(*orig_shape[:-1], out.size(-1))
+
+ bias = None
+ if self.linear_fc2.te_return_bias:
+ bias = self.linear_fc2.bias
+ if isinstance(bias, torch.Tensor) and bias.numel() == 0:
+ bias = None
+
+ return out, bias
+
else:
TEFusedMLP = None # type: ignore[assignment, misc]
+ TEFusedMLPWithGroupedLinear = None # type: ignore[assignment, misc]
class TEDelayedScaling(te.common.recipe.DelayedScaling):
@@ -2826,8 +3393,8 @@ def get_cpu_offload_context(
retain_pinned_cpu_buffers,
):
"""Get CPU offload context and sync function."""
- if is_te_min_version("2.5.0"):
- # Enables the additional double buffering switch for activations during LLM training
+ if is_te_min_version("2.10.0"):
+ # TE 2.10+ supports retain_pinned_cpu_buffers
context, sync_func = _get_cpu_offload_context(
enabled,
num_layers,
@@ -2837,6 +3404,16 @@ def get_cpu_offload_context(
double_buffering,
retain_pinned_cpu_buffers=retain_pinned_cpu_buffers,
)
+ elif is_te_min_version("2.5.0"):
+ # TE 2.5-2.9 supports double_buffering but not retain_pinned_cpu_buffers
+ context, sync_func = _get_cpu_offload_context(
+ enabled,
+ num_layers,
+ model_layers,
+ activation_offloading,
+ weight_offloading,
+ double_buffering,
+ )
elif is_te_min_version("1.10.0.dev0"):
context, sync_func = _get_cpu_offload_context(
enabled, num_layers, model_layers, activation_offloading, weight_offloading
diff --git a/megatron/core/extensions/transformer_engine_spec_provider.py b/megatron/core/extensions/transformer_engine_spec_provider.py
index 04228e02e88..352f3b15a8a 100644
--- a/megatron/core/extensions/transformer_engine_spec_provider.py
+++ b/megatron/core/extensions/transformer_engine_spec_provider.py
@@ -44,7 +44,7 @@ def column_parallel_linear(self) -> type:
"""Which column parallel linear module TE backend uses"""
return TEColumnParallelLinear
- def row_parallel_linear(self) -> type:
+ def row_parallel_linear(self) -> type[TERowParallelLinear]:
"""Which row parallel linear module TE backend uses"""
return TERowParallelLinear
diff --git a/megatron/core/fault_injector.py b/megatron/core/fault_injector.py
new file mode 100644
index 00000000000..68e0464fad7
--- /dev/null
+++ b/megatron/core/fault_injector.py
@@ -0,0 +1,233 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+import datetime
+import logging
+import math
+import random
+from dataclasses import dataclass
+from typing import Optional, Protocol, Sequence, TypeVar, cast
+
+import torch
+import torch.distributed as dist
+
+try:
+ from nvidia_resiliency_ext.shared_utils.inject_fault import ( # type: ignore[import-untyped]
+ Fault,
+ clear_workload_exception,
+ dispatch_fault_injection,
+ maybe_raise_workload_exception,
+ )
+
+ has_nvidia_resiliency_ext = True
+except ModuleNotFoundError:
+ has_nvidia_resiliency_ext = False
+
+ def maybe_raise_workload_exception(): # pylint: disable=missing-function-docstring
+ raise ModuleNotFoundError(
+ "nvidia_resiliency_ext is required for fault injection. "
+ "Please install it or disable fault injection."
+ )
+
+
+__all__ = ["FaultInjectorConfig", "setup_fault_injection", "maybe_raise_workload_exception"]
+
+
+def _require_nvidia_resiliency_ext():
+ if not has_nvidia_resiliency_ext:
+ raise ModuleNotFoundError(
+ "nvidia_resiliency_ext is required for fault injection. "
+ "Please install it or disable fault injection."
+ )
+
+
+logger = logging.getLogger(__name__)
+
+_T = TypeVar("_T")
+
+
+@dataclass(kw_only=True)
+class FaultInjectorConfig:
+ """Configuration for fault injection testing via nvidia_resiliency_ext."""
+
+ fault_injector_ranks: Optional[str] = None
+ """Comma-separated list of ranks to inject faults on."""
+
+ fault_injector_num_ranks: Optional[int] = None
+ """Number of ranks to inject faults on (random selection)."""
+
+ fault_injector_fault_types: Optional[str] = None
+ """Comma-separated list of fault types to inject (e.g. 'hang,crash')."""
+
+ fault_injector_fault_probabilities: Optional[str] = None
+ """Comma-separated list of fault probabilities (normalized at runtime)."""
+
+ fault_injector_fault_delay: Optional[float] = None
+ """Force a specific fault delay in seconds from training start or delay_start_iteration."""
+
+ fault_injector_delay_start_iteration: Optional[int] = None
+ """Start the fault delay timer after iteration N completes.
+ If unset, fault delay timing starts from the beginning of training."""
+
+ fault_injector_mtti_seconds: Optional[float] = None
+ """Mean time to inject (MTTI) in seconds; used when fault_delay is None."""
+
+ fault_injector_offset_seconds: Optional[float] = None
+ """Offset seconds added to the sampled fault delay."""
+
+ fault_injector_seed: Optional[int] = None
+ """RNG seed for the fault injector."""
+
+
+class _FaultInjectorRNG(Protocol):
+ """Minimal RNG interface used by fault injector helper functions."""
+
+ def sample(self, population: Sequence[int], k: int) -> list[int]:
+ """Return ``k`` sampled items from the given population."""
+ ...
+
+ def choices(self, population: Sequence[_T], weights: Sequence[float], k: int) -> list[_T]:
+ """Return ``k`` weighted samples from the given population."""
+ ...
+
+ def random(self) -> float:
+ """Return a floating-point value in the half-open interval [0.0, 1.0)."""
+ ...
+
+
+rng: _FaultInjectorRNG | None = None
+
+
+def _require_rng() -> _FaultInjectorRNG:
+ assert rng is not None, "fault injector rng must be initialized"
+ return rng
+
+
+def get_fault_ranks(config: FaultInjectorConfig):
+ """Return list of ranks to inject faults on, from explicit list or random sample."""
+ global rng
+
+ force_ranks = config.fault_injector_ranks
+ world_size = dist.get_world_size()
+
+ if force_ranks is not None:
+ assert (
+ config.fault_injector_num_ranks is None
+ ), "Cannot specify both force_ranks and num_ranks"
+ if ',' in force_ranks:
+ fault_ranks = [int(r) for r in force_ranks.split(",")]
+ else:
+ fault_ranks = [int(force_ranks)]
+ assert all(
+ 0 <= r < world_size for r in fault_ranks
+ ), f"Fault ranks must be between 0 and {world_size - 1}"
+ assert len(fault_ranks) > 0, "Must specify at least one fault rank"
+ else:
+ assert (
+ config.fault_injector_num_ranks is not None
+ ), "Must specify either force_ranks or num_ranks"
+ fault_ranks = _require_rng().sample(range(1, world_size), k=config.fault_injector_num_ranks)
+
+ return fault_ranks
+
+
+def get_fault(config: FaultInjectorConfig):
+ """Sample a fault type according to the configured types and probabilities."""
+ _require_nvidia_resiliency_ext()
+ global rng
+
+ fault_types_config = config.fault_injector_fault_types
+ fault_probabilities_config = config.fault_injector_fault_probabilities
+ assert fault_types_config is not None, "fault_injector_fault_types must be specified"
+
+ if ',' in fault_types_config:
+ fault_types = [Fault[t.upper()] for t in fault_types_config.split(",")]
+ else:
+ fault_types = [Fault[fault_types_config.upper()]]
+
+ if fault_probabilities_config is not None:
+ if ',' in fault_probabilities_config:
+ fault_probabilities = [float(p) for p in fault_probabilities_config.split(",")]
+ else:
+ fault_probabilities = [float(fault_probabilities_config)]
+ fault_probabilities = [p / sum(fault_probabilities) for p in fault_probabilities]
+ else:
+ fault_probabilities = [1 / len(fault_types) for _ in fault_types]
+
+ assert len(fault_types) > 0, "Must specify at least one fault type"
+ assert len(fault_types) == len(
+ fault_probabilities
+ ), "Number of fault types and fault probabilities must match"
+
+ return _require_rng().choices(fault_types, fault_probabilities, k=1)[0]
+
+
+def should_setup_fault_injection_at_start(config: FaultInjectorConfig):
+ """Return True when fault timing is anchored to training start."""
+ return config.fault_injector_delay_start_iteration is None
+
+
+def should_setup_fault_injection_at_iteration(config: FaultInjectorConfig, iteration):
+ """Return True when fault timing should start from the given iteration."""
+ delay_start_iteration = config.fault_injector_delay_start_iteration
+ return delay_start_iteration is not None and delay_start_iteration == iteration
+
+
+def get_fault_delay(config: FaultInjectorConfig):
+ """Return fault delay in seconds from the configured scheduling anchor."""
+ global rng
+
+ fault_delay = config.fault_injector_fault_delay
+ assert (
+ fault_delay is not None or config.fault_injector_mtti_seconds is not None
+ ), "fault_injector_fault_delay or fault_injector_mtti_seconds must be specified"
+ if fault_delay is None:
+ mtti_seconds = config.fault_injector_mtti_seconds
+ assert mtti_seconds is not None, "fault_injector_mtti_seconds must be specified"
+ offset_seconds = config.fault_injector_offset_seconds or 0.0
+ lambda_inj = 1.0 / mtti_seconds
+ fault_delay = offset_seconds + (-math.log(1.0 - _require_rng().random()) / lambda_inj)
+
+ return fault_delay
+
+
+def setup_fault_injection(config: FaultInjectorConfig):
+ """Broadcast fault plan across ranks and dispatch injection on target ranks."""
+ _require_nvidia_resiliency_ext()
+ global rng
+
+ my_rank = dist.get_rank()
+ world_size = dist.get_world_size()
+
+ device = torch.device("cuda", torch.cuda.current_device())
+ plan_tensor = torch.full((world_size + 1,), float("nan"), dtype=torch.float64, device=device)
+
+ clear_workload_exception()
+
+ if my_rank == 0:
+ if rng is None:
+ rng = cast(_FaultInjectorRNG, random.Random(config.fault_injector_seed))
+
+ fault_ranks = get_fault_ranks(config)
+ fault = get_fault(config)
+ fault_delay = get_fault_delay(config)
+
+ for rank in fault_ranks:
+ plan_tensor[rank] = float(fault.value)
+ plan_tensor[world_size] = fault_delay
+
+ dist.broadcast(plan_tensor, src=0)
+
+ planned_fault = float(plan_tensor[my_rank].item())
+ is_target_rank = not math.isnan(planned_fault)
+
+ if is_target_rank:
+ fault = Fault(int(planned_fault))
+ fault_delay = float(plan_tensor[world_size].item())
+ current_time = datetime.datetime.now()
+ fault_time = current_time + datetime.timedelta(seconds=fault_delay)
+ timestamp = current_time.strftime("%Y-%m-%d %H:%M:%S.%f")
+ fault_timestamp = fault_time.strftime("%Y-%m-%d %H:%M:%S.%f")
+ logger.warning(
+ f"[{timestamp}] FAULT INJECTION: Rank {my_rank} will inject fault "
+ f"{fault.name} at {fault_timestamp}"
+ )
+ dispatch_fault_injection(fault=fault, delay=fault_delay, callback=None)
diff --git a/megatron/core/fp4_utils.py b/megatron/core/fp4_utils.py
index cc67855180e..be02914ce26 100644
--- a/megatron/core/fp4_utils.py
+++ b/megatron/core/fp4_utils.py
@@ -62,6 +62,13 @@
HAVE_TE_MXFP4_TENSOR_CLASS = False
MXFP4_TENSOR_CLASS = None
+try:
+ from transformer_engine.pytorch.tensor.utils import (
+ post_all_gather_processing as te_post_all_gather_processing,
+ )
+except ImportError:
+ te_post_all_gather_processing = None
+
def is_nvfp4tensor(tensor: torch.Tensor) -> bool:
"""Check if a tensor is a Transformer Engine NVFP4Tensor."""
@@ -71,6 +78,77 @@ def is_mxfp4tensor(tensor: torch.Tensor) -> bool:
"""Check if a tensor is a Transformer Engine MXFP4Tensor."""
return HAVE_TE_MXFP4_TENSOR_CLASS and isinstance(tensor, MXFP4_TENSOR_CLASS)
+def get_nvfp4_rowwise_packed_shape(shape: torch.Size) -> torch.Size:
+ """Return packed byte shape for NVFP4 rowwise storage (last dim // 2)."""
+ if len(shape) == 0:
+ return shape
+ assert shape[-1] % 2 == 0, "NVFP4 requires inner dimension divisible by 2"
+ packed = list(shape)
+ packed[-1] = packed[-1] // 2
+ return torch.Size(packed)
+
+
+def modify_nvfp4_rowwise_storage(fp4_tensor: torch.Tensor, new_rowwise_data: torch.Tensor) -> None:
+ """Replace NVFP4 tensor's rowwise raw data with a new uint8 storage view.
+
+ Copies existing bytes into the new buffer, then swaps the underlying pointer.
+ """
+ if not is_nvfp4tensor(fp4_tensor):
+ raise ValueError("modify_nvfp4_rowwise_storage expects an NVFP4 tensor")
+ # Access TE's internal storage fields
+ old_rowwise = getattr(fp4_tensor, "_rowwise_data", None)
+ if old_rowwise is None:
+ raise RuntimeError("NVFP4 tensor is missing rowwise data to replace")
+ assert (
+ old_rowwise.dtype == new_rowwise_data.dtype == torch.uint8
+ ), "Rowwise NVFP4 storage must be uint8"
+ # Preserve existing values and then swap storage
+ new_rowwise_data.detach().copy_(old_rowwise)
+ fp4_tensor._rowwise_data = new_rowwise_data
+ del old_rowwise
+
+
+def quantize_nvfp4_param_shard(
+ model_params, main_params, start_offsets, data_parallel_group, fsdp_shard_model_params=None
+):
+ """Cast shard FP32 master weights to NVFP4 model params (rowwise/columnwise).
+
+ This function wraps Transformer Engine's quantize_master_weights, which handles:
+ - Two-level NVFP4 scaling (global FP32 scale + per-block FP8 E4M3 scale)
+ - Partial casting with nibble-accurate updates
+ - Coordinated amax reduction across data parallel group
+
+ Args:
+ model_params: List of NVFP4 model parameters (NVFP4Tensor).
+ main_params: List of FP32 master weights (shards).
+ start_offsets: List of starting offsets in the full model weight for each shard.
+ data_parallel_group: Distributed group for amax reduction.
+ fsdp_shard_model_params: Optional list of FSDP sharded model params.
+ """
+ if not HAVE_TE_FP4_TENSOR_CLASS:
+ raise RuntimeError("NVFP4 shard quantization requires Transformer Engine >= 2.7.0.dev0")
+
+ try:
+ from transformer_engine.pytorch.tensor.utils import quantize_master_weights
+ except ImportError:
+ raise RuntimeError(
+ "quantize_master_weights not available in this Transformer Engine version"
+ )
+
+ if len(model_params) == 0:
+ return
+
+ args = [model_params, main_params, start_offsets, data_parallel_group]
+ if fsdp_shard_model_params is not None:
+ args.append(fsdp_shard_model_params)
+
+ kwargs = {}
+ if te_post_all_gather_processing is not None:
+ kwargs["manual_post_all_gather_processing"] = True
+
+ quantize_master_weights(*args, **kwargs)
+
+
def get_fp4_align_size(fp4_recipe: Fp4Recipe) -> int:
"""
Get the alignment size required for FP4 GEMM.
@@ -199,6 +277,10 @@ def get_fp4_context(config: TransformerConfig, layer_no: int = -1, is_init: bool
in inspect.signature(transformer_engine.pytorch.fp8_model_init).parameters
):
context_args["recipe"] = fp4_recipe
+ if "preserve_high_precision_init_val" in (
+ inspect.signature(transformer_engine.pytorch.fp8_model_init).parameters
+ ):
+ context_args["preserve_high_precision_init_val"] = torch.is_grad_enabled()
fp4_context = transformer_engine.pytorch.fp8_model_init(**context_args)
return fp4_context
diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py
index d2ba7b8c2f4..1034059da6d 100644
--- a/megatron/core/fp8_utils.py
+++ b/megatron/core/fp8_utils.py
@@ -530,6 +530,24 @@ def is_first_last_bf16_layer(config: TransformerConfig, layer_no: int):
return False
+def is_mxfp8_output_proj_active(config) -> bool:
+ """Return True when the LM-head output projection should run under MXFP8.
+
+ Active when ``fp8_output_proj=True``, ``fp8=True``, ``fp8_recipe='mxfp8'``,
+ and Transformer Engine is installed.
+ """
+ if not HAVE_TE:
+ return False
+ if not getattr(config, "fp8_output_proj", False):
+ return False
+ if not getattr(config, "fp8", False):
+ return False
+
+ fp8_recipe = getattr(config, "fp8_recipe", None)
+ recipe_value = getattr(fp8_recipe, "value", fp8_recipe)
+ return str(recipe_value).lower() == "mxfp8" or str(fp8_recipe).lower().endswith(".mxfp8")
+
+
if HAVE_TE:
from megatron.core import parallel_state
from megatron.core.extensions.transformer_engine import TEDelayedScaling
@@ -569,7 +587,7 @@ def get_fp8_recipe(config: TransformerConfig):
)
elif config.fp8_recipe == Fp8Recipe.mxfp8:
fp8_recipe = transformer_engine.common.recipe.MXFP8BlockScaling(
- fp8_format=fp8_format
+ fp8_format=fp8_format, fp8_dpa=config.fp8_dot_product_attention
)
elif config.fp8_recipe == Fp8Recipe.custom:
assert config.fp8_quantizer_factory is not None
diff --git a/megatron/core/full_cuda_graph.py b/megatron/core/full_cuda_graph.py
index 7c11195f33b..abee2bf811e 100644
--- a/megatron/core/full_cuda_graph.py
+++ b/megatron/core/full_cuda_graph.py
@@ -2,6 +2,7 @@
"""Full iteration CUDA graph for training."""
+import gc
import logging
import torch
@@ -10,6 +11,47 @@
logger = logging.getLogger(__name__)
+# Process-wide handle so full-iter and optimizer graph captures share one pool and one
+# non-default stream (per-stream alloc segments can inflate memory_reserved; see
+# tools/debug_cuda_graph_pool_memory*.py).
+_shared_graph_pool = None
+_shared_capture_stream = None
+
+
+def get_shared_capture_stream():
+ """Return one `torch.cuda.Stream` for all full-iter and optimizer graph captures.
+
+ Call after the target CUDA device is selected.
+ """
+ global _shared_capture_stream
+ if _shared_capture_stream is None:
+ _shared_capture_stream = torch.cuda.Stream()
+ return _shared_capture_stream
+
+
+def get_shared_graph_pool():
+ """Return a process-wide handle so all call sites share one graph memory pool.
+
+ `torch.cuda.graph_pool_handle()` returns a new pool each time; this lazy singleton
+ ensures e.g. full-iteration and optimizer captures reuse the same pool.
+ """
+ global _shared_graph_pool
+ if _shared_graph_pool is None:
+ _shared_graph_pool = torch.cuda.graph_pool_handle()
+ return _shared_graph_pool
+
+
+def get_graph_pool(use_single_mempool):
+ """Return graph pool handle for full-iter/optimizer graph capture.
+
+ When `use_single_mempool` is True, train/eval and optimizer captures reuse one
+ process-wide pool. Otherwise, each capture call gets a new pool handle.
+ """
+ if use_single_mempool:
+ return get_shared_graph_pool()
+ return torch.cuda.graph_pool_handle()
+
+
# The below functions traverse through nested data structures (tuples, lists, dicts)
# present in src and creates a deep copy where all PyTorch tensors are cloned,
# detached from the computation graph, and moved to CUDA device. Non-tensor objects
@@ -70,6 +112,7 @@ def __call__(self, inputs, stage, microbatch):
assert isinstance(inputs, dict)
if microbatch == len(StaticBufferLoader.static_buffers[stage]):
+ self.stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(self.stream):
StaticBufferLoader.static_buffers[stage].append(copy_tensors_in_struct(inputs))
else:
@@ -83,6 +126,7 @@ def __call__(self, inputs, stage, microbatch):
else:
StaticBufferLoader.static_buffers[stage][microbatch][k] = inputs[k]
+ self.stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(self.stream):
clone_tensors_in_struct(
StaticBufferLoader.static_buffers[stage][microbatch], inputs
@@ -98,10 +142,11 @@ class FullCudaGraphWrapper:
cuda_graph = {'training': None, 'validation': None}
result = {'training': None, 'validation': None}
- def __init__(self, forward_backward_func, cuda_graph_warmup_steps=1):
+ def __init__(self, forward_backward_func, cuda_graph_warmup_steps=1, use_single_mempool=False):
self.forward_backward_func = forward_backward_func
self.static_loader = StaticBufferLoader()
self.cuda_graph_warmup_steps = cuda_graph_warmup_steps
+ self.use_single_mempool = use_single_mempool
def data_read(self, data_iterator, model, training, num_microbatches):
"""Read all microbatch inputs from Dataloader and copy to static buffers."""
@@ -168,10 +213,11 @@ def __call__(self, *args, **kwargs):
for _, state in get_all_rng_states().items():
FullCudaGraphWrapper.cuda_graph[training_str].register_generator_state(state)
torch.cuda.synchronize()
- capture_stream = torch.cuda.Stream()
+ capture_stream = get_shared_capture_stream()
with torch.cuda.graph(
FullCudaGraphWrapper.cuda_graph[training_str],
stream=capture_stream,
+ pool=get_graph_pool(self.use_single_mempool),
capture_error_mode="thread_local",
):
FullCudaGraphWrapper.result[training_str] = self.forward_backward_func(
@@ -180,12 +226,10 @@ def __call__(self, *args, **kwargs):
torch.cuda.synchronize()
torch.distributed.barrier()
logger.info(f'CUDA graph capture done for {training_str}!!!')
-
if FullCudaGraphWrapper.cuda_graph[training_str] is None:
FullCudaGraphWrapper.result[training_str] = self.forward_backward_func(*args, **kwargs)
else:
FullCudaGraphWrapper.cuda_graph[training_str].replay()
-
self.next_iter(training_str)
return FullCudaGraphWrapper.result[training_str]
@@ -196,3 +240,19 @@ def curr_iter(self, stage):
def next_iter(self, stage):
"""Increment current training/validation iteration."""
FullCudaGraphWrapper.curr_iteration[stage] += 1
+
+ def reset_cuda_graph(self, stage=None):
+ """Reset CUDA graph."""
+ if stage is None or stage == 'training':
+ if FullCudaGraphWrapper.cuda_graph['training'] is not None:
+ del FullCudaGraphWrapper.cuda_graph['training']
+ FullCudaGraphWrapper.cuda_graph['training'] = None
+ FullCudaGraphWrapper.result['training'] = None
+ FullCudaGraphWrapper.curr_iteration['training'] = 0
+ if stage is None or stage == 'validation':
+ if FullCudaGraphWrapper.cuda_graph['validation'] is not None:
+ del FullCudaGraphWrapper.cuda_graph['validation']
+ FullCudaGraphWrapper.cuda_graph['validation'] = None
+ FullCudaGraphWrapper.result['validation'] = None
+ FullCudaGraphWrapper.curr_iteration['validation'] = 0
+ gc.collect()
diff --git a/megatron/core/inference/README.md b/megatron/core/inference/README.md
new file mode 100644
index 00000000000..1c133349445
--- /dev/null
+++ b/megatron/core/inference/README.md
@@ -0,0 +1,92 @@
+# Megatron Inference
+
+Use `MegatronLLM` (sync) or `MegatronAsyncLLM` (async, with HTTP serving via `serve()`) for typical inference workflows. Both classes hide the underlying engine pipeline (`DynamicInferenceContext` + `GPTInferenceWrapper` + `TextGenerationController` + `DynamicInferenceEngine`) and provide a vLLM-style `generate(prompts, sampling_params)` API. Choose **direct mode** (`use_coordinator=False`) when you manage data sharding yourself; **coordinator mode** (`use_coordinator=True`) when you want the engine to route requests across data-parallel replicas (required for HTTP serving).
+
+## Quickstart
+
+### Offline batch (sync)
+
+```python
+from megatron.core.inference.apis import MegatronLLM, SamplingParams
+
+# Caller owns initialize_megatron(...), model construction, and model.eval().
+# See examples/inference/offline_inference.py for a runnable end-to-end script.
+with MegatronLLM(
+ model=model,
+ tokenizer=tokenizer,
+ inference_config=inference_config,
+ use_coordinator=False,
+) as llm:
+ results = llm.generate(
+ ["Megatron inference is", "Hello, world"],
+ SamplingParams(num_tokens_to_generate=64),
+ )
+ for r in results:
+ print(r.generated_text)
+```
+
+### OpenAI-compatible HTTP server
+
+```python
+import asyncio
+from megatron.core.inference.apis import MegatronAsyncLLM, ServeConfig
+
+async def main():
+ async with MegatronAsyncLLM(
+ model=model,
+ tokenizer=tokenizer,
+ inference_config=inference_config,
+ use_coordinator=True, # serve() requires coordinator mode
+ ) as llm:
+ await llm.serve(ServeConfig(host="0.0.0.0", port=5000)) # blocks until shutdown
+
+asyncio.run(main())
+```
+
+## Public API
+
+| Symbol | Purpose |
+|---|---|
+| `MegatronLLM` | Sync entry. Methods: `generate`, `pause`/`unpause`/`suspend`/`resume`, `shutdown`/`wait_for_shutdown`. Properties: `engine`, `context`, `controller`, `is_primary_rank`. Context-manager protocol. |
+| `MegatronAsyncLLM` | Async-flavored equivalent. Adds `serve(serve_config, blocking=True)` for HTTP. |
+| `ServeConfig` | Dataclass for the HTTP frontend. Fields: `host` (`"0.0.0.0"`), `port` (`5000`), `parsers` (`[]`), `verbose` (`False`), `frontend_replicas` (`4`). |
+| `SamplingParams`, `DynamicInferenceRequest`, `DynamicInferenceRequestRecord` | Re-exports from `megatron.core.inference`. |
+
+## Caller responsibilities
+
+- Call `initialize_megatron(...)` (full Megatron distributed setup) BEFORE construction.
+- Call `model.eval()` BEFORE construction. The class does not toggle model state.
+- Lifecycle methods (`pause`/`unpause`/`suspend`/`resume`) require `use_coordinator=True`; they raise `RuntimeError` in direct mode.
+
+## Future roadmap
+
+Planned new features:
+
+- **Dynamic streaming.** Offline streaming via `engine.async_step()`; HTTP streaming requires extending the coordinator / `InferenceClient` protocol to carry partial outputs (not just final request records).
+
+- **Weight update APIs.** `suspend_for_refit()`, `update_weights_from_collective()`, `resume_after_refit()` wrapping the existing resharding/refit primitives for RL workflows where weights swap between rollout steps.
+
+- **`megatron serve` CLI.** Single-binary launcher reusing `MegatronAsyncLLM.serve(...)`, with single-node and multi-node / headless modes — mirrors `vllm serve`.
+
+- **Config-based model construction.** `MegatronLLM(model="...")` style with model recipes and checkpoint resolution, removing manual model building from caller responsibilities.
+
+## Known limitations
+
+- **`MegatronAsyncLLM` requires `use_coordinator=True`** -- constructing with `use_coordinator=False` raises `ValueError` at `__init__`. The underlying `DynamicInferenceEngine` caches its loop reference at construction time and binds internal asyncio primitives (`_cond`, `_state_events`) to it. Coordinator mode rebinds those to a dedicated daemon-thread loop via `start_listening_to_data_parallel_coordinator`; direct mode has no such rebinding, so the synchronous `engine.generate()` path collides with the caller's running asyncio loop and raises `RuntimeError: This event loop is already running`. Use `MegatronLLM` for sync direct/coordinator workflows. Tracked for an upstream `engine.async_generate(...)` (or engine loop-rebinding) fix that would let `MegatronAsyncLLM` support direct mode.
+
+- **`llm.engine.reset()` is unsafe in coordinator mode.** Two failure modes, both upstream in `dynamic_engine.py`:
+ - *Deadlock*: `reset()` *rebinds* (does not mutate in-place) `_cond` / `_state_events`. Any coroutine on the engine-loop task that is `await`ing one of those primitives holds a reference to the OLD object in its suspended frame. Subsequent `notify_all()` / `set()` calls hit the NEW objects, leaving the suspended waiter stranded; the next `generate()` hangs.
+ - *Silent corruption*: `reset()` also sets `self.use_coordinator = False`, which silently re-routes failed-request handling, scheduling notification, and `suspend()`'s state machine to direct-mode branches. Outcome: not-a-hang but wrong behavior, harder to diagnose.
+ - The example `offline_inference.py` blocks `--inference-repeat-n > 1` with `--use-coordinator` for these reasons. Direct-mode reset is safe.
+
+- **HTTP frontend is fixed to global rank 0.** There is no per-rank `role` override on `ServeConfig` to host the HTTP server on a non-rank-0 rank or to opt a rank out of HTTP. Control placement via the launcher (e.g., torchrun rank-0 placement), mirroring how vLLM's `--headless` is invoked today.
+
+- **Server returns `"model": "EMPTY"`.** The HTTP frontend doesn't expose a `ServeConfig.model_name` to echo in `/v1/completions` / `/v1/chat/completions` responses, doesn't validate the request `model` field against a configured name, and exposes no `GET /v1/models` discovery endpoint. Clients can still pass any `model` in their request body — the dynamic server ignores it.
+
+## Low-level APIs
+
+For step-level control, custom forward-step integration, or migration from existing pipelines, drop down to the building blocks in this directory: `DynamicInferenceEngine` (manual `add_request` / `step_modern` stepping), `DynamicInferenceContext`, `TextGenerationController`, and the model inference wrappers under `model_inference_wrappers/`. Runnable examples live in [`examples/inference/advanced/`](../../examples/inference/advanced/): `gpt_dynamic_inference.py` (manual stepping), `gpt_dynamic_inference_with_coordinator.py` (explicit coordinator + `InferenceClient` lifecycle), `gpt_static_inference.py` (static engine), and `simple_t5_batch_inference.py` (T5).
+
+## See also
+
+- Examples: [`examples/inference/offline_inference.py`](../../examples/inference/offline_inference.py) (4 modes via `--mode` / `--use-coordinator`), [`examples/inference/launch_inference_server.py`](../../examples/inference/launch_inference_server.py) (HTTP server).
diff --git a/megatron/core/inference/apis/__init__.py b/megatron/core/inference/apis/__init__.py
new file mode 100644
index 00000000000..19b27250406
--- /dev/null
+++ b/megatron/core/inference/apis/__init__.py
@@ -0,0 +1,19 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.inference.apis.async_llm import MegatronAsyncLLM
+from megatron.core.inference.apis.llm import MegatronLLM
+from megatron.core.inference.apis.serve_config import ServeConfig
+from megatron.core.inference.inference_request import (
+ DynamicInferenceRequest,
+ DynamicInferenceRequestRecord,
+)
+from megatron.core.inference.sampling_params import SamplingParams
+
+__all__ = [
+ "DynamicInferenceRequest",
+ "DynamicInferenceRequestRecord",
+ "MegatronAsyncLLM",
+ "MegatronLLM",
+ "SamplingParams",
+ "ServeConfig",
+]
diff --git a/megatron/core/inference/apis/_llm_base.py b/megatron/core/inference/apis/_llm_base.py
new file mode 100644
index 00000000000..0c0f9881b11
--- /dev/null
+++ b/megatron/core/inference/apis/_llm_base.py
@@ -0,0 +1,462 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Internal building blocks for the Megatron inference high-level API.
+
+This module hosts private helpers shared by ``MegatronLLM`` and
+``MegatronAsyncLLM``: ``_EventLoopManager``, ``_CoordinatorRuntime``, and
+``_MegatronLLMBase``. The public sync/async wrappers live on the subclasses;
+this base only exposes shared engine state, runtime spawn, validation
+helpers, and the private ``__impl`` coroutines.
+"""
+
+import asyncio
+import concurrent.futures
+import threading
+from typing import Coroutine, List, Optional, Tuple, Union
+
+import torch.distributed as dist
+
+from megatron.core.inference.config import InferenceConfig
+from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext
+from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine, EngineState
+from megatron.core.inference.inference_request import DynamicInferenceRequest
+from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
+ GPTInferenceWrapper,
+)
+from megatron.core.inference.sampling_params import SamplingParams
+from megatron.core.inference.text_generation_controllers.text_generation_controller import (
+ TextGenerationController,
+)
+
+
+class _EventLoopManager:
+ """Per-instance background daemon thread + persistent asyncio event loop.
+
+ Bridges sync and async user-thread callers to coroutines that run on the
+ background loop via ``asyncio.run_coroutine_threadsafe``.
+ """
+
+ def __init__(self) -> None:
+ self._loop: Optional[asyncio.AbstractEventLoop] = None
+ self._thread: Optional[threading.Thread] = None
+ self._started: bool = False
+ self._stopped: bool = False
+
+ def start(self) -> None:
+ """Spawn the daemon thread and start the event loop. Idempotent."""
+ if self._started:
+ return
+
+ # PyTorch's CUDA current-device is thread-local and defaults to 0 on
+ # new threads. Capture the spawning thread's device so NCCL ops
+ # scheduled on the runtime loop (e.g. inside
+ # ``start_listening_to_data_parallel_coordinator``) hit the right GPU
+ # under torchrun, where every process sees all GPUs and rank-to-device
+ # mapping is set on the main thread only.
+ import torch
+
+ parent_device = torch.cuda.current_device() if torch.cuda.is_available() else None
+
+ loop_ready = threading.Event()
+
+ def _run_loop() -> None:
+ if parent_device is not None:
+ torch.cuda.set_device(parent_device)
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ self._loop = loop
+ # Fires once run_forever() starts dispatching callbacks, so
+ # callers blocked on loop_ready.wait() resume only after the
+ # loop is actually running.
+ loop.call_soon(loop_ready.set)
+ loop.run_forever()
+
+ self._thread = threading.Thread(target=_run_loop, daemon=True)
+ self._thread.start()
+ loop_ready.wait()
+ self._started = True
+
+ @property
+ def loop(self) -> asyncio.AbstractEventLoop:
+ """The background asyncio loop. Raises if ``start()`` has not been called."""
+ if not self._started or self._loop is None:
+ raise RuntimeError("_EventLoopManager.start() must be called before accessing loop.")
+ return self._loop
+
+ def submit(self, coro: Coroutine) -> concurrent.futures.Future:
+ """Schedule ``coro`` on the background loop and return its future.
+
+ The caller decides how to wait on the returned future (e.g.
+ ``.result()`` for blocking sync, ``asyncio.wrap_future(...)`` for
+ awaiting from another loop).
+ """
+ if not self._started or self._loop is None:
+ raise RuntimeError("_EventLoopManager.start() must be called before submit().")
+ return asyncio.run_coroutine_threadsafe(coro, self._loop)
+
+ def run_sync(self, coro: Coroutine):
+ """Schedule ``coro`` on the background loop and block on its result.
+
+ Must not be called from a coroutine running on ``self._loop`` itself
+ -- that would deadlock, since the only loop that could dispatch
+ ``coro`` would be the one already blocked waiting for the caller.
+ Calling from a different loop (e.g., the user's main-thread asyncio
+ loop) is allowed: ``coro`` runs on the background loop while the
+ caller's loop is stalled until ``.result()`` returns.
+ """
+ try:
+ running = asyncio.get_running_loop()
+ except RuntimeError:
+ running = None # no loop on this thread, safe
+ if running is self._loop:
+ raise RuntimeError(
+ "run_sync called from a coroutine running on the background "
+ "loop -- would deadlock waiting for the same loop."
+ )
+ return self.submit(coro).result()
+
+ async def run_async(self, coro: Coroutine):
+ """Schedule ``coro`` on the background loop and await it from any loop."""
+ return await asyncio.wrap_future(self.submit(coro))
+
+ def stop(self) -> None:
+ """Stop the event loop and join the background thread. Idempotent."""
+ if not self._started or self._stopped:
+ return
+ assert self._loop is not None
+ assert self._thread is not None
+ self._loop.call_soon_threadsafe(self._loop.stop)
+ self._thread.join()
+ self._stopped = True
+ self._started = False
+
+
+class _CoordinatorRuntime:
+ """Owns the dynamic-inference coordinator and ``InferenceClient`` lifecycle.
+
+ Async-native: :meth:`setup` and :meth:`teardown` are coroutines meant to
+ run on a background loop owned by :class:`_EventLoopManager`. The primary
+ rank additionally holds an :class:`InferenceClient` used by the high-level
+ API to submit requests and send control signals.
+ """
+
+ def __init__(
+ self,
+ engine: "DynamicInferenceEngine",
+ *,
+ is_primary: bool,
+ coordinator_host: Optional[str],
+ coordinator_port: Optional[int],
+ ) -> None:
+ self._engine = engine
+ self._is_primary = is_primary
+ self._coordinator_host = coordinator_host
+ self._coordinator_port = coordinator_port
+ self._client: "Optional[InferenceClient]" = None
+ self._coord_addr: Optional[str] = None
+
+ async def setup(self, *, loop: asyncio.AbstractEventLoop) -> None:
+ """Bring the coordinator and (on primary) the ``InferenceClient`` up.
+
+ Calls ``engine.start_listening_to_data_parallel_coordinator(loop=loop)``
+ on every rank. Only host/port kwargs that the caller actually supplied
+ are forwarded so the engine can auto-bind when both are ``None``.
+ """
+ kwargs = {"loop": loop}
+ if self._coordinator_host is not None:
+ kwargs["hostname"] = self._coordinator_host
+ if self._coordinator_port is not None:
+ kwargs["inference_coordinator_port"] = self._coordinator_port
+
+ coord_addr = await self._engine.start_listening_to_data_parallel_coordinator(**kwargs)
+ self._coord_addr = coord_addr
+
+ if self._is_primary:
+ # Lazy import: keep this module importable without pyzmq/msgpack
+ # installed when the user only needs direct mode.
+ from megatron.core.inference.inference_client import InferenceClient
+
+ # deserialize=True returns DynamicInferenceRequest objects from
+ # add_request futures, matching the high-level API contract.
+ client = InferenceClient(coord_addr, deserialize=True)
+ client.start(loop=loop)
+ self._client = client
+
+ async def teardown(self) -> None:
+ """Idempotent best-effort shutdown of the coordinator + client.
+
+ Safe to call from partial-setup state (e.g., when :meth:`setup` raised
+ after the coordinator subprocess spawned but before the client opened).
+ Worker ranks are always no-op; their ``engine_loop_task`` is awaited by
+ :meth:`_MegatronLLMBase._shutdown_impl` after the primary has issued
+ the STOP signal.
+ """
+ if not self._is_primary:
+ return
+
+ # Happy path: client open -> graceful protocol shutdown.
+ if self._client is not None:
+ try:
+ self._client.shutdown_coordinator()
+ self._client.stop()
+ finally:
+ self._client = None
+ return
+
+ # Partial-setup path: client never opened. If the coordinator
+ # subprocess was spawned, kill it via the engine's process handle.
+ proc = getattr(self._engine, "inference_coordinator_process", None)
+ if proc is not None and proc.is_alive():
+ proc.terminate()
+ proc.join(timeout=5)
+ if proc.is_alive():
+ proc.kill()
+ proc.join(timeout=2)
+
+ @property
+ def client(self) -> "Optional[InferenceClient]":
+ """The :class:`InferenceClient` on the primary rank; ``None`` on workers."""
+ return self._client
+
+ @property
+ def coord_addr(self) -> Optional[str]:
+ """Address returned by ``start_listening_to_data_parallel_coordinator``."""
+ return self._coord_addr
+
+
+class _MegatronLLMBase:
+ """Private base shared by ``MegatronLLM`` and ``MegatronAsyncLLM``.
+
+ This base intentionally exposes no public ``generate`` / lifecycle
+ methods -- those live on the subclasses, which call into the private
+ ``__impl`` coroutines defined here. The base owns:
+
+ - the engine pipeline (engine, context, controller),
+ - the per-instance background runtime (``_loop_manager``,
+ ``_coord_runtime``) when ``use_coordinator=True``,
+ - validation helpers (``_assert_primary``, ``_assert_coordinator``) and
+ the input shape helper (``_normalize_prompts``).
+
+ Two execution modes are supported:
+
+ - **Direct mode** (``use_coordinator=False``): every rank is treated as
+ primary and ``generate`` runs the engine synchronously (offloaded to a
+ thread when called from an event loop). Lifecycle methods are invalid
+ and raise :class:`RuntimeError` via ``_assert_coordinator``.
+ - **Coordinator mode** (``use_coordinator=True``): a background event loop
+ hosts the engine pipeline and an :class:`InferenceClient` (on global
+ rank 0). Only the primary rank may submit requests via ``generate``.
+
+ ``model`` must be in eval mode before construction; this class does not
+ modify the model state.
+ """
+
+ def __init__(
+ self,
+ *,
+ model,
+ tokenizer,
+ inference_config: Optional[InferenceConfig] = None,
+ use_coordinator: bool = False,
+ coordinator_host: Optional[str] = None,
+ coordinator_port: Optional[int] = None,
+ ) -> None:
+ if (coordinator_host is not None or coordinator_port is not None) and not use_coordinator:
+ raise ValueError("coordinator_host/port require use_coordinator=True")
+
+ if not use_coordinator:
+ from megatron.core import parallel_state
+
+ ep_size = parallel_state.get_expert_model_parallel_world_size()
+ if ep_size > 1:
+ raise ValueError(
+ f"use_coordinator=True is required when expert_model_parallel_size > 1 "
+ f"(got EP={ep_size}). Use coordinator mode to handle EP routing."
+ )
+
+ if inference_config is None:
+ inference_config = InferenceConfig()
+
+ # Build the engine pipeline. Mirrors examples/inference/gpt/gpt_dynamic_inference.py.
+ context = DynamicInferenceContext(model.config, inference_config)
+ wrapper = GPTInferenceWrapper(model, context)
+ controller = TextGenerationController(inference_wrapped_model=wrapper, tokenizer=tokenizer)
+ engine = DynamicInferenceEngine(controller=controller, context=context)
+
+ if use_coordinator:
+ is_primary_rank = dist.get_rank() == 0
+ else:
+ is_primary_rank = True
+
+ self._engine = engine
+ self._context = context
+ self._controller = controller
+ self._use_coordinator = use_coordinator
+ self._is_primary_rank = is_primary_rank
+ self._loop_manager: "Optional[_EventLoopManager]" = None
+ self._coord_runtime: "Optional[_CoordinatorRuntime]" = None
+ self._shutdown_called: bool = False
+
+ if use_coordinator:
+ loop_manager = _EventLoopManager()
+ loop_manager.start()
+ coord_runtime: "Optional[_CoordinatorRuntime]" = None
+ try:
+ coord_runtime = _CoordinatorRuntime(
+ engine,
+ is_primary=is_primary_rank,
+ coordinator_host=coordinator_host,
+ coordinator_port=coordinator_port,
+ )
+ loop_manager.run_sync(coord_runtime.setup(loop=loop_manager.loop))
+ except BaseException:
+ if coord_runtime is not None:
+ try:
+ loop_manager.run_sync(coord_runtime.teardown())
+ except Exception:
+ pass # best-effort; don't mask the original failure
+ loop_manager.stop()
+ raise
+ self._loop_manager = loop_manager
+ self._coord_runtime = coord_runtime
+
+ # ---- properties ----
+
+ @property
+ def is_primary_rank(self) -> bool:
+ """Whether ``generate`` may be called on this rank."""
+ return self._is_primary_rank
+
+ @property
+ def engine(self) -> "DynamicInferenceEngine":
+ """The underlying :class:`DynamicInferenceEngine`."""
+ return self._engine
+
+ @property
+ def context(self) -> "DynamicInferenceContext":
+ """The underlying :class:`DynamicInferenceContext`."""
+ return self._context
+
+ @property
+ def controller(self) -> "TextGenerationController":
+ """The underlying :class:`TextGenerationController`."""
+ return self._controller
+
+ # ---- internal helpers ----
+
+ def _assert_primary(self) -> None:
+ if not self._is_primary_rank:
+ raise RuntimeError(
+ "generate(...) is only valid on the primary rank in coordinator mode"
+ )
+
+ def _assert_coordinator(self) -> None:
+ if not self._use_coordinator:
+ raise RuntimeError("This method requires use_coordinator=True")
+
+ def _normalize_prompts(
+ self, prompts: Union[str, List[int], List[str], List[List[int]]]
+ ) -> Tuple[Union[List[str], List[List[int]]], bool]:
+ """Return ``(normalized_list, is_batch_input)``.
+
+ - ``"abc"`` -> ``(["abc"], False)``
+ - ``[1, 2, 3]`` -> ``([[1, 2, 3]], False)`` (single token-id prompt)
+ - ``["abc", "def"]`` -> ``(["abc", "def"], True)``
+ - ``[[1, 2], [3, 4]]`` -> ``([[1, 2], [3, 4]], True)``
+ - ``[]`` -> ``([], True)``
+
+ Only the first element is inspected to distinguish single vs batch;
+ per-element type validation is left to the engine.
+ """
+ if isinstance(prompts, str):
+ return [prompts], False
+ if isinstance(prompts, list):
+ if not prompts:
+ return [], True
+ first = prompts[0]
+ if isinstance(first, int):
+ return [prompts], False
+ if isinstance(first, (str, list)):
+ return prompts, True
+ raise TypeError(
+ f"Unsupported prompt element type: {type(first)}; "
+ "expected str, list[int], list[str], or list[list[int]]."
+ )
+ raise TypeError(
+ f"prompts must be str, list[int], list[str], or list[list[int]]; "
+ f"got {type(prompts)}"
+ )
+
+ # ---- private impl coroutines ----
+ # Subclasses' public methods bridge to these via ``_EventLoopManager``
+ # (coordinator mode, on the runtime loop) or await them directly
+ # (direct mode, on the caller's event loop).
+ # We need this bridge in coordinator mode because the coordinator requires
+ # a long running event loop, so we need to route the user's event
+ # loop to our runtime loop
+
+ async def _generate_impl(
+ self, prompts: Union[List[str], List[List[int]]], sp: SamplingParams
+ ) -> List["DynamicInferenceRequest"]:
+ """Run inference for a non-empty list of prompts; returns input-ordered list.
+
+ - Coordinator mode: must run on the runtime loop (via
+ ``_loop_manager.run_async``); enqueues requests through
+ ``client.add_request`` and gathers all futures.
+ - Direct mode: runs on the caller's event loop; offloads the synchronous
+ ``engine.generate`` to a thread.
+ """
+ if self._use_coordinator:
+ # ``add_request`` calls ``asyncio.get_running_loop().create_future()``
+ # so it must be invoked from a coroutine on the runtime loop. This
+ # coroutine runs on that same loop, so ``asyncio.gather`` over the
+ # returned futures is safe.
+ assert self._coord_runtime is not None and self._coord_runtime.client is not None
+ futures = [self._coord_runtime.client.add_request(p, sp) for p in prompts]
+ return list(await asyncio.gather(*futures))
+ # TODO: replace with an upstream ``engine.async_generate`` so direct-mode
+ # async generate doesn't block the caller's event loop.
+ records = self._engine.generate(prompts, sp)
+ return [r.merge() for r in records]
+
+ async def _pause_impl(self) -> None:
+ if self._is_primary_rank:
+ assert self._coord_runtime is not None and self._coord_runtime.client is not None
+ self._coord_runtime.client.pause_engines()
+ await self._engine.wait_until(EngineState.PAUSED)
+
+ async def _unpause_impl(self) -> None:
+ if self._is_primary_rank:
+ assert self._coord_runtime is not None and self._coord_runtime.client is not None
+ self._coord_runtime.client.unpause_engines()
+ await self._engine.wait_until(EngineState.RUNNING)
+
+ async def _suspend_impl(self) -> None:
+ if self._is_primary_rank:
+ assert self._coord_runtime is not None and self._coord_runtime.client is not None
+ self._coord_runtime.client.suspend_engines()
+ await self._engine.wait_until(EngineState.SUSPENDED)
+
+ async def _resume_impl(self) -> None:
+ if self._is_primary_rank:
+ assert self._coord_runtime is not None and self._coord_runtime.client is not None
+ self._coord_runtime.client.resume_engines()
+ await self._engine.wait_until(EngineState.RESUMED)
+
+ async def _shutdown_impl(self) -> None:
+ if self._is_primary_rank:
+ assert self._coord_runtime is not None and self._coord_runtime.client is not None
+ # The coordinator only honors STOP from PAUSED or SUSPENDED. If
+ # the engine is RUNNING (the typical state at shutdown), pause
+ # first so the STOP isn't ignored.
+ if self._engine.state == EngineState.RUNNING:
+ self._coord_runtime.client.pause_engines()
+ await self._engine.wait_until(EngineState.PAUSED)
+ self._coord_runtime.client.stop_engines()
+ await self._engine.wait_until(EngineState.STOPPED)
+ await self._coord_runtime.teardown()
+ else:
+ await self._engine.engine_loop_task
+
+ async def _wait_for_shutdown_impl(self) -> None:
+ await self._engine.engine_loop_task
diff --git a/megatron/core/inference/apis/async_llm.py b/megatron/core/inference/apis/async_llm.py
new file mode 100644
index 00000000000..f2cea47b848
--- /dev/null
+++ b/megatron/core/inference/apis/async_llm.py
@@ -0,0 +1,231 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Async high-level inference API for Megatron (``MegatronAsyncLLM``)."""
+
+from typing import List, Optional, Union
+
+from megatron.core.inference.apis._llm_base import _MegatronLLMBase
+from megatron.core.inference.apis.serve_config import ServeConfig
+from megatron.core.inference.config import InferenceConfig
+from megatron.core.inference.inference_request import DynamicInferenceRequest
+from megatron.core.inference.sampling_params import SamplingParams
+
+
+class MegatronAsyncLLM(_MegatronLLMBase):
+ """Async high-level inference API for Megatron.
+
+ Asyncio-native wrapper over the shared engine + runtime managed by
+ :class:`_MegatronLLMBase` -- see that class for caller responsibilities
+ and the ``model.eval()`` contract. Requires ``use_coordinator=True``;
+ direct mode is rejected at ``__init__`` (see Known Limitations in the
+ package README).
+
+ On top of the base, this class provides:
+
+ - ``async generate`` accepting single or batched prompts.
+ - ``async`` lifecycle controls: ``pause`` / ``unpause`` / ``suspend`` /
+ ``resume`` / ``shutdown`` / ``wait_for_shutdown``.
+ - :meth:`serve` for OpenAI-compatible HTTP serving on the primary rank.
+ - ``async with`` context-manager protocol; exit calls :meth:`shutdown`.
+ """
+
+ def __init__(
+ self,
+ *,
+ model,
+ tokenizer,
+ inference_config: Optional[InferenceConfig] = None,
+ use_coordinator: bool = False,
+ coordinator_host: Optional[str] = None,
+ coordinator_port: Optional[int] = None,
+ ) -> None:
+ # MegatronAsyncLLM requires coordinator mode: direct mode invokes the
+ # synchronous ``engine.generate()`` from inside the caller's asyncio
+ # loop, which collides with the engine's loop-bound internal state
+ # (``_cond``, ``_state_events``). Coordinator mode rebinds those to a
+ # daemon-thread loop via ``start_listening_to_data_parallel_coordinator``
+ # and avoids the conflict.
+ if not use_coordinator:
+ raise ValueError(
+ "MegatronAsyncLLM requires use_coordinator=True. Direct mode is "
+ "not supported in async because the underlying engine's "
+ "asyncio primitives bind to the caller's loop and collide with "
+ "the synchronous engine.generate() path. Use MegatronLLM for "
+ "sync direct/coordinator workflows."
+ )
+ super().__init__(
+ model=model,
+ tokenizer=tokenizer,
+ inference_config=inference_config,
+ use_coordinator=use_coordinator,
+ coordinator_host=coordinator_host,
+ coordinator_port=coordinator_port,
+ )
+ # Set in serve() when this rank starts the HTTP frontend; consulted by shutdown().
+ self._serve_started: bool = False
+
+ async def generate(
+ self,
+ prompts: Union[str, List[int], List[str], List[List[int]]],
+ sampling_params: Optional[SamplingParams] = None,
+ ) -> Union["DynamicInferenceRequest", List["DynamicInferenceRequest"]]:
+ """Run inference for one prompt or a batch of prompts.
+
+ Single input (``str`` or ``list[int]``) returns a single
+ ``DynamicInferenceRequest``; batched input (``list[str]`` or
+ ``list[list[int]]``) returns ``list[DynamicInferenceRequest]`` in
+ input order.
+
+ Raises:
+ RuntimeError: if called on a non-primary rank.
+ """
+ self._assert_primary()
+ if sampling_params is None:
+ sampling_params = SamplingParams()
+
+ normalized, is_batch = self._normalize_prompts(prompts)
+
+ if not normalized:
+ # Empty batch: nothing to schedule. ``is_batch`` is always True
+ # here since single input is wrapped to a one-element list.
+ return []
+
+ assert self._loop_manager is not None
+ results = await self._loop_manager.run_async(
+ self._generate_impl(normalized, sampling_params)
+ )
+ return results if is_batch else results[0]
+
+ async def pause(self) -> None:
+ """Transition the engine to ``PAUSED``.
+
+ Raises:
+ RuntimeError: in direct mode (``use_coordinator=False``).
+ """
+ self._assert_coordinator()
+ assert self._loop_manager is not None
+ await self._loop_manager.run_async(self._pause_impl())
+
+ async def unpause(self) -> None:
+ """Transition the engine from ``PAUSED`` back to ``RUNNING``.
+
+ Raises:
+ RuntimeError: in direct mode (``use_coordinator=False``).
+ """
+ self._assert_coordinator()
+ assert self._loop_manager is not None
+ await self._loop_manager.run_async(self._unpause_impl())
+
+ async def suspend(self) -> None:
+ """Transition the engine to ``SUSPENDED`` (offloads GPU buffers).
+
+ The caller must ``pause()`` first; this method does not enforce that.
+
+ Raises:
+ RuntimeError: in direct mode (``use_coordinator=False``).
+ """
+ self._assert_coordinator()
+ assert self._loop_manager is not None
+ await self._loop_manager.run_async(self._suspend_impl())
+
+ async def resume(self) -> None:
+ """Transition the engine from ``SUSPENDED`` to ``RESUMED``.
+
+ Raises:
+ RuntimeError: in direct mode (``use_coordinator=False``).
+ """
+ self._assert_coordinator()
+ assert self._loop_manager is not None
+ await self._loop_manager.run_async(self._resume_impl())
+
+ async def shutdown(self) -> None:
+ """Stop the engine, tear down the coordinator, and join the runtime thread.
+
+ Idempotent. No-op in direct mode.
+ """
+ if self._shutdown_called:
+ return
+ self._shutdown_called = True
+
+ # If we started an HTTP frontend, stop it first so no new requests
+ # arrive while we tear down the coordinator. Invariant:
+ # ``_serve_started`` can only be True when ``use_coordinator=True``
+ # because ``serve()`` raises otherwise.
+ if self._serve_started:
+ from megatron.core.inference.text_generation_server.dynamic_text_gen_server.text_generation_server import ( # pylint: disable=line-too-long
+ stop_text_gen_server,
+ )
+
+ stop_text_gen_server()
+ self._serve_started = False
+
+ if not self._use_coordinator:
+ return
+ assert self._loop_manager is not None
+ await self._loop_manager.run_async(self._shutdown_impl())
+ self._loop_manager.stop()
+
+ async def serve(self, serve_config: ServeConfig, *, blocking: bool = True) -> None:
+ """Start the OpenAI-compatible HTTP frontend.
+
+ Coordinator mode only. The HTTP frontend runs only on the primary
+ rank (global rank 0); other ranks no-op the HTTP setup but still
+ respect ``blocking`` (so all ranks return together).
+
+ With ``blocking=True`` (default), this awaits the engine loop until
+ :meth:`shutdown` is called -- suitable for standalone serving scripts.
+ With ``blocking=False``, this returns once the HTTP frontend is up
+ (primary) or immediately (workers); the engine loop continues in the
+ background runtime, and the user can call :meth:`generate` /
+ :meth:`shutdown` afterward.
+
+ Raises:
+ ValueError: if ``use_coordinator=False`` (HTTP serving requires
+ the coordinator path).
+ """
+ if not self._use_coordinator:
+ raise ValueError("MegatronAsyncLLM.serve() requires use_coordinator=True")
+
+ if self._is_primary_rank:
+ # Lazy import: keep the module importable in environments where
+ # the HTTP server backend (Quart/Hypercorn) isn't installed.
+ import torch.distributed as dist
+
+ from megatron.core.inference.text_generation_server.dynamic_text_gen_server.text_generation_server import ( # pylint: disable=line-too-long
+ start_text_gen_server,
+ )
+
+ assert self._coord_runtime is not None
+ start_text_gen_server(
+ coordinator_addr=self._coord_runtime.coord_addr,
+ tokenizer=self._controller.tokenizer,
+ rank=dist.get_rank(),
+ server_port=serve_config.port,
+ parsers=serve_config.parsers,
+ verbose=serve_config.verbose,
+ num_replicas=serve_config.frontend_replicas,
+ hostname=serve_config.host,
+ )
+ self._serve_started = True
+
+ if blocking:
+ # Block until the engine loop terminates (shutdown was invoked
+ # somewhere in this process; for serve(blocking=True) typically by
+ # SIGINT or out-of-band orchestration).
+ await self.wait_for_shutdown()
+
+ async def wait_for_shutdown(self) -> None:
+ """Block until the engine's background loop task terminates.
+
+ No-op in direct mode.
+ """
+ if not self._use_coordinator:
+ return
+ assert self._loop_manager is not None
+ await self._loop_manager.run_async(self._wait_for_shutdown_impl())
+
+ async def __aenter__(self) -> "MegatronAsyncLLM":
+ return self
+
+ async def __aexit__(self, exc_type, exc, tb) -> None:
+ await self.shutdown()
diff --git a/megatron/core/inference/apis/llm.py b/megatron/core/inference/apis/llm.py
new file mode 100644
index 00000000000..7179bafa427
--- /dev/null
+++ b/megatron/core/inference/apis/llm.py
@@ -0,0 +1,153 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Sync high-level inference API for Megatron (``MegatronLLM``)."""
+
+from typing import List, Optional, Union
+
+from megatron.core.inference.apis._llm_base import _MegatronLLMBase
+from megatron.core.inference.config import InferenceConfig
+from megatron.core.inference.inference_request import DynamicInferenceRequest
+from megatron.core.inference.sampling_params import SamplingParams
+
+
+class MegatronLLM(_MegatronLLMBase):
+ """Sync high-level inference API for Megatron.
+
+ See :class:`_MegatronLLMBase` for execution modes (direct vs
+ coordinator), caller responsibilities, and the ``model.eval()`` contract.
+
+ On top of the base, this class provides:
+
+ - :meth:`generate` accepting one prompt or a batch; **always returns a
+ ``list[DynamicInferenceRequest]``** (single-prompt input returns a
+ one-element list -- deliberate asymmetry vs the async API).
+ - Sync lifecycle controls: :meth:`pause` / :meth:`unpause` /
+ :meth:`suspend` / :meth:`resume` / :meth:`shutdown` /
+ :meth:`wait_for_shutdown`.
+ - Context-manager protocol: ``with MegatronLLM(...) as llm:``; exit
+ calls :meth:`shutdown`.
+
+ Note:
+ ``serve()`` (online HTTP serving) is async-only by design; use
+ :class:`MegatronAsyncLLM` for serving.
+ """
+
+ def __init__(
+ self,
+ *,
+ model,
+ tokenizer,
+ inference_config: Optional[InferenceConfig] = None,
+ use_coordinator: bool = False,
+ coordinator_host: Optional[str] = None,
+ coordinator_port: Optional[int] = None,
+ ) -> None:
+ super().__init__(
+ model=model,
+ tokenizer=tokenizer,
+ inference_config=inference_config,
+ use_coordinator=use_coordinator,
+ coordinator_host=coordinator_host,
+ coordinator_port=coordinator_port,
+ )
+
+ def generate(
+ self,
+ prompts: Union[str, List[int], List[str], List[List[int]]],
+ sampling_params: Optional[SamplingParams] = None,
+ ) -> List["DynamicInferenceRequest"]:
+ """Run inference for one prompt or a batch.
+
+ Returns ``list[DynamicInferenceRequest]`` in input order. Single-prompt
+ input returns a one-element list -- the always-list shape is the
+ deliberate sync-vs-async asymmetry.
+
+ No concurrency guard: sync is single-caller by Python's GIL. If you
+ need to call ``generate`` concurrently from multiple threads, callers
+ must serialize externally.
+
+ Raises:
+ RuntimeError: if called on a non-primary rank in coordinator mode.
+ """
+ self._assert_primary()
+ if sampling_params is None:
+ sampling_params = SamplingParams()
+
+ normalized, _is_batch = self._normalize_prompts(prompts)
+ if not normalized:
+ return []
+
+ if self._use_coordinator:
+ assert self._loop_manager is not None
+ return self._loop_manager.run_sync(self._generate_impl(normalized, sampling_params))
+ # Direct mode: bypass _generate_impl (which would use to_thread,
+ # pointless for sync). Call the engine directly and merge.
+ records = self._engine.generate(normalized, sampling_params)
+ return [r.merge() for r in records]
+
+ def pause(self) -> None:
+ """Transition the engine to ``PAUSED``. Coordinator mode only.
+
+ Raises:
+ RuntimeError: in direct mode (``use_coordinator=False``).
+ """
+ self._assert_coordinator()
+ assert self._loop_manager is not None
+ self._loop_manager.run_sync(self._pause_impl())
+
+ def unpause(self) -> None:
+ """Transition the engine from ``PAUSED`` back to ``RUNNING``.
+
+ Raises:
+ RuntimeError: in direct mode (``use_coordinator=False``).
+ """
+ self._assert_coordinator()
+ assert self._loop_manager is not None
+ self._loop_manager.run_sync(self._unpause_impl())
+
+ def suspend(self) -> None:
+ """Transition the engine to ``SUSPENDED`` (offloads GPU buffers).
+
+ The caller must ``pause()`` first; this method does not enforce that.
+
+ Raises:
+ RuntimeError: in direct mode (``use_coordinator=False``).
+ """
+ self._assert_coordinator()
+ assert self._loop_manager is not None
+ self._loop_manager.run_sync(self._suspend_impl())
+
+ def resume(self) -> None:
+ """Transition the engine from ``SUSPENDED`` to ``RESUMED``.
+
+ Raises:
+ RuntimeError: in direct mode (``use_coordinator=False``).
+ """
+ self._assert_coordinator()
+ assert self._loop_manager is not None
+ self._loop_manager.run_sync(self._resume_impl())
+
+ def shutdown(self) -> None:
+ """Tear down the engine and runtime. Idempotent. Direct mode is a no-op."""
+ if self._shutdown_called:
+ return
+ self._shutdown_called = True
+ if not self._use_coordinator:
+ return # direct mode: nothing to tear down
+ assert self._loop_manager is not None
+ self._loop_manager.run_sync(self._shutdown_impl())
+ # Sync caller already on its own thread; no need for to_thread.
+ self._loop_manager.stop()
+
+ def wait_for_shutdown(self) -> None:
+ """Block until the engine loop terminates. Direct mode no-op."""
+ if not self._use_coordinator:
+ return
+ assert self._loop_manager is not None
+ self._loop_manager.run_sync(self._wait_for_shutdown_impl())
+
+ def __enter__(self) -> "MegatronLLM":
+ return self
+
+ def __exit__(self, exc_type, exc, tb) -> None:
+ self.shutdown()
diff --git a/megatron/core/inference/apis/serve_config.py b/megatron/core/inference/apis/serve_config.py
new file mode 100644
index 00000000000..aa7c6afe8fd
--- /dev/null
+++ b/megatron/core/inference/apis/serve_config.py
@@ -0,0 +1,44 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class ServeConfig:
+ """Programmatic configuration for ``MegatronAsyncLLM.serve(...)``.
+
+ This dataclass also serves as the future source of truth for a
+ ``megatron serve`` CLI. It controls only the HTTP serving surface; engine
+ construction and coordinator addressing are configured separately via the
+ ``MegatronLLM`` / ``MegatronAsyncLLM`` constructor.
+ """
+
+ host: str = "0.0.0.0"
+ """HTTP bind host for the OpenAI-compatible frontend.
+
+ Distinct from the ``MegatronLLM`` / ``MegatronAsyncLLM`` constructor's
+ ``coordinator_host`` argument: ``coordinator_host`` is the internal/routable
+ address used for coordinator ZMQ traffic, whereas ``host`` is the
+ externally-visible interface where the HTTP server accepts client
+ connections.
+ """
+
+ port: int = 5000
+ """HTTP bind port for the OpenAI-compatible frontend."""
+
+ parsers: list[str] = field(default_factory=list)
+ """Response parser names to enable on the HTTP frontend.
+
+ Examples include ``["json", "tool_use"]``. Values are passed through to the
+ underlying text-generation server unchanged.
+ """
+
+ verbose: bool = False
+ """Whether the HTTP frontend should log per-request detail."""
+
+ frontend_replicas: int = 4
+ """Number of HTTP frontend processes spawned on the primary rank.
+
+ The default of 4 matches the existing ``start_text_gen_server`` default of
+ ``num_replicas=4``.
+ """
diff --git a/megatron/core/inference/batch_dimensions_utils.py b/megatron/core/inference/batch_dimensions_utils.py
index e27438e63d0..d8793f01d67 100644
--- a/megatron/core/inference/batch_dimensions_utils.py
+++ b/megatron/core/inference/batch_dimensions_utils.py
@@ -14,7 +14,7 @@
import torch
-from megatron.core.utils import get_pg_size
+from megatron.core.utils import get_pg_size, round_up_to_nearest_multiple
@dataclass(order=True, frozen=True)
@@ -85,6 +85,10 @@ def is_valid(
Returns:
True if the config is valid, False otherwise
"""
+ # A dimension with no tokens serves no requests.
+ if self.token_count <= 0:
+ return False
+
# Check if total requests exceed maximum
if self.prefill_req_count + self.decode_req_count > max_requests:
return False
@@ -138,81 +142,64 @@ def req_count(self) -> int:
@staticmethod
def adjust_batch_dims_for_expert_parallelism(
local_batch_dims,
- strict: bool,
- decode_only_cuda_graphs: bool,
- smallest_non_decode_cuda_graph_size: int,
ep_group: Optional[torch.distributed.ProcessGroup] = None,
+ ep_zmq_communicator=None,
) -> Optional["InferenceBatchDimensions"]:
- """Adjusted cuda graph batch dimensions for expert parallelism.
- We take the max token count across expert model parallel group.
+ """Adjust CUDA graph batch dimensions for expert parallelism.
+
+ All-reduce-max the token count and non-decode flag across the EP group.
+ If any rank has a prefill (non-decode) step, all ranks fall back to eager
+ mode (return None) — the non-CG path handles variable token counts via
+ use_allgather_v. Otherwise return adjusted dims with the max token count.
Args:
local_batch_dims: The local batch dimensions to adjust.
- strict: Whether to use strict matching for batch dimensions.
- decode_only_cuda_graphs: Whether CUDA graphs are only used for decode steps.
ep_group: Optional expert parallel process group. If None, uses global parallel state.
When using different EP sizes for inference vs training, pass the
inference EP group explicitly.
+ ep_zmq_communicator: Optional AsyncZMQCommunicator over the EP group. When
+ provided, the cross-rank MAX reduction runs on the CPU via ZMQ
+ (no GPU kernel, no H2D/D2H), avoiding a per-step NCCL AllReduce
+ on the compute stream. When absent, falls back to
+ torch.distributed.all_reduce on a GPU tensor.
- Return:
- (InferenceBatchDimensions) A new InferenceBatchDimensions object with
- adjusted dimensions, or None if eager mode should be used.
+ Returns:
+ InferenceBatchDimensions with max token count, or None for eager mode.
"""
ep_size = get_pg_size(ep_group)
if ep_size <= 1:
return local_batch_dims
- # all reduce local work across expert model parallel group
is_non_decode = local_batch_dims.prefill_req_count > 0
- sync_tensor = torch.tensor(
- [
- local_batch_dims.token_count,
- int(is_non_decode),
- local_batch_dims.prefill_req_count,
- local_batch_dims.decode_req_count,
- ],
- dtype=torch.int32,
- device=torch.cuda.current_device(),
- )
+ if ep_zmq_communicator is not None:
+ # CPU-only sync via ZMQ: avoids a NCCL AllReduce kernel on the
+ # compute stream plus the H2D/D2H pair that sandwiches it.
+ (max_token_count, max_is_non_decode) = ep_zmq_communicator.sync_all_reduce_max(
+ local_batch_dims.token_count, int(is_non_decode)
+ )
+ else:
+ sync_tensor = torch.tensor(
+ [local_batch_dims.token_count, int(is_non_decode)],
+ dtype=torch.int32,
+ device=torch.cuda.current_device(),
+ )
+ torch.distributed.all_reduce(
+ sync_tensor, op=torch.distributed.ReduceOp.MAX, group=ep_group
+ )
+ sync_tensor = sync_tensor.cpu()
+ max_token_count = int(sync_tensor[0].item())
+ max_is_non_decode = int(sync_tensor[1].item())
- torch.distributed.all_reduce(sync_tensor, op=torch.distributed.ReduceOp.MAX, group=ep_group)
-
- sync_tensor = sync_tensor.cpu()
- is_any_ep_rank_in_non_decode = sync_tensor[1].item() == 1
-
- # We force eager mode for scenarios where some ranks will run with CUDA graphs
- # while others will not. Without this check, communication in the
- # expert routing layer would pad up to the maximum capacity only for the ranks that
- # are using CUDA graphs in this step, leading to a hang.
- # This can happen if we only allow decode CUDA graphs but some ranks are running
- # non-decode batches.
- if is_any_ep_rank_in_non_decode and decode_only_cuda_graphs:
- return None # indicate no match, run in eager mode
-
- # If strict matching is enabled, we sync the request counts across EP ranks
- # to ensure the graph captures the maximum needed capacity.
- # TODO(ksanthanam): Add functional test for this scenario
- adjusted_prefill_req_count = (
- int(sync_tensor[2].item()) if strict else local_batch_dims.prefill_req_count
- )
- adjusted_decode_req_count = (
- int(sync_tensor[3].item()) if strict else local_batch_dims.decode_req_count
- )
- adjusted_token_count = int(sync_tensor[0].item())
+ is_any_ep_rank_in_non_decode = max_is_non_decode == 1
- # When any EP rank has prefill requests (non-strict mode), elevate
- # the token count to be >= the smallest prefill/mixed cuda graph.
- # This ensures decode-only ranks don't match a fine-grained decode
- # graph while prefill ranks match a coarser mixed graph, which would
- # produce inconsistent token counts across EP ranks.
- if is_any_ep_rank_in_non_decode and not strict:
- adjusted_token_count = max(adjusted_token_count, smallest_non_decode_cuda_graph_size)
+ if is_any_ep_rank_in_non_decode:
+ return None # any rank has prefill → eager mode
adjusted_batch_dim = InferenceBatchDimensions(
- token_count=adjusted_token_count,
- prefill_req_count=adjusted_prefill_req_count,
- decode_req_count=adjusted_decode_req_count,
+ token_count=max_token_count,
+ prefill_req_count=local_batch_dims.prefill_req_count,
+ decode_req_count=local_batch_dims.decode_req_count,
)
return adjusted_batch_dim
@@ -259,7 +246,9 @@ def _calculate_cuda_graph_token_counts(
)
# Align each entry to TP size
cuda_graph_token_counts = list(
- dict.fromkeys(math.ceil(s / tp_size) * tp_size for s in cuda_graph_token_counts)
+ dict.fromkeys(
+ round_up_to_nearest_multiple(s, tp_size) for s in cuda_graph_token_counts
+ )
)
# Clamp to max tokens
cuda_graph_token_counts = [
@@ -281,7 +270,9 @@ def _calculate_cuda_graph_token_counts(
math.ceil(int(cuda_graph_step_size) / CUDAGraphBatchDimensionBuilder.CUDA_GRAPH_ROUNDER)
)
# Make sure divisible by TP size
- cuda_graph_step_size = math.ceil(cuda_graph_step_size / tp_size) * tp_size
+ cuda_graph_step_size = round_up_to_nearest_multiple(cuda_graph_step_size, tp_size)
+ # Ensure non-zero step size (can happen when max_tokens < num_cuda_graphs).
+ cuda_graph_step_size = max(cuda_graph_step_size, tp_size)
# round down cuda graph max tokens to be multiple of TP size
cuda_graph_max_tokens = (cuda_graph_max_tokens // tp_size) * tp_size
@@ -378,11 +369,9 @@ def add_if_valid(token_count: int, prefill_req_count: int, decode_req_count: int
):
cuda_graph_max_tokens = max_tokens
- assert cuda_graph_max_tokens == max_requests * (num_speculative_tokens + 1), (
- f"cuda_graph_max_tokens ({cuda_graph_max_tokens}) must equal max_requests *"
- f"(num_speculative_tokens + 1) ({max_requests * (num_speculative_tokens + 1)}). "
- "This is required for correctly syncing EP ranks: "
- f"prefill and decode graph pools must have the same token count granularity."
+ assert cuda_graph_max_tokens >= max_requests * (num_speculative_tokens + 1), (
+ f"cuda_graph_max_tokens ({cuda_graph_max_tokens}) must be >= max_requests * "
+ f"(num_speculative_tokens + 1) ({max_requests * (num_speculative_tokens + 1)})."
)
if num_cuda_graphs != -1:
@@ -496,10 +485,10 @@ def add_if_valid(token_count: int, prefill_req_count: int, decode_req_count: int
def match_graph_config(
real_batch_dim: InferenceBatchDimensions,
cuda_graph_batch_dimensions_list: List[InferenceBatchDimensions],
- smallest_non_decode_cuda_graph_size: int,
strict: bool = False,
- decode_only_cuda_graphs: bool = False,
ep_group: Optional[torch.distributed.ProcessGroup] = None,
+ ep_zmq_communicator=None,
+ match_ep_token_counts: bool = True,
) -> Optional[InferenceBatchDimensions]:
"""
Matches the best CUDA graph batch dimension for the given real batch dimension.
@@ -515,6 +504,14 @@ def match_graph_config(
ep_group: Optional expert parallel process group. If None, uses global parallel state.
When using different EP sizes for inference vs training, pass the
inference EP group explicitly.
+ ep_zmq_communicator: Optional AsyncZMQCommunicator over the EP group. When
+ provided, batch-dimension MAX reduction uses a CPU-only ZMQ sync
+ instead of a GPU NCCL AllReduce. Forwarded to
+ adjust_batch_dims_for_expert_parallelism.
+ match_ep_token_counts: If True (default), token counts are synced across EP ranks via
+ all-reduce-max so all ranks select the same CUDA graph. Set to False when the
+ dispatcher handles per-rank token variation internally (e.g. AGV/RSV in the NVLS
+ path) and external EP sync is not needed.
Returns:
The best matching CUDA graph batch dimension, or None if no applicable match is found
"""
@@ -523,19 +520,20 @@ def match_graph_config(
# no need to match if no cuda graph batch dimensions are provided
return None
- adjusted_batch_dim = InferenceBatchDimensions.adjust_batch_dims_for_expert_parallelism(
- real_batch_dim,
- strict=strict,
- decode_only_cuda_graphs=decode_only_cuda_graphs,
- ep_group=ep_group,
- smallest_non_decode_cuda_graph_size=smallest_non_decode_cuda_graph_size,
- )
+ if match_ep_token_counts:
+ # NCCL dispatcher: all EP ranks must select the same CUDA graph. Sync batch dims
+ # across the EP group so graph selection is consistent.
+ adjusted_batch_dim = InferenceBatchDimensions.adjust_batch_dims_for_expert_parallelism(
+ real_batch_dim, ep_group=ep_group, ep_zmq_communicator=ep_zmq_communicator
+ )
- if adjusted_batch_dim is None:
- # we hit this scenario if decode_only_cuda_graphs is true,
- # and one of the EP ranks is running a non-decode step
- # in that case, all ranks have to run in eager mode
- return None
+ if adjusted_batch_dim is None:
+ # we hit this scenario if decode_only_cuda_graphs is true,
+ # and one of the EP ranks is running a non-decode step
+ # in that case, all ranks have to run in eager mode
+ return None
+ else:
+ adjusted_batch_dim = real_batch_dim
# first filter out batch dimensions with smaller token count, prefill req count,
# or decode req count, as they are not applicable
diff --git a/megatron/core/inference/communication/torch_symm_triton/__init__.py b/megatron/core/inference/communication/torch_symm_triton/__init__.py
index 967dc8329f1..75da02eaf4b 100644
--- a/megatron/core/inference/communication/torch_symm_triton/__init__.py
+++ b/megatron/core/inference/communication/torch_symm_triton/__init__.py
@@ -3,3 +3,8 @@
from .collectives import multimem_all_gather, multimem_all_gather_fused, multimem_reduce_scatter
from .fused_collectives import fused_multimem_rs_add_norm_ag
from .utils import are_tensors_nvls_eligible, is_device_nvls_capable
+from .variable_collectives import (
+ multimem_all_gather_v,
+ multimem_all_gatherv_3tensor,
+ multimem_reduce_scatter_v,
+)
diff --git a/megatron/core/inference/communication/torch_symm_triton/multimem_asm.py b/megatron/core/inference/communication/torch_symm_triton/multimem_asm.py
index 859b9010aea..eace10ff167 100644
--- a/megatron/core/inference/communication/torch_symm_triton/multimem_asm.py
+++ b/megatron/core/inference/communication/torch_symm_triton/multimem_asm.py
@@ -211,6 +211,182 @@ def add_v8_bf16_from_u32(
)
+@triton.jit
+def ld_64(ptr, mask):
+ """
+ Loads 64 bits from local global memory into two 32-bit registers.
+
+ Uses `ld.global.v2.u32`. Mirrors the non-multicast path of ld_128.
+
+ Args:
+ ptr: source pointer typed as uint64 (8-byte aligned).
+ mask: boolean predicate — if False, the load is skipped.
+
+ Returns:
+ (x, y): two tl.uint32 registers containing 64 bits of loaded data.
+ """
+ return tl.inline_asm_elementwise(
+ """
+ {
+ .reg .pred %p0;
+ setp.ne.s32 %p0, $3, 1;
+ @%p0 bra end;
+ ld.global.v2.u32 {$0, $1}, [$2];
+ end:
+ }
+ """,
+ "=r,=r,l,r",
+ args=[ptr, mask.to(tl.int32)],
+ dtype=(tl.uint32, tl.uint32),
+ is_pure=True,
+ pack=1,
+ )
+
+
+@triton.jit
+def st_64(ptr, x, y, mask, multicast_op: tl.constexpr):
+ """
+ Stores 64 bits (two 32-bit registers) to memory.
+
+ Mirrors st_128 but operates on 64-bit (v2) quantities.
+
+ 1. **Standard Store (`multicast_op=False`)**:
+ - `st.global.v2.f32` — writes 64 bits to local global memory.
+
+ 2. **Multicast Store (`multicast_op=True`)**:
+ - `multimem.st.relaxed.sys.global.v2.f32` — broadcasts 64 bits to all
+ peers in the multicast group simultaneously.
+
+ Args:
+ ptr: destination pointer typed as uint64 (8-byte aligned).
+ x, y: two tl.uint32 registers containing the data to store.
+ mask: boolean predicate — if False, the store is skipped.
+ multicast_op (tl.constexpr): False = local store, True = multicast broadcast.
+ """
+ if multicast_op:
+ return tl.inline_asm_elementwise(
+ """
+ {
+ .reg .pred %p0;
+ setp.ne.s32 %p0, $4, 1;
+ @%p0 bra end;
+ multimem.st.relaxed.sys.global.v2.f32 [$1], {$2, $3};
+ end:
+ }
+ """,
+ "=r,l,r,r,r",
+ args=[ptr, x, y, mask.to(tl.int32)],
+ dtype=(tl.uint32),
+ is_pure=False,
+ pack=1,
+ )
+ else:
+ return tl.inline_asm_elementwise(
+ """
+ {
+ .reg .pred %p0;
+ setp.ne.s32 %p0, $4, 1;
+ @%p0 bra end;
+ st.global.v2.f32 [$1], {$2, $3};
+ end:
+ }
+ """,
+ "=r,l,r,r,r",
+ args=[ptr, x, y, mask.to(tl.int32)],
+ dtype=(tl.uint32),
+ is_pure=False,
+ pack=1,
+ )
+
+
+@triton.jit
+def ld_32(ptr, mask):
+ """
+ Loads 32 bits from local global memory into one 32-bit register.
+
+ Uses `ld.global.u32`. Scalar version of ld_64/ld_128.
+
+ Args:
+ ptr: source pointer typed as uint32 (4-byte aligned).
+ mask: boolean predicate — if False, the load is skipped.
+
+ Returns:
+ x: one tl.uint32 register containing 32 bits of loaded data.
+ """
+ return tl.inline_asm_elementwise(
+ """
+ {
+ .reg .pred %p0;
+ setp.ne.s32 %p0, $2, 1;
+ @%p0 bra end;
+ ld.global.u32 $0, [$1];
+ end:
+ }
+ """,
+ "=r,l,r",
+ args=[ptr, mask.to(tl.int32)],
+ dtype=(tl.uint32,),
+ is_pure=True,
+ pack=1,
+ )
+
+
+@triton.jit
+def st_32(ptr, x, mask, multicast_op: tl.constexpr):
+ """
+ Stores 32 bits (one 32-bit register) to memory.
+
+ Scalar version of st_64/st_128.
+
+ 1. **Standard Store (`multicast_op=False`)**:
+ - `st.global.f32` — writes 32 bits to local global memory.
+
+ 2. **Multicast Store (`multicast_op=True`)**:
+ - `multimem.st.relaxed.sys.global.f32` — broadcasts 32 bits to all
+ peers in the multicast group simultaneously.
+
+ Args:
+ ptr: destination pointer typed as uint32 (4-byte aligned).
+ x: one tl.uint32 register containing the data to store.
+ mask: boolean predicate — if False, the store is skipped.
+ multicast_op (tl.constexpr): False = local store, True = multicast broadcast.
+ """
+ if multicast_op:
+ return tl.inline_asm_elementwise(
+ """
+ {
+ .reg .pred %p0;
+ setp.ne.s32 %p0, $3, 1;
+ @%p0 bra end;
+ multimem.st.relaxed.sys.global.f32 [$1], $2;
+ end:
+ }
+ """,
+ "=r,l,r,r",
+ args=[ptr, x, mask.to(tl.int32)],
+ dtype=(tl.uint32),
+ is_pure=False,
+ pack=1,
+ )
+ else:
+ return tl.inline_asm_elementwise(
+ """
+ {
+ .reg .pred %p0;
+ setp.ne.s32 %p0, $3, 1;
+ @%p0 bra end;
+ st.global.f32 [$1], $2;
+ end:
+ }
+ """,
+ "=r,l,r,r",
+ args=[ptr, x, mask.to(tl.int32)],
+ dtype=(tl.uint32),
+ is_pure=False,
+ pack=1,
+ )
+
+
@triton.jit
def asm_rsqrt(x, eps):
"""
diff --git a/megatron/core/inference/communication/torch_symm_triton/variable_collectives.py b/megatron/core/inference/communication/torch_symm_triton/variable_collectives.py
new file mode 100644
index 00000000000..a32b20b9a14
--- /dev/null
+++ b/megatron/core/inference/communication/torch_symm_triton/variable_collectives.py
@@ -0,0 +1,776 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Variable-count NVLS collectives (AllGatherV / ReduceScatterV).
+
+Unlike the uniform collectives in collectives.py, each rank may contribute
+a different number of tokens. The caller provides:
+ - rank_token_offset: prefix sum of token counts for all lower-ranked ranks.
+ - local_tokens: this rank's token count.
+
+One CTA processes one token; the outer loop is persistent over local_tokens.
+"""
+
+from unittest.mock import MagicMock
+
+import torch
+
+from megatron.core.utils import null_decorator
+
+try:
+ import triton
+ import triton.language as tl
+
+ HAVE_TRITON = True
+except ImportError:
+ triton = MagicMock()
+ triton.jit = null_decorator
+ tl = MagicMock()
+ HAVE_TRITON = False
+
+try:
+ from torch._C._distributed_c10d import _SymmetricMemory
+except ImportError:
+ _SymmetricMemory = MagicMock()
+
+from .barrier import symm_mem_sync
+from .multimem_asm import ld_64, ld_128, st_64, st_128
+from .utils import is_device_nvls_capable, sync_threads
+
+
+@triton.jit
+def _multimem_all_gather_v_kernel(
+ local_ptr,
+ multicast_ptr,
+ signal_pad_ptrs,
+ local_tokens,
+ rank_token_offset_ptr,
+ ep_max_tokens_ptr,
+ output_byte_offset,
+ HIDDEN_SIZE: tl.constexpr,
+ BLOCK_SIZE: tl.constexpr,
+ NUMEL_PER_THREAD: tl.constexpr,
+ BITS: tl.constexpr,
+ RANK: tl.constexpr,
+ WORLD_SIZE: tl.constexpr,
+):
+ """Variable-count multicast all-gather kernel. One CTA processes one token.
+
+ Each rank contributes local_tokens tokens starting at rank_token_offset in
+ the global output. Ranks may have different local_tokens values.
+
+ Args:
+ local_ptr: pointer to this rank's local input, shape [local_tokens, hidden_size].
+ multicast_ptr: multicast pointer to the output symmetric memory buffer.
+ signal_pad_ptrs: signal pads for barrier synchronization.
+ local_tokens: number of tokens this rank contributes.
+ rank_token_offset_ptr: pointer to a scalar int32 CUDA tensor holding the index
+ of the first token this rank writes in the global output (prefix sum of
+ local_tokens for all lower-ranked ranks). Fixed address; value set each step.
+ ep_max_tokens_ptr: pointer to a scalar int32 CUDA tensor holding the
+ maximum local_tokens across all EP ranks for this iteration. Fixed address;
+ value set each step. CTAs with pid >= this value exit immediately. Safe
+ because the value is identical on all ranks, so paired CTAs on every rank
+ exit together — the barrier for those CTAs is never entered on any rank.
+ output_byte_offset: byte offset of this tensor within the symmetric memory buffer.
+ HIDDEN_SIZE: hidden dimension, i.e. number of elements per token row (constexpr).
+ BLOCK_SIZE: threads per block (constexpr, >= numel_per_token).
+ NUMEL_PER_THREAD: elements per thread per load/store, i.e. BITS / element_bits (constexpr).
+ BITS: width of each load/store in bits — 128 for activations (bf16) and expert
+ indices (int64, always 16-byte aligned for any topk); 64 for routing probs
+ (fp32 with topk=6 or topk=22 yields 24/88-byte rows, not 16-byte aligned
+ but 8-byte aligned) (constexpr).
+ RANK: this rank's index (constexpr).
+ WORLD_SIZE: total number of ranks (constexpr).
+ """
+ pid = tl.program_id(axis=0)
+
+ # Exit before the barrier if this CTA's pid exceeds the iteration maximum.
+ # ep_max_tokens is the max over all EP ranks, so all ranks agree on
+ # which CTAs exit — the barrier slots for those CTAs are never touched on any rank.
+ ep_max_tokens = tl.load(ep_max_tokens_ptr)
+ if pid >= ep_max_tokens:
+ return
+
+ tid = tl.arange(0, BLOCK_SIZE)
+ rank_token_offset = tl.load(rank_token_offset_ptr)
+
+ numel_per_token = tl.cdiv(HIDDEN_SIZE, NUMEL_PER_THREAD)
+ local_numel = local_tokens * numel_per_token
+ # BLOCK_SIZE is the next power of 2 >= numel_per_token, so it may be larger.
+ # channel_mask deactivates the extra padding threads (tid >= numel_per_token).
+ channel_mask = tid < numel_per_token
+
+ for token_offset in range(pid, local_tokens, tl.num_programs(axis=0)):
+ for channel_offset in range(0, numel_per_token, BLOCK_SIZE):
+ local_offsets = token_offset * numel_per_token + channel_offset + tid
+ # Two independent masks in orthogonal dimensions:
+ # channel_mask — deactivates power-of-2 padding threads (tid >= numel_per_token).
+ # token_mask — deactivates overflow threads in the last inner-loop chunk
+ # when numel_per_token > BLOCK_SIZE and the window
+ # [channel_offset, channel_offset+BLOCK_SIZE) extends past
+ # the final token row.
+ token_mask = local_offsets < local_numel
+ mask = token_mask & channel_mask
+
+ # This rank's tokens start at rank_token_offset in the global output.
+ global_offsets = rank_token_offset * numel_per_token + local_offsets
+
+ if BITS == 128:
+ # Each 128-bit pack occupies 2 uint64 units; output_byte_offset // 8 converts
+ # the tensor's byte offset within the symm-mem buffer to uint64 units.
+ # The global offset is multiplied by 2 to convert from 128-bit
+ # units to uint64 units.
+ multicast_ptrs = (
+ multicast_ptr.to(tl.pointer_type(tl.uint64))
+ + output_byte_offset // 8
+ + global_offsets * 2
+ )
+ local_ptrs = local_ptr.to(tl.pointer_type(tl.uint64)) + local_offsets * 2
+ (x, y, z, w) = ld_128(local_ptrs, mask=mask, multicast_op=False)
+ st_128(multicast_ptrs, x, y, z, w, mask=mask, multicast_op=True)
+ else:
+ # Each 64-bit pack is exactly 1 uint64, so offsets index directly (no * 2 stride).
+ multicast_ptrs = (
+ multicast_ptr.to(tl.pointer_type(tl.uint64))
+ + output_byte_offset // 8
+ + global_offsets
+ )
+ local_ptrs = local_ptr.to(tl.pointer_type(tl.uint64)) + local_offsets
+ (x, y) = ld_64(local_ptrs, mask=mask)
+ st_64(multicast_ptrs, x, y, mask=mask, multicast_op=True)
+
+ sync_threads()
+ symm_mem_sync(
+ signal_pad_ptrs,
+ None,
+ RANK,
+ WORLD_SIZE,
+ hasPreviousMemAccess=True,
+ hasSubsequentMemAccess=True,
+ )
+
+
+@triton.jit
+def _multimem_reduce_scatter_v_kernel(
+ local_ptr,
+ multicast_ptr,
+ signal_pad_ptrs,
+ local_tokens,
+ rank_token_offset_ptr,
+ ep_max_tokens_ptr,
+ input_byte_offset,
+ HIDDEN_SIZE: tl.constexpr,
+ BLOCK_SIZE: tl.constexpr,
+ NUMEL_PER_THREAD: tl.constexpr,
+ RANK: tl.constexpr,
+ WORLD_SIZE: tl.constexpr,
+ REDUCE_F32: tl.constexpr = False,
+):
+ """Variable-count multicast reduce-scatter kernel. One CTA processes one token.
+
+ Reads this rank's token shard from the symmetric buffer via multimem.ld_reduce
+ (which atomically sums contributions from all EP ranks) and writes the result
+ to local memory.
+
+ The barrier runs first — it waits for all ranks to have written their expert
+ GEMM outputs into the symmetric buffer before any rank starts reading.
+
+ Args:
+ local_ptr: output pointer to this rank's local buffer, shape [local_tokens, hidden_size].
+ multicast_ptr: multicast pointer to the symmetric memory buffer holding all expert outputs.
+ signal_pad_ptrs: signal pads for barrier synchronization.
+ local_tokens: number of tokens this rank owns.
+ rank_token_offset_ptr: pointer to a scalar int32 CUDA tensor holding the index of the
+ first token this rank owns in the global token sequence. Fixed address; set each step.
+ ep_max_tokens_ptr: pointer to a scalar int32 CUDA tensor holding the maximum local_tokens
+ across all EP ranks. Fixed address; set each step. CTAs with pid >= this value exit
+ immediately — safe because the value is identical on all ranks.
+ input_byte_offset: byte offset of the input tensor within the symmetric memory buffer.
+ HIDDEN_SIZE: number of elements per token row (constexpr).
+ BLOCK_SIZE: threads per block (constexpr, >= numel_per_token).
+ NUMEL_PER_THREAD: elements per thread per load/store, i.e. 128 / element_bits (constexpr).
+ RANK: this rank's index (constexpr).
+ WORLD_SIZE: total number of ranks (constexpr).
+ """
+ pid = tl.program_id(axis=0)
+
+ # Exit before the barrier if this CTA's pid exceeds the iteration maximum.
+ # ep_max_tokens is the max over all EP ranks, so all ranks agree on which
+ # CTAs exit — the barrier slots for those CTAs are never touched on any rank.
+ ep_max_tokens = tl.load(ep_max_tokens_ptr)
+ if pid >= ep_max_tokens:
+ return
+
+ # Wait for all ranks to have written their expert GEMM outputs to symm_mem
+ # before any rank starts the reduce-load.
+ symm_mem_sync(
+ signal_pad_ptrs,
+ None,
+ RANK,
+ WORLD_SIZE,
+ hasPreviousMemAccess=False,
+ hasSubsequentMemAccess=False,
+ )
+ sync_threads()
+
+ tid = tl.arange(0, BLOCK_SIZE)
+ rank_token_offset = tl.load(rank_token_offset_ptr)
+
+ numel_per_token = tl.cdiv(HIDDEN_SIZE, NUMEL_PER_THREAD)
+ local_numel = local_tokens * numel_per_token
+ # channel_mask: deactivates power-of-2 padding threads (tid >= numel_per_token).
+ channel_mask = tid < numel_per_token
+
+ for token_offset in range(pid, local_tokens, tl.num_programs(axis=0)):
+ program_offset = token_offset * numel_per_token
+
+ for channel_offset in range(0, numel_per_token, BLOCK_SIZE):
+ local_offsets = program_offset + channel_offset + tid
+ # Two independent masks in orthogonal dimensions:
+ # channel_mask — deactivates power-of-2 padding threads (tid >= numel_per_token).
+ # token_mask — deactivates overflow threads in the last inner-loop chunk
+ # when numel_per_token > BLOCK_SIZE and the window
+ # [channel_offset, channel_offset+BLOCK_SIZE) extends past
+ # the final token row.
+ token_mask = local_offsets < local_numel
+ mask = token_mask & channel_mask
+
+ # This rank's tokens start at rank_token_offset in the global input.
+ global_offsets = rank_token_offset * numel_per_token + local_offsets
+
+ # Each 128-bit pack occupies 2 uint64 units; input_byte_offset // 8 converts
+ # the tensor's byte offset within the symm-mem buffer to uint64 units.
+ multicast_ptrs = (
+ multicast_ptr.to(tl.pointer_type(tl.uint64))
+ + input_byte_offset // 8
+ + global_offsets * 2
+ )
+ local_ptrs = local_ptr.to(tl.pointer_type(tl.uint64)) + local_offsets * 2
+
+ (x, y, z, w) = ld_128(
+ multicast_ptrs, mask=mask, multicast_op=True, reduce_f32=REDUCE_F32
+ )
+ st_128(local_ptrs, x, y, z, w, mask=mask, multicast_op=False)
+
+
+def multimem_reduce_scatter_v(
+ output_tensor: torch.Tensor,
+ input_tensor: torch.Tensor,
+ symm_mem_hdl: _SymmetricMemory,
+ rank_token_offset: torch.Tensor,
+ ep_max_tokens: torch.Tensor,
+ per_rank_max_tokens: int,
+ input_byte_offset: int = 0,
+ **kwargs,
+) -> torch.Tensor:
+ """Variable-count multicast reduce-scatter for a single 2-D tensor.
+
+ Reduces expert GEMM outputs across all EP ranks. Each rank reads its owned
+ token shard [rank_token_offset : rank_token_offset + local_tokens] from the
+ symmetric buffer using multimem.ld_reduce (which atomically sums all ranks'
+ contributions), and writes the result to output_tensor.
+
+ Both tensors must be 2-D and 16-byte row-aligned (128-bit path only).
+ hidden_size is inferred from output_tensor.shape[1].
+
+ Args:
+ output_tensor: local output, shape [local_tokens, hidden_size].
+ input_tensor: symmetric memory buffer holding all expert outputs,
+ shape [global_tokens, hidden_size].
+ symm_mem_hdl: symmetric memory handle for input_tensor.
+ rank_token_offset: pre-allocated scalar int32 CUDA tensor. The dispatcher
+ writes this rank's token offset into it each step before kernel launch.
+ ep_max_tokens: pre-allocated scalar int32 CUDA tensor. The dispatcher writes
+ the maximum local_tokens across all EP ranks each step. CTAs with
+ pid >= ep_max_tokens exit immediately without entering the barrier.
+ per_rank_max_tokens: static int set at model init. Determines the CTA grid size
+ as min(per_rank_max_tokens, MAX_NUM_BLOCKS).
+ input_byte_offset: byte offset of input_tensor within the symmetric memory
+ buffer (for packing multiple tensors into one buffer; 0 otherwise).
+
+ Returns:
+ output_tensor populated with this rank's reduced token outputs.
+ """
+ assert HAVE_TRITON, "Triton is required for multimem reduce-scatter-v."
+ assert (
+ output_tensor.ndim == 2 and input_tensor.ndim == 2
+ ), "output_tensor and input_tensor must be 2-D [tokens, hidden_size]."
+ assert is_device_nvls_capable(
+ output_tensor.device
+ ), "multimem_reduce_scatter_v requires a Hopper+ GPU with NVLink (SM >= 9)."
+ assert (
+ rank_token_offset.numel() == 1
+ and rank_token_offset.dtype == torch.int32
+ and rank_token_offset.is_cuda
+ ), "rank_token_offset must be a scalar int32 CUDA tensor."
+ assert output_tensor.dtype in (
+ torch.bfloat16,
+ torch.float32,
+ ), f"Only bfloat16 and float32 are supported, got {output_tensor.dtype}"
+ assert (
+ output_tensor.dtype == input_tensor.dtype
+ ), f"output and input dtype mismatch: {output_tensor.dtype} vs {input_tensor.dtype}"
+
+ hidden_size = output_tensor.shape[1]
+ assert (
+ input_tensor.shape[1] == hidden_size
+ ), f"input and output hidden_size mismatch: {input_tensor.shape[1]} vs {hidden_size}"
+ row_bytes = hidden_size * output_tensor.element_size()
+ assert row_bytes % 16 == 0, (
+ f"Row size ({hidden_size} elements × {output_tensor.element_size()} bytes) = "
+ f"{row_bytes} bytes is not 16-byte aligned; RSV requires 128-bit alignment."
+ )
+
+ MAX_NUM_BLOCKS = kwargs.get("max_num_blocks", 128)
+ MAX_BLOCK_SIZE = 1024
+ WARP_SIZE = 32
+
+ local_tokens = output_tensor.shape[0]
+ numel_per_thread = 128 // (output_tensor.element_size() * 8)
+ numel_per_token = (hidden_size + numel_per_thread - 1) // numel_per_thread
+
+ block_size = min(triton.next_power_of_2(numel_per_token), MAX_BLOCK_SIZE)
+ num_warps = max(1, block_size // WARP_SIZE)
+ num_blocks = min(per_rank_max_tokens, MAX_NUM_BLOCKS)
+
+ reduce_f32 = output_tensor.dtype == torch.float32
+ _multimem_reduce_scatter_v_kernel[(num_blocks, 1, 1)](
+ output_tensor.data_ptr(),
+ symm_mem_hdl.multicast_ptr,
+ symm_mem_hdl.signal_pad_ptrs_dev,
+ local_tokens=local_tokens,
+ rank_token_offset_ptr=rank_token_offset,
+ ep_max_tokens_ptr=ep_max_tokens,
+ input_byte_offset=input_byte_offset,
+ HIDDEN_SIZE=hidden_size,
+ BLOCK_SIZE=block_size,
+ NUMEL_PER_THREAD=numel_per_thread,
+ RANK=symm_mem_hdl.rank,
+ WORLD_SIZE=symm_mem_hdl.world_size,
+ REDUCE_F32=reduce_f32,
+ num_warps=num_warps,
+ )
+
+ return output_tensor
+
+
+@triton.jit
+def _multimem_all_gatherv_3tensor_kernel(
+ local_ptr_0,
+ multicast_ptr_0,
+ output_byte_offset_0,
+ local_ptr_1,
+ multicast_ptr_1,
+ output_byte_offset_1,
+ local_ptr_2,
+ multicast_ptr_2,
+ output_byte_offset_2,
+ signal_pad_ptrs,
+ local_tokens,
+ rank_token_offset_ptr,
+ ep_max_tokens_ptr,
+ HIDDEN_SIZE_0: tl.constexpr,
+ HIDDEN_SIZE_1: tl.constexpr,
+ HIDDEN_SIZE_2: tl.constexpr,
+ BLOCK_SIZE: tl.constexpr,
+ NUMEL_PER_THREAD_0: tl.constexpr,
+ NUMEL_PER_THREAD_1: tl.constexpr,
+ NUMEL_PER_THREAD_2: tl.constexpr,
+ BITS_0: tl.constexpr,
+ BITS_1: tl.constexpr,
+ BITS_2: tl.constexpr,
+ RANK: tl.constexpr,
+ WORLD_SIZE: tl.constexpr,
+):
+ """Variable-count multicast all-gather for three tensors in a single kernel.
+
+ Identical semantics to _multimem_all_gather_v_kernel but processes three
+ tensors per CTA iteration, sharing a single barrier. This avoids launching
+ three separate kernels (and three separate barriers) for the common case
+ of gathering hidden states, routing probabilities, and expert indices together.
+
+ The outer token loop is shared across all three tensors; each tensor has its
+ own inner channel loop with independent masking. BLOCK_SIZE is the maximum
+ of the three per-tensor block sizes — smaller tensors mask out the extra threads
+ via channel_mask.
+
+ signal_pad_ptrs from the first output buffer's symmetric memory handle are used
+ for the single end-of-kernel barrier. Since all three writes complete before the
+ barrier, a single sync suffices for all three tensors.
+
+ Args:
+ local_ptr_0/1/2: pointers to each rank's local input for tensors 0/1/2.
+ multicast_ptr_0/1/2: multicast pointers to the output symmetric memory buffers.
+ output_byte_offset_0/1/2: byte offsets of each tensor within its symmetric
+ memory buffer (0 when the buffer holds only that tensor).
+ signal_pad_ptrs: signal pads from symm_mem_hdl_0, used for the single barrier.
+ local_tokens: number of tokens this rank contributes (shared across tensors).
+ rank_token_offset_ptr: pointer to a scalar int32 CUDA tensor holding this rank's
+ write offset in the global output (prefix sum over lower-ranked EP ranks).
+ ep_max_tokens_ptr: pointer to a scalar int32 CUDA tensor holding the maximum
+ local_tokens across all EP ranks. CTAs with pid >= this value exit immediately.
+ HIDDEN_SIZE_0/1/2: hidden dimension (elements per token row) for each tensor (constexpr).
+ BLOCK_SIZE: threads per block — max of the three per-tensor block sizes (constexpr).
+ NUMEL_PER_THREAD_0/1/2: elements per thread per load/store for each tensor (constexpr).
+ BITS_0/1/2: load/store width in bits (128 or 64) for each tensor (constexpr).
+ RANK: this rank's index (constexpr).
+ WORLD_SIZE: total number of ranks (constexpr).
+ """
+ pid = tl.program_id(axis=0)
+
+ ep_max_tokens = tl.load(ep_max_tokens_ptr)
+ if pid >= ep_max_tokens:
+ return
+
+ tid = tl.arange(0, BLOCK_SIZE)
+ rank_token_offset = tl.load(rank_token_offset_ptr)
+
+ numel_per_token_0 = tl.cdiv(HIDDEN_SIZE_0, NUMEL_PER_THREAD_0)
+ numel_per_token_1 = tl.cdiv(HIDDEN_SIZE_1, NUMEL_PER_THREAD_1)
+ numel_per_token_2 = tl.cdiv(HIDDEN_SIZE_2, NUMEL_PER_THREAD_2)
+
+ local_numel_0 = local_tokens * numel_per_token_0
+ local_numel_1 = local_tokens * numel_per_token_1
+ local_numel_2 = local_tokens * numel_per_token_2
+
+ # channel_mask: deactivates threads beyond each tensor's numel_per_token (power-of-2 padding).
+ channel_mask_0 = tid < numel_per_token_0
+ channel_mask_1 = tid < numel_per_token_1
+ channel_mask_2 = tid < numel_per_token_2
+
+ for token_offset in range(pid, local_tokens, tl.num_programs(axis=0)):
+ # --- Tensor 0 ---
+ for channel_offset in range(0, numel_per_token_0, BLOCK_SIZE):
+ local_offsets = token_offset * numel_per_token_0 + channel_offset + tid
+ token_mask = local_offsets < local_numel_0
+ mask = token_mask & channel_mask_0
+ global_offsets = rank_token_offset * numel_per_token_0 + local_offsets
+ if BITS_0 == 128:
+ multicast_ptrs = (
+ multicast_ptr_0.to(tl.pointer_type(tl.uint64))
+ + output_byte_offset_0 // 8
+ + global_offsets * 2
+ )
+ local_ptrs = local_ptr_0.to(tl.pointer_type(tl.uint64)) + local_offsets * 2
+ (x, y, z, w) = ld_128(local_ptrs, mask=mask, multicast_op=False)
+ st_128(multicast_ptrs, x, y, z, w, mask=mask, multicast_op=True)
+ else:
+ multicast_ptrs = (
+ multicast_ptr_0.to(tl.pointer_type(tl.uint64))
+ + output_byte_offset_0 // 8
+ + global_offsets
+ )
+ local_ptrs = local_ptr_0.to(tl.pointer_type(tl.uint64)) + local_offsets
+ (x, y) = ld_64(local_ptrs, mask=mask)
+ st_64(multicast_ptrs, x, y, mask=mask, multicast_op=True)
+
+ # --- Tensor 1 ---
+ for channel_offset in range(0, numel_per_token_1, BLOCK_SIZE):
+ local_offsets = token_offset * numel_per_token_1 + channel_offset + tid
+ token_mask = local_offsets < local_numel_1
+ mask = token_mask & channel_mask_1
+ global_offsets = rank_token_offset * numel_per_token_1 + local_offsets
+ if BITS_1 == 128:
+ multicast_ptrs = (
+ multicast_ptr_1.to(tl.pointer_type(tl.uint64))
+ + output_byte_offset_1 // 8
+ + global_offsets * 2
+ )
+ local_ptrs = local_ptr_1.to(tl.pointer_type(tl.uint64)) + local_offsets * 2
+ (x, y, z, w) = ld_128(local_ptrs, mask=mask, multicast_op=False)
+ st_128(multicast_ptrs, x, y, z, w, mask=mask, multicast_op=True)
+ else:
+ multicast_ptrs = (
+ multicast_ptr_1.to(tl.pointer_type(tl.uint64))
+ + output_byte_offset_1 // 8
+ + global_offsets
+ )
+ local_ptrs = local_ptr_1.to(tl.pointer_type(tl.uint64)) + local_offsets
+ (x, y) = ld_64(local_ptrs, mask=mask)
+ st_64(multicast_ptrs, x, y, mask=mask, multicast_op=True)
+
+ # --- Tensor 2 ---
+ for channel_offset in range(0, numel_per_token_2, BLOCK_SIZE):
+ local_offsets = token_offset * numel_per_token_2 + channel_offset + tid
+ token_mask = local_offsets < local_numel_2
+ mask = token_mask & channel_mask_2
+ global_offsets = rank_token_offset * numel_per_token_2 + local_offsets
+ if BITS_2 == 128:
+ multicast_ptrs = (
+ multicast_ptr_2.to(tl.pointer_type(tl.uint64))
+ + output_byte_offset_2 // 8
+ + global_offsets * 2
+ )
+ local_ptrs = local_ptr_2.to(tl.pointer_type(tl.uint64)) + local_offsets * 2
+ (x, y, z, w) = ld_128(local_ptrs, mask=mask, multicast_op=False)
+ st_128(multicast_ptrs, x, y, z, w, mask=mask, multicast_op=True)
+ else:
+ multicast_ptrs = (
+ multicast_ptr_2.to(tl.pointer_type(tl.uint64))
+ + output_byte_offset_2 // 8
+ + global_offsets
+ )
+ local_ptrs = local_ptr_2.to(tl.pointer_type(tl.uint64)) + local_offsets
+ (x, y) = ld_64(local_ptrs, mask=mask)
+ st_64(multicast_ptrs, x, y, mask=mask, multicast_op=True)
+
+ sync_threads()
+ symm_mem_sync(
+ signal_pad_ptrs,
+ None,
+ RANK,
+ WORLD_SIZE,
+ hasPreviousMemAccess=True,
+ hasSubsequentMemAccess=True,
+ )
+
+
+def multimem_all_gather_v(
+ output_tensor: torch.Tensor,
+ input_tensor: torch.Tensor,
+ symm_mem_hdl: _SymmetricMemory,
+ rank_token_offset: torch.Tensor,
+ ep_max_tokens: torch.Tensor,
+ per_rank_max_tokens: int,
+ output_byte_offset: int = 0,
+ **kwargs,
+) -> torch.Tensor:
+ """Variable-count multicast all-gather for a single 2-D tensor.
+
+ Gathers [local_tokens, hidden_size] from each EP rank into a shared
+ output_tensor of shape [global_tokens, hidden_size], where global_tokens is
+ the sum of all ranks' local_tokens. Each rank writes its slice starting at
+ rank_token_offset in the output.
+
+ Both tensors must be 2-D; hidden_size is inferred from input_tensor.shape[1].
+ The 128-bit or 64-bit NVLS path is selected automatically based on row alignment.
+
+ Args:
+ output_tensor: symmetric memory buffer, shape [global_tokens, hidden_size].
+ input_tensor: this rank's local input, shape [local_tokens, hidden_size].
+ symm_mem_hdl: symmetric memory handle for output_tensor.
+ rank_token_offset: pre-allocated scalar int32 CUDA tensor. The dispatcher
+ writes this rank's token offset (prefix sum over lower-ranked EP ranks)
+ into it each step before kernel launch.
+ ep_max_tokens: pre-allocated scalar int32 CUDA tensor. The dispatcher writes
+ the maximum local_tokens across all EP ranks into it each step. CTAs with
+ pid >= ep_max_tokens exit immediately — safe because all ranks agree on
+ this value, so the corresponding CTAs exit on every rank simultaneously.
+ per_rank_max_tokens: static int set at model init. Determines the CTA grid size
+ as min(per_rank_max_tokens, MAX_NUM_BLOCKS). Typically > MAX_NUM_BLOCKS so
+ we always launch MAX_NUM_BLOCKS CTAs.
+ output_byte_offset: byte offset of this tensor within the symmetric memory buffer
+ (for packing multiple tensors into one buffer; 0 if the buffer holds only
+ this tensor).
+
+ Returns:
+ output_tensor with all ranks' data written.
+ """
+ assert HAVE_TRITON, "Triton is required for multimem all-gather-v."
+ assert input_tensor.ndim == 2 and output_tensor.ndim == 2, (
+ f"input_tensor and output_tensor must be 2-D [tokens, hidden_size], "
+ f"got input_tensor.shape={input_tensor.shape}, output_tensor.shape={output_tensor.shape}."
+ )
+ assert is_device_nvls_capable(
+ input_tensor.device
+ ), "multimem_all_gather_v requires a Hopper+ GPU with NVLink (SM >= 9)."
+ assert (
+ rank_token_offset.numel() == 1
+ and rank_token_offset.dtype == torch.int32
+ and rank_token_offset.is_cuda
+ ), "rank_token_offset must be a scalar int32 CUDA tensor."
+
+ hidden_size = input_tensor.shape[1]
+ assert (
+ input_tensor.shape[1] == output_tensor.shape[1]
+ ), f"input and output hidden_size mismatch: {input_tensor.shape[1]} vs {output_tensor.shape[1]}"
+
+ row_bytes = hidden_size * input_tensor.element_size()
+ assert row_bytes % 8 == 0, (
+ f"Row size ({hidden_size} elements × {input_tensor.element_size()} bytes) = "
+ f"{row_bytes} bytes is not 8-byte aligned; cannot use NVLS."
+ )
+ bits = 128 if row_bytes % 16 == 0 else 64
+
+ MAX_NUM_BLOCKS = kwargs.get("max_num_blocks", 128)
+ MAX_BLOCK_SIZE = 1024
+ WARP_SIZE = 32
+
+ local_tokens = input_tensor.shape[0]
+ numel_per_thread = bits // (input_tensor.element_size() * 8)
+ numel_per_token = (hidden_size + numel_per_thread - 1) // numel_per_thread
+
+ # BLOCK_SIZE must be a constexpr and >= numel_per_token; round up to next power of 2.
+ block_size = min(triton.next_power_of_2(numel_per_token), MAX_BLOCK_SIZE)
+ num_warps = max(1, block_size // WARP_SIZE)
+
+ # All ranks launch the same fixed number of CTAs. CTAs with
+ # pid >= ep_max_tokens exit immediately at kernel entry.
+ num_blocks = min(per_rank_max_tokens, MAX_NUM_BLOCKS)
+
+ _multimem_all_gather_v_kernel[(num_blocks, 1, 1)](
+ input_tensor.data_ptr(),
+ symm_mem_hdl.multicast_ptr,
+ symm_mem_hdl.signal_pad_ptrs_dev,
+ local_tokens=local_tokens,
+ rank_token_offset_ptr=rank_token_offset,
+ ep_max_tokens_ptr=ep_max_tokens,
+ output_byte_offset=output_byte_offset,
+ HIDDEN_SIZE=hidden_size,
+ BLOCK_SIZE=block_size,
+ NUMEL_PER_THREAD=numel_per_thread,
+ BITS=bits,
+ RANK=symm_mem_hdl.rank,
+ WORLD_SIZE=symm_mem_hdl.world_size,
+ num_warps=num_warps,
+ )
+
+ return output_tensor
+
+
+def multimem_all_gatherv_3tensor(
+ output_tensor_0: torch.Tensor,
+ output_tensor_1: torch.Tensor,
+ output_tensor_2: torch.Tensor,
+ input_tensor_0: torch.Tensor,
+ input_tensor_1: torch.Tensor,
+ input_tensor_2: torch.Tensor,
+ symm_mem_hdl_0: _SymmetricMemory,
+ symm_mem_hdl_1: _SymmetricMemory,
+ symm_mem_hdl_2: _SymmetricMemory,
+ rank_token_offset: torch.Tensor,
+ ep_max_tokens: torch.Tensor,
+ per_rank_max_tokens: int,
+ output_byte_offset_0: int = 0,
+ output_byte_offset_1: int = 0,
+ output_byte_offset_2: int = 0,
+ **kwargs,
+) -> tuple:
+ """Variable-count multicast all-gather for three tensors in a single kernel launch.
+
+ Gathers three independent [local_tokens, hidden_size_i] tensors from every EP rank
+ into their respective output symmetric memory buffers in one fused kernel, sharing a
+ single end-of-kernel barrier. This is more efficient than calling multimem_all_gather_v
+ three times because the barrier cost (one per kernel) is paid only once.
+
+ All three input tensors must share the same local_tokens dimension (i.e. the same
+ number of token rows per rank). Each tensor may have a different hidden_size and dtype.
+ The 128-bit or 64-bit NVLS path is selected independently per tensor based on row
+ alignment.
+
+ The barrier at the end of the kernel uses signal_pad_ptrs from symm_mem_hdl_0. Since
+ all three multicast stores complete before the barrier, a single sync covers all three
+ tensors. All three handles must belong to the same EP group (identical rank/world_size).
+
+ Args:
+ output_tensor_0/1/2: symmetric memory buffers for each tensor,
+ shape [global_tokens, hidden_size_i].
+ input_tensor_0/1/2: this rank's local inputs, shape [local_tokens, hidden_size_i].
+ symm_mem_hdl_0/1/2: symmetric memory handles for each output buffer.
+ signal_pad_ptrs from hdl_0 are used for the single end-of-kernel barrier.
+ rank_token_offset: pre-allocated scalar int32 CUDA tensor. The dispatcher writes
+ this rank's token offset (prefix sum over lower-ranked EP ranks) each step.
+ ep_max_tokens: pre-allocated scalar int32 CUDA tensor. The dispatcher writes the
+ maximum local_tokens across all EP ranks each step. CTAs with
+ pid >= ep_max_tokens exit immediately — safe because all ranks agree.
+ per_rank_max_tokens: static int set at model init. Determines the CTA grid size as
+ min(per_rank_max_tokens, MAX_NUM_BLOCKS).
+ output_byte_offset_0/1/2: byte offset of each tensor within its symmetric memory
+ buffer (for packing multiple tensors into one buffer; 0 otherwise).
+
+ Returns:
+ Tuple of (output_tensor_0, output_tensor_1, output_tensor_2) with all ranks'
+ data written.
+ """
+ assert HAVE_TRITON, "Triton is required for multimem all-gather-v3."
+ for i, (inp, out) in enumerate(
+ zip(
+ (input_tensor_0, input_tensor_1, input_tensor_2),
+ (output_tensor_0, output_tensor_1, output_tensor_2),
+ )
+ ):
+ assert inp.ndim == 2 and out.ndim == 2, (
+ f"input_tensor_{i} and output_tensor_{i} must be 2-D [tokens, hidden_size], "
+ f"got input_tensor_{i}.shape={inp.shape}, output_tensor_{i}.shape={out.shape}."
+ )
+ assert inp.shape[1] == out.shape[1], (
+ f"input_tensor_{i} and output_tensor_{i} hidden_size mismatch: "
+ f"{inp.shape[1]} vs {out.shape[1]}."
+ )
+ assert (
+ input_tensor_0.shape[0] == input_tensor_1.shape[0] == input_tensor_2.shape[0]
+ ), "All three input tensors must have the same local_tokens (first dimension)."
+ assert is_device_nvls_capable(
+ input_tensor_0.device
+ ), "multimem_all_gatherv_3tensor requires a Hopper+ GPU with NVLink (SM >= 9)."
+ assert (
+ rank_token_offset.numel() == 1
+ and rank_token_offset.dtype == torch.int32
+ and rank_token_offset.is_cuda
+ ), "rank_token_offset must be a scalar int32 CUDA tensor."
+ assert (
+ symm_mem_hdl_0.rank == symm_mem_hdl_1.rank == symm_mem_hdl_2.rank
+ ), "All three symmetric memory handles must belong to the same EP group (rank mismatch)."
+ assert (
+ symm_mem_hdl_0.world_size == symm_mem_hdl_1.world_size == symm_mem_hdl_2.world_size
+ ), "All three symmetric memory handles must belong to the same EP group (world_size mismatch)."
+
+ MAX_NUM_BLOCKS = kwargs.get("max_num_blocks", 128)
+ MAX_BLOCK_SIZE = 1024
+ WARP_SIZE = 32
+
+ local_tokens = input_tensor_0.shape[0]
+
+ def _tensor_params(inp):
+ hidden_size = inp.shape[1]
+ row_bytes = hidden_size * inp.element_size()
+ assert row_bytes % 8 == 0, (
+ f"Row size ({hidden_size} elements × {inp.element_size()} bytes) = "
+ f"{row_bytes} bytes is not 8-byte aligned; cannot use NVLS."
+ )
+ bits = 128 if row_bytes % 16 == 0 else 64
+ numel_per_thread = bits // (inp.element_size() * 8)
+ numel_per_token = (hidden_size + numel_per_thread - 1) // numel_per_thread
+ block_size = min(triton.next_power_of_2(numel_per_token), MAX_BLOCK_SIZE)
+ return hidden_size, bits, numel_per_thread, block_size
+
+ hidden_size_0, bits_0, numel_per_thread_0, block_size_0 = _tensor_params(input_tensor_0)
+ hidden_size_1, bits_1, numel_per_thread_1, block_size_1 = _tensor_params(input_tensor_1)
+ hidden_size_2, bits_2, numel_per_thread_2, block_size_2 = _tensor_params(input_tensor_2)
+
+ # Use the largest block size so all threads are occupied for at least one tensor;
+ # smaller tensors mask out excess threads via channel_mask inside the kernel.
+ block_size = max(block_size_0, block_size_1, block_size_2)
+ num_warps = max(1, block_size // WARP_SIZE)
+ num_blocks = min(per_rank_max_tokens, MAX_NUM_BLOCKS)
+
+ _multimem_all_gatherv_3tensor_kernel[(num_blocks, 1, 1)](
+ input_tensor_0.data_ptr(),
+ symm_mem_hdl_0.multicast_ptr,
+ output_byte_offset_0,
+ input_tensor_1.data_ptr(),
+ symm_mem_hdl_1.multicast_ptr,
+ output_byte_offset_1,
+ input_tensor_2.data_ptr(),
+ symm_mem_hdl_2.multicast_ptr,
+ output_byte_offset_2,
+ symm_mem_hdl_0.signal_pad_ptrs_dev,
+ local_tokens=local_tokens,
+ rank_token_offset_ptr=rank_token_offset,
+ ep_max_tokens_ptr=ep_max_tokens,
+ HIDDEN_SIZE_0=hidden_size_0,
+ HIDDEN_SIZE_1=hidden_size_1,
+ HIDDEN_SIZE_2=hidden_size_2,
+ BLOCK_SIZE=block_size,
+ NUMEL_PER_THREAD_0=numel_per_thread_0,
+ NUMEL_PER_THREAD_1=numel_per_thread_1,
+ NUMEL_PER_THREAD_2=numel_per_thread_2,
+ BITS_0=bits_0,
+ BITS_1=bits_1,
+ BITS_2=bits_2,
+ RANK=symm_mem_hdl_0.rank,
+ WORLD_SIZE=symm_mem_hdl_0.world_size,
+ num_warps=num_warps,
+ )
+
+ return output_tensor_0, output_tensor_1, output_tensor_2
diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py
index 4063ffbc977..e8769f3d6e7 100644
--- a/megatron/core/inference/config.py
+++ b/megatron/core/inference/config.py
@@ -1,8 +1,8 @@
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-from dataclasses import dataclass
+from dataclasses import InitVar, dataclass
from enum import Enum
-from typing import List, Optional, Tuple
+from typing import List, Literal, Optional, Tuple
import torch
@@ -24,7 +24,7 @@ class MambaInferenceStateConfig:
layer_type_list: List[str]
"""
A list of strings that indicates the layer type (Mamba / Attention / MLP) for each layer.
- See `megatron/core/ssm/mamba_hybrid_layer_allocation.py` for the list of symbols.
+ See `megatron/core/models/hybrid/hybrid_layer_allocation.py` for the list of symbols.
"""
conv_states_shape: Tuple[int]
@@ -50,7 +50,7 @@ def from_model(
ssm_states_dtype: Optional[torch.dtype] = None,
) -> Optional["MambaInferenceStateConfig"]:
"""Returns Mamba inference state config from the model if it is a hybrid model."""
- from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols
+ from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols
decoder = get_attr_wrapped_model(model, "decoder")
layer_type_list = getattr(decoder, "layer_type_list", None)
@@ -188,8 +188,12 @@ class InferenceConfig:
# =================================
num_cuda_graphs: Optional[int] = None
"""
- Maximum number of cuda graphs to capture, where the cuda graph batch sizes range from 1 to
- `max_requests`. Due to rounding, the actual number of cuda graphs may not equal this argument.
+ Maximum number of cuda graphs to capture.
+ Graph token counts are spaced from 1 up to a per-graph-type budget:
+ - Decode-only graphs are always bounded by `max_requests * (num_speculative_tokens + 1)`.
+ - Prefill/mixed graphs share that same bound by default,
+ or extend up to `max_tokens` when `cuda_graph_all_prefills` is set.
+ Due to rounding, the actual number of cuda graphs may not equal this argument.
"""
cuda_graph_mixed_prefill_count: Optional[int] = 16
@@ -202,6 +206,14 @@ class InferenceConfig:
Whether to use CUDA graphs for non-decode steps.
"""
+ cuda_graph_all_prefills: bool = False
+ """
+ Whether prefill/mixed CUDA graphs should span up to `max_tokens`.
+ When False (default), prefill/mixed graphs are bounded by the same token limit as decode graphs:
+ `max_requests * (num_speculative_tokens + 1)`.
+ When True, prefill/mixed graph capture is extended to cover the full `max_tokens` budget.
+ """
+
static_kv_memory_pointers: bool = False
"""
Whether the KV cache (and Mamba states) will reside at the same memory addresses
@@ -297,10 +309,13 @@ class InferenceConfig:
Defaults to 0, which means no logging.
"""
- request_metadata_types: Optional[List[Tuple[str, torch.dtype, bool]]] = None
+ sampling_backend: Literal['torch', 'flashinfer'] = 'torch'
+ """Which sampling kernels to use during inference."""
+
+ request_metadata_types: Optional[List[Tuple[str, torch.dtype]]] = None
"""
A list of the per-request metadata types to track. Each entry is a tuple
- consisting of the string label, the target dtype, and whether to store the data on GPU.
+ consisting of the string label and the target dtype.
"""
use_synchronous_zmq_collectives: bool = False
@@ -309,9 +324,42 @@ class InferenceConfig:
performance variability for MoEs.
"""
- def __post_init__(self):
+ disable_ep_consensus: bool = False
+ """If True, the engine skips the EP-group consensus all-reduce in
+ `run_engine_with_coordinator` and decides whether to step based on local
+ state alone. The rank still calls `controller.dummy_forward()` whenever
+ `local_pending == 0`, so EP collectives (NCCL all-to-all, etc.) stay in
+ sync — without this, a peer running a real forward would deadlock waiting
+ on this rank's all-to-all participation. Trades off the consensus
+ all-reduce CPU cost for unconditional dummy_forwards on idle ranks.
+ """
+
+ ep_consensus_interval: int = 20
+ """How many steps to skip between EP-consensus all-reduces when the engine
+ has pending work. Consensus is always run immediately when there is no
+ global work (to detect new arrivals quickly); this interval only applies
+ to the busy case, where skipping avoids per-step all-reduce overhead.
+ In the worst case, pausing is delayed by this many steps (~10–20 ms per
+ step at typical decode throughput).
+ """
+
+ verbose: InitVar[bool] = False
+ """Whether to log detailed context configuration at initialization.
+ This is an InitVar and is not stored as a field on the config."""
+
+ def __post_init__(self, verbose: bool):
+ self._verbose = verbose
if not (0.0 <= self.prefix_caching_routing_alpha <= 1.0):
raise ValueError(
f"prefix_caching_routing_alpha must be in [0, 1], "
f"got {self.prefix_caching_routing_alpha}"
)
+
+ if self.sampling_backend == 'flashinfer':
+ try:
+ import flashinfer # noqa: F401
+ except ImportError as e:
+ raise ImportError(
+ "sampling_backend='flashinfer' requires the flashinfer package; "
+ "install it or set sampling_backend='torch'."
+ ) from e
diff --git a/megatron/core/inference/contexts/attention_context/mamba_metadata.py b/megatron/core/inference/contexts/attention_context/mamba_metadata.py
index 19091d35bfb..3e98f0324e6 100644
--- a/megatron/core/inference/contexts/attention_context/mamba_metadata.py
+++ b/megatron/core/inference/contexts/attention_context/mamba_metadata.py
@@ -35,14 +35,15 @@ def __init__(
# Maximum possible chunks across all batch configurations
self.max_chunks = max_tokens // mamba_chunk_size + max_requests
- # Map from requests to slots in the static Mamba state buffer
+ # Map from requests to slots in the static Mamba state buffer (CPU for bookkeeping).
self.request_to_mamba_state_idx = torch.full(
- (self.max_requests,), -1, dtype=torch.int32, device=torch.cuda.current_device()
+ (self.max_requests,), -1, dtype=torch.int32, device='cpu'
)
- # Map from requests to slots in the static Mamba state buffer for active decode requests
+ # Map from requests to slots in the static Mamba state buffer for active decode requests.
+ # int64 so selective_state_update can index directly without a per-layer upcast kernel;
self._batch_indices_decode_buffer = torch.full(
- (self.max_requests,), -1, dtype=torch.int32, device=self.device
+ (self.max_requests,), -1, dtype=torch.int64, device=self.device
)
# Map from requests to slots in the static Mamba state buffer for active prefill requests
@@ -84,9 +85,9 @@ def __init__(
self._conv_seq_idx_buffer = torch.zeros(max_tokens, dtype=torch.int32, device=self.device)
self._conv_seq_start_buffer = torch.zeros(max_tokens, dtype=torch.int32, device=self.device)
- # Allocator for Mamba state slots
+ # Allocator for Mamba state slots (CPU for bookkeeping).
self.mamba_state_free_slots = torch.arange(
- self.max_requests, dtype=torch.int32, device=torch.cuda.current_device()
+ self.max_requests, dtype=torch.int32, device='cpu'
)
self.mamba_state_free_slot_count = self.max_requests
@@ -107,8 +108,31 @@ def __init__(
else:
self.conv_gather_offsets = None
+ # Coalesced production path: pinned CPU views + shared GPU views bound
+ # by DynamicInferenceContext so that the per-step Mamba metadata fields
+ # ride along with the single coalesced H2D in transfer_bookkeeping_to_gpu.
+ # The legacy update() path above keeps using the standalone _*_buffer
+ # tensors (exercised only by unit tests that construct MambaMetadata
+ # without a context).
+ self._cpu_bufs = None
+ self._gpu_view = None
+
self.reset_varlen_metadata()
+ def bind_cpu_buffers(self, bufs: dict) -> None:
+ """Attach pinned CPU views from DynamicInferenceContext._cpu_bookkeeping_buf.
+
+ ``bufs`` maps field names to 1D (or (1, max_tokens) for ``seq_idx``)
+ pinned CPU views that compute_cpu_metadata writes into. The matching
+ GPU views on the other side of the H2D are exposed via
+ :meth:`bind_gpu_buffers`.
+ """
+ self._cpu_bufs = bufs
+
+ def bind_gpu_buffers(self, gpu_view) -> None:
+ """Attach shared GPU views from the context's :class:`ContextGPUView`."""
+ self._gpu_view = gpu_view
+
def reset(self) -> None:
"""
Resets all Mamba states and frees all allocated slots.
@@ -119,7 +143,7 @@ def reset(self) -> None:
# 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.max_requests, dtype=torch.int32, device='cpu'
)
self.mamba_state_free_slot_count = self.max_requests
@@ -324,7 +348,10 @@ def update(
# This converts per-request token offsets to chunk indices and
# absolute positions, padded to fixed size for CUDA graph compat.
self._update_intermediate_metadata(
- intermediate_offsets_gpu, intermediate_counts_gpu, real_prefill_count
+ intermediate_offsets_gpu,
+ intermediate_counts_gpu,
+ real_prefill_count,
+ padded_prefill_count,
)
if padded_decode_count > 0 and padded_prefill_count > 0:
@@ -339,6 +366,8 @@ def _update_intermediate_metadata(
intermediate_offsets_gpu: Optional[torch.Tensor],
intermediate_counts_gpu: Optional[torch.Tensor],
real_prefill_count: int,
+ padded_prefill_count: int,
+ cu_seqlens_gpu: Optional[torch.Tensor] = None,
) -> None:
"""Precompute intermediate extraction metadata for CUDA graph compatibility.
@@ -352,18 +381,32 @@ def _update_intermediate_metadata(
intermediate_counts_gpu: [real_prefill_count] int32 GPU tensor of
per-request offset counts (0-3), or None.
real_prefill_count: Number of real (non-padding) prefill requests.
+ cu_seqlens_gpu: GPU cu_seqlens tensor to read from. Defaults to
+ the legacy standalone ``_cu_seqlens_buffer`` used by
+ :meth:`update`; the coalesced production path passes the
+ shared ``ContextGPUView.mamba_cu_seqlens`` view.
"""
chunk_size = self.mamba_chunk_size
- max_count = self.max_intermediate_count
+ max_count = padded_prefill_count * MAX_INTERMEDIATE_OFFSETS_PER_REQUEST
+ if cu_seqlens_gpu is None:
+ cu_seqlens_gpu = self._cu_seqlens_buffer
if intermediate_offsets_gpu is not None and real_prefill_count > 0:
- # Transfer counts to CPU (single sync) for per_request_counts and total check
+ # counts_list is CPU-cheap (source is already CPU from MambaSlotAllocator).
counts_list = intermediate_counts_gpu.tolist()
total = sum(counts_list)
+ # Ensure GPU copies for vectorized GPU ops below.
+ if not intermediate_offsets_gpu.is_cuda:
+ intermediate_offsets_gpu = intermediate_offsets_gpu.to(
+ self.device, non_blocking=True
+ )
+ if not intermediate_counts_gpu.is_cuda:
+ intermediate_counts_gpu = intermediate_counts_gpu.to(self.device, non_blocking=True)
+
if total > 0:
# Compute cumulative chunk counts from cu_seqlens (already on GPU)
- cu = self._cu_seqlens_buffer[: real_prefill_count + 1]
+ cu = cu_seqlens_gpu[: real_prefill_count + 1]
seq_lens = (cu[1 : real_prefill_count + 1] - cu[:real_prefill_count]).to(
torch.int64
)
@@ -405,15 +448,15 @@ def _update_intermediate_metadata(
# - abs_positions=d_conv: conv gather reads tokens [0..d_conv-1],
# which are within bounds and produce a valid but unused state
if real_count < max_count:
- self._intermediate_chunk_indices_buffer[real_count:].fill_(0)
- self._intermediate_abs_positions_buffer[real_count:].fill_(self.d_conv)
+ self._intermediate_chunk_indices_buffer[real_count:max_count].fill_(0)
+ self._intermediate_abs_positions_buffer[real_count:max_count].fill_(self.d_conv)
self.intermediate_count = real_count
self.per_request_intermediate_counts = counts_list
else:
# All counts are 0
- self._intermediate_chunk_indices_buffer.fill_(0)
- self._intermediate_abs_positions_buffer.fill_(self.d_conv)
+ self._intermediate_chunk_indices_buffer[:max_count] = 0
+ self._intermediate_abs_positions_buffer[:max_count] = self.d_conv
self.intermediate_count = 0
self.per_request_intermediate_counts = counts_list
@@ -422,13 +465,231 @@ def _update_intermediate_metadata(
else:
# No extraction: fill with safe defaults for CUDA graph warmup
# (same rationale as padding comment above)
- self._intermediate_chunk_indices_buffer.fill_(0)
- self._intermediate_abs_positions_buffer.fill_(self.d_conv)
+ self._intermediate_chunk_indices_buffer[:max_count] = 0
+ self._intermediate_abs_positions_buffer[:max_count] = self.d_conv
self.intermediate_count = 0
self.per_request_intermediate_counts = []
self.intermediate_chunk_indices = self._intermediate_chunk_indices_buffer[:max_count]
self.intermediate_abs_positions = self._intermediate_abs_positions_buffer[:max_count]
+ def compute_cpu_metadata(
+ self,
+ active_mamba_indices: torch.Tensor,
+ token_to_request_idx: torch.Tensor,
+ cpu_cu_query: torch.Tensor,
+ batch_dimensions: InferenceBatchDimensions,
+ padded_batch_dimensions: InferenceBatchDimensions,
+ enable_chunked_prefill: bool,
+ intermediate_offsets_gpu: Optional[torch.Tensor] = None,
+ intermediate_counts_gpu: Optional[torch.Tensor] = None,
+ ) -> dict:
+ """Compute all Mamba metadata on CPU, writing directly into the bound
+ pinned CPU views.
+
+ The values written here are transferred to GPU by the single coalesced
+ H2D in :meth:`DynamicInferenceContext.transfer_bookkeeping_to_gpu`.
+ The returned dict contains only Python scalars + the intermediate GPU
+ tensors, which :meth:`load_from_cpu` consumes after the H2D.
+
+ Args:
+ active_mamba_indices: CPU tensor of Mamba slot indices for active requests.
+ token_to_request_idx: CPU tensor mapping tokens to request indices.
+ cpu_cu_query: CPU cumulative query lengths from MHA metadata computation.
+ batch_dimensions: Dimensions of the current batch.
+ padded_batch_dimensions: Dimensions of the padded batch.
+ enable_chunked_prefill: Whether chunked prefill is enabled.
+ intermediate_offsets_gpu: GPU tensor of per-request intermediate offsets, or None.
+ intermediate_counts_gpu: GPU tensor of per-request intermediate counts, or None.
+ """
+ assert self._cpu_bufs is not None, "bind_cpu_buffers() must be called first"
+ bufs = self._cpu_bufs
+
+ real_decode_count = batch_dimensions.decode_req_count
+ real_prefill_count = batch_dimensions.prefill_req_count
+ padded_decode_count = padded_batch_dimensions.decode_req_count
+ padded_prefill_count = padded_batch_dimensions.prefill_req_count
+ padded_token_count = padded_batch_dimensions.token_count
+ chunk_size = self.mamba_chunk_size
+
+ result = {
+ "padded_decode_count": padded_decode_count,
+ "padded_prefill_count": padded_prefill_count,
+ "padded_token_count": padded_token_count,
+ "real_decode_count": real_decode_count,
+ "real_prefill_count": real_prefill_count,
+ }
+
+ # Decode batch indices (write into pinned view; padded slots = -1).
+ if padded_decode_count > 0:
+ bufs['batch_indices_decode'][:real_decode_count] = active_mamba_indices[
+ :real_decode_count
+ ]
+ if padded_decode_count > real_decode_count:
+ bufs['batch_indices_decode'][real_decode_count:padded_decode_count] = -1
+
+ # Prefill batch indices, seq_idx, cu_seqlens, chunk/conv metadata.
+ if padded_prefill_count > 0:
+ if real_prefill_count > 0:
+ start = real_decode_count
+ bufs['batch_indices_prefill'][:real_prefill_count] = active_mamba_indices[
+ start : start + real_prefill_count
+ ]
+ if padded_prefill_count > real_prefill_count:
+ bufs['batch_indices_prefill'][real_prefill_count:padded_prefill_count] = -1
+
+ # seq_idx: normalized token-to-request mapping for prefill tokens.
+ prefill_start_req = real_decode_count
+ end_prefill_req = real_decode_count + real_prefill_count
+ start_token = cpu_cu_query[prefill_start_req].item()
+ end_token = cpu_cu_query[end_prefill_req].item()
+ seq_len = end_token - start_token
+
+ if seq_len > 0:
+ raw = token_to_request_idx[start_token:end_token]
+ bufs['seq_idx'][0, :seq_len] = raw - raw[0]
+ if padded_token_count > seq_len:
+ bufs['seq_idx'][0, seq_len:padded_token_count] = -1
+ result["seq_len"] = seq_len
+
+ # cu_seqlens for prefill.
+ cu_seqlens_view = bufs['cu_seqlens']
+ cu_seqlens_view[0] = 0
+ if real_prefill_count > 0:
+ cu_seqlens_view[1 : real_prefill_count + 1] = (
+ cpu_cu_query[prefill_start_req + 1 : end_prefill_req + 1]
+ - cpu_cu_query[prefill_start_req]
+ )
+ if real_prefill_count < padded_prefill_count:
+ last_val = cu_seqlens_view[real_prefill_count].item()
+ cu_seqlens_view[real_prefill_count + 1 : padded_prefill_count + 1] = last_val
+
+ cu_seqlens_list = cu_seqlens_view[: real_prefill_count + 1].tolist()
+ real_prefill_tokens = (
+ cu_seqlens_list[real_prefill_count] if real_prefill_count > 0 else 0
+ )
+ result["cu_seqlens_list"] = cu_seqlens_list
+ result["real_prefill_token_count"] = real_prefill_tokens
+
+ # Chunk metadata (Python loop, pure CPU).
+ cu_seqlens_all = cu_seqlens_view[: padded_prefill_count + 1].tolist()
+ chunk_boundaries = [0]
+ last_chunk_idx_list = []
+ chunk_to_seq_list = []
+
+ for i in range(padded_prefill_count):
+ start = cu_seqlens_all[i]
+ end = cu_seqlens_all[i + 1]
+ s_len = end - start
+ n_chunks = max(1, (s_len + chunk_size - 1) // chunk_size)
+ boundaries = [min(start + (k + 1) * chunk_size, end) for k in range(n_chunks)]
+ chunk_boundaries.extend(boundaries)
+ chunk_to_seq_list.extend([i] * n_chunks)
+ last_chunk_idx_list.append(len(chunk_boundaries) - 2)
+
+ padded_max_chunks = padded_token_count // chunk_size + padded_prefill_count
+ last_boundary = chunk_boundaries[-1]
+ pad_b = padded_max_chunks + 1 - len(chunk_boundaries)
+ if pad_b > 0:
+ chunk_boundaries.extend([last_boundary] * pad_b)
+ pad_s = padded_max_chunks - len(chunk_to_seq_list)
+ if pad_s > 0:
+ chunk_to_seq_list.extend([0] * pad_s)
+
+ n_cu = padded_max_chunks + 1
+ bufs['cu_chunk_seqlens'][:n_cu] = torch.tensor(
+ chunk_boundaries[:n_cu], dtype=torch.int32
+ )
+ bufs['last_chunk_indices'][:padded_prefill_count] = torch.tensor(
+ last_chunk_idx_list, dtype=torch.int32
+ )
+ bufs['seq_idx_for_varlen'][:padded_max_chunks] = torch.tensor(
+ chunk_to_seq_list[:padded_max_chunks], dtype=torch.int32
+ )
+ result["padded_max_chunks"] = padded_max_chunks
+
+ # Conv1d per-token metadata (CPU repeat_interleave).
+ conv_seq_idx_view = bufs['conv_seq_idx']
+ conv_seq_start_view = bufs['conv_seq_start']
+ if real_prefill_tokens > 0:
+ cu_t = cu_seqlens_view[: real_prefill_count + 1]
+ lengths = (cu_t[1:] - cu_t[:-1]).to(torch.int64)
+ seq_indices = torch.arange(real_prefill_count, dtype=torch.int32)
+ seq_starts = cu_t[:real_prefill_count].to(torch.int32)
+ conv_seq_idx_view[:real_prefill_tokens] = torch.repeat_interleave(
+ seq_indices, lengths
+ )
+ conv_seq_start_view[:real_prefill_tokens] = torch.repeat_interleave(
+ seq_starts, lengths
+ )
+ if padded_token_count > real_prefill_tokens:
+ conv_seq_idx_view[real_prefill_tokens:padded_token_count] = 0
+ conv_seq_start_view[real_prefill_tokens:padded_token_count] = 0
+
+ # Intermediate metadata still requires GPU data: defer to load_from_cpu.
+ result["intermediate_offsets_gpu"] = intermediate_offsets_gpu
+ result["intermediate_counts_gpu"] = intermediate_counts_gpu
+
+ # device_decode_prefill scalars.
+ if padded_decode_count > 0 and padded_prefill_count > 0:
+ result["decode_prefill_0"] = cpu_cu_query[real_decode_count].item()
+ result["decode_prefill_1"] = (
+ cpu_cu_query[real_decode_count + real_prefill_count].item()
+ - cpu_cu_query[real_decode_count].item()
+ )
+
+ return result
+
+ def load_from_cpu(self, d: dict) -> None:
+ """Point state attributes at the freshly-transferred shared GPU views.
+
+ No H2D copies happen here: the Mamba metadata fields were transferred
+ as part of the coalesced bookkeeping H2D. This method just slices the
+ bound GPU views to the per-step sizes and runs the intermediate
+ metadata computation (which reads from the now-valid GPU cu_seqlens).
+
+ Args:
+ d: Dict returned by compute_cpu_metadata().
+ """
+ assert self._gpu_view is not None, "bind_gpu_buffers() must be called first"
+ v = self._gpu_view
+
+ padded_decode_count = d["padded_decode_count"]
+ padded_prefill_count = d["padded_prefill_count"]
+ padded_token_count = d["padded_token_count"]
+ real_prefill_count = d["real_prefill_count"]
+
+ if padded_decode_count > 0:
+ self.batch_indices_decode = v.mamba_batch_indices_decode[:padded_decode_count]
+
+ if padded_prefill_count > 0:
+ self.batch_indices_prefill = v.mamba_batch_indices_prefill[:padded_prefill_count]
+ self.seq_idx = v.mamba_seq_idx[:, :padded_token_count]
+ self.cu_seqlens = v.mamba_cu_seqlens[: padded_prefill_count + 1]
+ self.cu_seqlens_list = d["cu_seqlens_list"]
+ self.real_prefill_token_count = d["real_prefill_token_count"]
+
+ padded_max_chunks = d["padded_max_chunks"]
+ self.cu_chunk_seqlens = v.mamba_cu_chunk_seqlens[: padded_max_chunks + 1]
+ self.last_chunk_indices = v.mamba_last_chunk_indices[:padded_prefill_count]
+ self.seq_idx_for_varlen = v.mamba_seq_idx_for_varlen[:padded_max_chunks]
+ self.conv_seq_idx = v.mamba_conv_seq_idx[:padded_token_count]
+ self.conv_seq_start = v.mamba_conv_seq_start[:padded_token_count]
+
+ # Intermediate metadata reads from the just-transferred cu_seqlens
+ # to compute chunk indices & absolute positions for state extraction.
+ self._update_intermediate_metadata(
+ d["intermediate_offsets_gpu"],
+ d["intermediate_counts_gpu"],
+ real_prefill_count,
+ padded_prefill_count,
+ cu_seqlens_gpu=v.mamba_cu_seqlens,
+ )
+
+ if padded_decode_count > 0 and padded_prefill_count > 0:
+ self._device_decode_prefill_buffer[0] = d["decode_prefill_0"]
+ self._device_decode_prefill_buffer[1] = d["decode_prefill_1"]
+ self.device_decode_prefill = self._device_decode_prefill_buffer
+
def allocate_slot(self) -> Optional[int]:
"""
Allocates a new slot for a request in the Mamba state buffers.
diff --git a/megatron/core/inference/contexts/attention_context/mha_metadata.py b/megatron/core/inference/contexts/attention_context/mha_metadata.py
index 07f8a349b51..a71da895ea5 100644
--- a/megatron/core/inference/contexts/attention_context/mha_metadata.py
+++ b/megatron/core/inference/contexts/attention_context/mha_metadata.py
@@ -1,215 +1,84 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
import torch
-from megatron.core.inference.batch_dimensions_utils import InferenceBatchDimensions
-
from .metadata_base import MetadataBase
class MHAMetadata(MetadataBase):
"""
Metadata for MHA layer using flash-attention.
+
+ GPU storage for the per-step fields (``query_lengths``,
+ ``cu_query_seq_lengths``, ``kv_seq_lengths``, ``cu_kv_seq_lengths``,
+ ``block_table``) lives inside the context's :class:`ContextGPUView`
+ unified buffer. Both :class:`GraphedMHAMetadata` and
+ :class:`NonGraphedMHAMetadata` bind to the same GPU views (only one is
+ active per step), so the single coalesced H2D in
+ :meth:`DynamicInferenceContext.transfer_bookkeeping_to_gpu` covers the
+ MHA fields along with the rest of the bookkeeping state.
"""
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.device = torch.cuda.current_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 = {}
+ # Set by bind_gpu_buffers(); references shared views in ContextGPUView._buf.
+ self._gpu_view = None
- def update(
- self,
- request_query_lengths: torch.Tensor,
- request_kv_length_offsets: torch.Tensor,
- request_to_kv_block_ids: torch.Tensor,
- batch_dimensions: InferenceBatchDimensions,
- padded_batch_dimensions: InferenceBatchDimensions,
- num_speculative_tokens: int = 0,
- ):
- """
- 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)
- batch_dimensions: Configuration object containing real batch settings
- padded_batch_dimensions: Configuration object containing padded batch settings
- num_speculative_tokens: Number of speculative tokens
- """
- # Extract values from configs
- real_batch_size = batch_dimensions.req_count
- padded_active_token_count = padded_batch_dimensions.token_count
- padded_active_request_count = padded_batch_dimensions.req_count
-
- 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
+ def bind_gpu_buffers(self, gpu_view) -> None:
+ """Attach shared GPU buffer views from the context's ContextGPUView.
- 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 padded_batch_dimensions.prefill_req_count == 0:
- self._max_seqlen_q = num_speculative_tokens + 1
- else:
- # Make sure we will launch the prefill kernel for prefill graphs
- self._max_seqlen_q = max(2, padded_batch_dimensions.token_count)
+ Called by :class:`DynamicInferenceContext` after ``self.gpu_view`` is
+ constructed. Both graphed and non-graphed MHA metadata bind to the
+ same views; only one is active per step, so sharing storage is safe.
+ """
+ self._gpu_view = gpu_view
- self._max_seqlen_k = self.max_seqlen
+ def set_state_data(
+ self, padded_active_request_count: int, max_seqlen_q: int, max_seqlen_k: int
+ ) -> None:
+ """Build ``state_data`` slices into the bound GPU buffers.
+ Called once per step from ``transfer_bookkeeping_to_gpu`` after the
+ coalesced H2D copy. No ``.copy_()`` calls, no kernel launches.
+ """
+ assert self._gpu_view is not None, "bind_gpu_buffers() must be called first"
+ n = padded_active_request_count
+ v = self._gpu_view
+ self._max_seqlen_q = max_seqlen_q
+ self._max_seqlen_k = max_seqlen_k
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,
+ "query_lengths": v.mha_query_lengths[:n],
+ "cu_query_seq_lengths": v.mha_cu_query_seq_lengths[: n + 1],
+ "cu_kv_seq_lengths": v.mha_cu_kv_seq_lengths[: n + 1],
+ "kv_seq_lengths": v.mha_kv_seq_lengths[:n],
+ "block_table": v.mha_block_table[:n, :],
+ "max_seqlen_q": max_seqlen_q,
+ "max_seqlen_k": max_seqlen_k,
}
def reset(self):
+ """Reset the metadata for the next batch.
+
+ The GPU buffers live in the context's unified buffer and are fully
+ overwritten by the next H2D copy; clearing them here would launch
+ redundant CUDA kernels with no correctness benefit.
"""
- 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,
- batch_dimensions: InferenceBatchDimensions,
- padded_batch_dimensions: InferenceBatchDimensions,
- num_speculative_tokens: int = 0,
- ):
- """
- 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)
- batch_dimensions: Configuration object containing real batch settings
- padded_batch_dimensions: Configuration object containing padded batch settings
- num_speculative_tokens: Number of speculative tokens
- """
- super().update(
- request_query_lengths,
- request_kv_length_offsets,
- request_to_kv_block_ids,
- batch_dimensions,
- padded_batch_dimensions,
- num_speculative_tokens,
- )
-
- def reset(self):
- super().reset()
+ """MHA metadata for CUDA-graphed execution."""
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,
- batch_dimensions: InferenceBatchDimensions,
- padded_batch_dimensions: InferenceBatchDimensions,
- num_speculative_tokens: int = 0,
- ):
- """
- 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)
- batch_dimensions: Configuration object containing real batch settings
- padded_batch_dimensions: Configuration object containing padded batch settings
- num_speculative_tokens: Number of speculative tokens
- """
- super().update(
- request_query_lengths,
- request_kv_length_offsets,
- request_to_kv_block_ids,
- batch_dimensions,
- padded_batch_dimensions,
- num_speculative_tokens,
- )
- 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"] = num_speculative_tokens + 1
- self.state_data["max_seqlen_k"] = 1
+ """MHA metadata for non-graphed (eager) execution."""
diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
index 1117f2b9c4b..4a0d0cba518 100644
--- a/megatron/core/inference/contexts/dynamic_context.py
+++ b/megatron/core/inference/contexts/dynamic_context.py
@@ -2,9 +2,10 @@
import logging
import math
+import operator
import warnings
from contextlib import nullcontext
-from typing import List, Optional, Sequence, Tuple
+from typing import Dict, List, Optional, Sequence, Tuple
import torch # type: ignore
import torch.nn.functional as F # type: ignore
@@ -28,16 +29,24 @@
)
from megatron.core.inference.utils import device_memory_summary, tensor_swap
from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb
+from megatron.core.models.hybrid.hybrid_layer_allocation import (
+ Symbols,
+ get_layer_maps_from_layer_type_list,
+)
from megatron.core.package_info import __version__ as mcore_version
-from megatron.core.ssm.mamba_hybrid_layer_allocation import get_layer_maps_from_layer_type_list
from megatron.core.transformer import MLATransformerConfig, TransformerConfig
+from megatron.core.transformer.moe.token_dispatcher_inference import (
+ NCCLAllGatherDispatcher,
+ NVLSAllGatherVDispatcher,
+)
from megatron.core.utils import deprecate_args
from megatron.core.utils import divide as core_divide
-from megatron.core.utils import get_pg_size, internal_api
+from megatron.core.utils import get_pg_rank, get_pg_size, internal_api
from .attention_context.mamba_metadata import MambaMetadata
from .attention_context.mha_metadata import GraphedMHAMetadata, NonGraphedMHAMetadata
from .base_context import BaseInferenceContext
+from .gpu_view import ContextGPUView
from .kv_block_allocator import KVBlockAllocator
from .mamba_slot_allocator import MambaSlotAllocator
from .routing_metadata import RoutingMetadata
@@ -205,6 +214,8 @@ def deserialize(cls, obj: dict) -> ContextOverflowError:
def get_mem_size_str(n_bytes: int) -> str:
"""Convert number of bytes to human-readable string."""
+ if n_bytes == 0:
+ return "0 bytes"
for exp, suffix in ((4, "TB"), (3, "GB"), (2, "MB"), (3, "KB"), (0, "bytes")):
nquery = int(1024**exp)
if round(n_bytes / nquery) >= 1:
@@ -317,6 +328,12 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
else:
self.expert_model_parallel_group = None
+ # Optional CPU-side collective for EP batch-dimension sync. Populated by
+ # the engine via set_ep_zmq_communicator() when available. When set,
+ # match_graph_config() uses this to perform the MAX reduction on the
+ # CPU, avoiding a per-step NCCL AllReduce kernel on the compute stream.
+ self._ep_zmq_communicator = None
+
# Mamba states.
mamba_inference_state_config = inference_config.mamba_inference_state_config
self.is_hybrid_model = mamba_inference_state_config is not None
@@ -330,19 +347,39 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
# For hybrid models, the layer map converts the global layer index to the
# corresponding attention layer index or Mamba layer index depending on the
# layer type.
- mamba_layer_map, gdn_layer_map, attention_layer_map, _, _ = (
- get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list)
+ attention_layer_map, dsa_layer_map, gdn_layer_map, mamba_layer_map = (
+ operator.itemgetter(
+ Symbols.ATTENTION, Symbols.DS_ATTENTION, Symbols.GDN, Symbols.MAMBA
+ )(get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list))
)
if len(gdn_layer_map) > 0:
raise NotImplementedError("GDN layers are not supported for inference.")
- self.num_attention_layers = len(attention_layer_map)
+ self.num_attention_layers = len(attention_layer_map) + len(dsa_layer_map)
self.num_mamba_layers = len(mamba_layer_map)
- self.layer_map = attention_layer_map | mamba_layer_map
+ self.layer_map = attention_layer_map | dsa_layer_map | mamba_layer_map
else:
# The layer map is the identity function for pure Transformer models.
- self.num_attention_layers = model_config.num_layers // pp_size
+ # Use the same per-PP-rank layer count as TransformerBlock (handles
+ # account_for_embedding_in_pipeline_split, account_for_loss_in_pipeline_split,
+ # uneven first/last PP stages, and pipeline_model_parallel_layout). Using
+ # num_layers // pp_size mis-sizes the KV layer_map and can raise KeyError in
+ # append_key_value_cache.
+ from megatron.core.transformer.transformer_block import get_num_layers_to_build
+
+ # Interleaved / virtual PP is not used for inference (see
+ # AbstractModelInferenceWrapper: Iterable models are rejected). Always pass
+ # vp_stage=None into get_num_layers_to_build, consistent with attention inference
+ # (e.g. get_transformer_layer_offset(..., vp_stage=None, pp_rank=...)).
+ # When pg_collection is set, use the PP group's rank (same as attention.py).
+ if pg_collection is not None:
+ pp_rank = get_pg_rank(pg_collection.pp)
+ else:
+ pp_rank = None
+ self.num_attention_layers = get_num_layers_to_build(
+ model_config, vp_stage=None, pp_rank=pp_rank
+ )
self.num_mamba_layers = 0
(self.mamba_conv_states_shape, self.mamba_ssm_states_shape) = (None, None)
self.layer_map = {i: i for i in range(self.num_attention_layers)}
@@ -449,6 +486,26 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
buffer_size_bytes = int(buffer_size_bytes * (1.0 - mamba_memory_ratio))
paused_buffer_size_bytes = int(paused_buffer_size_bytes * (1.0 - mamba_memory_ratio))
+ block_count = buffer_size_bytes // self.block_size_bytes
+ block_count = max(2, block_count) # need >= 1 active block + 1 dummy block
+ paused_block_count = paused_buffer_size_bytes // self.block_size_bytes
+ elif self.is_hybrid_model and inference_config.max_requests is not None:
+ # Auto-derive mamba/KV split from max_requests. Allocate exactly enough
+ # mamba memory for max_requests, and give the rest to KV cache blocks.
+ total_memory = buffer_size_bytes + paused_buffer_size_bytes
+ mamba_memory_needed = inference_config.max_requests * mamba_states_memory_per_request
+ assert mamba_memory_needed < total_memory, (
+ f"Not enough memory for {inference_config.max_requests} mamba requests. "
+ f"Need {mamba_memory_needed / 1024**3:.2f} GB for mamba states, "
+ f"but total buffer is {total_memory / 1024**3:.2f} GB."
+ )
+ mamba_max_requests = inference_config.max_requests
+
+ # Subtract mamba memory proportionally from active and paused buffers.
+ mamba_memory_ratio = mamba_memory_needed / total_memory
+ buffer_size_bytes = int(buffer_size_bytes * (1.0 - mamba_memory_ratio))
+ paused_buffer_size_bytes = int(paused_buffer_size_bytes * (1.0 - mamba_memory_ratio))
+
block_count = buffer_size_bytes // self.block_size_bytes
block_count = max(2, block_count) # need >= 1 active block + 1 dummy block
paused_block_count = paused_buffer_size_bytes // self.block_size_bytes
@@ -499,8 +556,12 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
self.params_dtype = model_config.params_dtype
self.max_sequence_length = inference_config.max_sequence_length
- # Block ids.
+ # Block ids. With speculative decoding, blocks are pre-allocated when the
+ # last block offset >= block_size - 1 - num_speculative_tokens, so we may
+ # need one extra block beyond what max_sequence_length alone requires.
self.max_kv_block_count = math.ceil(self.max_sequence_length / self.block_size_tokens)
+ if self.num_speculative_tokens > 0:
+ self.max_kv_block_count += 1
# Set max_requests, max_tokens.
if inference_config.max_requests is None:
@@ -557,15 +618,44 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
), "Router recording/replay requested but no MoE experts specified!"
self.moe_routing_metadata = RoutingMetadata(self, model_config.moe_router_topk)
- # CUDA graph config list
+ # are we using the inference_optimized nccl ep dispatcher for MoEs?
+ self._nccl_ep_dispatcher = (
+ get_pg_size(self.expert_model_parallel_group) > 1
+ and model_config.inference_moe_token_dispatcher_type == 'nccl'
+ )
+
+ # are we using the training a2a dispatcher for MoEs?
+ # Note that this is not optimal for speed.
+ self._training_ep_dispatcher = (
+ get_pg_size(self.expert_model_parallel_group) > 1
+ and model_config.transformer_impl == "transformer_engine"
+ )
+
+ # We only allow non-decode cuda graphs for the nvls dispatcher
+ force_disable_non_decode_cuda_graphs = (
+ self._nccl_ep_dispatcher or self._training_ep_dispatcher
+ )
+
self.use_cuda_graphs_for_non_decode_steps = (
inference_config.use_cuda_graphs_for_non_decode_steps
+ and not (force_disable_non_decode_cuda_graphs)
)
+
+ # CUDA graph token budget for prefill/mixed graphs. Decode graphs are always
+ # capped at max_requests * (num_speculative_tokens + 1) inside the helper; this
+ # only widens the prefill/mixed range when `cuda_graph_all_prefills` is set.
+ cuda_graph_max_tokens = (
+ self.max_tokens
+ if inference_config.cuda_graph_all_prefills
+ else self.max_requests * (self.num_speculative_tokens + 1)
+ )
+
+ # CUDA graph config list.
self.cuda_graph_batch_dimensions_list, self.cuda_graph_token_counts = (
CUDAGraphBatchDimensionBuilder.generate_cuda_graph_batch_dimensions_list(
tp_size=tp_size,
num_cuda_graphs=inference_config.num_cuda_graphs,
- cuda_graph_max_tokens=self.max_requests * (self.num_speculative_tokens + 1),
+ cuda_graph_max_tokens=cuda_graph_max_tokens,
cuda_graph_mixed_prefill_request_count=inference_config.cuda_graph_mixed_prefill_count,
max_requests=self.max_requests,
max_tokens=self.max_tokens,
@@ -575,9 +665,21 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
)
)
- self.smallest_non_decode_cuda_graph_size = min(
- inference_config.cuda_graph_mixed_prefill_count, self.max_requests
- )
+ # Allocate per-step dispatcher buffers upfront so update_metadata never
+ # triggers an allocation inside a captured CUDA graph.
+ if get_pg_size(self.expert_model_parallel_group) > 1:
+ if self._nccl_ep_dispatcher:
+ NCCLAllGatherDispatcher.allocate_buffers()
+ else:
+ # Use moe_latent_size if set (latent MoE: SuperV3, UltraV3), else hidden_size.
+ moe_hidden_size = model_config.moe_latent_size or model_config.hidden_size
+ NVLSAllGatherVDispatcher.allocate_buffers(
+ per_rank_worst_case_token_count=self.round_up_tokens(self.max_tokens)
+ // tp_size,
+ topk=model_config.moe_router_topk,
+ hidden_size=moe_hidden_size,
+ ep_group=self.expert_model_parallel_group,
+ )
# Deal with chunked prefill
self.enable_chunked_prefill = inference_config.enable_chunked_prefill
@@ -588,19 +690,83 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
elif inference_config.use_flashinfer_fused_rope is None:
inference_config.use_flashinfer_fused_rope = HAVE_FLASHINFER
self.use_flashinfer_fused_rope = inference_config.use_flashinfer_fused_rope
+ self.inference_grouped_gemm_backend = model_config.inference_grouped_gemm_backend
# Allocate GPU state.
self.is_tensor_state_allocated = False
self.initialize_all_tensors()
# Print info.
- logging.info(
- "DynamicInferenceContext: allocated context with active buffer size %s (%d blocks)."
- % (
- get_mem_size_str(self.kv_block_allocator.active_count * self.block_size_bytes),
- self.kv_block_allocator.active_count,
+ active_blocks = self.kv_block_allocator.active_count
+ total_blocks = self.kv_block_allocator.total_count
+ paused_blocks = self.kv_block_allocator.paused_count
+ active_kv_bytes = active_blocks * self.block_size_bytes
+ total_kv_bytes = total_blocks * self.block_size_bytes
+ paused_kv_bytes = paused_blocks * self.block_size_bytes
+
+ log_lines = [
+ "DynamicInferenceContext: configuration summary",
+ f" max_requests: {self.max_requests}",
+ f" max_tokens: {self.max_tokens}",
+ f" max_sequence_length: {self.max_sequence_length}",
+ f" block_size_tokens: {self.block_size_tokens}",
+ f" max_kv_blocks_per_req: {self.max_kv_block_count}",
+ f" KV cache:",
+ f" block_size_bytes: {get_mem_size_str(self.block_size_bytes)}",
+ f" active_blocks: {active_blocks} ({get_mem_size_str(active_kv_bytes)})",
+ f" paused_blocks: {paused_blocks} ({get_mem_size_str(paused_kv_bytes)})",
+ f" total_blocks: {total_blocks} ({get_mem_size_str(total_kv_bytes)})",
+ ]
+
+ if self.is_hybrid_model:
+ mamba_conv_bytes = (
+ math.prod(self.mamba_conv_states_shape)
+ * self.mamba_conv_states_dtype.itemsize
+ * self.num_mamba_layers
)
- )
+ mamba_ssm_bytes = (
+ math.prod(self.mamba_ssm_states_shape)
+ * self.mamba_ssm_states_dtype.itemsize
+ * self.num_mamba_layers
+ )
+ mamba_bytes_per_req = mamba_conv_bytes + mamba_ssm_bytes
+ mamba_total_bytes = mamba_bytes_per_req * self.max_requests
+ log_lines += [
+ f" Mamba states:",
+ f" num_mamba_layers: {self.num_mamba_layers}",
+ f" conv_state_shape: {self.mamba_conv_states_shape}",
+ f" ssm_state_shape: {self.mamba_ssm_states_shape}",
+ f" per_request: {get_mem_size_str(mamba_bytes_per_req)}",
+ f" total ({self.max_requests} requests): {get_mem_size_str(mamba_total_bytes)}",
+ ]
+
+ if self.num_speculative_tokens > 0:
+ spec_multiplier = self.num_speculative_tokens + 1
+ spec_bytes_per_req = mamba_bytes_per_req * spec_multiplier
+ spec_total_bytes = spec_bytes_per_req * self.max_requests
+ log_lines += [
+ f" Mamba speculative buffers (num_speculative_tokens={self.num_speculative_tokens}):",
+ f" per_request: {get_mem_size_str(spec_bytes_per_req)}",
+ f" total ({self.max_requests} requests): {get_mem_size_str(spec_total_bytes)}",
+ ]
+
+ prefix_caching_mamba_gb = inference_config.prefix_caching_mamba_gb
+ if (
+ inference_config.enable_prefix_caching
+ and prefix_caching_mamba_gb is not None
+ and prefix_caching_mamba_gb > 0
+ ):
+ prefix_cache_bytes = int(prefix_caching_mamba_gb * 1024**3)
+ prefix_cache_slots = prefix_cache_bytes // mamba_bytes_per_req
+ log_lines += [
+ f" Mamba prefix cache:",
+ f" budget: {get_mem_size_str(prefix_cache_bytes)}",
+ f" slots: {prefix_cache_slots}",
+ f" per_slot: {get_mem_size_str(mamba_bytes_per_req)}",
+ ]
+
+ if inference_config._verbose and torch.distributed.get_rank() == 0:
+ logging.info("\n".join(log_lines))
def _allocate_memory_buffer(self):
"""Allocate the KV cache memory buffer."""
@@ -644,8 +810,26 @@ def _allocate_mamba_states(self):
self.mamba_metadata = MambaMetadata(
max_requests=self.max_requests,
max_tokens=self.max_tokens,
+ mamba_chunk_size=self.mamba_chunk_size,
d_conv=self.mamba_conv_states_shape[-1],
)
+ # Bind the unified CPU/GPU buffers so the per-step Mamba metadata
+ # fields ride along with the single coalesced H2D in
+ # transfer_bookkeeping_to_gpu().
+ self.mamba_metadata.bind_cpu_buffers(
+ {
+ "batch_indices_decode": self._cpu_mamba_batch_indices_decode,
+ "batch_indices_prefill": self._cpu_mamba_batch_indices_prefill,
+ "seq_idx": self._cpu_mamba_seq_idx,
+ "cu_seqlens": self._cpu_mamba_cu_seqlens,
+ "cu_chunk_seqlens": self._cpu_mamba_cu_chunk_seqlens,
+ "last_chunk_indices": self._cpu_mamba_last_chunk_indices,
+ "seq_idx_for_varlen": self._cpu_mamba_seq_idx_for_varlen,
+ "conv_seq_idx": self._cpu_mamba_conv_seq_idx,
+ "conv_seq_start": self._cpu_mamba_conv_seq_start,
+ }
+ )
+ self.mamba_metadata.bind_gpu_buffers(self.gpu_view)
self.mamba_conv_states = torch.empty(
(self.num_mamba_layers, self.max_requests) + self.mamba_conv_states_shape,
dtype=self.mamba_conv_states_dtype,
@@ -721,58 +905,326 @@ def initialize_all_tensors(self) -> None:
f"Please move tensor '{key}'."
)
- # Per-request state.
+ # Per-request state (CPU, pinned memory for fast H2D transfer).
self.request_ids = torch.full(
- (self.max_requests,), -1, dtype=torch.int32, device=torch.cuda.current_device()
+ (self.max_requests,), -1, dtype=torch.int32, device='cpu', pin_memory=True
)
# request_query_lengths is the input prompt tokens length during prefill phase (1st step) and then 1 for the decode phase (i.e During generation)
- self.request_query_lengths = torch.empty_like(self.request_ids)
+ self.request_query_lengths = torch.empty(
+ self.max_requests, dtype=torch.int32, device='cpu', pin_memory=True
+ )
# True only for a new request , then after a forward pass it is set to False
- self.request_in_prefill_status_tensor = torch.empty_like(self.request_ids)
+ self.request_in_prefill_status_tensor = torch.empty(
+ self.max_requests, dtype=torch.int32, device='cpu', pin_memory=True
+ )
# request_output_lengths is len(input_prompt_tokens) + num_tokens_to_generate
- self.request_output_lengths = torch.empty_like(self.request_ids)
+ self.request_output_lengths = torch.empty(
+ self.max_requests, dtype=torch.int32, device='cpu', pin_memory=True
+ )
# request_kv_length_offsets is the same as query length during prefill phase (1st step) and then 1 for the decode phase (i.e During generation)
- self.request_kv_length_offsets = torch.empty_like(self.request_ids)
- self.request_kv_block_counts = torch.empty_like(self.request_ids)
- self.request_last_kv_block_id = torch.empty_like(self.request_ids)
+ self.request_kv_length_offsets = torch.empty(
+ self.max_requests, dtype=torch.int32, device='cpu', pin_memory=True
+ )
+ self.request_kv_block_counts = torch.empty(
+ self.max_requests, dtype=torch.int32, device='cpu', pin_memory=True
+ )
+ self.request_last_kv_block_id = torch.empty(
+ self.max_requests, dtype=torch.int32, device='cpu', pin_memory=True
+ )
# request_last_kv_block_offset represents number of tokens in the last kv block
- self.request_last_kv_block_offset = torch.empty_like(self.request_ids)
+ self.request_last_kv_block_offset = torch.empty(
+ self.max_requests, dtype=torch.int32, device='cpu', pin_memory=True
+ )
self.request_to_kv_block_ids = torch.full(
(self.max_requests, self.max_kv_block_count),
-1,
dtype=torch.int,
- device=torch.cuda.current_device(),
+ device='cpu',
+ pin_memory=True,
)
- # Track request metadata.
+ # Track request metadata. Backed by pinned CPU memory: bookkeeping is
+ # CPU-resident; GPU consumers read from the active-slice mirror in
+ # `active_request_metadata` (also CPU pinned, refreshed each step).
self.request_metadata = {
- label: torch.empty(
- (self.max_requests,), dtype=dtype, device=torch.cuda.current_device()
- )
- for label, dtype, _ in self.request_metadata_types
+ label: torch.empty((self.max_requests,), dtype=dtype, device='cpu', pin_memory=True)
+ for label, dtype in self.request_metadata_types
}
- # Per-token state.
- self.token_to_input_ids = torch.full(
- (self.max_tokens,), 0, dtype=torch.long, device=torch.cuda.current_device()
- )
- self.token_to_pos_ids = torch.full_like(self.token_to_input_ids, 0)
- self.token_to_request_idx = torch.empty_like(self.token_to_input_ids)
- self.token_to_block_idx = torch.empty_like(self.token_to_input_ids)
+ # Static tensor addresses of active slices to enable fast inference
+ # kernels. Pinned CPU mirrors of `request_metadata`, refreshed each
+ # step by `build_active_slices()` from the active subrange.
+ self.active_request_metadata = {
+ label: torch.empty_like(tensor, pin_memory=True)
+ for label, tensor in self.request_metadata.items()
+ }
+
+ # Coalesced pinned CPU buffer for the bookkeeping fields that get
+ # transferred to GPU each step via transfer_bookkeeping_to_gpu().
+ # Layout matches ContextGPUView._buf so a single cudaMemcpyAsync
+ # suffices. Int64 token fields come first (8-byte aligned automatically),
+ # then int32 token fields, then int32/float32 request-staging fields.
+ # token_to_input_ids (int64, max_tokens)
+ # token_to_pos_ids (int64, max_tokens)
+ # token_to_block_idx (int32, max_tokens)
+ # token_to_local_position_within_kv_block (int32, max_tokens)
+ # token_to_request_idx (int32, max_tokens)
+ # token_to_position_in_request (int32, max_tokens)
+ # request_in_prefill_status (staging) (int32, max_requests)
+ # request_query_lengths (staging) (int32, max_requests)
+ # request_kv_length_offsets (staging) (int32, max_requests)
+ # temperature (staging) (float32, max_requests)
+ # top_k (staging) (int32, max_requests)
+ # top_p (staging) (float32, max_requests)
+ # active_request_last_token_idxs (alias) (int32, max_requests)
+ #
+ # Token fields are aliased with the source-of-truth attributes
+ # (`self.token_to_input_ids`, etc.) because the forward pass reads
+ # `gpu_view.token_to_input_ids[:n_tok]` which matches the CPU slot
+ # layout `[0, n_tok)`. Request fields, however, are read on GPU at
+ # `[:n_active]` but on CPU at `[paused_count:total_count)` — so the
+ # staging slots here are refreshed each step by copying the active
+ # slice from the persistent `request_*` tensors above.
+ _tok_int64_bytes = self.max_tokens * 8
+ _tok_int32_bytes = self.max_tokens * 4
+ # Request-level fields are all 4 bytes wide (5 int32 + 2 float32 = 7 fields).
+ _req_4byte_bytes = self.max_requests * 4
+ # MHA section: 5 fields (int32) shared between GraphedMHAMetadata and
+ # NonGraphedMHAMetadata. max_bs == max_requests.
+ _mha_query_lengths_bytes = self.max_requests * 4
+ _mha_cu_query_seq_lengths_bytes = (self.max_requests + 1) * 4
+ _mha_kv_seq_lengths_bytes = self.max_requests * 4
+ _mha_cu_kv_seq_lengths_bytes = (self.max_requests + 1) * 4
+ _mha_block_table_bytes = self.max_requests * self.max_kv_block_count * 4
+ # Mamba section: 9 int32 fields (hybrid models only). Must match the
+ # MambaMetadata shapes (mirrors the layout documented in ContextGPUView).
+ if self.is_hybrid_model:
+ self._max_mamba_chunks = self.max_tokens // self.mamba_chunk_size + self.max_requests
+ _mamba_batch_indices_decode_bytes = self.max_requests * 4
+ _mamba_batch_indices_prefill_bytes = self.max_requests * 4
+ _mamba_seq_idx_bytes = self.max_tokens * 4
+ _mamba_cu_seqlens_bytes = (self.max_requests + 1) * 4
+ _mamba_cu_chunk_seqlens_bytes = (self._max_mamba_chunks + 1) * 4
+ _mamba_last_chunk_indices_bytes = self.max_requests * 4
+ _mamba_seq_idx_for_varlen_bytes = self._max_mamba_chunks * 4
+ _mamba_conv_seq_idx_bytes = self.max_tokens * 4
+ _mamba_conv_seq_start_bytes = self.max_tokens * 4
+ else:
+ self._max_mamba_chunks = 0
+ _mamba_batch_indices_decode_bytes = 0
+ _mamba_batch_indices_prefill_bytes = 0
+ _mamba_seq_idx_bytes = 0
+ _mamba_cu_seqlens_bytes = 0
+ _mamba_cu_chunk_seqlens_bytes = 0
+ _mamba_last_chunk_indices_bytes = 0
+ _mamba_seq_idx_for_varlen_bytes = 0
+ _mamba_conv_seq_idx_bytes = 0
+ _mamba_conv_seq_start_bytes = 0
+ _total_bytes = (
+ 2 * _tok_int64_bytes
+ + 4 * _tok_int32_bytes
+ + 7 * _req_4byte_bytes
+ + _mha_query_lengths_bytes
+ + _mha_cu_query_seq_lengths_bytes
+ + _mha_kv_seq_lengths_bytes
+ + _mha_cu_kv_seq_lengths_bytes
+ + _mha_block_table_bytes
+ + _mamba_batch_indices_decode_bytes
+ + _mamba_batch_indices_prefill_bytes
+ + _mamba_seq_idx_bytes
+ + _mamba_cu_seqlens_bytes
+ + _mamba_cu_chunk_seqlens_bytes
+ + _mamba_last_chunk_indices_bytes
+ + _mamba_seq_idx_for_varlen_bytes
+ + _mamba_conv_seq_idx_bytes
+ + _mamba_conv_seq_start_bytes
+ )
+ self._cpu_bookkeeping_buf = torch.empty(
+ _total_bytes, dtype=torch.uint8, device='cpu', pin_memory=True
+ )
+ # token_to_input_ids and token_to_pos_ids were previously torch.full(0);
+ # zero the whole buffer so their views start at 0 too, and so the
+ # request staging slots start with a deterministic value.
+ self._cpu_bookkeeping_buf.fill_(0)
+
+ _off = 0
+ # Per-token state (source-of-truth lives in the coalesced buffer since
+ # the CPU-side bookkeeping and the GPU forward pass use the same
+ # `[:n_tok]` slice).
+ self.token_to_input_ids = self._cpu_bookkeeping_buf[_off : _off + _tok_int64_bytes].view(
+ torch.long
+ )
+ _off += _tok_int64_bytes
+ self.token_to_pos_ids = self._cpu_bookkeeping_buf[_off : _off + _tok_int64_bytes].view(
+ torch.long
+ )
+ _off += _tok_int64_bytes
+ self.token_to_block_idx = self._cpu_bookkeeping_buf[_off : _off + _tok_int32_bytes].view(
+ torch.int32
+ )
+ _off += _tok_int32_bytes
# i.e For a set of tokens A B C D E F .. and block_size 4:
# token_to_position_in_request is [0, 1, 2, 3, 4, 5]
# token_to_local_position_within_kv_block is [0 , 1, 2, 3, 0, 1, 2]
- self.token_to_position_in_request = torch.empty_like(self.token_to_input_ids)
- self.token_to_local_position_within_kv_block = torch.empty_like(self.token_to_input_ids)
-
- # NOTE: Need to build this outside the UVM / TMS context to avoid IMA.
+ self.token_to_local_position_within_kv_block = self._cpu_bookkeeping_buf[
+ _off : _off + _tok_int32_bytes
+ ].view(torch.int32)
+ _off += _tok_int32_bytes
+ self.token_to_request_idx = self._cpu_bookkeeping_buf[_off : _off + _tok_int32_bytes].view(
+ torch.int32
+ )
+ _off += _tok_int32_bytes
+ self.token_to_position_in_request = self._cpu_bookkeeping_buf[
+ _off : _off + _tok_int32_bytes
+ ].view(torch.int32)
+ _off += _tok_int32_bytes
+
+ # Request-level staging views into the coalesced buffer. Write-only on
+ # CPU (refreshed from persistent tensors in transfer_bookkeeping_to_gpu);
+ # read-only on GPU via matching slots in ContextGPUView._buf.
+ self._staging_request_in_prefill_status = self._cpu_bookkeeping_buf[
+ _off : _off + _req_4byte_bytes
+ ].view(torch.int32)
+ _off += _req_4byte_bytes
+ self._staging_request_query_lengths = self._cpu_bookkeeping_buf[
+ _off : _off + _req_4byte_bytes
+ ].view(torch.int32)
+ _off += _req_4byte_bytes
+ self._staging_request_kv_length_offsets = self._cpu_bookkeeping_buf[
+ _off : _off + _req_4byte_bytes
+ ].view(torch.int32)
+ _off += _req_4byte_bytes
+
+ # Sampling-parameter staging slots, refreshed from `active_request_metadata`
+ # in transfer_bookkeeping_to_gpu(). FlashInfer reads these via
+ # `gpu_view.{temperature, top_k, top_p}`.
+ self._staging_temperature = self._cpu_bookkeeping_buf[_off : _off + _req_4byte_bytes].view(
+ torch.float32
+ )
+ _off += _req_4byte_bytes
+ self._staging_top_k = self._cpu_bookkeeping_buf[_off : _off + _req_4byte_bytes].view(
+ torch.int32
+ )
+ _off += _req_4byte_bytes
+ self._staging_top_p = self._cpu_bookkeeping_buf[_off : _off + _req_4byte_bytes].view(
+ torch.float32
+ )
+ _off += _req_4byte_bytes
+
+ # Per-request last-token row indices. Aliased with the matching gpu_view slot:
+ # build_active_slices/pad_active_slices populate this CPU view.
+ self.active_request_last_token_idxs = self._cpu_bookkeeping_buf[
+ _off : _off + _req_4byte_bytes
+ ].view(torch.int32)
+ _off += _req_4byte_bytes
+
+ # Static tensor addresses to make `last_token_logits` graphable with speculative decoding.
+ max_logit_idxs = self.max_requests * (self.num_speculative_tokens + 1)
+ self.active_logit_idxs = torch.zeros(
+ max_logit_idxs, dtype=torch.int32, device=torch.cuda.current_device()
+ )
+ self._decode_logit_idxs = torch.arange(
+ max_logit_idxs, dtype=torch.int32, device=torch.cuda.current_device()
+ )
+
+ # MHA flash-attention metadata views (write-only on CPU, read-only on
+ # GPU via the matching region of ContextGPUView._buf). Populated per
+ # step by initialize_attention_state(); transferred as part of the
+ # single coalesced H2D in transfer_bookkeeping_to_gpu().
+ self._cpu_mha_query_lengths = self._cpu_bookkeeping_buf[
+ _off : _off + _mha_query_lengths_bytes
+ ].view(torch.int32)
+ _off += _mha_query_lengths_bytes
+ self._cpu_mha_cu_query_seq_lengths = self._cpu_bookkeeping_buf[
+ _off : _off + _mha_cu_query_seq_lengths_bytes
+ ].view(torch.int32)
+ _off += _mha_cu_query_seq_lengths_bytes
+ self._cpu_mha_kv_seq_lengths = self._cpu_bookkeeping_buf[
+ _off : _off + _mha_kv_seq_lengths_bytes
+ ].view(torch.int32)
+ _off += _mha_kv_seq_lengths_bytes
+ self._cpu_mha_cu_kv_seq_lengths = self._cpu_bookkeeping_buf[
+ _off : _off + _mha_cu_kv_seq_lengths_bytes
+ ].view(torch.int32)
+ _off += _mha_cu_kv_seq_lengths_bytes
+ self._cpu_mha_block_table = (
+ self._cpu_bookkeeping_buf[_off : _off + _mha_block_table_bytes]
+ .view(torch.int32)
+ .view(self.max_requests, self.max_kv_block_count)
+ )
+ _off += _mha_block_table_bytes
+
+ # Mamba varlen metadata views (hybrid models only). Populated per step
+ # by MambaMetadata.compute_cpu_metadata(); transferred as part of the
+ # single coalesced H2D in transfer_bookkeeping_to_gpu().
if self.is_hybrid_model:
- self.mamba_metadata = MambaMetadata(
- max_requests=self.max_requests,
- max_tokens=self.max_tokens,
- mamba_chunk_size=self.mamba_chunk_size,
- d_conv=self.mamba_conv_states_shape[-1],
+ self._cpu_mamba_batch_indices_decode = self._cpu_bookkeeping_buf[
+ _off : _off + _mamba_batch_indices_decode_bytes
+ ].view(torch.int32)
+ _off += _mamba_batch_indices_decode_bytes
+ self._cpu_mamba_batch_indices_prefill = self._cpu_bookkeeping_buf[
+ _off : _off + _mamba_batch_indices_prefill_bytes
+ ].view(torch.int32)
+ _off += _mamba_batch_indices_prefill_bytes
+ self._cpu_mamba_seq_idx = (
+ self._cpu_bookkeeping_buf[_off : _off + _mamba_seq_idx_bytes]
+ .view(torch.int32)
+ .view(1, self.max_tokens)
)
+ _off += _mamba_seq_idx_bytes
+ self._cpu_mamba_cu_seqlens = self._cpu_bookkeeping_buf[
+ _off : _off + _mamba_cu_seqlens_bytes
+ ].view(torch.int32)
+ _off += _mamba_cu_seqlens_bytes
+ self._cpu_mamba_cu_chunk_seqlens = self._cpu_bookkeeping_buf[
+ _off : _off + _mamba_cu_chunk_seqlens_bytes
+ ].view(torch.int32)
+ _off += _mamba_cu_chunk_seqlens_bytes
+ self._cpu_mamba_last_chunk_indices = self._cpu_bookkeeping_buf[
+ _off : _off + _mamba_last_chunk_indices_bytes
+ ].view(torch.int32)
+ _off += _mamba_last_chunk_indices_bytes
+ self._cpu_mamba_seq_idx_for_varlen = self._cpu_bookkeeping_buf[
+ _off : _off + _mamba_seq_idx_for_varlen_bytes
+ ].view(torch.int32)
+ _off += _mamba_seq_idx_for_varlen_bytes
+ self._cpu_mamba_conv_seq_idx = self._cpu_bookkeeping_buf[
+ _off : _off + _mamba_conv_seq_idx_bytes
+ ].view(torch.int32)
+ _off += _mamba_conv_seq_idx_bytes
+ self._cpu_mamba_conv_seq_start = self._cpu_bookkeeping_buf[
+ _off : _off + _mamba_conv_seq_start_bytes
+ ].view(torch.int32)
+ _off += _mamba_conv_seq_start_bytes
+
+ assert _off == _total_bytes, f"layout bug: wrote {_off} of {_total_bytes} bytes"
+
+ # GPU view: the single interface for GPU code to read context state.
+ # Populated per-step by transfer_bookkeeping_to_gpu().
+ self.gpu_view = ContextGPUView(
+ max_requests=self.max_requests,
+ max_tokens=self.max_tokens,
+ max_kv_blocks=self.max_kv_block_count,
+ device=torch.cuda.current_device(),
+ max_mamba_chunks=self._max_mamba_chunks,
+ )
+
+ # Cache of (input_ids_view, pos_ids_view) keyed by num_tokens. Instead of slicing and
+ # unsqueezing on every new inference step (constructing new TensorImpls at 30-60 us),
+ # we fix the underlying storage so views are reusable across steps. The number of entries
+ # is bounded by the graph sizes plus eager-mode token counts, which are rounded up to
+ # multiples of TOKEN_ROUNDER and capped at max_tokens / TOKEN_ROUNDER distinct values.
+ self._input_position_views: Dict[int, Tuple[Tensor, Tensor]] = {}
+
+ # Bind the shared MHA GPU views to both graph and non-graph metadata;
+ # only one is active per step, so sharing storage is safe.
+ self.graph_attn_metadata["mha_metadata"].bind_gpu_buffers(self.gpu_view)
+ self.non_graph_attn_metadata["mha_metadata"].bind_gpu_buffers(self.gpu_view)
+
+ # Deferred Mamba GPU operations. Populated by add_request() /
+ # update_requests() (CPU phase), executed by transfer_bookkeeping_to_gpu().
+ self._pending_mamba_zeros: list = []
+ self._pending_mamba_restores: list = []
# Allocate large non-graphed buffers.
need_static_addr = (
@@ -918,7 +1370,14 @@ def is_static_batching(self) -> bool:
def is_decode_only(self) -> bool:
"""
Return if this iteration we run decode only implementation.
+
+ When CUDA graphs are active, uses padded_batch_dimensions because it
+ reflects the post-expert-parallel sync state. Otherwise falls back to
+ num_prefill_requests which is always up-to-date regardless of where we
+ are in the step lifecycle.
"""
+ if self._using_cuda_graph_this_step:
+ return self.padded_batch_dimensions.prefill_req_count == 0
return self.num_prefill_requests == 0
def using_cuda_graph_this_step(self) -> bool:
@@ -960,6 +1419,61 @@ def get_active_request_count(self):
"""Returns the current number of active requests."""
return self.total_request_count - self.paused_request_count
+ def build_active_slices(self, batch_size: int):
+ """Build the active slices of specific tensors. This is run on every forward step.
+
+ If the context is reordered to active -> paused -> finished, this can be graphed.
+ """
+ padded_slice = slice(self.paused_request_count, self.paused_request_count + batch_size)
+
+ # Request metadata all needs to be sliced.
+ for label in self.request_metadata:
+ self.active_request_metadata[label][:batch_size].copy_(
+ self.request_metadata[label][padded_slice], non_blocking=True
+ )
+
+ torch.cumsum(
+ self.request_query_lengths[padded_slice],
+ dim=0,
+ out=self.active_request_last_token_idxs[:batch_size],
+ )
+ self.active_request_last_token_idxs[:batch_size].sub_(1)
+
+ def pad_active_slices(self):
+ """Pad the active slices of specific tensors."""
+ active_request_count = self.total_request_count - self.paused_request_count
+ active_decode_count = self.num_decode_requests
+ active_prefill_count = active_request_count - active_decode_count
+ active_decode_token_count = active_decode_count * (self.num_speculative_tokens + 1)
+
+ # Decode prefix: positions [0, 1, ..., active_decode_token_count - 1].
+ self.active_logit_idxs[:active_decode_token_count].copy_(
+ self._decode_logit_idxs[:active_decode_token_count]
+ )
+
+ # Prefill last-token positions: cumsum the prefill query lengths in place,
+ # then shift by (active_decode_token_count - 1) to get absolute positions.
+ prefill_dst = self.active_logit_idxs[
+ active_decode_token_count : active_decode_token_count + active_prefill_count
+ ]
+ prefill_idxs = self.paused_request_count + active_decode_count
+ prefill_lengths = self.request_query_lengths[prefill_idxs : self.total_request_count]
+ if active_prefill_count > 0:
+ prefill_cumsum = torch.cumsum(prefill_lengths, dim=0, dtype=torch.int32)
+ prefill_cumsum.add_(active_decode_token_count - 1)
+ prefill_dst.copy_(prefill_cumsum, non_blocking=True)
+
+ self.active_logit_idxs[active_decode_token_count + active_prefill_count :].zero_()
+
+ padding_request_slice = slice(active_request_count, self.padded_active_request_count)
+
+ # Sampling metadata: pad with neutral defaults, so that the kernel early-exits.
+ self.active_request_metadata["temperature"][padding_request_slice].fill_(1.0)
+ self.active_request_metadata["top_k"][padding_request_slice].fill_(0)
+ self.active_request_metadata["top_p"][padding_request_slice].fill_(0.0)
+ # Padded gather indices fan in to row 0 harmlessly when used by FlashInfer.
+ self.active_request_last_token_idxs[padding_request_slice].fill_(0)
+
def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None:
"""Append to KV cache.
@@ -978,12 +1492,12 @@ def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor)
value=value,
memory_buffer=self.memory_buffer,
padded_active_token_count=self.padded_active_token_count,
- token_to_block_idx=self.token_to_block_idx,
- token_to_local_position_within_kv_block=self.token_to_local_position_within_kv_block,
+ token_to_block_idx=self.gpu_view.token_to_block_idx,
+ token_to_local_position_within_kv_block=self.gpu_view.token_to_local_position_within_kv_block,
)
- block_idx = self.token_to_block_idx[: self.padded_active_token_count]
- local_kv_seq_idx = self.token_to_local_position_within_kv_block[
+ block_idx = self.gpu_view.token_to_block_idx[: self.padded_active_token_count]
+ local_kv_seq_idx = self.gpu_view.token_to_local_position_within_kv_block[
: self.padded_active_token_count
]
@@ -1119,7 +1633,7 @@ def apply_fused_qk_rotary_emb(
# use .view instead of .reshape to avoid extra transpose operations
query_rope, key_rope = flashinfer.rope.apply_rope_with_cos_sin_cache(
- positions=self.token_to_pos_ids[:n],
+ positions=self.gpu_view.token_to_pos_ids[:n],
query=query[:n].reshape(n, num_q_heads * head_size),
key=key[:n].reshape(n, num_k_heads * head_size),
head_size=head_size,
@@ -1152,7 +1666,7 @@ def apply_rotary_emb_query(
(Tensor) Query tensor after applying rotary embeddings.
"""
n = self.padded_active_token_count
- query_seq_idx = self.token_to_pos_ids[:n]
+ query_seq_idx = self.gpu_view.token_to_pos_ids[:n]
query_emb = query_emb[query_seq_idx]
query[:n] = apply_rotary_pos_emb(
t=query[:n],
@@ -1184,7 +1698,7 @@ def apply_rotary_emb_key(
(Tensor) Key tensor after applying rotary embeddings.
"""
n = self.padded_active_token_count
- key_seq_idx = self.token_to_position_in_request[:n]
+ key_seq_idx = self.gpu_view.token_to_position_in_request[:n]
key_emb = key_emb[key_seq_idx]
if self.is_decode_only():
if key.shape[0] != n:
@@ -1203,6 +1717,20 @@ def apply_rotary_emb_key(
)
return key
+ def set_ep_zmq_communicator(self, communicator) -> None:
+ """Attach an EP-group ZMQ communicator for CPU-side sync collectives.
+
+ When set, match_graph_config() uses this communicator's
+ sync_all_reduce_max() to perform the EP batch-dimension MAX reduction on
+ the CPU instead of launching a NCCL AllReduce kernel on the compute
+ stream. Expected to be called once by the inference engine after both
+ the context and the communicator have been created.
+
+ Args:
+ communicator: AsyncZMQCommunicator over the EP process group.
+ """
+ self._ep_zmq_communicator = communicator
+
def reset_attention_state(self) -> None:
"""Reset state used within attention, after each step."""
# Attention metadata reset is now handled by MHAMetadata.reset()
@@ -1289,9 +1817,9 @@ def add_dummy_requests_parallel(
self.request_output_lengths[request_slice] = lengths_tensor + tokens_to_generate_tensor
self.request_kv_length_offsets[request_slice] = 0
self.request_kv_block_counts[request_slice] = block_counts
- for i, (label, dtype, _) in enumerate(self.request_metadata_types):
+ for i, (label, dtype) in enumerate(self.request_metadata_types):
self.request_metadata[label][request_slice] = torch.tensor(
- metadata_cols[i], dtype=dtype, device=torch.cuda.current_device()
+ metadata_cols[i], dtype=dtype, device='cpu'
)
dummy_block_idx = self.kv_block_allocator.dummy_block_idx
@@ -1350,8 +1878,7 @@ def add_dummy_requests_parallel(
raise ContextOverflowError(
requests[logical_idx].request_id, "No Mamba slots available"
)
- self.mamba_conv_states[:, mamba_idx] = 0.0
- self.mamba_ssm_states[:, mamba_idx] = 0.0
+ self._pending_mamba_zeros.append(mamba_idx)
self.mamba_metadata.request_to_mamba_state_idx[request_idx] = mamba_idx
self.active_token_count = token_end
@@ -1373,7 +1900,7 @@ def add_dummy_requests_for_cudagraph_capture(
# Pre-construct shared objects (safe due to deep copy in DynamicInferenceRequest.__post_init__)
shared_sampling_params = SamplingParams(num_tokens_to_generate=1, termination_id=-1)
shared_decode_tokens = torch.zeros(
- self.num_speculative_tokens + 1, dtype=torch.long, device=torch.cuda.current_device()
+ self.num_speculative_tokens + 1, dtype=torch.long, device='cpu'
)
decode_requests = [
@@ -1403,9 +1930,7 @@ def add_dummy_requests_for_cudagraph_capture(
assert per_prefill_tokens > 0
# Create a single large tensor and slice from it for each prefill request
max_prefill_tokens = per_prefill_tokens + (1 if rem_prefill_tokens > 0 else 0)
- shared_prefill_tokens = torch.zeros(
- max_prefill_tokens, dtype=torch.long, device=torch.cuda.current_device()
- )
+ shared_prefill_tokens = torch.zeros(max_prefill_tokens, dtype=torch.long, device='cpu')
prefill_requests = [
DynamicInferenceRequest(
@@ -1425,43 +1950,56 @@ def num_decode_requests(self) -> int:
"""
return self.total_request_count - self.paused_request_count - self.num_prefill_requests
- def add_dummy_requests_for_expert_parallel_step(self) -> None:
+ def add_dummy_requests_for_expert_parallel_step(
+ self, graph_dimensions: InferenceBatchDimensions
+ ) -> None:
"""Minimal context setup so an EP rank with no real requests can replay
an already-captured cuda graph without crashing or corrupting memory.
This is the fast alternative to add_dummy_requests_for_cudagraph_capture
(which goes through the heavyweight add_dummy_requests_parallel path).
- We setup minimal state such the initialize_attention_state and the forward
+ We setup minimal state such that initialize_attention_state and the forward
pass can run without error.
+ Called AFTER the EP sync so graph_dimensions reflects the agreed-upon graph.
"""
- smallest_cuda_graph_dimensions = min(
- [x for x in self.cuda_graph_batch_dimensions_list if x.prefill_req_count == 0]
- )
- # the smallest cuda graph is decode only.
- assert smallest_cuda_graph_dimensions.prefill_req_count == 0
-
- N = smallest_cuda_graph_dimensions.decode_req_count
- tokens_per_request = self.num_speculative_tokens + 1
- T = smallest_cuda_graph_dimensions.token_count # N * tokens_per_request
+ N_decode = graph_dimensions.decode_req_count
+ N_prefill = graph_dimensions.prefill_req_count
+ N = N_decode + N_prefill
+ tokens_per_decode_request = self.num_speculative_tokens + 1
+ T = graph_dimensions.token_count
dummy_block_idx = self.kv_block_allocator.dummy_block_idx
# 1. Request counts and token count.
- # With speculative decoding each decode request has (num_speculative_tokens + 1) tokens.
self.total_request_count = N
self.active_token_count = T
- self.num_prefill_requests = 0
+ self.num_prefill_requests = N_prefill
+
+ # 2. Per-request state consumed by initialize_attention_state().
+ # Decode requests come first, followed by prefill requests.
+ self.request_query_lengths[0:N_decode].fill_(tokens_per_decode_request)
+ if N_prefill > 0:
+ prefill_tokens = T - N_decode * tokens_per_decode_request
+ per_prefill_tokens = prefill_tokens // N_prefill
+ rem_prefill_tokens = prefill_tokens % N_prefill
+ self.request_query_lengths[N_decode:N].fill_(per_prefill_tokens)
+ if rem_prefill_tokens > 0:
+ self.request_query_lengths[N_decode : N_decode + rem_prefill_tokens] += 1
- # 2. Per-request state consumed by mha_metadata.update().
- self.request_query_lengths[0:N].fill_(tokens_per_request)
self.request_kv_length_offsets[0:N].fill_(0)
self.request_to_kv_block_ids[0:N, 0] = dummy_block_idx
# 3. Token-level state consumed by the triton KV append kernel.
self.token_to_block_idx[0:T] = dummy_block_idx
- self.token_to_local_position_within_kv_block[0:T] = (
- torch.arange(T, device=self.token_to_block_idx.device) % tokens_per_request
+ # Compute per-request token positions: e.g. query_lengths [3,2] -> [0,1,2,0,1]
+ query_lengths = self.request_query_lengths[0:N]
+ starts = torch.cumsum(query_lengths, dim=0) - query_lengths
+ # Per-token start offset: e.g. starts [0,3], query_lengths [3,2] -> [0,0,0,3,3]
+ per_token_start = torch.repeat_interleave(starts, query_lengths)
+ positions = torch.arange(T, device=query_lengths.device) - per_token_start
+ self.token_to_local_position_within_kv_block[0:T] = torch.remainder(
+ positions, self.block_size_tokens
)
if self.is_hybrid_model:
@@ -1473,7 +2011,7 @@ def add_dummy_requests_for_expert_parallel_step(self) -> None:
device=self.token_to_request_idx.device,
dtype=self.token_to_request_idx.dtype,
),
- tokens_per_request,
+ self.request_query_lengths[0:N],
)
# 5. Mamba state: allocate slots for dummy requests.
@@ -1497,16 +2035,27 @@ def initialize_attention_state(
Return:
None.
"""
+ # Launch deferred Mamba GPU ops first (state zeroing/restore) so they
+ # overlap with the CPU work below. These are non-blocking GPU kernels.
+ self._execute_pending_mamba_ops()
+
self.is_creating_cuda_graphs = construct_graph_dimensions is not None
assert not (
self.is_creating_cuda_graphs and is_expert_parallel_dummy_cuda_graph_step
), "Dummy expert model parallel steps should not be creating cuda graphs."
- # If in CUDA graph creation mode, add dummy requests for CUDA graph capture
- if is_expert_parallel_dummy_cuda_graph_step:
- self.add_dummy_requests_for_expert_parallel_step()
- elif self.is_creating_cuda_graphs:
+ # If in CUDA graph creation mode, add dummy requests for CUDA graph capture.
+ # EP dummy requests are added AFTER the EP sync below.
+ if self.is_creating_cuda_graphs:
self.add_dummy_requests_for_cudagraph_capture(construct_graph_dimensions)
+ elif is_expert_parallel_dummy_cuda_graph_step:
+ self.add_dummy_requests_for_expert_parallel_step(
+ InferenceBatchDimensions(
+ token_count=self.num_speculative_tokens + 1,
+ prefill_req_count=0,
+ decode_req_count=1,
+ )
+ )
batch_dimensions = InferenceBatchDimensions(
token_count=self.active_token_count,
@@ -1519,23 +2068,16 @@ def initialize_attention_state(
best_graph = CUDAGraphBatchDimensionBuilder.match_graph_config(
batch_dimensions,
self.cuda_graph_batch_dimensions_list,
- smallest_non_decode_cuda_graph_size=self.smallest_non_decode_cuda_graph_size,
strict=self.is_hybrid_model,
- decode_only_cuda_graphs=(not self.use_cuda_graphs_for_non_decode_steps),
ep_group=self.expert_model_parallel_group,
+ match_ep_token_counts=self._nccl_ep_dispatcher or self._training_ep_dispatcher,
+ ep_zmq_communicator=self._ep_zmq_communicator,
)
self._using_cuda_graph_this_step = best_graph is not None
if construct_graph_dimensions is not None:
assert self._using_cuda_graph_this_step
- if is_expert_parallel_dummy_cuda_graph_step and not self.using_cuda_graph_this_step():
- # If we are here, this means that CUDAGraphBatchDimensionBuilder.match_graph_config
- # could not find a compatible cuda graph for the dummy forward step.
- # Now, we need not do the remaining setup. The controller
- # will directly call the model forward pass with a single token.
- return
-
if self.using_cuda_graph_this_step():
self.padded_batch_dimensions = best_graph
else:
@@ -1570,6 +2112,11 @@ def initialize_attention_state(
self.padded_active_request_count = self.padded_batch_dimensions.req_count
self.padding_slice = slice(self.active_token_count, self.padded_active_token_count)
+ self.build_active_slices(
+ min(self.padded_active_request_count, self.max_requests - self.paused_request_count)
+ )
+ self.pad_active_slices()
+
# Update token position indexes.
self.token_to_block_idx[self.active_token_count : self.padded_active_token_count] = (
self.kv_block_allocator.dummy_block_idx
@@ -1607,31 +2154,98 @@ def initialize_attention_state(
)
assert self.active_attn_metadata is not None
- 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,
- batch_dimensions=attn_dimensions,
- padded_batch_dimensions=self.padded_batch_dimensions,
- num_speculative_tokens=self.num_speculative_tokens,
+
+ # Compute MHA metadata directly into the pinned CPU section of
+ # _cpu_bookkeeping_buf. The single coalesced H2D in
+ # transfer_bookkeeping_to_gpu() covers these fields along with the rest
+ # of the bookkeeping state, so no ephemeral tensors and no per-field
+ # cudaMemcpyAsyncs.
+ real_bs = attn_dimensions.req_count
+ padded_bs = self.padded_batch_dimensions.req_count
+ mha = self.active_attn_metadata["mha_metadata"]
+
+ # Query lengths: [0:real_bs] real data, [real_bs:padded_bs] zero pad.
+ self._cpu_mha_query_lengths[:real_bs] = query_lengths_view[:real_bs]
+ if real_bs < padded_bs:
+ self._cpu_mha_query_lengths[real_bs:padded_bs] = 0
+
+ # Cumulative query lengths (padded slots repeat cu[real_bs]).
+ self._cpu_mha_cu_query_seq_lengths[0] = 0
+ if real_bs > 0:
+ self._cpu_mha_cu_query_seq_lengths[1 : real_bs + 1] = torch.cumsum(
+ query_lengths_view[:real_bs], dim=0
+ )
+ if real_bs < padded_bs:
+ self._cpu_mha_cu_query_seq_lengths[real_bs + 1 : padded_bs + 1] = (
+ self._cpu_mha_cu_query_seq_lengths[real_bs]
+ )
+
+ # KV sequence lengths: [0:real_bs] = kv_offsets + query_lengths.
+ self._cpu_mha_kv_seq_lengths[:real_bs] = (
+ request_kv_length_offsets_view[:real_bs] + query_lengths_view[:real_bs]
+ )
+ if real_bs < padded_bs:
+ self._cpu_mha_kv_seq_lengths[real_bs:padded_bs] = 0
+
+ # Cumulative KV lengths.
+ self._cpu_mha_cu_kv_seq_lengths[0] = 0
+ if real_bs > 0:
+ self._cpu_mha_cu_kv_seq_lengths[1 : real_bs + 1] = torch.cumsum(
+ self._cpu_mha_kv_seq_lengths[:real_bs], dim=0
+ )
+ if real_bs < padded_bs:
+ self._cpu_mha_cu_kv_seq_lengths[real_bs + 1 : padded_bs + 1] = (
+ self._cpu_mha_cu_kv_seq_lengths[real_bs]
+ )
+
+ # Block table: [0:real_bs] real, [real_bs:padded_bs] = -1 sentinel.
+ self._cpu_mha_block_table[:real_bs] = request_to_kv_block_ids_view[:real_bs]
+ if real_bs < padded_bs:
+ self._cpu_mha_block_table[real_bs:padded_bs] = -1
+
+ # Max sequence lengths (Python scalars; consumed as kernel launch args).
+ if not self.using_cuda_graph_this_step() and real_bs > 0:
+ # NonGraphedMHAMetadata: use actual max values.
+ max_seqlen_q = self._cpu_mha_query_lengths[:real_bs].max().item()
+ max_seqlen_k = self._cpu_mha_kv_seq_lengths[:real_bs].max().item()
+ else:
+ # GraphedMHAMetadata: use conservative bounds.
+ if self.padded_batch_dimensions.prefill_req_count == 0:
+ max_seqlen_q = self.num_speculative_tokens + 1
+ else:
+ max_seqlen_q = max(2, self.padded_batch_dimensions.token_count)
+ max_seqlen_k = mha.max_seqlen
+ if not self.using_cuda_graph_this_step() and real_bs == 0:
+ max_seqlen_q = self.num_speculative_tokens + 1
+ max_seqlen_k = 1
+
+ # Bind state_data to GPU views now. set_state_data() only creates Python
+ # slice references into the GPU buffer (no GPU reads), so it's safe to
+ # call before the H2D in transfer_bookkeeping_to_gpu(). This guarantees
+ # that callers reading state_data["block_table"] etc. between
+ # initialize_attention_state() and transfer_bookkeeping_to_gpu() see
+ # populated entries (the actual data fill happens at the H2D).
+ mha.set_state_data(
+ padded_active_request_count=padded_bs,
+ max_seqlen_q=max_seqlen_q,
+ max_seqlen_k=max_seqlen_k,
)
if self.is_hybrid_model:
- active_mamba_indices_view = self.mamba_metadata.request_to_mamba_state_idx[active_slice]
- token_to_request_idx_view = self.token_to_request_idx[: self.active_token_count]
- cu_seqlens = self.active_attn_metadata["mha_metadata"].state_data[
- "cu_query_seq_lengths"
- ]
+ # Mamba metadata update is deferred to transfer_bookkeeping_to_gpu()
+ # because it writes to GPU buffers. Store the parameters here.
+ # intermediate_offsets_gpu / intermediate_counts_gpu get the CPU-side
+ # slices here; H2D transfer happens in transfer_bookkeeping_to_gpu().
intermediate_offsets_gpu = None
intermediate_counts_gpu = None
if self.mamba_slot_allocator is not None:
intermediate_offsets_gpu, intermediate_counts_gpu = (
- self.mamba_slot_allocator.get_intermediate_gpu_data()
+ self.mamba_slot_allocator.get_intermediate_cpu_data()
)
- self.mamba_metadata.update(
- active_mamba_indices_view,
- token_to_request_idx_view,
- cu_seqlens,
+ self._pending_mamba_transfer = self.mamba_metadata.compute_cpu_metadata(
+ active_mamba_indices=self.mamba_metadata.request_to_mamba_state_idx[active_slice],
+ token_to_request_idx=self.token_to_request_idx[: self.active_token_count],
+ cpu_cu_query=self._cpu_mha_cu_query_seq_lengths,
batch_dimensions=attn_dimensions,
padded_batch_dimensions=self.padded_batch_dimensions,
enable_chunked_prefill=self.is_chunked_prefill_enabled(),
@@ -1645,8 +2259,111 @@ def initialize_attention_state(
else:
self.moe_routing_metadata.disable_static_buffer_recording()
+ # Flip NCCLAllGather dispatcher's path selector to not use allgathers.
+ # _nccl_ep_dispatcher already implies ep_size > 1, so no extra EP guard.
+ if self._nccl_ep_dispatcher:
+ NCCLAllGatherDispatcher._use_allgather_v = not self.using_cuda_graph_this_step()
+
+ # Flush any Mamba ops queued by add_dummy_requests_for_cudagraph_capture
+ # (warmup) or add_dummy_requests_for_expert_parallel_step (EP dummy step).
+ # The earlier call at the top drained ops queued by add_request() before
+ # this function ran; this call covers ops queued during the function.
+ # No-op when the queue is already empty (regular non-warmup steps).
+ self._execute_pending_mamba_ops()
+
+ # Run the H2D transfer here so callers that bypass the controller
+ # (e.g. unit tests that call `model.forward()` directly after
+ # `initialize_attention_state()`) see populated GPU bookkeeping. The
+ # text-generation controller still calls `transfer_bookkeeping_to_gpu`
+ # explicitly; that second call is a cheap idempotent re-copy.
+ self.transfer_bookkeeping_to_gpu()
+
+ def _execute_pending_mamba_ops(self) -> None:
+ """Execute Mamba GPU operations deferred from add_request() / update_requests().
+
+ This runs at the start of initialize_attention_state() so that all GPU
+ Mamba state is correct before the forward pass.
+ """
+ if not (self._pending_mamba_restores or self._pending_mamba_zeros):
+ return
+
+ # Restore cached Mamba state to live buffers. On failure, fall back to zeroing.
+ for request_idx, block_id, mamba_idx in self._pending_mamba_restores:
+ restored = self.mamba_slot_allocator.restore_to_live(request_idx, block_id)
+ if not restored:
+ self._pending_mamba_zeros.append(mamba_idx)
+ self._pending_mamba_restores.clear()
+
+ # Batch-zero newly allocated Mamba slots.
+ if self._pending_mamba_zeros:
+ device = self.mamba_conv_states.device
+ indices = torch.tensor(self._pending_mamba_zeros, dtype=torch.long, device=device)
+ self.mamba_conv_states[:, indices] = 0.0
+ self.mamba_ssm_states[:, indices] = 0.0
+ self._pending_mamba_zeros.clear()
+
+ def transfer_bookkeeping_to_gpu(self) -> None:
+ """Batch transfer CPU bookkeeping state to GPU staging buffers.
+
+ Called after initialize_attention_state() and before the forward pass.
+ All copies use non_blocking=True with pinned CPU memory. CUDA stream
+ ordering guarantees the forward pass sees completed transfers.
+
+ The bookkeeping fields are backed by one contiguous pinned CPU buffer
+ and one contiguous GPU buffer; a single cudaMemcpyAsync suffices.
+ Request-level staging slots are refreshed from the persistent CPU
+ tensors immediately before the H2D (GPU reads them at `[:n_active]`
+ while CPU bookkeeping keeps them at `[paused_count:total_count)`).
+ """
+ n_active = self.total_request_count - self.paused_request_count
+ active_slice = slice(self.paused_request_count, self.total_request_count)
+ padded_active = max(n_active, self.padded_active_request_count)
+
+ # Refresh request-level staging slots from the persistent CPU source.
+ # CPU-to-CPU slice assignment on pinned memory (~15 KB total for 6
+ # 4-byte fields at max_requests=624). Negligible vs. the launch overhead
+ # we save by merging the H2D memcpys into 1.
+ self._staging_request_in_prefill_status[:n_active] = self.request_in_prefill_status_tensor[
+ active_slice
+ ]
+ self._staging_request_query_lengths[:n_active] = self.request_query_lengths[active_slice]
+ self._staging_request_kv_length_offsets[:n_active] = self.request_kv_length_offsets[
+ active_slice
+ ]
+ # Sampling-parameter staging slots: read from `active_request_metadata`,
+ # which `build_active_slices` + `pad_active_slices` already populated for
+ # `[:padded_active]` (active values + neutral padding defaults).
+ self._staging_temperature[:padded_active] = self.active_request_metadata["temperature"][
+ :padded_active
+ ]
+ self._staging_top_k[:padded_active] = self.active_request_metadata["top_k"][:padded_active]
+ self._staging_top_p[:padded_active] = self.active_request_metadata["top_p"][:padded_active]
+
+ # Full-iteration CUDA graphs may have captured GPU consumers with the
+ # padded graph request count. Keep those padded staging rows bounded so
+ # graph replay never builds indices from stale request lengths.
+ if n_active < padded_active:
+ self._staging_request_in_prefill_status[n_active:padded_active] = 0
+ self._staging_request_query_lengths[n_active:padded_active] = 0
+ self._staging_request_kv_length_offsets[n_active:padded_active] = 0
+
+ # Coalesced H2D: one cudaMemcpyAsync for the entire bookkeeping buffer.
+ # Copying the whole (max_tokens + max_requests)-sized buffer including
+ # unused slots is cheap (~71 KB total, ~3-5 us on PCIe Gen4) and saves
+ # 8 redundant launch overheads vs. the prior per-field copies.
+ self.gpu_view._buf.copy_(self._cpu_bookkeeping_buf, non_blocking=True)
+
+ # MHA metadata GPU views were already bound to state_data in
+ # initialize_attention_state(); the H2D above populates the underlying
+ # bytes. Nothing else to do here for MHA.
+
+ # Mamba metadata: copy pre-computed CPU tensors to GPU buffers.
+ if hasattr(self, '_pending_mamba_transfer') and self._pending_mamba_transfer is not None:
+ self.mamba_metadata.load_from_cpu(self._pending_mamba_transfer)
+ self._pending_mamba_transfer = None
+
def reset_tensors(self) -> None:
- """Fill all GPU tensors with sentinel values."""
+ """Fill all bookkeeping tensors with sentinel values."""
# Reset request indexes.
self.request_ids.fill_(-1)
@@ -1749,33 +2466,75 @@ def current_input_and_position_ids(
assert num_tokens >= self.padded_batch_dimensions.decode_req_count * (
self.num_speculative_tokens + 1
)
- return (
- self.token_to_input_ids[:num_tokens].unsqueeze(0),
- self.token_to_pos_ids[:num_tokens].unsqueeze(0),
- )
+ cached = self._input_position_views.get(num_tokens)
+ if cached is not None:
+ return cached
+ input_ids = self.gpu_view.token_to_input_ids[:num_tokens].unsqueeze(0)
+ pos_ids = self.gpu_view.token_to_pos_ids[:num_tokens].unsqueeze(0)
+ cached = (input_ids, pos_ids)
+ self._input_position_views[num_tokens] = cached
+ return cached
+
+ def speculative_required_logit_indices(self) -> Tensor:
+ """Token-level indices needed for speculative decode verification.
+
+ Returns all decode token positions (base + speculative) concatenated
+ with the last token position of each prefill request.
+
+ Return:
+ (Tensor) 1-D indices into the packed token sequence, length
+ ``num_decode_requests * (num_speculative_tokens + 1) + num_prefill_requests``
+ in eager, or the equivalent padded count under non-eager.
+ """
+ return self.active_logit_idxs[: self.num_last_token_logits]
+
+ @property
+ def num_last_token_logits(self) -> int:
+ """Number of rows produced by `last_token_logits` for the current step.
+
+ Single source of truth for the bound: one row per request, with
+ `(num_speculative_tokens + 1)` rows per decode request when MTP is active.
+ """
+ if self.num_speculative_tokens > 0:
+ if self._using_cuda_graph_this_step:
+ return (
+ self.padded_batch_dimensions.decode_req_count
+ * (self.num_speculative_tokens + 1)
+ + self.padded_batch_dimensions.prefill_req_count
+ )
+ else:
+ return (
+ self.num_decode_requests * (self.num_speculative_tokens + 1)
+ + self.num_prefill_requests
+ )
+ else:
+ if self._using_cuda_graph_this_step:
+ return self.padded_active_request_count
+ else:
+ return self.total_request_count - self.paused_request_count
def last_token_logits(self, logits: Tensor) -> Tensor:
- """Last tokens of logits.
+ """Select the logit positions needed for token generation.
+
+ When speculative decoding is active, decode requests need logits for all
+ their tokens (base + speculative) for verification, while prefill requests
+ only need the last token logit. This avoids materializing the full
+ vocab-sized logits for every prefill token, which causes large memory
+ spikes during prefill-heavy batches.
Args:
- logits (Tensor): Output logits of forward pass.
+ logits (Tensor): Output logits of forward pass, shape [1, S, H].
Return:
- (Tensor) Last token logits.
+ (Tensor) Selected logits, shape [N, H], where N == num_last_token_logits.
"""
- paused = self.paused_request_count
- total = self.total_request_count
- query_lengths = self.request_query_lengths[paused:total]
-
# todo: @lmcafee, remove these asserts?
assert logits.size(0) == 1, f"logits.size(0) ({tuple(logits.shape)}) != 1"
assert logits.size(1) == self.padded_active_token_count, (
f"logits.size(1) ({tuple(logits.shape)}) != "
f"padded_active_token_count ({self.padded_active_token_count})."
)
- logits_2d = logits.squeeze(0)
- last_token_idxs = torch.cumsum(query_lengths, dim=0) - 1
- return logits_2d[last_token_idxs, :]
+ return logits.squeeze(0)[self.active_logit_idxs[: self.num_last_token_logits], :]
def _compute_prefix_match(
self, req: DynamicInferenceRequest, prefill_chunk_length: int
@@ -1844,6 +2603,14 @@ def _compute_prefix_match(
elif self.is_hybrid_model and finished == 0:
prefix_skip_tokens = 0
+ # Clamp so that effective_prefill_chunk_length >= 2 when possible.
+ # A single-token prefill chunk (effective == 1) causes max_seqlen_q == 1,
+ # which routes the batch into the flash-attention decode kernel and crashes.
+ # Round down to a block boundary to keep block-table indexing consistent.
+ if prefill_chunk_length - prefix_skip_tokens < 2 and prefill_chunk_length >= 2:
+ max_skip = prefill_chunk_length - 2
+ prefix_skip_tokens = (max_skip // self.block_size_tokens) * self.block_size_tokens
+
effective_prefill_chunk_length = prefill_chunk_length - prefix_skip_tokens
num_blocks_from_pool = max(
0, overall_required_blocks - already_allocated_blocks - num_matched
@@ -1980,9 +2747,7 @@ def add_request(
# Increment ref counts and update timestamps for matched (shared) blocks
if num_matched_blocks > 0:
- matched_tensor = torch.tensor(
- matched_block_ids, dtype=torch.int32, device=torch.cuda.current_device()
- )
+ matched_tensor = torch.tensor(matched_block_ids, dtype=torch.int32, device='cpu')
self.kv_block_allocator.block_ref_counts[matched_tensor] += 1
if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU:
self.kv_block_allocator.update_timestamps(matched_tensor)
@@ -2007,7 +2772,7 @@ def add_request(
metadata = req.tracked_metadata
metadata_types = req.get_metadata_types()
for m, m_type in zip(metadata, metadata_types):
- label, _, _ = m_type
+ label, _ = m_type
if not isinstance(m, torch.Tensor):
m = torch.as_tensor(
m,
@@ -2106,17 +2871,18 @@ def _register_range(start: int, end: int):
# Restore Mamba state from the block corresponding to prefix_skip_tokens
restore_block_count = prefix_skip_tokens // self.block_size_tokens
- restored = False
if restore_block_count > 0 and self.mamba_slot_allocator is not None:
restore_block_id = matched_block_ids[restore_block_count - 1]
- restored = self.mamba_slot_allocator.restore_to_live(
- self.total_request_count, restore_block_id
+ self._pending_mamba_restores.append(
+ (self.total_request_count, restore_block_id, mamba_idx)
)
- if not restored:
- self.mamba_conv_states[:, mamba_idx] = 0.0
- self.mamba_ssm_states[:, mamba_idx] = 0.0
+ else:
+ self._pending_mamba_zeros.append(mamba_idx)
- # Compute intermediate offsets for state extraction during forward pass
+ # compute_and_store_offsets sets both CPU state (hash_to_block_id,
+ # _eos_cache_block_id_gpu) and GPU staging buffers. Runs immediately
+ # because commit_intermediate_states() reads the CPU state after the
+ # forward pass.
if self.mamba_slot_allocator is not None:
self.mamba_slot_allocator.compute_and_store_offsets(
req,
@@ -2240,13 +3006,13 @@ def release_memory_blocks_from_request_indexes(self, request_indexes) -> None:
if self.is_hybrid_model:
self.mamba_metadata.free_slots(request_indexes)
- # Clear intermediate offset entries for released requests
+ # Clear intermediate offset entries for released requests (CPU writes).
if self.mamba_slot_allocator is not None:
sa = self.mamba_slot_allocator
- sa._intermediate_counts_gpu[request_indexes] = 0
- sa._intermediate_offsets_gpu[request_indexes] = 0
- sa._intermediate_block_ids_gpu[request_indexes] = -1
- sa._eos_cache_block_id_gpu[request_indexes] = -1
+ sa._intermediate_counts_cpu[request_indexes] = 0
+ sa._intermediate_offsets_cpu[request_indexes] = 0
+ sa._intermediate_block_ids_cpu[request_indexes] = -1
+ sa._eos_cache_block_id_cpu[request_indexes] = -1
def resume_paused_requests(
self, active_request_count: int, newly_paused_request_ids: torch.Tensor
@@ -2380,7 +3146,7 @@ def evict_overflow_paused_requests(
-1,
-1,
dtype=paused_block_counts_cumsum.dtype,
- device=torch.cuda.current_device(),
+ device='cpu',
)
net_block_counts = paused_block_counts_cumsum - remaining_paused_request_counts
evict_request_count = torch.nonzero(net_block_counts >= 0)[0].item() + 1
@@ -2388,9 +3154,7 @@ def evict_overflow_paused_requests(
# Eviction index range.
evict_start_idx = self.paused_request_count - evict_request_count
evict_end_idx = self.paused_request_count
- evict_request_idxs = torch.arange(
- evict_start_idx, evict_end_idx, device=torch.cuda.current_device()
- )
+ evict_request_idxs = torch.arange(evict_start_idx, evict_end_idx, device='cpu')
# Clone needed: subsequent release_memory_blocks_from_request_indexes and
# _swap_book_keeping_tensors calls mutate self.request_ids in place.
evict_request_ids = self.request_ids[evict_start_idx:evict_end_idx].clone()
@@ -2405,24 +3169,24 @@ def evict_overflow_paused_requests(
src_idxs = torch.arange(
self.paused_request_count - evict_request_count,
self.paused_request_count,
- device=torch.cuda.current_device(),
+ device='cpu',
)
dst_idxs = torch.arange(
self.total_request_count - evict_request_count,
self.total_request_count,
- device=torch.cuda.current_device(),
+ device='cpu',
)
else:
# Swap all active requests with left-most evicted requests.
src_idxs = torch.arange(
self.paused_request_count - evict_request_count,
self.paused_request_count - evict_request_count + active_request_count,
- device=torch.cuda.current_device(),
+ device='cpu',
)
dst_idxs = torch.arange(
self.paused_request_count,
self.paused_request_count + active_request_count,
- device=torch.cuda.current_device(),
+ device='cpu',
)
# Swap evicted and active requests.
@@ -2498,6 +3262,14 @@ def update_requests(
# active_request_count -> This corresponds to requests that have not reached EOD or max length
# finished_request_count are requests that have reached the termination criterion
+ # Ensure all inputs are on CPU for bookkeeping operations.
+ if active_requests_mask.is_cuda:
+ active_requests_mask = active_requests_mask.cpu()
+ if new_tokens.is_cuda:
+ new_tokens = new_tokens.cpu()
+ if new_speculative_tokens is not None and new_speculative_tokens.is_cuda:
+ new_speculative_tokens = new_speculative_tokens.cpu()
+
self.num_prefill_requests = 0 # all turns to decode
# All request that were in prefill become decode requests.
# For the chunked prefill request we will overwrite this the next time add_request
@@ -2802,14 +3574,14 @@ def update_requests(
self.token_to_pos_ids[: self.active_token_count] = self.request_kv_length_offsets[
self.paused_request_count : self.total_request_count
].repeat_interleave(num_generated_tokens) + torch.arange(
- num_generated_tokens, device=torch.cuda.current_device()
+ num_generated_tokens, device='cpu'
).repeat(
active_request_count
)
#
# Token to request idx : [0, 0, 0, 1, 1, 1, 2, 2, 2 ...]
self.token_to_request_idx[: self.active_token_count] = torch.arange(
- self.paused_request_count, self.total_request_count, device=torch.cuda.current_device()
+ self.paused_request_count, self.total_request_count, device='cpu'
).repeat_interleave(num_generated_tokens)
self.token_to_position_in_request[: self.active_token_count] = self.token_to_pos_ids[
@@ -2831,7 +3603,7 @@ def update_requests(
raw_positions = (
old_offsets[:, None]
+ 1 # Offset by 1 because old_offsets points to the LAST token
- + torch.arange(num_generated_tokens, device=torch.cuda.current_device())[None, :]
+ + torch.arange(num_generated_tokens, device='cpu')[None, :]
)
#
# A token crosses to the next block if its raw_position >= block_size
@@ -2947,10 +3719,9 @@ def calculate_log_probs(
#
# active_token_ids[new_token_idx] = new_tokens
# : [ 52 | 12 | 16 3 | 12 72 24 88 86 ]
- active_token_ids = self.token_to_input_ids[: self.active_token_count].roll(-1, 0)
- active_query_lengths = self.request_query_lengths[
- self.paused_request_count : self.total_request_count
- ]
+ n_active = self.total_request_count - self.paused_request_count
+ active_token_ids = self.gpu_view.token_to_input_ids[: self.active_token_count].roll(-1, 0)
+ active_query_lengths = self.gpu_view.request_query_lengths[:n_active]
new_token_idx = active_query_lengths.cumsum(0) - 1
active_token_ids[new_token_idx] = new_tokens
diff --git a/megatron/core/inference/contexts/gpu_view.py b/megatron/core/inference/contexts/gpu_view.py
new file mode 100644
index 00000000000..65c401163b0
--- /dev/null
+++ b/megatron/core/inference/contexts/gpu_view.py
@@ -0,0 +1,228 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+import torch
+
+
+class ContextGPUView:
+ """GPU-resident snapshot of context bookkeeping data for the forward pass.
+
+ This is the ONLY interface GPU code (attention kernels, KV append, RoPE,
+ sampling, log-probs, speculative verification) uses to read context state.
+ CPU bookkeeping code accesses context tensors directly.
+
+ Populated once per step by ``DynamicInferenceContext.transfer_bookkeeping_to_gpu()``.
+ All tensors have fixed addresses for CUDA graph compatibility.
+
+ Convention:
+ ``context.foo`` -> CPU (source of truth, used by bookkeeping)
+ ``context.gpu_view.foo`` -> GPU (snapshot, used by forward pass)
+
+ Layout note: the bookkeeping fields are backed by a single contiguous
+ ``uint8`` buffer (``self._buf``). Each field is a ``view(dtype)`` onto a
+ slice of that buffer. This matches the pinned-CPU-buffer layout in
+ :class:`DynamicInferenceContext` so that the per-step H2D transfer is a
+ single ``cudaMemcpyAsync`` instead of one per field.
+ """
+
+ def __init__(
+ self,
+ max_requests: int,
+ max_tokens: int,
+ max_kv_blocks: int,
+ device: torch.device,
+ max_mamba_chunks: int = 0,
+ ):
+ # Field layout (must match DynamicInferenceContext's CPU buffer layout):
+ # int64 token fields first (auto 8-byte alignment), then int32 token
+ # fields, then int32 request fields, then int32 MHA fields, then
+ # int32 Mamba fields (hybrid models only; omitted when
+ # max_mamba_chunks == 0).
+ tok_int64_bytes = max_tokens * 8 # 2 fields of int64 = 8 bytes/elem
+ tok_int32_bytes = max_tokens * 4 # 4 fields of int32 = 4 bytes/elem
+ # Request-level fields are all 4 bytes wide. 3 int32 (in_prefill_status,
+ # query_lengths, kv_length_offsets) + 1 int32 (top_k) + 2 float32
+ # (temperature, top_p) + 1 int32 (active_request_last_token_idxs) = 7 fields.
+ req_4byte_bytes = max_requests * 4
+
+ # MHA section: 5 fields shared by both graphed and non-graphed MHAMetadata
+ # (only one is active per step, so sharing storage is fine).
+ # mha_query_lengths int32 (max_bs,) = max_bs * 4
+ # mha_cu_query_seq_lengths int32 (max_bs + 1,) = (max_bs+1) * 4
+ # mha_kv_seq_lengths int32 (max_bs,) = max_bs * 4
+ # mha_cu_kv_seq_lengths int32 (max_bs + 1,) = (max_bs+1) * 4
+ # mha_block_table int32 (max_bs, max_kv_blocks)
+ # max_bs == max_requests in DynamicInferenceContext.
+ max_bs = max_requests
+ mha_query_lengths_bytes = max_bs * 4
+ mha_cu_query_seq_lengths_bytes = (max_bs + 1) * 4
+ mha_kv_seq_lengths_bytes = max_bs * 4
+ mha_cu_kv_seq_lengths_bytes = (max_bs + 1) * 4
+ mha_block_table_bytes = max_bs * max_kv_blocks * 4
+
+ # Mamba section: 9 int32 fields, only present for hybrid models.
+ # mamba_batch_indices_decode int32 (max_bs,)
+ # mamba_batch_indices_prefill int32 (max_bs,)
+ # mamba_seq_idx int32 (1, max_tokens)
+ # mamba_cu_seqlens int32 (max_bs + 1,)
+ # mamba_cu_chunk_seqlens int32 (max_mamba_chunks + 1,)
+ # mamba_last_chunk_indices int32 (max_bs,)
+ # mamba_seq_idx_for_varlen int32 (max_mamba_chunks,)
+ # mamba_conv_seq_idx int32 (max_tokens,)
+ # mamba_conv_seq_start int32 (max_tokens,)
+ if max_mamba_chunks > 0:
+ mamba_batch_indices_decode_bytes = max_bs * 4
+ mamba_batch_indices_prefill_bytes = max_bs * 4
+ mamba_seq_idx_bytes = max_tokens * 4
+ mamba_cu_seqlens_bytes = (max_bs + 1) * 4
+ mamba_cu_chunk_seqlens_bytes = (max_mamba_chunks + 1) * 4
+ mamba_last_chunk_indices_bytes = max_bs * 4
+ mamba_seq_idx_for_varlen_bytes = max_mamba_chunks * 4
+ mamba_conv_seq_idx_bytes = max_tokens * 4
+ mamba_conv_seq_start_bytes = max_tokens * 4
+ else:
+ mamba_batch_indices_decode_bytes = 0
+ mamba_batch_indices_prefill_bytes = 0
+ mamba_seq_idx_bytes = 0
+ mamba_cu_seqlens_bytes = 0
+ mamba_cu_chunk_seqlens_bytes = 0
+ mamba_last_chunk_indices_bytes = 0
+ mamba_seq_idx_for_varlen_bytes = 0
+ mamba_conv_seq_idx_bytes = 0
+ mamba_conv_seq_start_bytes = 0
+
+ total_bytes = (
+ 2 * tok_int64_bytes
+ + 4 * tok_int32_bytes
+ + 7 * req_4byte_bytes
+ + mha_query_lengths_bytes
+ + mha_cu_query_seq_lengths_bytes
+ + mha_kv_seq_lengths_bytes
+ + mha_cu_kv_seq_lengths_bytes
+ + mha_block_table_bytes
+ + mamba_batch_indices_decode_bytes
+ + mamba_batch_indices_prefill_bytes
+ + mamba_seq_idx_bytes
+ + mamba_cu_seqlens_bytes
+ + mamba_cu_chunk_seqlens_bytes
+ + mamba_last_chunk_indices_bytes
+ + mamba_seq_idx_for_varlen_bytes
+ + mamba_conv_seq_idx_bytes
+ + mamba_conv_seq_start_bytes
+ )
+
+ # Zero-initialized so pre-transfer reads see zeros (matches prior semantics).
+ self._buf = torch.zeros(total_bytes, dtype=torch.uint8, device=device)
+
+ # Token-level tensors (consumed by embedding, RoPE, KV append, Mamba).
+ off = 0
+ self.token_to_input_ids = self._buf[off : off + tok_int64_bytes].view(torch.long)
+ off += tok_int64_bytes
+ self.token_to_pos_ids = self._buf[off : off + tok_int64_bytes].view(torch.long)
+ off += tok_int64_bytes
+ self.token_to_block_idx = self._buf[off : off + tok_int32_bytes].view(torch.int32)
+ off += tok_int32_bytes
+ self.token_to_local_position_within_kv_block = self._buf[off : off + tok_int32_bytes].view(
+ torch.int32
+ )
+ off += tok_int32_bytes
+ self.token_to_request_idx = self._buf[off : off + tok_int32_bytes].view(torch.int32)
+ off += tok_int32_bytes
+ self.token_to_position_in_request = self._buf[off : off + tok_int32_bytes].view(torch.int32)
+ off += tok_int32_bytes
+
+ # Request-level tensors (consumed by sampling, log-probs, speculative verification, MTP).
+ self.request_in_prefill_status = self._buf[off : off + req_4byte_bytes].view(torch.int32)
+ off += req_4byte_bytes
+ self.request_query_lengths = self._buf[off : off + req_4byte_bytes].view(torch.int32)
+ off += req_4byte_bytes
+ self.request_kv_length_offsets = self._buf[off : off + req_4byte_bytes].view(torch.int32)
+ off += req_4byte_bytes
+ # Sampling parameters (consumed by FlashInfer sampling).
+ # Mirror the active slice of `active_request_metadata[{label}]`;
+ # padded slots get neutral defaults from `pad_active_slices` (T=1.0, top_k=0, top_p=0.0).
+ self.temperature = self._buf[off : off + req_4byte_bytes].view(torch.float32)
+ off += req_4byte_bytes
+ self.top_k = self._buf[off : off + req_4byte_bytes].view(torch.int32)
+ off += req_4byte_bytes
+ self.top_p = self._buf[off : off + req_4byte_bytes].view(torch.float32)
+ off += req_4byte_bytes
+ # Per-request last-token row indices (consumed by sampling kernels as `gather_indices`).
+ # The CPU side of this slot IS `context.active_request_last_token_idxs`,
+ # populated by `build_active_slices` and `pad_active_slices`.
+ self.active_request_last_token_idxs = self._buf[off : off + req_4byte_bytes].view(
+ torch.int32
+ )
+ off += req_4byte_bytes
+
+ # MHA flash-attention metadata (shared between GraphedMHAMetadata and
+ # NonGraphedMHAMetadata — only one is active per step).
+ self.mha_query_lengths = self._buf[off : off + mha_query_lengths_bytes].view(torch.int32)
+ off += mha_query_lengths_bytes
+ self.mha_cu_query_seq_lengths = self._buf[off : off + mha_cu_query_seq_lengths_bytes].view(
+ torch.int32
+ )
+ off += mha_cu_query_seq_lengths_bytes
+ self.mha_kv_seq_lengths = self._buf[off : off + mha_kv_seq_lengths_bytes].view(torch.int32)
+ off += mha_kv_seq_lengths_bytes
+ self.mha_cu_kv_seq_lengths = self._buf[off : off + mha_cu_kv_seq_lengths_bytes].view(
+ torch.int32
+ )
+ off += mha_cu_kv_seq_lengths_bytes
+ self.mha_block_table = (
+ self._buf[off : off + mha_block_table_bytes]
+ .view(torch.int32)
+ .view(max_bs, max_kv_blocks)
+ )
+ off += mha_block_table_bytes
+
+ # Mamba varlen metadata (hybrid models only). Each GPU view matches a
+ # pinned CPU view in DynamicInferenceContext._cpu_bookkeeping_buf; the
+ # per-step coalesced H2D copy covers both MHA and Mamba alongside the
+ # token/request bookkeeping.
+ if max_mamba_chunks > 0:
+ self.mamba_batch_indices_decode = self._buf[
+ off : off + mamba_batch_indices_decode_bytes
+ ].view(torch.int32)
+ off += mamba_batch_indices_decode_bytes
+ self.mamba_batch_indices_prefill = self._buf[
+ off : off + mamba_batch_indices_prefill_bytes
+ ].view(torch.int32)
+ off += mamba_batch_indices_prefill_bytes
+ self.mamba_seq_idx = (
+ self._buf[off : off + mamba_seq_idx_bytes].view(torch.int32).view(1, max_tokens)
+ )
+ off += mamba_seq_idx_bytes
+ self.mamba_cu_seqlens = self._buf[off : off + mamba_cu_seqlens_bytes].view(torch.int32)
+ off += mamba_cu_seqlens_bytes
+ self.mamba_cu_chunk_seqlens = self._buf[off : off + mamba_cu_chunk_seqlens_bytes].view(
+ torch.int32
+ )
+ off += mamba_cu_chunk_seqlens_bytes
+ self.mamba_last_chunk_indices = self._buf[
+ off : off + mamba_last_chunk_indices_bytes
+ ].view(torch.int32)
+ off += mamba_last_chunk_indices_bytes
+ self.mamba_seq_idx_for_varlen = self._buf[
+ off : off + mamba_seq_idx_for_varlen_bytes
+ ].view(torch.int32)
+ off += mamba_seq_idx_for_varlen_bytes
+ self.mamba_conv_seq_idx = self._buf[off : off + mamba_conv_seq_idx_bytes].view(
+ torch.int32
+ )
+ off += mamba_conv_seq_idx_bytes
+ self.mamba_conv_seq_start = self._buf[off : off + mamba_conv_seq_start_bytes].view(
+ torch.int32
+ )
+ off += mamba_conv_seq_start_bytes
+ else:
+ self.mamba_batch_indices_decode = None
+ self.mamba_batch_indices_prefill = None
+ self.mamba_seq_idx = None
+ self.mamba_cu_seqlens = None
+ self.mamba_cu_chunk_seqlens = None
+ self.mamba_last_chunk_indices = None
+ self.mamba_seq_idx_for_varlen = None
+ self.mamba_conv_seq_idx = None
+ self.mamba_conv_seq_start = None
+
+ assert off == total_bytes, f"layout bug: wrote {off} of {total_bytes} bytes"
diff --git a/megatron/core/inference/contexts/kv_block_allocator.py b/megatron/core/inference/contexts/kv_block_allocator.py
index 87039835c7f..d555c925c93 100644
--- a/megatron/core/inference/contexts/kv_block_allocator.py
+++ b/megatron/core/inference/contexts/kv_block_allocator.py
@@ -3,6 +3,7 @@
from collections import deque
from typing import Callable, Dict, Optional
+import numpy as np
import torch
from torch import Tensor
@@ -47,32 +48,31 @@ def __init__(
assert self.active_count >= 1 # ensures paused_count < total_count - 1
self.dummy_block_idx = self.total_count - 1
- # Initialize block pool as a "stack" data structure
- self.block_bag = torch.arange(
- self.total_count, dtype=torch.int32, device=torch.cuda.current_device()
- )
+ # Initialize block pool as a "stack" data structure (CPU for bookkeeping).
+ self.block_bag = torch.arange(self.total_count, dtype=torch.int32, device='cpu')
if self.enable_prefix_caching:
# Block hash tracking for prefix caching: -1 = uncomputed, positive = valid hash
- self.block_hashes = torch.full(
- (self.total_count,), -1, dtype=torch.int64, device=torch.cuda.current_device()
- )
+ self.block_hashes = torch.full((self.total_count,), -1, dtype=torch.int64, device='cpu')
# Hash-to-block mapping for O(1) prefix lookup
self.kv_hash_to_block_id: Dict[int, int] = {}
# Reference count per block: 0 = cached (evictable), >0 = actively used
self.block_ref_counts = torch.zeros(
- (self.total_count,), dtype=torch.int32, device=torch.cuda.current_device()
+ (self.total_count,), dtype=torch.int32, device='cpu'
)
# LRU timestamps for eviction ordering (higher = more recently used)
# Only needed in LRU mode; RZ mode evicts immediately on ref_count==0
if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU:
self.block_timestamps = torch.zeros(
- (self.total_count,), dtype=torch.int64, device=torch.cuda.current_device()
+ (self.total_count,), dtype=torch.int64, device='cpu'
)
+ # Per-block MoE routing storage (populated when routing replay is enabled)
+ self.block_routing: Dict[int, np.ndarray] = {}
+
def __str__(self):
return (
f"using: total {self.get_total_used()}/{self.total_count - 1}"
@@ -183,6 +183,10 @@ def allocate_memory_blocks(self, num_blocks: int) -> Optional[Tensor]:
if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU:
self.update_timestamps(block_ids)
+ # Clear stale routing data for re-allocated blocks
+ for bid in block_ids.tolist():
+ self.block_routing.pop(bid, None)
+
return block_ids
def release_memory_blocks(self, blocks: Tensor) -> None:
@@ -239,9 +243,7 @@ def reset(self) -> None:
# Without resetting the block bag, context request memory will clash and
# requests will point to each other's memory blocks, resulting in faulty
# generations.
- self.block_bag = torch.arange(
- self.total_count, dtype=torch.int32, device=torch.cuda.current_device()
- )
+ self.block_bag = torch.arange(self.total_count, dtype=torch.int32, device='cpu')
self.total_avail = self.total_count - 1
@@ -255,6 +257,9 @@ def reset(self) -> None:
if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU:
self.block_timestamps.fill_(0)
+ # Clear per-block routing storage
+ self.block_routing.clear()
+
# =========================================================================
# Prefix caching methods
# =========================================================================
@@ -358,3 +363,123 @@ def evict_lru_blocks(self, num_blocks_needed: int) -> bool:
self._deregister_blocks(blocks_to_evict)
return True
+
+ # =========================================================================
+ # Per-block routing storage methods (for MoE routing replay)
+ # =========================================================================
+
+ def store_routing_per_block(self, flat_routing: Optional[np.ndarray]) -> None:
+ """Scatter flat routing indices into per-block storage.
+
+ Uses the context's token-to-block mapping to distribute each token's
+ routing data into the appropriate block. Matched (prefix-cached) blocks
+ already have routing from the original request and are not overwritten
+ here since their tokens are not in the active token layout.
+
+ Args:
+ flat_routing: ndarray of shape [active_token_count, num_layers, topk]
+ aligned with the context's active-token layout, or None.
+ """
+ if flat_routing is None:
+ return
+
+ context = self.context
+ token_count = context.active_token_count
+ if token_count == 0:
+ return
+
+ assert (
+ flat_routing.shape[0] == token_count
+ ), f"Routing token count {flat_routing.shape[0]} != active token count {token_count}"
+
+ # Token-to-block mapping for all active tokens
+ block_ids_np = context.token_to_block_idx[:token_count].cpu().numpy()
+ positions_np = context.token_to_local_position_within_kv_block[:token_count].cpu().numpy()
+
+ dummy = self.dummy_block_idx
+
+ # Group tokens by block_id using sort for efficient scatter
+ unique_blocks, inverse, counts = np.unique(
+ block_ids_np, return_inverse=True, return_counts=True
+ )
+ sorted_indices = np.argsort(inverse, kind='stable')
+ sorted_positions = positions_np[sorted_indices]
+ sorted_routing = flat_routing[sorted_indices]
+
+ offset = 0
+ for bid, count in zip(unique_blocks, counts):
+ bid = int(bid)
+ count = int(count)
+ if bid == dummy:
+ offset += count
+ continue
+ block_pos = sorted_positions[offset : offset + count]
+ block_rout = sorted_routing[offset : offset + count]
+ self.store_block_routing(bid, block_pos, block_rout)
+ offset += count
+
+ def reconstruct_routing_from_blocks(
+ self, block_ids: list[int], total_routing_tokens: int
+ ) -> Optional[np.ndarray]:
+ """Reconstruct routing indices from per-block storage.
+
+ Concatenates per-block routing ndarrays in block order, trimming the
+ last block to exactly ``total_routing_tokens`` entries.
+
+ Args:
+ block_ids: Ordered list of block IDs for the request.
+ total_routing_tokens: Expected number of routing tokens
+ (total_tokens - 1, since the last generated token has no
+ forward-pass routing).
+
+ Returns:
+ ndarray [total_routing_tokens, num_layers, topk] or None if any
+ block is missing routing data.
+ """
+ block_size = self.context.block_size_tokens
+ routing_parts = []
+ tokens_collected = 0
+
+ for bid in block_ids:
+ routing = self.get_block_routing(bid)
+ if routing is None:
+ return None # Missing routing data for this block
+ remaining = total_routing_tokens - tokens_collected
+ if remaining <= 0:
+ break
+ take = min(block_size, remaining)
+ routing_parts.append(routing[:take])
+ tokens_collected += take
+
+ if not routing_parts or tokens_collected != total_routing_tokens:
+ return None
+
+ return np.concatenate(routing_parts, axis=0)
+
+ def store_block_routing(
+ self, block_id: int, positions: np.ndarray, routing: np.ndarray
+ ) -> None:
+ """Store routing indices for specific token positions in a block.
+
+ Args:
+ block_id: The block ID.
+ positions: ndarray of token positions within the block (1D, int).
+ routing: ndarray of routing data [num_positions, num_layers, topk].
+ """
+ if block_id not in self.block_routing:
+ self.block_routing[block_id] = np.zeros(
+ (self.context.block_size_tokens, routing.shape[-2], routing.shape[-1]),
+ dtype=routing.dtype,
+ )
+ self.block_routing[block_id][positions] = routing
+
+ def get_block_routing(self, block_id: int) -> Optional[np.ndarray]:
+ """Get routing indices for a block.
+
+ Args:
+ block_id: The block ID.
+
+ Returns:
+ ndarray [block_size_tokens, num_layers, topk] or None if not stored.
+ """
+ return self.block_routing.get(block_id)
diff --git a/megatron/core/inference/contexts/mamba_slot_allocator.py b/megatron/core/inference/contexts/mamba_slot_allocator.py
index d7c57046c8a..60c8dd3416b 100644
--- a/megatron/core/inference/contexts/mamba_slot_allocator.py
+++ b/megatron/core/inference/contexts/mamba_slot_allocator.py
@@ -47,59 +47,70 @@ def __init__(
self.max_slots = max_slots
self.num_mamba_layers = num_mamba_layers
- device = torch.cuda.current_device()
+ gpu_device = torch.cuda.current_device()
num_blocks = context.kv_block_allocator.total_count
- # Block <-> slot mappings
- self.block_to_slot = torch.full((num_blocks,), -1, dtype=torch.int32, device=device)
- self.slot_to_block = torch.full((max_slots,), -1, dtype=torch.int32, device=device)
+ # Block <-> slot mappings (CPU for bookkeeping).
+ self.block_to_slot = torch.full((num_blocks,), -1, dtype=torch.int32, device='cpu')
+ self.slot_to_block = torch.full((max_slots,), -1, dtype=torch.int32, device='cpu')
- # Free slot pool (stack)
- self.free_slots = torch.arange(max_slots, dtype=torch.int32, device=device)
+ # Free slot pool (stack, CPU).
+ self.free_slots = torch.arange(max_slots, dtype=torch.int32, device='cpu')
self.free_count = max_slots
- # State tensors
+ # State tensors (GPU - accessed by Mamba CUDA kernels).
self.conv_states = torch.zeros(
(num_mamba_layers, max_slots) + conv_states_shape,
dtype=conv_states_dtype,
- device=device,
+ device=gpu_device,
)
self.ssm_states = torch.zeros(
- (num_mamba_layers, max_slots) + ssm_states_shape, dtype=ssm_states_dtype, device=device
+ (num_mamba_layers, max_slots) + ssm_states_shape,
+ dtype=ssm_states_dtype,
+ device=gpu_device,
)
# Hash-to-block mapping: only blocks with cached Mamba state
self.hash_to_block_id: Dict[int, int] = {}
- # Per-request intermediate state storage (GPU tensors, fixed-size per request)
- # 0 = no offset, -1 = no block
+ # Per-request intermediate state storage.
+ # offsets_cpu and counts_cpu: CPU source of truth. GPU copies are
+ # populated by transfer_bookkeeping_to_gpu() since Triton kernels read them.
+ # block_ids and eos_cache_block_id: CPU only (consumed by CPU code).
k = MAX_INTERMEDIATE_OFFSETS_PER_REQUEST
- self._intermediate_offsets_gpu = torch.zeros(
- (context.max_requests, k), dtype=torch.int32, device=device
+ self._intermediate_offsets_cpu = torch.zeros(
+ (context.max_requests, k), dtype=torch.int32, device='cpu'
)
- self._intermediate_block_ids_gpu = torch.full(
- (context.max_requests, k), -1, dtype=torch.int32, device=device
+ self._intermediate_counts_cpu = torch.zeros(
+ context.max_requests, dtype=torch.int32, device='cpu'
+ )
+ self._intermediate_offsets_gpu = torch.zeros(
+ (context.max_requests, k), dtype=torch.int32, device=gpu_device
)
self._intermediate_counts_gpu = torch.zeros(
- context.max_requests, dtype=torch.int32, device=device
+ context.max_requests, dtype=torch.int32, device=gpu_device
)
- self._eos_cache_block_id_gpu = torch.full(
- (context.max_requests,), -1, dtype=torch.int32, device=device
+ # CPU-only: consumed by _collect_commit_data() which needs .tolist() anyway.
+ self._intermediate_block_ids_cpu = torch.full(
+ (context.max_requests, k), -1, dtype=torch.int32, device='cpu'
+ )
+ self._eos_cache_block_id_cpu = torch.full(
+ (context.max_requests,), -1, dtype=torch.int32, device='cpu'
)
# CPU flag to skip GPU sync when no intermediates exist
self._has_intermediates = False
- # Pre-allocated output buffers for CUDA graph compatible extraction
+ # Pre-allocated output buffers for CUDA graph compatible extraction (GPU).
self.max_intermediate_count = MAX_INTERMEDIATE_OFFSETS_PER_REQUEST * context.max_requests
self.intermediate_ssm_out = torch.zeros(
(num_mamba_layers, self.max_intermediate_count) + ssm_states_shape,
dtype=ssm_states_dtype,
- device=device,
+ device=gpu_device,
)
self.intermediate_conv_out = torch.zeros(
(num_mamba_layers, self.max_intermediate_count) + conv_states_shape,
dtype=conv_states_dtype,
- device=device,
+ device=gpu_device,
)
# =========================================================================
@@ -320,9 +331,11 @@ def store_from_live_batch(self, slots: list, request_indices: list) -> None:
return
device = self.conv_states.device
slot_tensor = torch.tensor(slots, dtype=torch.int64, device=device)
- req_tensor = torch.tensor(request_indices, dtype=torch.int64, device=device)
- # Batch lookup mamba state indices (1 GPU sync)
- mamba_indices = self.context.mamba_metadata.request_to_mamba_state_idx[req_tensor].tolist()
+ # Lookup mamba indices from CPU bookkeeping, then move to GPU for state copy.
+ req_tensor_cpu = torch.tensor(request_indices, dtype=torch.int64)
+ mamba_indices = self.context.mamba_metadata.request_to_mamba_state_idx[
+ req_tensor_cpu
+ ].tolist()
mamba_idx_tensor = torch.tensor(mamba_indices, dtype=torch.int64, device=device)
# Fancy-indexed copy (2 kernel launches instead of 2E)
self.conv_states[:, slot_tensor] = self.context.mamba_conv_states[:, mamba_idx_tensor]
@@ -413,42 +426,39 @@ def compute_and_store_offsets(
offsets = sorted(offsets_set)
count = len(offsets)
- # Vectorized block ID lookup: GPU gather avoids per-block .item() syncs
+ # CPU bookkeeping writes (no GPU kernel launches).
if count > 0:
- device = self._intermediate_offsets_gpu.device
- abs_tokens = torch.tensor(
- [skip_tokens + o for o in offsets], dtype=torch.int64, device=device
- )
- block_indices = abs_tokens // ctx.block_size_tokens - 1
- bids = ctx.request_to_kv_block_ids[current_id][block_indices]
+ abs_tokens_cpu = torch.tensor([skip_tokens + o for o in offsets], dtype=torch.int64)
+ block_indices_cpu = abs_tokens_cpu // ctx.block_size_tokens - 1
+ bids_cpu = ctx.request_to_kv_block_ids[current_id][block_indices_cpu]
- self._intermediate_offsets_gpu[current_id, :count] = torch.tensor(
- offsets, dtype=torch.int32, device=device
+ self._intermediate_offsets_cpu[current_id, :count] = torch.tensor(
+ offsets, dtype=torch.int32
)
- self._intermediate_block_ids_gpu[current_id, :count] = bids.to(torch.int32)
+ self._intermediate_block_ids_cpu[current_id, :count] = bids_cpu.to(torch.int32)
self._has_intermediates = True
- self._intermediate_counts_gpu[current_id] = count
+ self._intermediate_counts_cpu[current_id] = count
# Block-aligned EOS: prompt_len is exactly block-aligned
if last_aligned_abs == prompt_len and prompt_len > 0:
last_block_idx = prompt_len // ctx.block_size_tokens - 1
if last_block_idx >= 0:
- self._eos_cache_block_id_gpu[current_id] = ctx.request_to_kv_block_ids[current_id][
+ self._eos_cache_block_id_cpu[current_id] = ctx.request_to_kv_block_ids[current_id][
last_block_idx
]
self._has_intermediates = True
else:
- self._eos_cache_block_id_gpu[current_id] = -1
+ self._eos_cache_block_id_cpu[current_id] = -1
else:
- self._eos_cache_block_id_gpu[current_id] = -1
+ self._eos_cache_block_id_cpu[current_id] = -1
- def get_intermediate_gpu_data(self):
- """Get intermediate offsets and counts as GPU tensor slices for current prefill batch.
+ def get_intermediate_cpu_data(self):
+ """Get intermediate offsets and counts as CPU tensor slices for current prefill batch.
Returns:
- Tuple of (offsets_gpu, counts_gpu) where:
- offsets_gpu: [prefill_count, 3] int32 GPU tensor
- counts_gpu: [prefill_count] int32 GPU tensor
+ Tuple of (offsets_cpu, counts_cpu) where:
+ offsets_cpu: [prefill_count, 3] int32 CPU tensor
+ counts_cpu: [prefill_count] int32 CPU tensor
Returns (None, None) if no prefill requests or no intermediates.
"""
if not self._has_intermediates:
@@ -463,10 +473,25 @@ def get_intermediate_gpu_data(self):
decode_count = ctx.batch_dimensions.decode_req_count
prefill_start = active_start + decode_count
- offsets = self._intermediate_offsets_gpu[prefill_start : prefill_start + prefill_count]
- counts = self._intermediate_counts_gpu[prefill_start : prefill_start + prefill_count]
+ offsets = self._intermediate_offsets_cpu[prefill_start : prefill_start + prefill_count]
+ counts = self._intermediate_counts_cpu[prefill_start : prefill_start + prefill_count]
return offsets, counts
+ def transfer_intermediate_to_gpu(self, prefill_start: int, prefill_count: int):
+ """Copy intermediate offsets/counts slice from CPU to GPU for Mamba kernels.
+
+ Returns the GPU tensor views for the forward-pass kernels to consume.
+ """
+ if prefill_count == 0:
+ return None, None
+ offsets_cpu = self._intermediate_offsets_cpu[prefill_start : prefill_start + prefill_count]
+ counts_cpu = self._intermediate_counts_cpu[prefill_start : prefill_start + prefill_count]
+ offsets_gpu = self._intermediate_offsets_gpu[prefill_start : prefill_start + prefill_count]
+ counts_gpu = self._intermediate_counts_gpu[prefill_start : prefill_start + prefill_count]
+ offsets_gpu.copy_(offsets_cpu, non_blocking=True)
+ counts_gpu.copy_(counts_cpu, non_blocking=True)
+ return offsets_gpu, counts_gpu
+
# =========================================================================
# Intermediate state commit
# =========================================================================
@@ -517,14 +542,14 @@ def _collect_commit_data(self):
decode_count = ctx.batch_dimensions.decode_req_count
prefill_start = active_start + decode_count
- # Batch-transfer block IDs and EOS block IDs from GPU (2 GPU syncs)
+ # Block IDs and EOS block IDs live on CPU (no GPU sync needed).
intermediate_count = metadata.intermediate_count
per_request_counts = metadata.per_request_intermediate_counts
- all_block_ids_cpu = self._intermediate_block_ids_gpu[
+ all_block_ids_cpu = self._intermediate_block_ids_cpu[
prefill_start : prefill_start + prefill_count
].tolist()
- eos_bids_cpu = self._eos_cache_block_id_gpu[
+ eos_bids_cpu = self._eos_cache_block_id_cpu[
prefill_start : prefill_start + prefill_count
].tolist()
@@ -586,10 +611,10 @@ def _clear_intermediate_state(self) -> None:
decode_count = ctx.batch_dimensions.decode_req_count
prefill_start = active_start + decode_count
end = prefill_start + prefill_count
- self._intermediate_counts_gpu[prefill_start:end].fill_(0)
- self._intermediate_offsets_gpu[prefill_start:end].fill_(0)
- self._intermediate_block_ids_gpu[prefill_start:end].fill_(-1)
- self._eos_cache_block_id_gpu[prefill_start:end].fill_(-1)
+ self._intermediate_counts_cpu[prefill_start:end].fill_(0)
+ self._intermediate_offsets_cpu[prefill_start:end].fill_(0)
+ self._intermediate_block_ids_cpu[prefill_start:end].fill_(-1)
+ self._eos_cache_block_id_cpu[prefill_start:end].fill_(-1)
self._has_intermediates = False
# =========================================================================
@@ -600,15 +625,13 @@ def reset(self) -> None:
"""Reset all state (mappings, free pool, cache, intermediate tracking)."""
self.block_to_slot.fill_(-1)
self.slot_to_block.fill_(-1)
- self.free_slots = torch.arange(
- self.max_slots, dtype=torch.int32, device=torch.cuda.current_device()
- )
+ self.free_slots = torch.arange(self.max_slots, dtype=torch.int32, device='cpu')
self.free_count = self.max_slots
self.hash_to_block_id.clear()
self.intermediate_ssm_out.zero_()
self.intermediate_conv_out.zero_()
- self._intermediate_offsets_gpu.fill_(0)
- self._intermediate_block_ids_gpu.fill_(-1)
- self._intermediate_counts_gpu.fill_(0)
- self._eos_cache_block_id_gpu.fill_(-1)
+ self._intermediate_offsets_cpu.fill_(0)
+ self._intermediate_counts_cpu.fill_(0)
+ self._intermediate_block_ids_cpu.fill_(-1)
+ self._eos_cache_block_id_cpu.fill_(-1)
self._has_intermediates = False
diff --git a/megatron/core/inference/engines/async_zmq_communicator.py b/megatron/core/inference/engines/async_zmq_communicator.py
index 52570845d61..aa13f659d40 100644
--- a/megatron/core/inference/engines/async_zmq_communicator.py
+++ b/megatron/core/inference/engines/async_zmq_communicator.py
@@ -131,6 +131,45 @@ async def all_reduce_max(self, *local_vals: int, async_op=True) -> int | tuple[i
except zmq.Again:
await asyncio.sleep(0.001)
+ def sync_all_reduce_max(self, *local_vals: int) -> int | tuple[int, ...]:
+ """Synchronous (non-asyncio) variant of all_reduce_max.
+
+ Uses blocking ZMQ sends/recvs so it can be called from synchronous
+ call sites that need a CPU-only MAX reduction across the process
+ group. Intended for tiny payloads (e.g. a few integers) that would
+ otherwise force a NCCL AllReduce kernel on the compute stream.
+
+ Note: when called from inside a running asyncio event loop, the
+ blocking recv will pause other coroutines on this rank until all
+ peers respond. This is acceptable here because every rank reaches
+ the call simultaneously and the message size is trivial.
+
+ Returns a single int when called with one argument, otherwise a tuple.
+ """
+ n = len(local_vals)
+ if n == 0:
+ raise ValueError("sync_all_reduce_max requires at least one value")
+
+ if self.world_size <= 1:
+ return local_vals[0] if n == 1 else local_vals
+
+ fmt = f'!{n}i'
+ payload = struct.pack(fmt, *local_vals)
+
+ if self.is_leader:
+ rows = [local_vals]
+ while len(rows) < self.world_size:
+ msg = self.gather_sock.recv()
+ rows.append(struct.unpack(fmt, msg))
+ maxes = tuple(max(row[i] for row in rows) for i in range(n))
+ self.bcast_sock.send(struct.pack(fmt, *maxes))
+ return maxes[0] if n == 1 else maxes
+ else:
+ self.gather_sock.send(payload)
+ msg = self.bcast_sock.recv()
+ result = struct.unpack(fmt, msg)
+ return result[0] if n == 1 else result
+
def close(self):
"""
Close the ZMQ sockets.
diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py
index a9c8337271f..92efff36073 100644
--- a/megatron/core/inference/engines/dynamic_engine.py
+++ b/megatron/core/inference/engines/dynamic_engine.py
@@ -3,9 +3,9 @@
import asyncio
import concurrent.futures
import logging
+import math
import multiprocessing
import socket
-import struct
import time
import warnings
from collections import deque
@@ -18,10 +18,10 @@
import torch
from torch import Tensor
-from torch.cuda.nvtx import range_pop, range_push
from megatron.core.inference.config import KVCacheManagementMode
from megatron.core.inference.contexts.dynamic_context import (
+ BlockOverflowError,
DynamicInferenceContext,
MaxSequenceLengthOverflowError,
TokenOverflowError,
@@ -42,15 +42,10 @@
from megatron.core.inference.text_generation_controllers.text_generation_controller import (
TextGenerationController,
)
-from megatron.core.inference.utils import (
- Counter,
- await_process_call,
- set_inference_cuda_graphed_iteration_for_ep_inference,
- unset_inference_cuda_graphed_iteration_for_ep_inference,
-)
+from megatron.core.inference.utils import Counter, InferenceMode, await_process_call
from megatron.core.process_groups_config import ProcessGroupCollection
from megatron.core.transformer.cuda_graphs import delete_cuda_graphs
-from megatron.core.transformer.enums import CudaGraphScope
+from megatron.core.transformer.enums import InferenceCudaGraphScope
from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
from megatron.core.utils import (
deprecate_args,
@@ -60,7 +55,11 @@
get_pg_size,
get_pg_src_rank,
internal_api,
+ nvtx_range_pop,
+ nvtx_range_push,
+ round_up_to_nearest_multiple,
trace_async_exceptions,
+ unwrap_model,
)
from .async_zmq_communicator import AsyncZMQCommunicator
@@ -212,12 +211,9 @@ def __init__(self, controller: TextGenerationController, context: DynamicInferen
if self.num_speculative_tokens > 0:
assert (
- self.num_speculative_tokens <= self.controller.num_mtp_heads
+ model_config.mtp_use_repeated_layer
+ or self.num_speculative_tokens <= self.controller.num_mtp_heads
), f"Number of speculative tokens {self.num_speculative_tokens} must be less than or equal to number of MTP heads {self.controller.num_mtp_heads}"
- assert (
- not self.materialize_only_last_token_logits
- ), "materialize_only_last_token_logits must be False when num_speculative_tokens > 0"
-
self.track_paused_request_events = inference_config.track_paused_request_events
self.track_generated_token_events = inference_config.track_generated_token_events
self.enable_chunked_prefill = inference_config.enable_chunked_prefill
@@ -225,8 +221,11 @@ def __init__(self, controller: TextGenerationController, context: DynamicInferen
self.logging_step_interval = inference_config.logging_step_interval
self.unified_memory_level = inference_config.unified_memory_level
self.use_synchronous_zmq_collectives = inference_config.use_synchronous_zmq_collectives
+ self.disable_ep_consensus = inference_config.disable_ep_consensus
+ self.ep_consensus_interval = inference_config.ep_consensus_interval
self.cuda_graph_impl = model_config.cuda_graph_impl
- self.cuda_graph_scope = model_config.cuda_graph_scope
+ self.inference_cuda_graph_scope = model_config.inference_cuda_graph_scope
+ self.cuda_graph_modules = model_config.cuda_graph_modules
# Initialize engine.
self.reset()
@@ -260,6 +259,9 @@ def __init__(self, controller: TextGenerationController, context: DynamicInferen
max_step = int(val)
self.inference_step_offset = int(max_step)
+ # Mark the inference engine as active. Cleared in `suspend()` and re-set in `resume()`.
+ InferenceMode.set_active()
+
# Create cuda graphs.
self.create_cuda_graphs()
@@ -333,17 +335,11 @@ def create_cuda_graphs(self, reset_context: bool = True):
reset_context (bool): Whether to reset the context after building cuda graphs.
"""
- if self.cuda_graph_impl != "local":
+ if self.inference_cuda_graph_scope == InferenceCudaGraphScope.none:
return
- if (
- CudaGraphScope.full_iteration in self.cuda_graph_scope
- and CudaGraphScope.full_iteration_inference not in self.cuda_graph_scope
- ):
- warnings.warn(
- "\n\n*** WARNING: 'full_iteration' CUDA graph scope used during inference! "
- "This will not create inference CUDA graphs. Use '--cuda-graph-scope=full_iteration_inference' instead. ***\n"
- )
+ if self.cuda_graph_impl != "local":
+ return
context = self.context
controller = self.controller
@@ -357,13 +353,21 @@ def create_cuda_graphs(self, reset_context: bool = True):
# Enable inference dispatcher for EP during graph capture
model_config = controller.inference_wrapped_model.model.config
- is_inference_optimized_ep = (
- model_config.transformer_impl == "inference_optimized"
- and model_config.expert_model_parallel_size > 1
+
+ # MTP warmup preparation: capture MTP CUDA graphs alongside the
+ # decoder graphs within the same loop rather than in a separate pass.
+ unwrapped = unwrap_model(controller.inference_wrapped_model.model)
+ mtp_warmup_enabled = (
+ controller.num_mtp_heads > 0
+ and (controller.num_speculative_tokens or 0) > 0
+ and hasattr(unwrapped, 'mtp')
)
- if is_inference_optimized_ep:
- unwrapped_model = controller.inference_wrapped_model.model
- set_inference_cuda_graphed_iteration_for_ep_inference(unwrapped_model)
+ if mtp_warmup_enabled:
+ tp_size = get_pg_size(controller.inference_wrapped_model.tp_group)
+ sp_enabled = model_config.sequence_parallel and tp_size > 1
+ mtp_pass_depth = not unwrapped.mtp.mtp_use_repeated_layer
+ mtp_warmup_depths = range(controller._num_mtp_depths) if mtp_pass_depth else [None]
+ mtp_seen_batch_sizes = set()
tbar = enumerate(context.cuda_graph_batch_dimensions_list)
if HAVE_TQDM:
@@ -383,18 +387,48 @@ def create_cuda_graphs(self, reset_context: bool = True):
# Enable routing recording during warmup if routing replay is enabled.
# This ensures the record_indices copy operation is captured in the CUDA graph.
- model_config = controller.inference_wrapped_model.model.config
if model_config.moe_enable_routing_replay:
RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)
# Forward pass -> logits.
- controller._dynamic_step_forward_logits(input_ids, position_ids)
+ with torch.inference_mode():
+ controller._dynamic_step_forward_logits(input_ids, position_ids)
- context.reset()
+ if controller._sampling_backend == "flashinfer":
+ if controller.num_speculative_tokens > 0:
+ controller._dynamic_step_sample_logits_and_verify_tokens(input_ids)
+ else:
+ controller._dynamic_step_sample_logits()
- # Disable inference dispatcher after graph capture
- if is_inference_optimized_ep:
- unset_inference_cuda_graphed_iteration_for_ep_inference(unwrapped_model)
+ # MTP CUDA graph warmup for this batch dimension.
+ if mtp_warmup_enabled:
+ n = cuda_graph_batch_dimension.req_count
+ # pylint: disable-next=possibly-used-before-assignment
+ if sp_enabled:
+ n = round_up_to_nearest_multiple(n, tp_size)
+ # pylint: disable-next=possibly-used-before-assignment
+ if n > 0 and n not in mtp_seen_batch_sizes:
+ mtp_seen_batch_sizes.add(n)
+ device = torch.cuda.current_device()
+ batch_dim = n // tp_size if sp_enabled else n
+ # Use zeros (not empty) — garbage token IDs cause OOB embedding lookups during graph capture/replay.
+ for depth in mtp_warmup_depths:
+ unwrapped.compute_mtp_single_step(
+ hidden_states=torch.zeros(
+ (batch_dim, 1, model_config.hidden_size),
+ device=device,
+ dtype=model_config.params_dtype,
+ ),
+ next_token_ids=torch.zeros((1, n), device=device, dtype=torch.long),
+ position_ids=torch.zeros((1, n), device=device, dtype=torch.int64),
+ depth=depth,
+ cache_key=("mtp", n, depth),
+ )
+
+ context.reset()
+
+ if mtp_warmup_enabled and mtp_seen_batch_sizes:
+ logging.info("> MTP CUDA graph warmup: %d batch size(s)", len(mtp_seen_batch_sizes))
# Memory usage.
time_end = time.time()
@@ -549,20 +583,16 @@ async def start_listening_to_data_parallel_coordinator(
mp_req_sock.bind_to_random_port(f"tcp://{local_ip}")
mp_req_addr = mp_req_sock.getsockopt_string(zmq.LAST_ENDPOINT)
- mp_len_sock = self.zmq_context.socket(zmq.PUB)
- mp_len_sock.bind_to_random_port(f"tcp://{local_ip}")
- mp_len_addr = mp_len_sock.getsockopt_string(zmq.LAST_ENDPOINT)
else:
mp_req_addr = None
- mp_len_addr = None
# Broadcast addresses to respective ranks.
bcast = [dp_addr]
torch.distributed.broadcast_object_list(bcast, src=dp_src, group=dp_group)
[dp_addr] = bcast
- bcast = [mp_req_addr, mp_len_addr]
+ bcast = [mp_req_addr]
torch.distributed.broadcast_object_list(bcast, src=mp_src, group=mp_group)
- [mp_req_addr, mp_len_addr] = bcast
+ [mp_req_addr] = bcast
identity = f'mp-coord-{dp_rank}'
if self.is_mp_coordinator:
@@ -579,37 +609,32 @@ async def start_listening_to_data_parallel_coordinator(
# 2. Create a publisher socket. This is used to publish or broadcast
# requests within the model parallel group
self.model_parallel_publisher_socket = mp_req_sock
-
- # 3. Create another publisher socket to broadcast the number of messages to receive.
- self.model_parallel_num_msgs_publisher_socket = mp_len_sock
self.zmq_sockets += [
self.socket_for_receiving_requests,
- self.model_parallel_num_msgs_publisher_socket,
self.model_parallel_publisher_socket,
]
- # All MP ranks subscribe to the two publisher sockets
+ # All MP ranks subscribe to the publisher socket
self.model_parallel_subscriber_socket = self.zmq_context.socket(zmq.SUB)
self.model_parallel_subscriber_socket.connect(mp_req_addr)
self.model_parallel_subscriber_socket.setsockopt_string(zmq.SUBSCRIBE, "")
- self.model_parallel_num_msgs_subscriber_socket = self.zmq_context.socket(zmq.SUB)
- self.model_parallel_num_msgs_subscriber_socket.connect(mp_len_addr)
- self.model_parallel_num_msgs_subscriber_socket.setsockopt_string(zmq.SUBSCRIBE, "")
-
- self.zmq_sockets += [
- self.model_parallel_subscriber_socket,
- self.model_parallel_num_msgs_subscriber_socket,
- ]
+ self.zmq_sockets += [self.model_parallel_subscriber_socket]
torch.distributed.barrier(mp_group)
# initialize zmq-based EP communicator
self.ep_rank = get_pg_rank(self.pg_collection.ep)
self.ep_world_size = get_pg_size(self.pg_collection.ep)
+ self._ep_consensus_loop_counter = 0
+ self._last_ep_consensus: tuple[int, bool] = (0, False)
if self.ep_world_size > 1:
self.expert_parallel_zmq_communicator = AsyncZMQCommunicator(
self.zmq_context, process_group=self.pg_collection.ep, hostname=hostname
)
+ # Give the context a CPU-side MAX-reduction primitive so
+ # match_graph_config() can avoid a per-step NCCL AllReduce kernel.
+ if hasattr(self.context, "set_ep_zmq_communicator"):
+ self.context.set_ep_zmq_communicator(self.expert_parallel_zmq_communicator)
# initialize zmq-based world communicator for consensus barriers
total_world_size = torch.distributed.get_world_size()
@@ -651,14 +676,14 @@ def suspend_resume_ctx(key: str, *, unified_memory_level: int) -> None:
start_mem = torch.cuda.memory_stats()
start_time = time.time()
- range_push(f"{key}-inference-context")
+ nvtx_range_push(f"{key}-inference-context")
torch.cuda.synchronize()
yield
finally:
- range_pop()
+ nvtx_range_pop(f"{key}-inference-context")
end_time = time.time()
end_mem = torch.cuda.memory_stats()
@@ -701,6 +726,8 @@ def suspend(self):
if self.state in (EngineState.SUSPENDED, EngineState.SUSPENDING):
return
+ InferenceMode.unset_active()
+
# Deallocate context tensors.
with self.__class__.suspend_resume_ctx(
"suspended", unified_memory_level=self.unified_memory_level
@@ -750,6 +777,8 @@ def resume(self):
if self.state not in (EngineState.SUSPENDED, EngineState.SUSPENDING):
return
+ InferenceMode.set_active()
+
# Resume.
with self.__class__.suspend_resume_ctx(
"resumed", unified_memory_level=self.unified_memory_level
@@ -820,8 +849,20 @@ def _handle_failed_request(self, request_id: int):
request = request_entry.record[-1]
if self.rank == 0:
+ errors = [
+ e.payload
+ for e in request.events
+ if e.type
+ in (
+ DynamicInferenceEventType.ERROR_NONTRANSIENT,
+ DynamicInferenceEventType.ERROR_TRANSIENT,
+ )
+ ]
+ errors_str = (
+ "; ".join(f"{type(e).__name__}: {e}" for e in errors) if errors else "unknown error"
+ )
warnings.warn(
- f"Request {request_id} failed to be added to the engine due to errors. "
+ f"Request {request_id} failed to be added to the engine ({errors_str}). "
f"Prompt Tokens: {len(request.prompt_tokens)} "
f"Tokens to generate: {request.sampling_params.num_tokens_to_generate} "
f"Max sequence length: {self.context.max_sequence_length} "
@@ -941,6 +982,16 @@ def _add_request(
request.status = Status.FAILED
request.add_event_error_nontransient(TokenOverflowError(request_id))
+ # Check that the KV cache has enough blocks for this request's max sequence length.
+ max_request_tokens = (
+ len(request.prompt_tokens) + request.sampling_params.num_tokens_to_generate
+ )
+ request_block_count = math.ceil(max_request_tokens / self.context.block_size_tokens)
+ total_blocks = self.context.kv_block_allocator.total_count - 1 # -1 for dummy block
+ if request_block_count > total_blocks:
+ request.status = Status.FAILED
+ request.add_event_error_nontransient(BlockOverflowError(request_id))
+
# Tokenize stop words if provided
if request.sampling_params.stop_words:
stop_word_ids = [
@@ -1025,9 +1076,9 @@ def post_process_requests(
accepted_tokens: torch.Tensor,
log_probs: torch.Tensor,
top_n_logprobs: Optional[Dict[int, List[Tuple[torch.Tensor, torch.Tensor]]]] = None,
- routing_indices_per_request: Optional[Dict[int, torch.Tensor]] = None,
pre_fwd_active_token_count: Optional[int] = None,
pre_fwd_step_count: Optional[int] = None,
+ finished_routing_block_ids: Optional[Dict[int, list[int]]] = None,
) -> Tuple[List[DynamicInferenceRequest], List[DynamicInferenceRequest]]:
"""
Handles post-processing for requests after a step.
@@ -1042,9 +1093,9 @@ def post_process_requests(
log_probs: (List): Log probs for each request
top_n_logprobs: (Dict): Top-n log probs for each request. Maps request_idx to
list of (top_n_logprobs, top_n_indices) tuples.
- routing_indices_per_request: (Dict[int, Tensor]): MoE routing indices
- pre-mapped by request_id. Each value is a tensor of shape
- [num_tokens_this_step, num_layers, topk].
+ finished_routing_block_ids: (Dict[int, List[int]]): Block IDs for
+ finished requests, saved before update_requests released them.
+ Used for per-block routing reconstruction.
Returns:
A list of active requests and completed requests as `DynamicInferenceRequest` objects
@@ -1105,10 +1156,15 @@ def post_process_requests(
len(request.generated_tokens) + len(tokens)
>= request.sampling_params.num_tokens_to_generate
):
- tokens = tokens[
- : request.sampling_params.num_tokens_to_generate
- - len(request.generated_tokens)
- ]
+ keep = request.sampling_params.num_tokens_to_generate - len(
+ request.generated_tokens
+ )
+ tokens = tokens[:keep]
+ # Trim log probs / top-n to match so the counts stay in sync.
+ if request_log_probs is not None:
+ request_log_probs = request_log_probs[:keep]
+ if top_n_logprobs is not None and req_idx in top_n_logprobs:
+ top_n_logprobs[req_idx] = top_n_logprobs[req_idx][:keep]
if request_id not in self.stop_word_being_finished_ids:
is_first_token = len(request.generated_tokens) == 0
request.generated_tokens += tokens
@@ -1145,10 +1201,13 @@ def post_process_requests(
request.ttft = (
first_token_event.timestamp - request.event_add_engine.timestamp
)
- if request.tpot is None:
- request.tpot = []
- per_token_step_time = step_time / len(tokens)
- request.tpot.extend([per_token_step_time] * len(tokens))
+ # TPOT is observability-only. step_time is 0.0 on
+ # non-logging steps (async_forward skips the event sync),
+ # so gate the update to keep the metric a truthful sparse
+ # sample instead of polluting it with zeros.
+ if step_time > 0:
+ per_token_step_time = step_time / len(tokens)
+ request.tpot.extend([per_token_step_time] * len(tokens))
# Check for stop words (after token is appended).
# With speculative decoding, a stop word may end before the last
@@ -1168,6 +1227,20 @@ def post_process_requests(
self._spec_tokens_accepted += actual_accepted
if request_id in finished_request_ids:
+ # Reconstruct routing from per-block storage before popping.
+ if (
+ finished_routing_block_ids
+ and request_id in finished_routing_block_ids
+ and len(self.requests[request_id].record.requests) == 1
+ ):
+ block_ids = finished_routing_block_ids[request_id]
+ total_tokens = len(request.prompt_tokens) + len(request.generated_tokens)
+ request.routing_indices = (
+ self.context.kv_block_allocator.reconstruct_routing_from_blocks(
+ block_ids, total_tokens - 1
+ )
+ )
+
# Request finished by normal means (termination_id, max_length, or stop word from previous step)
request.generated_length = len(request.generated_tokens)
request.status = Status.COMPLETED
@@ -1199,7 +1272,13 @@ def post_process_requests(
top_n_logprobs[req_idx] = top_n_logprobs[req_idx][:-num_stop_word_trim]
# Process log_probs if available (unified for both regular and chunked prefill)
- if request_log_probs is not None:
+ # Skip for requests being finished due to stop words — tokens are not
+ # appended for these requests, so log probs must also be skipped to keep
+ # the two lists in sync.
+ if (
+ request_log_probs is not None
+ and request_id not in self.stop_word_being_finished_ids
+ ):
# Initialize lists if they don't exist
if not request.prompt_log_probs:
request.prompt_log_probs = []
@@ -1232,7 +1311,12 @@ def post_process_requests(
request.generated_log_probs.extend(request_log_probs[split_idx:])
# Process top_n_logprobs if available (unified for both regular and chunked prefill)
- if top_n_logprobs is not None and req_idx in top_n_logprobs:
+ # Same stop-word guard as log probs above.
+ if (
+ top_n_logprobs is not None
+ and req_idx in top_n_logprobs
+ and request_id not in self.stop_word_being_finished_ids
+ ):
# Initialize lists if they don't exist
if request.prompt_top_n_logprobs is None:
request.prompt_top_n_logprobs = []
@@ -1266,23 +1350,6 @@ def post_process_requests(
else:
request.generated_top_n_logprobs.append(logit_dict)
- # Process routing indices if available (keyed by request_id)
- # Each step's routing is a tensor of shape [num_tokens_this_step, num_layers, topk]
- # We concatenate along dim=0 to accumulate: [total_tokens, num_layers, topk]
- if (
- routing_indices_per_request is not None
- and request_id in routing_indices_per_request
- ):
- step_routing = routing_indices_per_request[
- request_id
- ] # [num_tokens, num_layers, topk]
- if request.routing_indices is None:
- request.routing_indices = step_routing.clone()
- else:
- request.routing_indices = torch.cat(
- [request.routing_indices, step_routing], dim=0
- )
-
# Handle evicted requests.
if evict_request_ids is not None and evict_request_ids.numel() > 0:
@@ -1619,56 +1686,74 @@ async def async_forward(self) -> Tuple[Dict, Dict, float]:
# schedule requests
self.schedule_waiting_requests()
- # Saving pre-step state, for printing output below.
+ # The print block (async_bookkeep) and metrics block both fire on this
+ # condition after step_count is incremented. Predict it up-front so we
+ # can skip the GPU-timing sync and the context_state dict builds that
+ # only exist to feed those logging/metrics blocks.
+ will_log_this_step = (
+ self.logging_step_interval > 0
+ and (self.context.step_count + 1) % self.logging_step_interval == 0
+ )
+
is_decode_only = self.context.is_decode_only()
- pre_step_context_state = {
- "is_decode_only": is_decode_only,
- "max_requests": self.context.max_requests,
- "total_request_count": self.context.total_request_count,
- "paused_request_count": self.context.paused_request_count,
- "active_token_count": self.context.active_token_count,
- "step_count": self.context.step_count,
- }
+ if will_log_this_step:
+ pre_step_context_state = {
+ "is_decode_only": is_decode_only,
+ "max_requests": self.context.max_requests,
+ "total_request_count": self.context.total_request_count,
+ "paused_request_count": self.context.paused_request_count,
+ "active_token_count": self.context.active_token_count,
+ "step_count": self.context.step_count,
+ }
+ else:
+ # active_token_count and step_count are still consumed by
+ # post_process_requests' pre_fwd_* args (for add_event_generated_token);
+ # the other four fields are only read in the gated print block.
+ pre_step_context_state = {
+ "active_token_count": self.context.active_token_count,
+ "step_count": self.context.step_count,
+ }
# Generate tokens.
- range_push("Prefill" if not is_decode_only else "Decode")
+ nvtx_range_push("Prefill" if not is_decode_only else "Decode")
# TODO @TDE: Account for this line when overlapping forward and bookkeep.
self.is_decode_only = is_decode_only
- self.step_start_event.record()
+ if will_log_this_step:
+ self.step_start_event.record()
result = await self.controller.async_generate_output_tokens_dynamic_batch()
- self.step_end_event.record()
- self.step_end_event.synchronize()
- step_time = self.step_start_event.elapsed_time(self.step_end_event) / 1e3
+ if will_log_this_step:
+ self.step_end_event.record()
+ self.step_end_event.synchronize()
+ step_time = self.step_start_event.elapsed_time(self.step_end_event) / 1e3
+ else:
+ step_time = 0.0
self.context.step_count += 1
self.context.prefix_cache_lru_clock += 1
- range_pop()
+ nvtx_range_pop("Prefill" if not is_decode_only else "Decode")
- if (
- self.logging_step_interval > 0
- and self.context.step_count > 0
- and self.context.step_count % self.logging_step_interval == 0
- and self.metrics_writer is not None
- ):
- kvcache_util_stats = self.context.get_kvcache_utilization_stats()
+ if will_log_this_step:
+ kvcache_util_stats = (
+ self.context.get_kvcache_utilization_stats()
+ if self.metrics_writer is not None
+ else None
+ )
+ post_step_context_state = {
+ "waiting_request_count": len(self.waiting_request_ids),
+ "finished_request_count": self.finished_request_count,
+ "evicted_request_count": self.evicted_request_count,
+ "kv_stats": kvcache_util_stats,
+ "total_active_block_count": self.context.kv_block_allocator.active_count,
+ "total_paused_block_count": self.context.kv_block_allocator.paused_count,
+ "total_active_used_blocks": self.context.kv_block_allocator.get_active_used(),
+ "total_paused_used_blocks": self.context.kv_block_allocator.get_paused_used(),
+ }
+ context_state = {**pre_step_context_state, **post_step_context_state}
else:
- kvcache_util_stats = None
-
- post_step_context_state = {
- "waiting_request_count": len(self.waiting_request_ids),
- "finished_request_count": self.finished_request_count,
- "evicted_request_count": self.evicted_request_count,
- "kv_stats": kvcache_util_stats,
- "padded_active_token_count": self.context.padded_active_token_count,
- "using_cuda_graph_this_step": self.context.using_cuda_graph_this_step(),
- "total_active_block_count": self.context.kv_block_allocator.active_count,
- "total_paused_block_count": self.context.kv_block_allocator.paused_count,
- "total_active_used_blocks": self.context.kv_block_allocator.get_active_used(),
- "total_paused_used_blocks": self.context.kv_block_allocator.get_paused_used(),
- }
-
- context_state = {**pre_step_context_state, **post_step_context_state}
+ # Keep kv_stats=None so the metrics-block gate at `async_bookkeep`
+ # (`if context_state["kv_stats"] is not None`) remains well-typed.
+ context_state = {**pre_step_context_state, "kv_stats": None}
return result, context_state, step_time
@@ -1690,7 +1775,7 @@ async def async_bookkeep(
cuda_graph_request_count (int): The CUDA graph batch size matching this step.
"""
# Increment finished_request_count.
- range_push("bookkeeping")
+ nvtx_range_push("bookkeeping")
cuda_graph_request_count = None
if step_result is not None:
@@ -1702,7 +1787,7 @@ async def async_bookkeep(
accepted_tokens = step_result["accepted_tokens"]
log_probs = step_result["log_probs"]
top_n_logprobs = step_result.get("top_n_logprobs", None)
- routing_indices_per_request = step_result.get("routing_indices_per_request", None)
+ finished_routing_block_ids = step_result.get("finished_routing_block_ids", None)
cuda_graph_request_count = step_result["cuda_graph_request_count"]
# Add paused events.
@@ -1720,9 +1805,9 @@ async def async_bookkeep(
accepted_tokens,
log_probs,
top_n_logprobs,
- routing_indices_per_request,
pre_fwd_active_token_count=context_state.get("active_token_count"),
pre_fwd_step_count=context_state.get("step_count"),
+ finished_routing_block_ids=finished_routing_block_ids,
)
else:
@@ -1739,13 +1824,13 @@ async def async_bookkeep(
), f"Failed request {failed_request_id} future has not been properly resolved."
self.failed_request_ids.clear()
- range_pop()
+ nvtx_range_pop("bookkeeping")
# Detokenize all finished requests if not using
# the coordinator. Otherwise, the coordinator will
# overlap detokenization with the engine.
if not self.use_coordinator:
- range_push("detokenization")
+ nvtx_range_push("detokenization")
for record in finished_request_records:
for request in record.requests:
if request.prompt is None:
@@ -1759,7 +1844,7 @@ async def async_bookkeep(
request.generated_tokens,
remove_EOD=not request.sampling_params.detokenize_stop_sequence,
)
- range_pop()
+ nvtx_range_pop("detokenization")
# Handle necessary ZMQ DP coordinator communication.
# Failed request replies were already sent in _handle_failed_request,
@@ -1769,13 +1854,13 @@ async def async_bookkeep(
r for r in finished_request_records if r.requests[-1].status != Status.FAILED
]
if records_to_send:
- range_push("coordinator_communication")
+ nvtx_range_push("coordinator_communication")
payload = msgpack.packb(
[Headers.ENGINE_REPLY.value, [r.merge().serialize() for r in records_to_send]],
use_bin_type=True,
)
self.socket_for_receiving_requests.send(payload)
- range_pop()
+ nvtx_range_pop("coordinator_communication")
# Drain prefix cache hit counters from context into engine accumulators.
if self.context.enable_prefix_caching:
@@ -1785,6 +1870,7 @@ async def async_bookkeep(
self.context.prefix_cache_blocks_matched = 0
# Log KV cache utilization stats to W&B
+ nvtx_range_push("wandb_logging")
if context_state["kv_stats"] is not None:
# Prepare metrics dictionary with all stats
# Use 'inference/' prefix for all metrics to separate from training metrics
@@ -1824,13 +1910,17 @@ async def async_bookkeep(
self.metrics_writer.log(metrics, commit=True)
else:
raise ValueError(f"Unsupported metrics writer type: {type(self.metrics_writer)}")
+ nvtx_range_pop("wandb_logging")
# Print context state.
+ nvtx_range_push("console_logging")
if (
self.logging_step_interval > 0
and self.context.step_count % self.logging_step_interval == 0
):
+ nvtx_range_push("cuda_memory_stats")
mem = torch.cuda.memory_stats()
+ nvtx_range_pop("cuda_memory_stats")
step_type = "decode" if context_state["is_decode_only"] else "non-decode"
output_str = (
"* rank %d | step %d | %s ... time: %.3f ms%s ... "
@@ -1897,6 +1987,8 @@ async def async_bookkeep(
self._prefix_cache_hits = 0
self._prefix_cache_blocks_matched = 0
+ nvtx_range_pop("console_logging")
+
return {
"active_request_ids": active_request_ids,
"finished_request_records": finished_request_records,
@@ -2015,7 +2107,7 @@ def schedule_requests(self) -> int:
int: The number of messages that were received and processed in this batch.
"""
- range_push("drain_zmq_socket")
+ nvtx_range_push("drain_zmq_socket")
all_messages = []
if self.is_mp_coordinator:
while True:
@@ -2025,30 +2117,14 @@ def schedule_requests(self) -> int:
except zmq.Again:
# This exception is hit as soon as the socket is empty.
break
- messages_to_dequeue = len(all_messages)
- # First publish the number of messages to dequeue.
- # This is important because we want all tensor parallel ranks
- # to dequeue the same number of messages.
- self.model_parallel_num_msgs_publisher_socket.send(
- struct.pack('!i', messages_to_dequeue)
+ self.model_parallel_publisher_socket.send_multipart(
+ [bytes([Headers.TP_BROADCAST.value])] + all_messages
)
- # Now publish the actual messages to all model parallel ranks
- if messages_to_dequeue > 0:
- self.model_parallel_publisher_socket.send_multipart(all_messages)
else:
- # First, receive the number of messages to dequeue from mp-rank 0
- messages_to_dequeue = struct.unpack(
- '!i', self.model_parallel_num_msgs_subscriber_socket.recv()
- )[0]
- # Now, dequeue the same number of messages from the subscriber socket.
- # Note that these receives are blocking, because the messages
- # are guaranteed to be available after the tp-rank 0 has sent them.
- if messages_to_dequeue > 0:
- all_messages = self.model_parallel_subscriber_socket.recv_multipart()
- else:
- all_messages = []
+ frames = self.model_parallel_subscriber_socket.recv_multipart()
+ all_messages = frames[1:]
- range_pop()
+ nvtx_range_pop("drain_zmq_socket")
# First pass: add requests.
# Control signals are queued for the second pass.
@@ -2059,9 +2135,9 @@ def schedule_requests(self) -> int:
if header == Headers.SUBMIT_REQUEST:
request_id, prompt, sampling_params = data[1:]
sampling_params = SamplingParams.deserialize(sampling_params)
- range_push("add_request")
+ nvtx_range_push("add_request")
self.add_request(request_id, prompt, sampling_params)
- range_pop()
+ nvtx_range_pop("add_request")
elif header == Headers.SET_GENERATION_EPOCH:
new_generation_epoch = data[1]
else:
@@ -2205,7 +2281,7 @@ async def _ep_establish_consensus(
(global_work, all_pausing): max work across EP, and whether
all peers signaled consensus.
"""
- range_push("_ep_establish_consensus")
+ nvtx_range_push("_ep_establish_consensus")
consensus_val = -1 if signal_consensus else 0
@@ -2230,7 +2306,7 @@ async def _ep_establish_consensus(
else:
global_work, global_consensus = local_work, consensus_val
- range_pop()
+ nvtx_range_pop("_ep_establish_consensus")
return global_work, global_consensus == -1
async def _world_barrier(self):
@@ -2242,12 +2318,12 @@ async def _world_barrier(self):
No-op when world_size == 1 (communicator is not created).
"""
- range_push("world_barrier")
+ nvtx_range_push("world_barrier")
if hasattr(self, 'world_zmq_communicator'):
await self.world_zmq_communicator.all_reduce_max(
1, async_op=(not self.use_synchronous_zmq_collectives)
)
- range_pop()
+ nvtx_range_pop("world_barrier")
@trace_async_exceptions
async def run_engine_with_coordinator(
@@ -2273,9 +2349,50 @@ async def run_engine_with_coordinator(
local_pending = self.context.get_active_request_count() + len(
self.waiting_request_ids
)
- global_work, all_pausing = await self._ep_establish_consensus(
- local_pending, signal_consensus=(self.state == EngineState.PAUSING)
- )
+ if self.disable_ep_consensus:
+ # Skip the EP consensus all-reduce; act on local state only.
+ # NOTE: even with no consensus we must still participate in EP
+ # collectives (NCCL all-to-all, etc.) every iteration. A peer with
+ # real work will block at its all-to-all kernel waiting for this
+ # rank, so when there is no local work we run dummy_forward()
+ # rather than sleeping. Sleeping here would deadlock EP > 1.
+ if self.state == EngineState.PAUSING:
+ await self._world_barrier()
+ self.state = EngineState.PAUSED
+ self._state_events[EngineState.PAUSED].set()
+ elif local_pending > 0:
+ await self.async_step()
+ else:
+ self.step_start_event.record()
+ nvtx_range_push("EP-dummy-forward")
+ self.controller.dummy_forward()
+ self.step_end_event.record()
+ self.step_end_event.synchronize()
+ nvtx_range_pop("EP-dummy-forward")
+ self.context.step_count += 1
+ self.context.prefix_cache_lru_clock += 1
+ # The consensus path yields via _ep_establish_consensus;
+ # without it we must still let other coroutines (signal
+ # delivery, request scheduling) run between steps.
+ await asyncio.sleep(0)
+ continue
+ global_work_from_last_consensus, _ = self._last_ep_consensus
+ if (
+ global_work_from_last_consensus == 0
+ or self._ep_consensus_loop_counter % self.ep_consensus_interval == 0
+ ):
+ # selectively enter ep_establish_consensus if
+ # 1. there is no global work -> engine is idle. At any step in the future
+ # one of the ranks can receive work. So we should be eagerly checking for that
+ # 2. it has been 20 steps since we last established consensus, and that consensus
+ # had some work.
+ # In the worst case, this delays pausing by 20 steps which is around
+ # 200-400 milliseconds.
+ self._last_ep_consensus = await self._ep_establish_consensus(
+ local_pending, signal_consensus=(self.state == EngineState.PAUSING)
+ )
+ global_work, all_pausing = self._last_ep_consensus
+ self._ep_consensus_loop_counter += 1
if all_pausing:
# All EP peers are PAUSING: pause immediately.
@@ -2289,9 +2406,11 @@ async def run_engine_with_coordinator(
else:
# Dummy forward to participate in the EP collective.
self.step_start_event.record()
+ nvtx_range_push("EP-dummy-forward")
self.controller.dummy_forward()
self.step_end_event.record()
self.step_end_event.synchronize()
+ nvtx_range_pop("EP-dummy-forward")
self.context.step_count += 1
self.context.prefix_cache_lru_clock += 1
else:
@@ -2306,6 +2425,10 @@ async def run_engine_with_coordinator(
self.state = EngineState.RUNNING
self._state_events[EngineState.PAUSED].clear()
self._state_events[EngineState.RUNNING].set()
+ # The cache from the PAUSING phase still has all_pausing=True;
+ # without this reset the next RUNNING iteration would skip
+ # consensus, read the stale flag, and immediately re-pause.
+ self._last_ep_consensus = (0, False)
elif self.state == EngineState.SUSPENDING:
await self._world_barrier()
diff --git a/megatron/core/inference/engines/static_engine.py b/megatron/core/inference/engines/static_engine.py
index 0b3b9c1b856..c079921a271 100644
--- a/megatron/core/inference/engines/static_engine.py
+++ b/megatron/core/inference/engines/static_engine.py
@@ -18,6 +18,7 @@
from megatron.core.inference.text_generation_controllers.text_generation_controller import (
TextGenerationController,
)
+from megatron.core.inference.utils import InferenceMode
from megatron.core.utils import get_asyncio_loop
try:
@@ -129,6 +130,8 @@ def __init__(
self.controller.inference_wrapped_model.inference_context = original_context
self.legacy = True
+ InferenceMode.set_active()
+
def get_new_request_id(self) -> str:
"""Gets a new request id from the scheduler"""
return self.scheduler.get_new_request_id()
diff --git a/megatron/core/inference/headers.py b/megatron/core/inference/headers.py
index aa2f0568975..8ad1913e6b1 100644
--- a/megatron/core/inference/headers.py
+++ b/megatron/core/inference/headers.py
@@ -20,6 +20,7 @@ class Headers(Enum):
STOP = auto()
DISCONNECT = auto()
SHUTDOWN = auto()
+ TP_BROADCAST = auto()
class UnknownHeaderError(Exception):
diff --git a/megatron/core/inference/inference_request.py b/megatron/core/inference/inference_request.py
index 27580f4b830..d6e7c67a959 100644
--- a/megatron/core/inference/inference_request.py
+++ b/megatron/core/inference/inference_request.py
@@ -1,18 +1,19 @@
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
import copy
+import hashlib
import time
import warnings
from dataclasses import asdict, dataclass, field
from enum import Enum, auto
-from itertools import accumulate
from typing import Any, Dict, List, Optional, Tuple
+import numpy as np
import torch
from megatron.core.inference.sampling_params import SamplingParams
from megatron.core.tokenizers import MegatronTokenizer
-from megatron.core.utils import experimental_api
+from megatron.core.utils import experimental_api, nvtx_range_pop, nvtx_range_push
def serialize_tensor(tensor: torch.Tensor) -> List:
@@ -24,12 +25,12 @@ def serialize_tensor(tensor: torch.Tensor) -> List:
Returns:
(List) Tensor as a list
"""
- torch.cuda.nvtx.range_push("serialize_tensor")
+ nvtx_range_push("serialize_tensor")
# simply convert tensor into a list
tensor = tensor.cpu().tolist()
- torch.cuda.nvtx.range_pop()
+ nvtx_range_pop("serialize_tensor")
return tensor
@@ -46,6 +47,16 @@ def deserialize_tensor(tensor_as_list: List) -> torch.Tensor:
return tensor
+def serialize_ndarray(arr: np.ndarray) -> dict:
+ """Serialize numpy array to a JSON-compatible dict."""
+ return {"data": arr.tolist(), "dtype": str(arr.dtype)}
+
+
+def deserialize_ndarray(obj: dict) -> np.ndarray:
+ """Deserialize numpy array from dict."""
+ return np.array(obj["data"], dtype=np.dtype(obj["dtype"]))
+
+
def unwrap_serialized_tensors(serialized_request: dict) -> dict:
"""Unwrap ("tensor", [...]) tuples produced by serialize() into plain lists.
@@ -76,53 +87,44 @@ class Status(Enum):
# Hash computation for prefix caching
# =========================================================================
-# Constants for hash computation
-# Using 2^61 - 1 (Mersenne prime) for ~10^18 hash space, reducing collision probability
-# from ~10^-9 to ~10^-18 compared to the previous prime (1000000007).
-HASH_PRIME = 2305843009213693951
-HASH_BASE = 31
-
-_hash_powers: Optional[torch.Tensor] = None
-
def compute_block_hashes_batched(prompt_tokens: torch.Tensor, block_size: int) -> List[int]:
- """Compute hashes for all complete blocks in a prompt in one batched operation.
+ """Compute SHA-256 based hashes for all complete blocks in a prompt.
- Reshapes prompt tokens into [num_blocks, block_size], computes all per-block
- token hashes via a single GPU matmul, transfers results with one .tolist() call,
- and chains parent hashes on CPU.
+ Each block hash is computed as SHA-256(parent_digest || block_bytes), where
+ parent_digest chains from the previous block (starting from a zero digest).
+ This provides cryptographic collision resistance with no exploitable algebraic
+ structure.
Args:
prompt_tokens: All prompt token IDs, shape [seq_len].
block_size: Number of tokens per block.
Returns:
- List of positive integer hash values (1 to HASH_PRIME), one per complete block.
+ List of positive integer hash values in [1, 2^63-1], one per complete block.
"""
num_complete_blocks = len(prompt_tokens) // block_size
if num_complete_blocks == 0:
return []
- global _hash_powers
- if _hash_powers is None or _hash_powers.shape[0] != block_size:
- positions = torch.arange(block_size, device=prompt_tokens.device, dtype=torch.int64)
- _hash_powers = torch.pow(HASH_BASE, positions).to(torch.int64) % HASH_PRIME
+ # Single GPU->CPU transfer, get contiguous bytes
+ tokens_cpu = prompt_tokens[: num_complete_blocks * block_size].to(torch.int64).cpu()
+ tokens_bytes = tokens_cpu.numpy().tobytes()
+ block_byte_size = block_size * tokens_cpu.element_size() # 8 bytes per int64
- # Reshape to [num_blocks, block_size] (zero-copy view) and compute all token hashes
- blocks = prompt_tokens[: num_complete_blocks * block_size].view(num_complete_blocks, block_size)
- token_hashes = (blocks.to(torch.int64) * _hash_powers).sum(dim=1) % HASH_PRIME
+ hashes = []
+ parent_digest = b'\x00' * 32 # SHA-256 digest size
- # Single GPU→CPU transfer
- token_hashes_list = token_hashes.tolist()
+ for i in range(num_complete_blocks):
+ block_bytes = tokens_bytes[i * block_byte_size : (i + 1) * block_byte_size]
+ digest = hashlib.sha256(parent_digest + block_bytes).digest()
- # Chain parent hashes on CPU (C-level accumulate, no Python loop)
- hashes = list(
- accumulate(
- token_hashes_list,
- lambda parent, th: (parent * HASH_BASE + th) % HASH_PRIME + 1,
- initial=0,
- )
- )[1:]
+ # Map to positive int64 range [1, 2^63-1], avoiding sentinels -1 and 0
+ raw = int.from_bytes(digest[:8], byteorder='little', signed=False)
+ hash_val = (raw % (2**63 - 1)) + 1
+
+ hashes.append(hash_val)
+ parent_digest = digest # Full 32-byte digest chains into next block
return hashes
@@ -153,7 +155,7 @@ class InferenceRequest:
prompt_top_n_logprobs: Optional[List[Dict[str, float]]] = None
generated_top_n_logprobs: Optional[List[Dict[str, float]]] = None
generated_length: Optional[int] = None
- tpot: Optional[List[int]] = None
+ tpot: List[float] = field(default_factory=list)
def __post_init__(self):
if self.sampling_params is None and self.inference_parameters is not None:
@@ -180,9 +182,13 @@ def serialize(self) -> dict:
self.inference_parameters.serialize() if self.inference_parameters else None
)
- # Serialize tensors.
+ # Serialize tensors and numpy arrays.
obj = {
- k: (("tensor", serialize_tensor(v)) if isinstance(v, torch.Tensor) else v)
+ k: (
+ ("tensor", serialize_tensor(v))
+ if isinstance(v, torch.Tensor)
+ else ("ndarray", serialize_ndarray(v)) if isinstance(v, np.ndarray) else v
+ )
for k, v in obj.items()
}
return obj
@@ -221,10 +227,12 @@ def _post_deserialize(self, obj: dict):
else SamplingParams.deserialize(obj["inference_parameters"])
)
- # Deserialize tensors and sampling params.
+ # Deserialize tensors, numpy arrays, and sampling params.
for k, v in obj.items():
if isinstance(v, list) and len(v) == 2 and v[0] == "tensor":
setattr(self, k, deserialize_tensor(v[1]))
+ elif isinstance(v, list) and len(v) == 2 and v[0] == "ndarray":
+ setattr(self, k, deserialize_ndarray(v[1]))
class DynamicInferenceEventType(Enum):
@@ -299,7 +307,7 @@ def serialize(self) -> dict:
Returns:
dict: Full event dict.
"""
- torch.cuda.nvtx.range_push("DynamicInferenceEvent.serialize")
+ nvtx_range_push("DynamicInferenceEvent.serialize")
# do not use asdict(self) - it has very high CPU overheads
# and if there are tensors, it will try to deepcopy them
obj = self.__dict__.copy()
@@ -315,7 +323,7 @@ def serialize(self) -> dict:
obj["payload"] = ContextErrorFactory.serialize(self.payload)
- torch.cuda.nvtx.range_pop()
+ nvtx_range_pop("DynamicInferenceEvent.serialize")
return obj
@classmethod
@@ -361,9 +369,8 @@ class DynamicInferenceRequest(InferenceRequest):
policy_epoch: Optional[list[tuple[int, int]]] = None
kv_cache_epoch: Optional[list[tuple[int, int]]] = None
latency: Optional[float] = None
- # routing_indices stores MoE routing decisions for all tokens generated so far.
- # Shape: [total_tokens, num_layers, topk] - accumulated across all generation steps
- routing_indices: Optional[torch.Tensor] = None
+ # routing_indices is reconstructed from per-block storage when a request finishes.
+ routing_indices: Optional[np.ndarray] = None
finished_chunk_token_count: int = 0
stop_word_ids: Optional[List[List[int]]] = None # Tokenized stop words (populated internally)
@@ -429,12 +436,12 @@ def serialize(self):
(dict) A dictionary representation of the instance suitable for
serialization.
"""
- torch.cuda.nvtx.range_push("DynamicInferenceRequest.serialize")
+ nvtx_range_push("DynamicInferenceRequest.serialize")
obj = super().serialize()
obj["events"] = [e.serialize() for e in self.events]
obj.pop("event_add_engine", None)
- # Sanity check routing_indices: Tensor [total_tokens - 1, num_layers, topk]
+ # Sanity check routing_indices: ndarray [total_tokens - 1, num_layers, topk]
if self.routing_indices is not None:
total_tokens = len(self.prompt_tokens) + len(self.generated_tokens)
# the last generated token does not undergo a forward pass
@@ -444,7 +451,7 @@ def serialize(self):
f"total tokens {total_tokens-1}."
)
- torch.cuda.nvtx.range_pop()
+ nvtx_range_pop("DynamicInferenceRequest.serialize")
return obj
def _post_deserialize(self, obj):
@@ -469,26 +476,25 @@ def tracked_metadata(self) -> List[Any]:
"in its sampling_params. Defaulting to -1."
)
sp.termination_id = -1
- return [getattr(sp, field) for field, _, _ in self.get_metadata_types()]
+ return [getattr(sp, field) for field, _ in self.get_metadata_types()]
@staticmethod
- def get_metadata_types() -> List[Tuple[str, torch.dtype, bool]]:
- """Keeps track of all request metadata names, dtypes, and target device.
+ def get_metadata_types() -> List[Tuple[str, torch.dtype]]:
+ """Keeps track of all request metadata names and dtypes.
Returns:
- List[Tuple[str, torch.dtype, bool]]: Mapping from metadata name to:
+ List[Tuple[str, torch.dtype]]: Mapping from metadata name to:
name (str) - The name of the metadata field.
dtype (torch.dtype) - The datatype of the metadata.
- on_device (bool) - Whether the metadata lives on GPU (True) or CPU (False).
"""
return [
- ("temperature", torch.float32, False), # CPU for torch sampling
- ("top_k", torch.int32, False), # CPU for torch sampling
- ("top_p", torch.float32, False), # CPU for torch sampling
- ("termination_id", torch.int64, True),
- ("return_log_probs", torch.bool, False), # CPU for non-selective logprobs
- ("skip_prompt_log_probs", torch.bool, False), # CPU for non-selective logprobs
- ("top_n_logprobs", torch.int32, False), # CPU for torch sampling
+ ("temperature", torch.float32),
+ ("top_k", torch.int32),
+ ("top_p", torch.float32),
+ ("termination_id", torch.int64),
+ ("return_log_probs", torch.bool),
+ ("skip_prompt_log_probs", torch.bool),
+ ("top_n_logprobs", torch.int32),
]
def add_event(
@@ -695,8 +701,9 @@ def merge_lists(key):
prompt_tokens = self.requests[0].prompt_tokens
prompt_text = self.requests[0].prompt
routing_indices = None
- if self.requests[0].routing_indices is not None:
- routing_indices = torch.cat([r.routing_indices for r in self.requests])
+ routing_parts = [r.routing_indices for r in self.requests if r.routing_indices is not None]
+ if routing_parts:
+ routing_indices = np.concatenate(routing_parts)
generated_tokens = merge_lists("generated_tokens")
try:
generated_text = "".join(r.generated_text for r in self.requests)
@@ -741,10 +748,10 @@ def serialize(self) -> dict:
(dict) A dictionary representation of the instance suitable for
serialization.
"""
- torch.cuda.nvtx.range_push("DynamicInferenceRequestRecord.serialize")
+ nvtx_range_push("DynamicInferenceRequestRecord.serialize")
obj = self.__dict__.copy() # shallow dict copy
obj["requests"] = [r.serialize() for r in obj["requests"]]
- torch.cuda.nvtx.range_pop()
+ nvtx_range_pop("DynamicInferenceRequestRecord.serialize")
return obj
@classmethod
diff --git a/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py b/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py
index 55efb24cb08..5fbbcc376f3 100644
--- a/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py
+++ b/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py
@@ -126,34 +126,6 @@ def _forward(self, inference_input):
runtime_gather_output=True, # Inference should always gather the logits
)
- @torch.inference_mode()
- def dummy_forward(self):
- """Run a dummy forward pass through the model, with a single token.
- Use-case: Used in EP on ranks which do not have any work, but are needed
- for the all-to-all communication.
- Runs under inference_mode so that transformer layers can distinguish this eager
- dummy_forward from training/validation passes and skip matching on CUDA graphs."""
-
- # we use num_dummy_tokens equal to tensor model parallel size
- # so that the dummy forward pass will work with sequence parallel
- num_dummy_tokens = self.tp_size
- tokens = torch.zeros(
- (1, num_dummy_tokens), dtype=torch.long, device=torch.cuda.current_device()
- )
- position_ids = torch.zeros(
- (1, num_dummy_tokens), dtype=torch.long, device=torch.cuda.current_device()
- )
- attention_mask = None
- # Always skip MTP during dummy forwards. When num_speculative_tokens > 0
- # the serial MTP path handles MTP separately (with its own dummy forward).
- # When num_speculative_tokens == 0 MTP is not needed at all. In both
- # cases, running MTP here would issue MoE all-to-all collectives that the
- # real EP ranks do not execute, causing a hang.
- is_spec_decode = (
- self.inference_context.is_dynamic_batching() and self.config.mtp_num_layers is not None
- )
- return self.model(tokens, position_ids, attention_mask, is_spec_decode=is_spec_decode)
-
def _get_batch_size_and_seq_len(
self, tokens: torch.Tensor, recv_buffer_seq_len: Optional[int] = None
):
diff --git a/megatron/core/inference/moe/__init__.py b/megatron/core/inference/moe/__init__.py
index dbbb24f07bf..cc64fb65110 100644
--- a/megatron/core/inference/moe/__init__.py
+++ b/megatron/core/inference/moe/__init__.py
@@ -2,55 +2,17 @@
import enum
-import torch
-
from .fused_moe import ActivationType, mcore_fused_moe
+from .vllm_fused_moe import vllm_fused_moe
class InferenceGroupedGemmBackend(enum.Enum):
- """Resolved backend for grouped GEMM operations during inference."""
+ """Backend for grouped GEMM operations during inference.
+
+ The string value matches the inference_grouped_gemm_backend config field so
+ TransformerConfig.__post_init__ can convert via InferenceGroupedGemmBackend(str).
+ """
FLASHINFER = "flashinfer"
TORCH = "torch"
- TE = "te"
-
-
-def resolve_inference_grouped_gemm_backend(
- backend: str, is_cuda_graphed: bool, is_mxfp8: bool = False
-) -> InferenceGroupedGemmBackend:
- """Resolve the grouped GEMM backend to use for the current iteration.
-
- Prerequisites are validated at init time in MoELayer; this function
- simply maps (backend, is_cuda_graphed) to the concrete backend enum.
-
- Args:
- backend: One of 'auto', 'torch', 'te'.
- is_cuda_graphed: Whether this is a CUDA-graphed iteration.
- is_mxfp8: Whether the model is using MXFP8 quantization (affects auto backend choice).
- Returns:
- An InferenceGroupedGemmBackend enum value.
- """
- if backend == 'auto':
- if is_mxfp8:
- assert hasattr(torch.nn.functional, 'scaled_grouped_mm'), (
- "Auto backend selection for MXFP8 requires "
- "torch.nn.functional.scaled_grouped_mm. "
- "Please install PyTorch 2.10+."
- )
- return InferenceGroupedGemmBackend.TORCH
- if is_cuda_graphed:
- return InferenceGroupedGemmBackend.FLASHINFER
- else:
- if hasattr(torch.nn.functional, 'grouped_mm'):
- return InferenceGroupedGemmBackend.TORCH
- else:
- return InferenceGroupedGemmBackend.TE
- elif backend == 'torch':
- return InferenceGroupedGemmBackend.TORCH
- elif backend == 'te':
- return InferenceGroupedGemmBackend.TE
- else:
- raise ValueError(
- f"Unknown inference_grouped_gemm_backend: '{backend}'. "
- "Must be 'auto', 'torch', or 'te'."
- )
+ VLLM = "vllm"
diff --git a/megatron/core/inference/moe/activations.py b/megatron/core/inference/moe/activations.py
index 169d8499116..ae5e4560ce3 100644
--- a/megatron/core/inference/moe/activations.py
+++ b/megatron/core/inference/moe/activations.py
@@ -30,25 +30,53 @@ def _ceil_div(a, b):
@triton.jit
-def _squared_relu_kernel(input_ptr, output_ptr, src_idx_ptr, M, N, BLOCK_N: tl.constexpr):
- """Squared ReLU that skips padding rows (permutation_map == -1)."""
- row = tl.program_id(0)
- if tl.load(src_idx_ptr + row) < 0:
- return
- for n in tl.range(0, N, BLOCK_N):
- o = n + tl.arange(0, BLOCK_N)
- m = o < N
- x = tl.load(input_ptr + row * N + o, mask=m).to(tl.float32)
- r = tl.maximum(x, 0.0)
- tl.store(output_ptr + row * N + o, (r * r).to(tl.bfloat16), mask=m)
+def _squared_relu_kernel(
+ input_ptr,
+ output_ptr,
+ src_idx_ptr,
+ n_used_ptr,
+ N,
+ max_rows, # output_size (fixed for CG)
+ BLOCK_N: tl.constexpr,
+ NUM_BLOCKS: tl.constexpr, # grid size (fixed for CG)
+):
+ """Squared ReLU that skips rows beyond n_used and alignment-padding rows (perm_map == -1).
+ Grid: fixed NUM_BLOCKS CTAs, each iterating over multiple rows.
+ n_used_ptr gates how many rows are processed — required for CUDA graph compatibility.
+ """
+ pid = tl.program_id(0)
+ n_used = tl.load(n_used_ptr)
+ if pid >= n_used:
+ return
+ for row in tl.range(pid, max_rows, NUM_BLOCKS):
+ if row < n_used:
+ if tl.load(src_idx_ptr + row) >= 0:
+ for n in tl.range(0, N, BLOCK_N):
+ o = n + tl.arange(0, BLOCK_N)
+ m = o < N
+ x = tl.load(input_ptr + row * N + o, mask=m).to(tl.float32)
+ r = tl.maximum(x, 0.0)
+ tl.store(output_ptr + row * N + o, (r * r).to(tl.bfloat16), mask=m)
+
+
+def padded_squared_relu(
+ x: torch.Tensor, permutation_map: torch.Tensor, n_used: torch.Tensor
+) -> torch.Tensor:
+ """Squared ReLU activation that skips rows beyond n_used and alignment-padding rows.
-def padded_squared_relu(x: torch.Tensor, permutation_map: torch.Tensor) -> torch.Tensor:
- """Squared ReLU activation that skips padding rows."""
+ Args:
+ x: [output_size, ffn_hidden] BF16 FC1 output.
+ permutation_map: [output_size] int32, original token index or -1 for padding.
+ n_used: scalar int32 CUDA tensor = inclusive_expert_offsets[-1].
+ """
M, N = x.shape
- out = torch.zeros(M, N, dtype=x.dtype, device=x.device)
+ out = torch.empty(M, N, dtype=x.dtype, device=x.device)
BLOCK_N = min(triton.next_power_of_2(N), 1024)
- _squared_relu_kernel[(M,)](x, out, permutation_map, M, N, BLOCK_N=BLOCK_N)
+ NUM_BLOCKS = min(M, 512)
+ _squared_relu_kernel[(NUM_BLOCKS,)](
+ x, out, permutation_map, n_used, N, M, BLOCK_N=BLOCK_N, NUM_BLOCKS=NUM_BLOCKS
+ )
return out
@@ -58,68 +86,71 @@ def _squared_relu_quantize_kernel(
out_fp8_ptr,
out_scale_ptr,
src_idx_ptr,
+ n_used_ptr, # pointer to inclusive_expert_offsets[-1]: number of used rows this iteration
K,
n_col_blocks,
- skip_padding: tl.constexpr,
+ max_rows, # output_size (fixed for CG)
REAL_GROUPS: tl.constexpr,
BLOCK_K: tl.constexpr,
BLOCK_GROUPS: tl.constexpr,
+ NUM_BLOCKS: tl.constexpr, # grid size (fixed for CG)
):
"""Fused squared ReLU + MXFP8 quantize + swizzle in one kernel.
- Grid: (M,) — one program per row.
- Reads BF16 FC1 output, applies squared ReLU, quantizes to FP8,
- writes FP8 data + swizzled scales in place.
+ Grid: fixed NUM_BLOCKS CTAs, each iterating over multiple rows.
+ Rows beyond n_used and alignment-padding rows (perm_map == -1) are skipped.
"""
- row = tl.program_id(0)
- if skip_padding:
- if tl.load(src_idx_ptr + row) < 0:
- return
-
- offs = tl.arange(0, BLOCK_K)
- mask = offs < K
-
- # Load and apply squared ReLU
- x = tl.load(input_ptr + row * K + offs, mask=mask, other=0.0).to(tl.float32)
- relu = tl.maximum(x, 0.0)
- activated = relu * relu
-
- # Per-group-of-32 quantization
- x_grouped = tl.reshape(activated, [BLOCK_GROUPS, 32])
- abs_grouped = tl.abs(x_grouped)
- max_vals = tl.max(abs_grouped, axis=1)
-
- dequant_scale = max_vals / 448.0
- dequant_exp = (dequant_scale.to(tl.uint32, bitcast=True) + 0x007FFFFF) & 0x7F800000
- dequant_rounded = dequant_exp.to(tl.float32, bitcast=True)
- quant_scale = tl.where(dequant_rounded == 0, 0.0, 1.0 / dequant_rounded)
-
- quantized = x_grouped * quant_scale[:, None]
- quantized_flat = tl.reshape(quantized, [BLOCK_K])
- out_fp8 = quantized_flat.to(tl.float8e4nv)
-
- # Store FP8 data
- tl.store(out_fp8_ptr + row * K + offs, out_fp8, mask=mask)
-
- # Store swizzled scales
- scale_exp = (dequant_exp >> 23).to(tl.uint8)
- col_offs = tl.arange(0, BLOCK_GROUPS)
- col_mask = col_offs < REAL_GROUPS
-
- macro_row_block = row // 128
- macro_col_block = col_offs // 4
- local_row = row % 128
- local_col = col_offs % 4
- group = local_row // 32
- sub_row = local_row % 32
- tile_idx = macro_row_block * n_col_blocks + macro_col_block
- swizzled_offs = tile_idx * 512 + sub_row * 16 + group * 4 + local_col
-
- tl.store(out_scale_ptr + swizzled_offs, scale_exp, mask=col_mask)
+ pid = tl.program_id(0)
+ n_used = tl.load(n_used_ptr)
+ if pid >= n_used:
+ return
+ for row in tl.range(pid, max_rows, NUM_BLOCKS):
+ if row < n_used:
+ if tl.load(src_idx_ptr + row) >= 0:
+ offs = tl.arange(0, BLOCK_K)
+ mask = offs < K
+
+ # Load and apply squared ReLU
+ x = tl.load(input_ptr + row * K + offs, mask=mask, other=0.0).to(tl.float32)
+ relu = tl.maximum(x, 0.0)
+ activated = relu * relu
+
+ # Per-group-of-32 quantization
+ x_grouped = tl.reshape(activated, [BLOCK_GROUPS, 32])
+ abs_grouped = tl.abs(x_grouped)
+ max_vals = tl.max(abs_grouped, axis=1)
+
+ dequant_scale = max_vals / 448.0
+ dequant_exp = (dequant_scale.to(tl.uint32, bitcast=True) + 0x007FFFFF) & 0x7F800000
+ dequant_rounded = dequant_exp.to(tl.float32, bitcast=True)
+ quant_scale = tl.where(dequant_rounded == 0, 0.0, 1.0 / dequant_rounded)
+
+ quantized = x_grouped * quant_scale[:, None]
+ quantized_flat = tl.reshape(quantized, [BLOCK_K])
+ out_fp8 = quantized_flat.to(tl.float8e4nv)
+
+ # Store FP8 data
+ tl.store(out_fp8_ptr + row * K + offs, out_fp8, mask=mask)
+
+ # Store swizzled scales
+ scale_exp = (dequant_exp >> 23).to(tl.uint8)
+ col_offs = tl.arange(0, BLOCK_GROUPS)
+ col_mask = col_offs < REAL_GROUPS
+
+ macro_row_block = row // 128
+ macro_col_block = col_offs // 4
+ local_row = row % 128
+ local_col = col_offs % 4
+ group = local_row // 32
+ sub_row = local_row % 32
+ tile_idx = macro_row_block * n_col_blocks + macro_col_block
+ swizzled_offs = tile_idx * 512 + sub_row * 16 + group * 4 + local_col
+
+ tl.store(out_scale_ptr + swizzled_offs, scale_exp, mask=col_mask)
def squared_relu_and_quantize_mxfp8(
- x: torch.Tensor, permutation_map: torch.Tensor, skip_padding: bool = True
+ x: torch.Tensor, permutation_map: torch.Tensor, n_used: torch.Tensor
):
"""Fused squared ReLU + MXFP8 quantize + swizzle.
@@ -127,12 +158,13 @@ def squared_relu_and_quantize_mxfp8(
swizzled scales. Single kernel replaces padded_squared_relu + mxfp8_quantize.
Args:
- x: [M, K] BF16 FC1 output.
- permutation_map: [M] int32, original token index or -1 for padding.
- skip_padding: if True, skip rows where permutation_map == -1.
+ x: [output_size, K] BF16 FC1 output.
+ permutation_map: [output_size] int32, original token index or -1 for padding.
+ n_used: scalar int32 CUDA tensor = inclusive_expert_offsets[-1]. Rows beyond
+ this are skipped before even checking the permutation_map.
Returns:
- MXFP8Tensor with .data [M, K] float8_e4m3fn and .scale (swizzled e8m0).
+ MXFP8Tensor with .data [output_size, K] float8_e4m3fn and .scale (swizzled e8m0).
"""
from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor
@@ -149,18 +181,21 @@ def squared_relu_and_quantize_mxfp8(
BLOCK_K = triton.next_power_of_2(K)
BLOCK_GROUPS = BLOCK_K // 32
+ NUM_BLOCKS = min(M, 512)
- _squared_relu_quantize_kernel[(M,)](
+ _squared_relu_quantize_kernel[(NUM_BLOCKS,)](
x,
out_fp8,
out_scale,
permutation_map,
+ n_used,
K,
n_col_blocks,
- skip_padding,
+ M,
REAL_GROUPS=scale_cols,
BLOCK_K=BLOCK_K,
BLOCK_GROUPS=BLOCK_GROUPS,
+ NUM_BLOCKS=NUM_BLOCKS,
)
return MXFP8Tensor(data=out_fp8, scale=out_scale.view(torch.float8_e8m0fnu), backend="triton")
diff --git a/megatron/core/inference/moe/fused_moe.py b/megatron/core/inference/moe/fused_moe.py
index 39382eee079..f6c0af4e94e 100644
--- a/megatron/core/inference/moe/fused_moe.py
+++ b/megatron/core/inference/moe/fused_moe.py
@@ -6,7 +6,7 @@
"""
from enum import Enum
-from typing import Callable, Optional
+from typing import Callable
import torch
@@ -14,7 +14,6 @@
padded_squared_relu,
squared_relu_and_quantize_mxfp8,
)
-from megatron.core.inference.moe.pad import pad_to_alignment, unpad_from_alignment
from megatron.core.inference.moe.permute import (
permute_and_quantize_mxfp8,
permute_tokens,
@@ -27,7 +26,9 @@
HAVE_GROUPED_MM = True
except ImportError:
- HAVE_GROUPED_MM = False
+ # Fallback to the private symbol for torch versions < 2.10.
+ grouped_mm = getattr(torch, "_grouped_mm", None)
+ HAVE_GROUPED_MM = grouped_mm is not None
try:
from torch.nn.functional import ScalingType, SwizzleType, scaled_grouped_mm
@@ -86,49 +87,46 @@ def mcore_fused_moe(
activation_type: ActivationType,
num_local_experts: int,
local_expert_start: int,
- routing_map: Optional[torch.Tensor] = None,
- tokens_per_expert: Optional[torch.Tensor] = None,
- skip_permute: bool = False,
+ valid_tokens: torch.Tensor,
+ routing_map: torch.Tensor,
disable_fused_quant_kernels: bool = False,
+ out: torch.Tensor = None,
) -> torch.Tensor:
- """Fused MoE: [permute ->] pad -> FC1 -> activation -> FC2 -> unpad [-> unpermute].
-
- Two modes:
- - skip_permute=False (default): tokens are unpermuted. Requires routing_map.
- Performs full permute -> compute -> unpermute.
- - skip_permute=True: tokens are already permuted by the dispatcher. Requires
- tokens_per_expert. Pads to alignment, computes, then unpads. Probs are
- applied during unpad.
+ """Fused MoE: permute -> pad -> FC1 -> activation -> FC2 -> unpad -> unpermute.
Unless disable_fused_quant_kernels=True, when weights are MXFP8, uses fused
kernels that combine permute/activation with MXFP8 quantization into single
kernel launches.
Args:
- hidden_states: [num_tokens, hidden_size] BF16 input.
- probs: routing probabilities. Shape is [num_tokens, topk] when
- skip_permute=False, or [num_tokens] (already gathered) when
- skip_permute=True.
+ hidden_states: [max_tokens, hidden_size] BF16 input. max_tokens =
+ max_local_tokens * ep_size; only the first valid_tokens rows are valid.
+ probs: [max_tokens, topk] routing probabilities.
fc1_weight: stacked weight for FC1 (torch.Tensor for BF16, MXFP8Tensor for MXFP8).
fc2_weight: stacked weight for FC2 (same type as fc1_weight).
activation_type: ActivationType enum (SQUARED_RELU).
num_local_experts: number of experts on this rank.
local_expert_start: first global expert index on this rank.
- routing_map: [num_tokens, topk] int expert assignments. Required when skip_permute=False.
- tokens_per_expert: [num_local_experts] int32 token counts. Required when skip_permute=True.
- skip_permute: if True, skip permute/unpermute (tokens already in expert order).
+ valid_tokens: scalar int32 CUDA tensor holding the number of valid tokens this
+ iteration. Kernels use this to ignore rows beyond the valid prefix — required
+ for CUDA graph compatibility since hidden_states is always max-sized.
+ routing_map: [max_tokens, topk] int expert assignments.
disable_fused_quant_kernels: if True, disable fused permute+quantize and
activation+quantize kernels for MXFP8, using separate launches instead.
Useful for debugging. Ignored when weights are BF16.
+ out: optional pre-allocated output buffer. If provided, unpermute writes
+ directly into this tensor (e.g. the RSV symmetric buffer), avoiding a
+ separate copy before reduce-scatter.
Returns:
- [num_tokens, hidden_size] BF16 output.
+ [max_tokens, hidden_size] BF16 output. Only the first valid_tokens rows are
+ meaningful; rows beyond that are undefined.
"""
assert (
hidden_states.dtype == torch.bfloat16
), f"mcore_fused_moe requires bf16 input, got {hidden_states.dtype}"
- num_tokens = hidden_states.shape[0]
+ max_tokens = hidden_states.shape[0]
use_mxfp8 = isinstance(fc1_weight, MXFP8Tensor)
# Fused quant kernels only apply to MXFP8 path
use_fused_quant = use_mxfp8 and not disable_fused_quant_kernels
@@ -151,54 +149,47 @@ def mcore_fused_moe(
activation_func = _get_activation_func(activation_type, fused_quant=use_fused_quant)
- # --- Pre-processing: permute or pad ---
- if skip_permute:
- assert tokens_per_expert is not None, "tokens_per_expert is required when skip_permute=True"
- tokens_per_expert = tokens_per_expert.cuda().int()
- assert routing_map is None, "routing_map must be None when skip_permute=True"
- hidden_states, permutation_map, offs = pad_to_alignment(
- hidden_states, tokens_per_expert, expert_alignment
+ # --- Pre-processing: permute ---
+ if use_fused_quant:
+ # Fused permute + MXFP8 quantize: single kernel produces MXFP8Tensor
+ hidden_states, permuted_probs, permutation_map, offs = permute_and_quantize_mxfp8(
+ hidden_states,
+ probs,
+ routing_map,
+ local_expert_start,
+ num_local_experts,
+ valid_tokens,
+ alignment=expert_alignment,
)
- permuted_probs = None
-
else:
- assert routing_map is not None, "routing_map is required when skip_permute=False"
- if use_fused_quant:
- # Fused permute + MXFP8 quantize: single kernel produces MXFP8Tensor
- hidden_states, permuted_probs, permutation_map, offs = permute_and_quantize_mxfp8(
- hidden_states,
- probs,
- routing_map,
- local_expert_start,
- num_local_experts,
- alignment=expert_alignment,
- )
- else:
- hidden_states, permuted_probs, permutation_map, offs = permute_tokens(
- hidden_states,
- probs,
- routing_map,
- local_expert_start,
- num_local_experts,
- alignment=expert_alignment,
- )
+ hidden_states, permuted_probs, permutation_map, offs = permute_tokens(
+ hidden_states,
+ probs,
+ routing_map,
+ local_expert_start,
+ num_local_experts,
+ valid_tokens,
+ alignment=expert_alignment,
+ )
# --- FC1 -> activation -> FC2 ---
# Quantize if MXFP8 path and hidden_states not already quantized (fused permute+quant
- # produces MXFP8Tensor directly; skip_permute path always needs separate quant).
- needs_quant = use_mxfp8 and not isinstance(hidden_states, MXFP8Tensor)
- if needs_quant:
+ # produces MXFP8Tensor directly).
+ if use_mxfp8 and not isinstance(hidden_states, MXFP8Tensor):
hidden_states = MXFP8Tensor.from_bf16(hidden_states, backend="triton")
fc1_output = mm_fn(hidden_states, fc1_weight, offs)
- activation_out = activation_func(fc1_output, permutation_map)
+ # offs[-1:] is a 1-element view pointing to inclusive_expert_offsets[-1] — the total
+ # number of rows actually used by experts this iteration (valid tokens + alignment
+ # padding within expert blocks). Passed to activation and unpermute to skip unused rows.
+ n_used = offs[-1:]
+ activation_out = activation_func(fc1_output, permutation_map, n_used)
# Fused activation+quant returns MXFP8Tensor; otherwise quantize separately.
if use_mxfp8 and not isinstance(activation_out, MXFP8Tensor):
activation_out = MXFP8Tensor.from_bf16(activation_out, backend="triton")
fc2_output = mm_fn(activation_out, fc2_weight, offs)
- # --- Post-processing: unpermute or unpad ---
- if skip_permute:
- probs_1d = probs.squeeze(-1) if probs.dim() > 1 else probs
- return unpad_from_alignment(fc2_output, permutation_map, num_tokens, probs=probs_1d)
- else:
- return unpermute_tokens(fc2_output, permuted_probs, permutation_map, num_tokens)
+
+ # --- Post-processing: unpermute ---
+ return unpermute_tokens(
+ fc2_output, permuted_probs, permutation_map, max_tokens, n_used, valid_tokens, out=out
+ )
diff --git a/megatron/core/inference/moe/metadata.py b/megatron/core/inference/moe/metadata.py
new file mode 100644
index 00000000000..8658ab9b42a
--- /dev/null
+++ b/megatron/core/inference/moe/metadata.py
@@ -0,0 +1,134 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Fused NVLS metadata update kernel for MoE expert parallelism.
+
+Replaces the multi-kernel sequence:
+ dist.all_gather_into_tensor(...) # NCCL
+ local_tokens_per_rank.sum() # kernel
+ local_tokens_per_rank[:rank].sum() # kernel
+ local_tokens_per_rank.max() # kernel
+ _step_metadata.copy_(...) # kernel
+
+with a single Triton kernel that:
+ 1. Multicast-stores this rank's local_tokens to the symmetric memory buffer.
+ 2. Barrier (all ranks have written).
+ 3. Reads all ranks' counts, computes sum / prefix-sum / max.
+ 4. Writes the 3-element step_metadata tensor in-place.
+"""
+
+from unittest.mock import MagicMock
+
+import torch
+
+from megatron.core.utils import null_decorator
+
+try:
+ import triton
+ import triton.language as tl
+
+ HAVE_TRITON = True
+except ImportError:
+ triton = MagicMock()
+ triton.jit = null_decorator
+ tl = MagicMock()
+ HAVE_TRITON = False
+
+try:
+ from torch._C._distributed_c10d import _SymmetricMemory
+except ImportError:
+ _SymmetricMemory = MagicMock()
+
+from megatron.core.inference.communication.torch_symm_triton.barrier import symm_mem_sync
+from megatron.core.inference.communication.torch_symm_triton.multimem_asm import st_32
+from megatron.core.inference.communication.torch_symm_triton.utils import sync_threads
+
+
+@triton.jit
+def _fused_metadata_kernel(
+ local_tokens,
+ local_buf_ptr,
+ multicast_ptr,
+ signal_pad_ptrs,
+ step_metadata_ptr,
+ RANK: tl.constexpr,
+ WORLD_SIZE: tl.constexpr,
+):
+ """Fused allgather + reduce kernel for MoE step metadata.
+
+ Single CTA. Writes this rank's local_tokens to the symmetric buffer
+ via multicast store, barriers, then reads all ranks' values from the
+ local buffer and computes [valid_tokens, rank_token_offset, ep_max_tokens].
+
+ Args:
+ local_tokens: scalar int32, this rank's token count.
+ local_buf_ptr: pointer to the local symmetric memory buffer (for reads).
+ multicast_ptr: multicast pointer to the symmetric memory buffer (for writes).
+ signal_pad_ptrs: signal pads for barrier synchronization.
+ step_metadata_ptr: pointer to the 3-element int32 output tensor.
+ RANK: this rank's index (constexpr).
+ WORLD_SIZE: total number of ranks (constexpr).
+ """
+
+ tid = tl.program_id(0)
+ if tid > 0:
+ return
+
+ # 1. Multicast-store local_tokens to buffer[RANK].
+ mc_ptr = multicast_ptr.to(tl.pointer_type(tl.uint32)) + RANK
+ mask = tl.full([], 1, dtype=tl.int1)
+ val = tl.full([], local_tokens, dtype=tl.uint32)
+ st_32(mc_ptr, val, mask, multicast_op=True)
+
+ # 2. Barrier — wait for all ranks to have written.
+ sync_threads()
+ symm_mem_sync(
+ signal_pad_ptrs,
+ None,
+ RANK,
+ WORLD_SIZE,
+ hasPreviousMemAccess=True,
+ hasSubsequentMemAccess=True,
+ )
+
+ # 3. Load all ranks' values, reduce, and write metadata.
+ offsets = tl.arange(0, WORLD_SIZE)
+ vals = tl.load(local_buf_ptr + offsets)
+
+ total = tl.sum(vals)
+ prefix = tl.sum(tl.where(offsets < RANK, vals, tl.zeros_like(vals)))
+ max_val = tl.max(vals)
+
+ tl.store(step_metadata_ptr, total)
+ tl.store(step_metadata_ptr + 1, prefix)
+ tl.store(step_metadata_ptr + 2, max_val)
+
+
+def fused_metadata_update(
+ local_tokens: int,
+ local_buf: torch.Tensor,
+ symm_mem_hdl: _SymmetricMemory,
+ step_metadata: torch.Tensor,
+) -> None:
+ """Fused NVLS allgather + reduce for MoE step metadata.
+
+ Args:
+ local_tokens: number of tokens on this rank this step.
+ local_buf: the local symmetric memory buffer tensor ([WORLD_SIZE] int32).
+ Used for reads after the barrier.
+ symm_mem_hdl: symmetric memory handle for the metadata buffer.
+ Provides the multicast pointer for writes and signal pads for barrier.
+ step_metadata: [3] int32 CUDA tensor to write
+ [valid_tokens, rank_token_offset, ep_max_tokens] into.
+ """
+ assert HAVE_TRITON, "Triton is required for fused_metadata_update."
+
+ _fused_metadata_kernel[(1, 1, 1)](
+ local_tokens,
+ local_buf,
+ symm_mem_hdl.multicast_ptr,
+ symm_mem_hdl.signal_pad_ptrs_dev,
+ step_metadata,
+ RANK=symm_mem_hdl.rank,
+ WORLD_SIZE=symm_mem_hdl.world_size,
+ num_warps=min(max(1, (symm_mem_hdl.world_size + 31) // 32), 8),
+ )
diff --git a/megatron/core/inference/moe/pad.py b/megatron/core/inference/moe/pad.py
deleted file mode 100644
index 656953b691c..00000000000
--- a/megatron/core/inference/moe/pad.py
+++ /dev/null
@@ -1,201 +0,0 @@
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-"""Pad / unpad utilities for already-permuted expert tokens.
-
-When the token dispatcher has already permuted tokens into expert-grouped
-order, these functions insert/remove alignment padding so that each expert's
-token block satisfies the alignment requirements of grouped_mm /
-scaled_grouped_mm.
-"""
-
-from unittest.mock import MagicMock
-
-import torch
-from packaging import version
-
-from megatron.core.utils import null_decorator
-
-try:
- import triton
- import triton.language as tl
-
- if version.parse(triton.__version__) < version.parse("3.4.0") and not torch.cuda.is_available():
- HAVE_TRITON = False
- else:
- HAVE_TRITON = tl.constexpr(version.parse(triton.__version__) >= version.parse("2.0.0"))
-except ImportError:
- HAVE_TRITON = False
-
-if not HAVE_TRITON:
- triton = MagicMock()
- triton.jit = null_decorator
- tl = MagicMock()
-
-from megatron.core.inference.moe.permute import compute_expert_offsets
-
-
-@triton.jit
-def _pad_tokens_kernel(
- src_ptr,
- dst_ptr,
- perm_map_ptr,
- tpe_ptr, # tokens_per_expert [num_experts]
- hidden_dim,
- num_experts: tl.constexpr,
- alignment: tl.constexpr,
- BLOCK_H: tl.constexpr,
-):
- """Copy one input row into the padded output buffer.
-
- Computes unpadded and padded cumulative offsets inline from
- tokens_per_expert, avoiding a separate cumsum kernel launch.
- """
- row = tl.program_id(0)
-
- # Walk tokens_per_expert to find which expert this row belongs to
- # and compute both unpadded and padded start offsets on the fly.
- unpadded_start = tl.zeros([], dtype=tl.int32)
- padded_start = tl.zeros([], dtype=tl.int32)
- expert_id = -1
- for e in tl.static_range(0, num_experts):
- count = tl.load(tpe_ptr + e).to(tl.int32)
- if expert_id < 0 and row < unpadded_start + count:
- expert_id = e
- if expert_id < 0:
- unpadded_start += count
- aligned = tl.where(
- count > 0,
- ((count + alignment - 1) // alignment) * alignment,
- tl.zeros([], dtype=tl.int32),
- )
- padded_start += aligned
-
- if expert_id < 0:
- return
-
- local_idx = row - unpadded_start
- dst_row = padded_start + local_idx
-
- # Write permutation_map: padded row → original unpadded row
- tl.store(perm_map_ptr + dst_row, row)
-
- # Copy hidden state
- for h in tl.range(0, hidden_dim, BLOCK_H):
- o = h + tl.arange(0, BLOCK_H)
- m = o < hidden_dim
- tl.store(
- dst_ptr + dst_row * hidden_dim + o,
- tl.load(src_ptr + row * hidden_dim + o, mask=m),
- mask=m,
- )
-
-
-def pad_to_alignment(
- hidden_states: torch.Tensor, tokens_per_expert: torch.Tensor, alignment: int
-) -> tuple:
- """Pad already-permuted tokens so each expert's block is aligned.
-
- Args:
- hidden_states: [total_tokens, hidden_size] already permuted by dispatcher.
- tokens_per_expert: [num_local_experts] int32 token counts.
- alignment: per-expert alignment.
-
- Returns:
- (padded_hidden, permutation_map, inclusive_offsets)
- - padded_hidden: [padded_total, hidden_size]
- - permutation_map: [padded_total] int32, original row index or -1 for padding.
- - inclusive_offsets: [num_local_experts] int32 cumulative aligned offsets for grouped_mm.
- """
- num_experts = tokens_per_expert.shape[0]
- total_tokens = hidden_states.shape[0]
- hidden_dim = hidden_states.shape[1]
-
- # We still need padded_inc for the return value (used as offs by grouped_mm)
- _, padded_inc = compute_expert_offsets(tokens_per_expert, alignment=alignment)
- padded_total = int(padded_inc[-1].item())
-
- padded_hidden = torch.zeros(
- padded_total, hidden_dim, dtype=hidden_states.dtype, device=hidden_states.device
- )
- permutation_map = torch.full(
- (padded_total,), -1, dtype=torch.int32, device=hidden_states.device
- )
-
- if total_tokens > 0:
- BLOCK_H = min(triton.next_power_of_2(hidden_dim), 1024)
- _pad_tokens_kernel[(total_tokens,)](
- hidden_states,
- padded_hidden,
- permutation_map,
- tokens_per_expert,
- hidden_dim,
- num_experts,
- alignment,
- BLOCK_H=BLOCK_H,
- )
-
- return padded_hidden, permutation_map, padded_inc
-
-
-@triton.jit
-def _unpad_tokens_kernel(
- src_ptr,
- dst_ptr,
- perm_map_ptr,
- probs_ptr,
- hidden_dim,
- has_probs: tl.constexpr,
- BLOCK_H: tl.constexpr,
-):
- """Copy one real (non-padding) row from padded to unpadded layout.
-
- Optionally multiplies each row by its routing probability.
- """
- row = tl.program_id(0)
- dst_row = tl.load(perm_map_ptr + row)
- if dst_row < 0:
- return
- if has_probs:
- prob = tl.load(probs_ptr + dst_row)
- for h in tl.range(0, hidden_dim, BLOCK_H):
- o = h + tl.arange(0, BLOCK_H)
- m = o < hidden_dim
- v = tl.load(src_ptr + row * hidden_dim + o, mask=m)
- if has_probs:
- v = v * prob
- tl.store(dst_ptr + dst_row * hidden_dim + o, v, mask=m)
-
-
-def unpad_from_alignment(
- padded_output: torch.Tensor,
- permutation_map: torch.Tensor,
- original_size: int,
- probs: torch.Tensor = None,
-) -> torch.Tensor:
- """Remove alignment padding, scattering results back to original positions.
-
- Args:
- padded_output: [padded_total, hidden_size] output from expert computation.
- permutation_map: [padded_total] int32, original row index or -1 for padding.
- original_size: number of rows in the unpadded output.
- probs: optional [original_size] routing probabilities to multiply during unpad.
-
- Returns:
- [original_size, hidden_size] unpadded output.
- """
- hidden_dim = padded_output.shape[1]
- output = torch.zeros(
- original_size, hidden_dim, dtype=padded_output.dtype, device=padded_output.device
- )
- has_probs = probs is not None
- if padded_output.shape[0] > 0:
- BLOCK_H = min(triton.next_power_of_2(hidden_dim), 1024)
- _unpad_tokens_kernel[(padded_output.shape[0],)](
- padded_output,
- output,
- permutation_map,
- probs if has_probs else padded_output, # dummy pointer when no probs
- hidden_dim,
- has_probs,
- BLOCK_H=BLOCK_H,
- )
- return output
diff --git a/megatron/core/inference/moe/permute.py b/megatron/core/inference/moe/permute.py
index b14d0b3dbd0..6906c877061 100644
--- a/megatron/core/inference/moe/permute.py
+++ b/megatron/core/inference/moe/permute.py
@@ -8,6 +8,7 @@
- Unpermute expert outputs back to original token order
"""
+from typing import Optional
from unittest.mock import MagicMock
import torch
@@ -28,15 +29,26 @@
tl = MagicMock()
+_NUM_SMS: Optional[int] = None
+
+
+def _get_num_sms(device: torch.device) -> int:
+ global _NUM_SMS
+ if _NUM_SMS is None:
+ _NUM_SMS = torch.cuda.get_device_properties(device).multi_processor_count
+ return _NUM_SMS
+
+
def _ceil_div(a, b):
return (a + b - 1) // b
@triton.jit
def _count_local_tokens_kernel(
- routing_map_ptr, # [num_tokens * topk] flattened expert assignments
+ routing_map_ptr, # [max_tokens, topk] flattened expert assignments
tokens_per_expert_ptr, # [num_local_experts] output counters (zeroed by caller)
- total_pairs, # num_tokens * topk — total (token, topk) pairs
+ valid_tokens_ptr, # scalar int32 CUDA tensor: number of valid tokens this iteration
+ topk, # number of expert choices per token
local_expert_start, # first global expert index owned by this rank
num_local_experts: tl.constexpr, # number of experts on this rank
BLOCK_SIZE: tl.constexpr, # number of pairs processed per program
@@ -45,33 +57,102 @@ def _count_local_tokens_kernel(
Each program processes BLOCK_SIZE (token, topk) pairs. Tokens assigned to
experts outside [local_expert_start, local_expert_start + num_local_experts)
- are silently skipped.
+ or beyond valid_tokens are silently skipped.
+
+ Grid is launched at max size (max_tokens * topk); valid_tokens gates which
+ pairs are actually processed — required for CUDA graph compatibility.
"""
pid = tl.program_id(0)
+ valid_tokens = tl.load(valid_tokens_ptr)
+ valid_pairs = valid_tokens * topk
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
- mask = offsets < total_pairs
+ mask = offsets < valid_pairs
expert_ids = tl.load(routing_map_ptr + offsets, mask=mask, other=-1)
- # Map global expert IDs to local indices; non-local experts become negative
local_ids = expert_ids - local_expert_start
is_local = (local_ids >= 0) & (local_ids < num_local_experts) & mask
tl.atomic_add(tokens_per_expert_ptr + local_ids, 1, mask=is_local)
+@triton.jit
+def _count_local_tokens_kernel_persistent(
+ routing_map_ptr, # [max_tokens, topk] flattened expert assignments
+ tokens_per_expert_ptr, # [num_local_experts] output counters (zeroed by caller)
+ valid_tokens_ptr, # scalar int32 CUDA tensor: number of valid tokens this iteration
+ topk, # number of expert choices per token
+ local_expert_start, # first global expert index owned by this rank
+ num_local_experts: tl.constexpr, # number of experts on this rank
+ num_sms, # number of SMs (grid size for persistent kernel)
+ BLOCK_SIZE: tl.constexpr, # number of pairs processed per iteration
+):
+ """Count tokens routed to local experts using a persistent grid.
+
+ Launches num_sms CTAs. Each CTA loops over its share of BLOCK_SIZE-sized
+ chunks, with total work determined device-side from valid_tokens.
+ """
+ pid = tl.program_id(0)
+ valid_tokens = tl.load(valid_tokens_ptr)
+ valid_pairs = valid_tokens * topk
+
+ total_blocks = tl.cdiv(valid_pairs, BLOCK_SIZE)
+ blocks_per_cta = tl.cdiv(total_blocks, num_sms)
+ block_start = pid * blocks_per_cta
+
+ if block_start < total_blocks:
+ block_end = tl.minimum(block_start + blocks_per_cta, total_blocks)
+
+ for block_id in tl.range(block_start, block_end):
+ offsets = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+ mask = offsets < valid_pairs
+ expert_ids = tl.load(routing_map_ptr + offsets, mask=mask, other=-1)
+ local_ids = expert_ids - local_expert_start
+ is_local = (local_ids >= 0) & (local_ids < num_local_experts) & mask
+ tl.atomic_add(tokens_per_expert_ptr + local_ids, 1, mask=is_local)
+
+
def compute_local_tokens_per_expert(
- routing_map: torch.Tensor, local_expert_start: int, num_local_experts: int
+ routing_map: torch.Tensor,
+ local_expert_start: int,
+ num_local_experts: int,
+ valid_tokens: torch.Tensor,
+ persistent: bool = False,
) -> torch.Tensor:
- """Count tokens routed to each local expert."""
- total_pairs = routing_map.numel()
+ """Count tokens routed to each local expert.
+
+ Args:
+ routing_map: [max_tokens, topk] expert assignments. Only the first
+ valid_tokens rows are processed; the rest are ignored.
+ local_expert_start: first global expert index on this rank.
+ num_local_experts: number of experts on this rank.
+ valid_tokens: scalar int32 CUDA tensor with the number of valid tokens
+ this iteration. Fixed address; value updated each step before graph replay.
+ persistent: use persistent-grid kernel variant (fewer CTAs, looped).
+ """
+ max_pairs = routing_map.numel()
+ topk = routing_map.shape[1]
tokens_per_expert = torch.zeros(num_local_experts, dtype=torch.int32, device=routing_map.device)
- BLOCK = 256
- _count_local_tokens_kernel[(_ceil_div(total_pairs, BLOCK),)](
- routing_map,
- tokens_per_expert,
- total_pairs,
- local_expert_start,
- num_local_experts,
- BLOCK_SIZE=BLOCK,
- )
+ BLOCK = 1024
+ if persistent:
+ num_sms = _get_num_sms(routing_map.device)
+ _count_local_tokens_kernel_persistent[(num_sms,)](
+ routing_map,
+ tokens_per_expert,
+ valid_tokens,
+ topk,
+ local_expert_start,
+ num_local_experts,
+ num_sms,
+ BLOCK_SIZE=BLOCK,
+ )
+ else:
+ _count_local_tokens_kernel[(_ceil_div(max_pairs, BLOCK),)](
+ routing_map,
+ tokens_per_expert,
+ valid_tokens,
+ topk,
+ local_expert_start,
+ num_local_experts,
+ BLOCK_SIZE=BLOCK,
+ )
return tokens_per_expert
@@ -101,6 +182,39 @@ def _prefix_sum_kernel(
tl.store(inclusive_offsets_ptr + r, inc, mask=mask)
+@triton.jit
+def _init_permutation_map_kernel(
+ perm_map_ptr,
+ n_used_ptr, # pointer to inclusive_expert_offsets[-1]: total used rows this iteration
+ BLOCK_SIZE: tl.constexpr,
+):
+ """Initialize permutation_map entries to -1 up to n_used rows.
+
+ Grid is launched at max size; entries beyond n_used are left untouched —
+ the activation and unpermute kernels are gated by the same n_used pointer
+ so they never read those entries.
+ """
+ pid = tl.program_id(0)
+ n_used = tl.load(n_used_ptr)
+ offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+ mask = offsets < n_used
+ tl.store(perm_map_ptr + offsets, tl.full([BLOCK_SIZE], -1, tl.int32), mask=mask)
+
+
+def init_permutation_map(permutation_map: torch.Tensor, n_used: torch.Tensor) -> None:
+ """Fill permutation_map[0:n_used] with -1.
+
+ Args:
+ permutation_map: [output_size] int32 buffer (pre-allocated at max size).
+ n_used: scalar int32 CUDA tensor = inclusive_expert_offsets[-1].
+ """
+ output_size = permutation_map.shape[0]
+ BLOCK_SIZE = 1024
+ _init_permutation_map_kernel[(_ceil_div(output_size, BLOCK_SIZE),)](
+ permutation_map, n_used, BLOCK_SIZE=BLOCK_SIZE
+ )
+
+
def compute_expert_offsets(tokens_per_expert: torch.Tensor, alignment: int = 1) -> tuple:
"""Compute exclusive and inclusive prefix sums of aligned token counts."""
n = tokens_per_expert.shape[0]
@@ -119,52 +233,55 @@ def compute_expert_offsets(tokens_per_expert: torch.Tensor, alignment: int = 1)
@triton.jit
def _permute_tokens_kernel(
- hidden_ptr, # [num_tokens, hidden_dim] input hidden states
- probs_ptr, # [num_tokens, topk] routing probabilities
- routing_map_ptr, # [num_tokens, topk] expert assignments (global IDs)
+ hidden_ptr, # [max_tokens, hidden_dim] input hidden states
+ probs_ptr, # [max_tokens, topk] routing probabilities
+ routing_map_ptr, # [max_tokens, topk] expert assignments (global IDs)
out_hidden_ptr, # [output_size, hidden_dim] output: permuted hidden states
out_probs_ptr, # [output_size] output: permuted probabilities
out_src_idx_ptr, # [output_size] output: permutation_map (original token index, -1 for padding)
- counters_ptr, # [num_local_experts] exclusive offsets,
- # atomically incremented to assign positions
- num_tokens, # number of input tokens
+ counters_ptr, # [num_local_experts] exclusive offsets, atomically incremented
+ valid_tokens_ptr, # scalar int32 CUDA tensor: number of valid tokens this iteration
hidden_dim, # hidden dimension
+ max_pairs, # max_tokens * topk (fixed for CG)
topk: tl.constexpr, # number of expert choices per token
local_expert_start, # first global expert index on this rank
num_local_experts: tl.constexpr, # number of experts on this rank
BLOCK_H: tl.constexpr, # tile size for copying hidden_dim
+ NUM_BLOCKS: tl.constexpr, # grid size (fixed for CG)
):
"""Permute tokens into expert-grouped order.
- Grid: one program per (token, topk) pair. Each program looks up the assigned
- expert, skips non-local experts, then atomically claims a position within
- that expert's block and copies the hidden state + prob + source index.
+ Grid: fixed NUM_BLOCKS CTAs, each iterating over multiple (token, topk) pairs.
+ valid_tokens gates which pairs are actually processed — required for CUDA graph
+ compatibility since the grid size never changes across steps.
"""
- # Each program handles one (token, topk) pair
- pair = tl.program_id(0)
- tok = pair // topk
- k = pair % topk
- if tok >= num_tokens:
- return
- eid = tl.load(routing_map_ptr + tok * topk + k)
- lid = eid - local_expert_start
- # Skip tokens routed to non-local experts
- if lid < 0 or lid >= num_local_experts:
+ pid = tl.program_id(0)
+ valid_tokens = tl.load(valid_tokens_ptr)
+ valid_pairs = valid_tokens * topk
+ if pid >= valid_pairs:
return
- # Atomically claim a position within this expert's aligned block
- pos = tl.atomic_add(counters_ptr + lid, 1)
- # Copy hidden state row
- for h in tl.range(0, hidden_dim, BLOCK_H):
- o = h + tl.arange(0, BLOCK_H)
- m = o < hidden_dim
- tl.store(
- out_hidden_ptr + pos * hidden_dim + o,
- tl.load(hidden_ptr + tok * hidden_dim + o, mask=m),
- mask=m,
- )
- tl.store(out_probs_ptr + pos, tl.load(probs_ptr + tok * topk + k))
- # Record source token index for unpermute
- tl.store(out_src_idx_ptr + pos, tok)
+ for pair in tl.range(pid, max_pairs, NUM_BLOCKS):
+ tok = pair // topk
+ if tok < valid_tokens:
+ k = pair % topk
+ eid = tl.load(routing_map_ptr + tok * topk + k)
+ lid = eid - local_expert_start
+ # Skip tokens routed to non-local experts
+ if lid >= 0 and lid < num_local_experts:
+ # Atomically claim a position within this expert's aligned block
+ pos = tl.atomic_add(counters_ptr + lid, 1)
+ # Copy hidden state row
+ for h in tl.range(0, hidden_dim, BLOCK_H):
+ o = h + tl.arange(0, BLOCK_H)
+ m = o < hidden_dim
+ tl.store(
+ out_hidden_ptr + pos * hidden_dim + o,
+ tl.load(hidden_ptr + tok * hidden_dim + o, mask=m),
+ mask=m,
+ )
+ tl.store(out_probs_ptr + pos, tl.load(probs_ptr + tok * topk + k))
+ # Record source token index for unpermute
+ tl.store(out_src_idx_ptr + pos, tok)
def permute_tokens(
@@ -173,6 +290,7 @@ def permute_tokens(
routing_map: torch.Tensor,
local_expert_start: int,
num_local_experts: int,
+ valid_tokens: torch.Tensor,
alignment: int = 1,
) -> tuple:
"""Permute tokens into expert-grouped order.
@@ -181,11 +299,14 @@ def permute_tokens(
permutation in a single call.
Args:
- hidden_states: [num_tokens, hidden_size] input.
- probs: [num_tokens, topk] routing probabilities.
- routing_map: [num_tokens, topk] expert assignments.
+ hidden_states: [max_tokens, hidden_size] input. Only the first valid_tokens
+ rows are valid; the rest are ignored.
+ probs: [max_tokens, topk] routing probabilities.
+ routing_map: [max_tokens, topk] expert assignments.
local_expert_start: first global expert index on this rank.
num_local_experts: number of experts on this rank.
+ valid_tokens: scalar int32 CUDA tensor with the number of valid tokens this
+ iteration. Fixed address; value updated each step before graph replay.
alignment: per-expert token alignment (default 1).
Returns:
@@ -197,13 +318,13 @@ def permute_tokens(
outputs back and by activation kernels to skip padding rows (-1).
- inclusive_offsets: [num_local_experts] int32 cumulative offsets for grouped_mm
"""
- num_tokens, hidden_dim = hidden_states.shape
+ max_tokens, hidden_dim = hidden_states.shape
topk = probs.shape[1]
# Count how many (token, topk) pairs are routed to each local expert.
- # Non-local experts are ignored. Result is [num_local_experts] int32.
+ # Non-local experts and rows beyond valid_tokens are ignored.
tokens_per_expert = compute_local_tokens_per_expert(
- routing_map, local_expert_start, num_local_experts
+ routing_map, local_expert_start, num_local_experts, valid_tokens
)
# exclusive_expert_offsets[i] = start of expert i's block in the padded output.
@@ -213,15 +334,21 @@ def permute_tokens(
exclusive_expert_offsets, inclusive_expert_offsets = compute_expert_offsets(
tokens_per_expert, alignment=alignment
)
- output_size = num_tokens * min(topk, num_local_experts) + alignment * num_local_experts
+ # Output sized at max to keep allocations fixed across steps (CUDA graph compatible).
+ output_size = max_tokens * min(topk, num_local_experts) + alignment * num_local_experts
permuted_hidden = torch.empty(
output_size, hidden_dim, dtype=hidden_states.dtype, device=hidden_states.device
)
permuted_probs = torch.empty(output_size, dtype=probs.dtype, device=probs.device)
- permutation_map = torch.full((output_size,), -1, dtype=torch.int32, device=probs.device)
+ permutation_map = torch.empty(output_size, dtype=torch.int32, device=probs.device)
+ # Only initialize [0, n_used) to -1; activation and unpermute kernels are gated
+ # by the same inclusive_expert_offsets[-1] pointer so they never read beyond n_used.
+ init_permutation_map(permutation_map, inclusive_expert_offsets[-1:])
BLOCK_H = min(triton.next_power_of_2(hidden_dim), 1024)
- _permute_tokens_kernel[(num_tokens * topk,)](
+ max_pairs = max_tokens * topk
+ NUM_BLOCKS = min(max_pairs, 512)
+ _permute_tokens_kernel[(NUM_BLOCKS,)](
hidden_states,
probs,
routing_map,
@@ -229,43 +356,80 @@ def permute_tokens(
permuted_probs,
permutation_map,
exclusive_expert_offsets,
- num_tokens,
+ valid_tokens,
hidden_dim,
+ max_pairs,
topk,
local_expert_start,
num_local_experts,
BLOCK_H=BLOCK_H,
+ NUM_BLOCKS=NUM_BLOCKS,
)
return permuted_hidden, permuted_probs, permutation_map, inclusive_expert_offsets
+@triton.jit
+def _zero_output_rows_kernel(
+ output_ptr, # [num_tokens, hidden_dim] fp32 buffer to partially zero
+ valid_tokens_ptr, # scalar int32 CUDA tensor: number of rows to zero
+ hidden_dim, # hidden dimension
+ num_tokens, # max token count (fixed for CG)
+ BLOCK_H: tl.constexpr,
+ NUM_BLOCKS: tl.constexpr, # grid size (fixed for CG)
+):
+ """Zero rows [0, valid_tokens) of the fp32 output buffer.
+
+ Grid: fixed NUM_BLOCKS CTAs, each iterating over multiple rows.
+ valid_tokens gates which rows are zeroed — required for CUDA graph compatibility.
+ """
+ pid = tl.program_id(0)
+ valid_tokens = tl.load(valid_tokens_ptr)
+ if pid >= valid_tokens:
+ return
+ zero = tl.zeros([BLOCK_H], dtype=tl.float32)
+ for row in tl.range(pid, num_tokens, NUM_BLOCKS):
+ if row < valid_tokens:
+ for h in tl.range(0, hidden_dim, BLOCK_H):
+ o = h + tl.arange(0, BLOCK_H)
+ m = o < hidden_dim
+ tl.store(output_ptr + row * hidden_dim + o, zero, mask=m)
+
+
@triton.jit
def _unpermute_tokens_kernel(
expert_out_ptr, # [output_size, hidden_dim] expert outputs in permuted order
probs_ptr, # [output_size] fp32 routing probabilities (permuted)
src_idx_ptr, # [output_size] permutation_map: original token index, or -1 for padding
- output_ptr, # [num_tokens, hidden_dim] fp32 output buffer (zeroed by caller)
+ output_ptr, # [max_tokens, hidden_dim] fp32 output buffer (zeroed by caller)
+ n_used_ptr, # pointer to inclusive_expert_offsets[-1]: number of used rows this iteration
hidden_dim, # hidden dimension
+ max_rows, # output_size (fixed for CG)
BLOCK_H: tl.constexpr, # tile size for processing hidden_dim
+ NUM_BLOCKS: tl.constexpr, # grid size (fixed for CG)
):
"""Scatter weighted expert outputs back to original token positions.
- Grid: one program per row of expert_out. Padding rows (src_idx == -1) are
- skipped. Multiple topk selections for the same token are accumulated via
- atomic adds. All arithmetic is in fp32 to avoid precision loss.
+ Grid: fixed NUM_BLOCKS CTAs, each iterating over multiple rows.
+ Rows beyond n_used and alignment-padding rows (src_idx == -1) are skipped.
+ Multiple topk selections for the same token are accumulated via atomic adds.
+ All arithmetic is in fp32 to avoid precision loss.
"""
- row = tl.program_id(0)
- source_idx = tl.load(src_idx_ptr + row)
- # Skip padding rows
- if source_idx < 0:
+ pid = tl.program_id(0)
+ n_used = tl.load(n_used_ptr)
+ if pid >= n_used:
return
- prob = tl.load(probs_ptr + row) # fp32
- for h in tl.range(0, hidden_dim, BLOCK_H):
- offsets = h + tl.arange(0, BLOCK_H)
- m = offsets < hidden_dim
- # Upcast bf16 expert output to fp32 before multiply + accumulate
- v = tl.load(expert_out_ptr + row * hidden_dim + offsets, mask=m).to(tl.float32)
- tl.atomic_add(output_ptr + source_idx * hidden_dim + offsets, v * prob, mask=m)
+ for row in tl.range(pid, max_rows, NUM_BLOCKS):
+ if row < n_used:
+ source_idx = tl.load(src_idx_ptr + row)
+ # Skip alignment-padding rows within the used range
+ if source_idx >= 0:
+ prob = tl.load(probs_ptr + row) # fp32
+ for h in tl.range(0, hidden_dim, BLOCK_H):
+ offsets = h + tl.arange(0, BLOCK_H)
+ m = offsets < hidden_dim
+ # Upcast bf16 expert output to fp32 before multiply + accumulate
+ v = tl.load(expert_out_ptr + row * hidden_dim + offsets, mask=m).to(tl.float32)
+ tl.atomic_add(output_ptr + source_idx * hidden_dim + offsets, v * prob, mask=m)
def unpermute_tokens(
@@ -273,22 +437,53 @@ def unpermute_tokens(
permuted_probs: torch.Tensor,
permutation_map: torch.Tensor,
num_tokens: int,
+ n_used: torch.Tensor,
+ valid_tokens: torch.Tensor,
+ out: torch.Tensor = None,
) -> torch.Tensor:
"""Unpermute expert outputs back to original token order.
Accumulates in fp32 to avoid precision loss from multiple topk atomic adds.
Returns fp32 output.
+
+ Args:
+ expert_output: [output_size, hidden_dim] expert outputs in permuted order.
+ permuted_probs: [output_size] fp32 routing probabilities.
+ permutation_map: [output_size] int32, original token index or -1 for padding.
+ num_tokens: max token count (output buffer height); always fixed for CG.
+ n_used: scalar int32 CUDA tensor = inclusive_expert_offsets[-1]. Rows
+ beyond this are skipped without reading permutation_map.
+ valid_tokens: scalar int32 CUDA tensor = number of valid input tokens.
+ Only rows [0, valid_tokens) are zeroed; all atomic_adds target
+ source_idx < valid_tokens so rows beyond are never written.
+ out: optional pre-allocated [num_tokens, hidden_dim] fp32 output buffer.
+ Pass a symmetric memory tensor to scatter directly into it, avoiding
+ a separate copy before RSV. If None, a local buffer is allocated.
"""
assert (
permuted_probs.dtype == torch.float32
), f"permuted_probs must be fp32, got {permuted_probs.dtype}"
output_size, hidden_dim = expert_output.shape
- output = torch.zeros(num_tokens, hidden_dim, dtype=torch.float32, device=expert_output.device)
BLOCK_H = min(triton.next_power_of_2(hidden_dim), 1024)
- _unpermute_tokens_kernel[(output_size,)](
- expert_output, permuted_probs, permutation_map, output, hidden_dim, BLOCK_H=BLOCK_H
+ if out is None:
+ out = torch.empty(num_tokens, hidden_dim, dtype=torch.float32, device=expert_output.device)
+ NUM_BLOCKS_ZERO = min(num_tokens, 512)
+ _zero_output_rows_kernel[(NUM_BLOCKS_ZERO,)](
+ out, valid_tokens, hidden_dim, num_tokens, BLOCK_H=BLOCK_H, NUM_BLOCKS=NUM_BLOCKS_ZERO
+ )
+ NUM_BLOCKS = min(output_size, 512)
+ _unpermute_tokens_kernel[(NUM_BLOCKS,)](
+ expert_output,
+ permuted_probs,
+ permutation_map,
+ out,
+ n_used,
+ hidden_dim,
+ output_size,
+ BLOCK_H=BLOCK_H,
+ NUM_BLOCKS=NUM_BLOCKS,
)
- return output
+ return out
@triton.jit
@@ -301,75 +496,80 @@ def _permute_quantize_mxfp8_kernel(
out_probs_ptr,
out_src_idx_ptr,
counters_ptr,
- num_tokens,
+ valid_tokens_ptr, # scalar int32 CUDA tensor: number of valid tokens this iteration
K,
n_col_blocks,
+ max_pairs, # max_tokens * topk (fixed for CG)
topk: tl.constexpr,
local_expert_start,
num_local_experts: tl.constexpr,
REAL_GROUPS: tl.constexpr,
BLOCK_K: tl.constexpr,
BLOCK_GROUPS: tl.constexpr,
+ NUM_BLOCKS: tl.constexpr, # grid size (fixed for CG)
):
"""Fused permute + MXFP8 quantize + swizzle in one kernel.
- Grid: (num_tokens * topk,) — one program per (token, k) pair.
- Reads BF16 from source token, quantizes to FP8 e4m3, writes FP8 data +
- swizzled e8m0 scales to the permuted write position.
+ Grid: fixed NUM_BLOCKS CTAs, each iterating over multiple (token, topk) pairs.
+ valid_tokens gates which pairs are actually processed — required for CUDA graph
+ compatibility since the grid size never changes across steps.
"""
- pair = tl.program_id(0)
- tok = pair // topk
- k = pair % topk
- if tok >= num_tokens:
- return
- eid = tl.load(routing_map_ptr + tok * topk + k)
- lid = eid - local_expert_start
- if lid < 0 or lid >= num_local_experts:
+ pid = tl.program_id(0)
+ valid_tokens = tl.load(valid_tokens_ptr)
+ valid_pairs = valid_tokens * topk
+ if pid >= valid_pairs:
return
- pos = tl.atomic_add(counters_ptr + lid, 1)
-
- # Load full row from source token
- offs = tl.arange(0, BLOCK_K)
- mask = offs < K
- x = tl.load(hidden_ptr + tok * K + offs, mask=mask, other=0.0).to(tl.float32)
-
- # Per-group-of-32 quantization
- x_grouped = tl.reshape(x, [BLOCK_GROUPS, 32])
- abs_grouped = tl.abs(x_grouped)
- max_vals = tl.max(abs_grouped, axis=1)
-
- dequant_scale = max_vals / 448.0
- dequant_exp = (dequant_scale.to(tl.uint32, bitcast=True) + 0x007FFFFF) & 0x7F800000
- dequant_rounded = dequant_exp.to(tl.float32, bitcast=True)
- quant_scale = tl.where(dequant_rounded == 0, 0.0, 1.0 / dequant_rounded)
-
- quantized = x_grouped * quant_scale[:, None]
- quantized_flat = tl.reshape(quantized, [BLOCK_K])
- out_fp8 = quantized_flat.to(tl.float8e4nv)
-
- # Store FP8 data at permuted position
- tl.store(out_fp8_ptr + pos * K + offs, out_fp8, mask=mask)
-
- # Store swizzled scales at permuted position
- scale_exp = (dequant_exp >> 23).to(tl.uint8)
- col_offs = tl.arange(0, BLOCK_GROUPS)
- col_mask = col_offs < REAL_GROUPS
-
- macro_row_block = pos // 128
- macro_col_block = col_offs // 4
- local_row = pos % 128
- local_col = col_offs % 4
- group = local_row // 32
- sub_row = local_row % 32
- tile_idx = macro_row_block * n_col_blocks + macro_col_block
- swizzled_offs = tile_idx * 512 + sub_row * 16 + group * 4 + local_col
-
- tl.store(out_scale_ptr + swizzled_offs, scale_exp, mask=col_mask)
-
- # Store prob and source index
- tl.store(out_probs_ptr + pos, tl.load(probs_ptr + tok * topk + k))
- tl.store(out_src_idx_ptr + pos, tok)
+ for pair in tl.range(pid, max_pairs, NUM_BLOCKS):
+ tok = pair // topk
+ if tok < valid_tokens:
+ k = pair % topk
+ eid = tl.load(routing_map_ptr + tok * topk + k)
+ lid = eid - local_expert_start
+ if lid >= 0 and lid < num_local_experts:
+ pos = tl.atomic_add(counters_ptr + lid, 1)
+
+ # Load full row from source token
+ offs = tl.arange(0, BLOCK_K)
+ mask = offs < K
+ x = tl.load(hidden_ptr + tok * K + offs, mask=mask, other=0.0).to(tl.float32)
+
+ # Per-group-of-32 quantization
+ x_grouped = tl.reshape(x, [BLOCK_GROUPS, 32])
+ abs_grouped = tl.abs(x_grouped)
+ max_vals = tl.max(abs_grouped, axis=1)
+
+ dequant_scale = max_vals / 448.0
+ dequant_exp = (dequant_scale.to(tl.uint32, bitcast=True) + 0x007FFFFF) & 0x7F800000
+ dequant_rounded = dequant_exp.to(tl.float32, bitcast=True)
+ quant_scale = tl.where(dequant_rounded == 0, 0.0, 1.0 / dequant_rounded)
+
+ quantized = x_grouped * quant_scale[:, None]
+ quantized_flat = tl.reshape(quantized, [BLOCK_K])
+ out_fp8 = quantized_flat.to(tl.float8e4nv)
+
+ # Store FP8 data at permuted position
+ tl.store(out_fp8_ptr + pos * K + offs, out_fp8, mask=mask)
+
+ # Store swizzled scales at permuted position
+ scale_exp = (dequant_exp >> 23).to(tl.uint8)
+ col_offs = tl.arange(0, BLOCK_GROUPS)
+ col_mask = col_offs < REAL_GROUPS
+
+ macro_row_block = pos // 128
+ macro_col_block = col_offs // 4
+ local_row = pos % 128
+ local_col = col_offs % 4
+ group = local_row // 32
+ sub_row = local_row % 32
+ tile_idx = macro_row_block * n_col_blocks + macro_col_block
+ swizzled_offs = tile_idx * 512 + sub_row * 16 + group * 4 + local_col
+
+ tl.store(out_scale_ptr + swizzled_offs, scale_exp, mask=col_mask)
+
+ # Store prob and source index
+ tl.store(out_probs_ptr + pos, tl.load(probs_ptr + tok * topk + k))
+ tl.store(out_src_idx_ptr + pos, tok)
def permute_and_quantize_mxfp8(
@@ -378,6 +578,7 @@ def permute_and_quantize_mxfp8(
routing_map: torch.Tensor,
local_expert_start: int,
num_local_experts: int,
+ valid_tokens: torch.Tensor,
alignment: int = 128,
) -> tuple:
"""Fused permute + MXFP8 quantize + swizzle.
@@ -387,11 +588,14 @@ def permute_and_quantize_mxfp8(
single kernel launch.
Args:
- hidden_states: [num_tokens, hidden_size] BF16 input.
- probs: [num_tokens, topk] routing probabilities.
- routing_map: [num_tokens, topk] expert assignments.
+ hidden_states: [max_tokens, hidden_size] BF16 input. Only the first
+ valid_tokens rows are valid; the rest are ignored.
+ probs: [max_tokens, topk] routing probabilities.
+ routing_map: [max_tokens, topk] expert assignments.
local_expert_start: first global expert index on this rank.
num_local_experts: number of experts on this rank.
+ valid_tokens: scalar int32 CUDA tensor with the number of valid tokens this
+ iteration. Fixed address; value updated each step before graph replay.
alignment: per-expert token alignment (default 128, required for MXFP8 swizzle).
Returns:
@@ -403,13 +607,14 @@ def permute_and_quantize_mxfp8(
"""
from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor
- num_tokens, K = hidden_states.shape
+ max_tokens, K = hidden_states.shape
topk = probs.shape[1]
assert K % 32 == 0
# Count how many (token, topk) pairs are routed to each local expert.
+ # Rows beyond valid_tokens are ignored.
tokens_per_expert = compute_local_tokens_per_expert(
- routing_map, local_expert_start, num_local_experts
+ routing_map, local_expert_start, num_local_experts, valid_tokens
)
# exclusive_expert_offsets[i] = start of expert i's block in the padded output.
@@ -417,7 +622,8 @@ def permute_and_quantize_mxfp8(
exclusive_expert_offsets, inclusive_expert_offsets = compute_expert_offsets(
tokens_per_expert, alignment=alignment
)
- output_size = num_tokens * min(topk, num_local_experts) + alignment * num_local_experts
+ # Output sized at max to keep allocations fixed across steps (CUDA graph compatible).
+ output_size = max_tokens * min(topk, num_local_experts) + alignment * num_local_experts
scale_cols = K // 32
n_row_blocks = _ceil_div(output_size, 128)
@@ -427,12 +633,14 @@ def permute_and_quantize_mxfp8(
out_fp8 = torch.empty(output_size, K, dtype=torch.float8_e4m3fn, device=hidden_states.device)
out_scale = torch.zeros(total_scale_bytes, dtype=torch.uint8, device=hidden_states.device)
permuted_probs = torch.empty(output_size, dtype=probs.dtype, device=probs.device)
- permutation_map = torch.full((output_size,), -1, dtype=torch.int32, device=probs.device)
+ permutation_map = torch.empty(output_size, dtype=torch.int32, device=probs.device)
+ init_permutation_map(permutation_map, inclusive_expert_offsets[-1:])
BLOCK_K = triton.next_power_of_2(K)
BLOCK_GROUPS = BLOCK_K // 32
-
- _permute_quantize_mxfp8_kernel[(num_tokens * topk,)](
+ max_pairs = max_tokens * topk
+ NUM_BLOCKS = min(max_pairs, 512)
+ _permute_quantize_mxfp8_kernel[(NUM_BLOCKS,)](
hidden_states,
probs,
routing_map,
@@ -441,15 +649,17 @@ def permute_and_quantize_mxfp8(
permuted_probs,
permutation_map,
exclusive_expert_offsets,
- num_tokens,
+ valid_tokens,
K,
n_col_blocks,
+ max_pairs,
topk,
local_expert_start,
num_local_experts,
REAL_GROUPS=scale_cols,
BLOCK_K=BLOCK_K,
BLOCK_GROUPS=BLOCK_GROUPS,
+ NUM_BLOCKS=NUM_BLOCKS,
)
permuted_mxfp8 = MXFP8Tensor(
diff --git a/megatron/core/inference/moe/vllm_fused_moe.py b/megatron/core/inference/moe/vllm_fused_moe.py
new file mode 100644
index 00000000000..287d5f2828e
--- /dev/null
+++ b/megatron/core/inference/moe/vllm_fused_moe.py
@@ -0,0 +1,680 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+# Some of this code was adopted from https://github.com/vllm-project/vllm.
+# This source code is licensed under the Apache license found in the
+# LICENSE file in the root directory of this source tree.
+"""vLLM-style Triton fused MoE kernel (BF16) for Megatron inference.
+
+CUDA-graph compatible: all indirection table construction happens on-device
+via Triton kernels with fixed-size buffers and valid_tokens gating.
+"""
+
+from typing import Optional
+from unittest.mock import MagicMock
+
+import torch
+
+from megatron.core.utils import null_decorator
+
+try:
+ import triton
+ import triton.language as tl
+
+ HAVE_TRITON = True
+except ImportError:
+ HAVE_TRITON = False
+
+if not HAVE_TRITON:
+ triton = MagicMock()
+ triton.jit = null_decorator
+ tl = MagicMock()
+
+from megatron.core.inference.moe.fused_moe import ActivationType
+from megatron.core.inference.moe.permute import (
+ _get_num_sms,
+ compute_expert_offsets,
+ compute_local_tokens_per_expert,
+)
+
+# ---------------------------------------------------------------------------
+# Triton kernel – BF16 grouped GEMM with indirect token addressing
+# ---------------------------------------------------------------------------
+
+
+def _get_default_config(M: int, E: int, top_k: int) -> dict:
+ """Pick BLOCK_SIZE_*, GROUP_SIZE_M, num_warps, num_stages from M, E, top_k.
+
+ Mirrors vLLM's ``get_default_config`` (bf16/fp16 branch) verbatim:
+ https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/fused_moe/fused_moe.py
+
+ M here is the host-side token-count hint (``num_tokens_hint`` in
+ ``vllm_fused_moe``), NOT ``hidden_states.size(0)``. The hint is the
+ expected per-step token count; the worst-case buffer size would over-tune
+ for prefill on every decode step.
+
+ Two intuitions drive the choices:
+ 1. Small M is memory-bound (favor tall/narrow tiles, more pipeline
+ stages); large M is compute-bound (favor short/wide tiles, more warps).
+ 2. Padding tax dominates at small M — the indirection table pads M-tiles
+ per expert, so small M-tiles minimize wasted rows.
+ """
+ # BLOCK_SIZE_M: shrink at small M to limit per-expert padding waste.
+ if M <= 32:
+ block_m = 16
+ elif M <= 96:
+ block_m = 32
+ elif M <= 512:
+ block_m = 64
+ else:
+ block_m = 128
+
+ # BLOCK_SIZE_N: small M is memory-bound on weights, narrow N keeps weight
+ # traffic in check; large M has enough FMAs per weight load for wider N.
+ block_n = 64 if M <= 64 else 128
+
+ # BLOCK_SIZE_K: small M needs depth in K to keep tensor cores fed; large M
+ # already has enough M*N work, so shorter K reduces accumulator stall.
+ block_k = 128 if M <= 64 else 64
+
+ # GROUP_SIZE_M: tile-grouping for L2 reuse on weight tiles. Only profitable
+ # when each expert sees enough adjacent M-tiles.
+ tokens_per_expert = M // max(E, 1)
+ group_m = 16 if tokens_per_expert > 128 else 1
+
+ # num_warps: small M doesn't justify register pressure of more warps;
+ # large M is compute-bound and feeds an MMA pipeline that wants more.
+ num_warps = 4 if M <= 128 else 8
+
+ # num_stages: extra prefetch only pays off when memory-bound (very small M).
+ num_stages = 4 if M <= 32 else 3
+
+ return {
+ 'BLOCK_SIZE_M': block_m,
+ 'BLOCK_SIZE_N': block_n,
+ 'BLOCK_SIZE_K': block_k,
+ 'GROUP_SIZE_M': group_m,
+ 'num_warps': num_warps,
+ 'num_stages': num_stages,
+ }
+
+
+@triton.jit
+def _fused_moe_kernel(
+ # Pointers
+ a_ptr,
+ b_ptr,
+ c_ptr,
+ topk_weights_ptr,
+ sorted_token_ids_ptr,
+ expert_ids_ptr,
+ num_tokens_post_padded_ptr,
+ # Dimensions
+ N,
+ K,
+ num_valid_tokens,
+ # Strides
+ stride_am,
+ stride_ak,
+ stride_be,
+ stride_bk,
+ stride_bn,
+ stride_cm,
+ stride_cn,
+ # Flags / constexprs
+ MUL_ROUTED_WEIGHT: tl.constexpr,
+ FUSE_SQUARED_RELU: tl.constexpr,
+ top_k: tl.constexpr,
+ BLOCK_SIZE_M: tl.constexpr,
+ BLOCK_SIZE_N: tl.constexpr,
+ BLOCK_SIZE_K: tl.constexpr,
+ GROUP_SIZE_M: tl.constexpr,
+):
+ """Fused MoE grouped GEMM with indirect token addressing.
+
+ Body mirrors vLLM's `fused_moe_kernel` verbatim except for the
+ `FUSE_SQUARED_RELU` branch (Megatron applies relu+square in fp32 on
+ the accumulator before the bf16 cast — strictly more accurate than
+ upstream's separate post-FC1 activation kernel).
+
+ Grid is sized host-side from `num_tokens_hint` (the typical-case token
+ count), not the worst-case buffer length, so launch overhead at decode
+ stays small. When the actual padded length exceeds the hinted grid
+ size (rare prefill spikes), each CTA strides over multiple tiles via
+ the outer `tl.range` loop.
+ """
+ num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+
+ num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+ num_pid_m = tl.cdiv(num_tokens_post_padded, BLOCK_SIZE_M)
+ total_tiles = num_pid_m * num_pid_n
+ num_pid_in_group = GROUP_SIZE_M * num_pid_n
+
+ pid_init = tl.program_id(axis=0)
+ grid_size = tl.num_programs(axis=0)
+
+ offs_k = tl.arange(0, BLOCK_SIZE_K)
+
+ for pid in tl.range(pid_init, total_tiles, grid_size):
+ # GROUP_SIZE_M swizzle: pid → (pid_m, pid_n). Mirrors upstream vLLM.
+ group_id = pid // num_pid_in_group
+ first_pid_m = group_id * GROUP_SIZE_M
+ group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+ pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+ pid_n = (pid % num_pid_in_group) // group_size_m
+
+ # Skip padding tiles whose expert slot was never assigned. In
+ # vLLM this also handles non-local experts via `write_zeros_to_output`;
+ # our scatter excludes non-local pairs from `sorted_token_ids` entirely,
+ # so `expert_id == -1` only fires on tail padding and we just skip.
+ # (Triton's JIT does not support `continue`, so we gate the body.)
+ off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
+ if off_experts != -1:
+ offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+ offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64)
+ token_mask = offs_token < num_valid_tokens
+
+ # `% N` keeps overflow lanes in-bounds; matching C-store mask drops
+ # their contribution. Saves a 2-D bounds check inside the K loop.
+ offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N
+
+ a_ptrs = a_ptr + (
+ offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak
+ )
+ b_ptrs = (
+ b_ptr
+ + off_experts * stride_be
+ + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+ )
+
+ accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+ for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+ a = tl.load(
+ a_ptrs,
+ mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
+ other=0.0,
+ )
+ b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+ accumulator += tl.dot(a, b)
+ a_ptrs += BLOCK_SIZE_K * stride_ak
+ b_ptrs += BLOCK_SIZE_K * stride_bk
+
+ # Megatron-only: squared-relu fused on the fp32 accumulator before
+ # the bf16 cast. Upstream runs relu+square as a separate bf16 kernel.
+ if FUSE_SQUARED_RELU:
+ accumulator = tl.maximum(accumulator, 0.0)
+ accumulator *= accumulator
+
+ if MUL_ROUTED_WEIGHT:
+ moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
+ accumulator *= moe_weight[:, None]
+
+ accumulator = accumulator.to(tl.bfloat16)
+ offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+ c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
+ c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
+ tl.store(c_ptrs, accumulator, mask=c_mask)
+
+
+# ---------------------------------------------------------------------------
+# Indirection table construction (CUDA-graph safe, fully on-device)
+# ---------------------------------------------------------------------------
+
+
+def _ceil_div(a, b):
+ return (a + b - 1) // b
+
+
+@triton.jit
+def _init_sorted_ids_kernel(
+ sorted_token_ids_ptr,
+ expert_ids_ptr,
+ max_sorted,
+ max_blocks,
+ SENTINEL: tl.constexpr,
+ BLOCK: tl.constexpr,
+):
+ """Initialize sorted_token_ids to SENTINEL and expert_ids to -1."""
+ pid = tl.program_id(0)
+ block_start = pid * BLOCK
+ if block_start < max_sorted or block_start < max_blocks:
+ offs = block_start + tl.arange(0, BLOCK)
+ tl.store(sorted_token_ids_ptr + offs, SENTINEL, mask=offs < max_sorted)
+ tl.store(expert_ids_ptr + offs, -1, mask=offs < max_blocks)
+
+
+@triton.jit
+def _scatter_token_indices_kernel(
+ routing_map_ptr,
+ sorted_token_ids_ptr,
+ counters_ptr,
+ valid_tokens_ptr,
+ topk: tl.constexpr,
+ local_expert_start,
+ num_local_experts: tl.constexpr,
+ max_pairs,
+ BLOCK_SIZE: tl.constexpr,
+):
+ """Scatter local-expert pair indices into the padded indirection table.
+
+ Only local expert pairs are written; non-local pairs are skipped (the
+ _moe_sum kernel handles them by checking the routing map directly).
+ """
+ pid = tl.program_id(0)
+ valid_tokens = tl.load(valid_tokens_ptr)
+ valid_pairs = valid_tokens * topk
+ if pid * BLOCK_SIZE >= valid_pairs:
+ return
+ offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+ mask = offs < valid_pairs
+
+ eids = tl.load(routing_map_ptr + offs, mask=mask, other=-1)
+ lids = eids - local_expert_start
+ is_local = (lids >= 0) & (lids < num_local_experts) & mask
+
+ local_pos = tl.atomic_add(counters_ptr + lids, 1, mask=is_local)
+ tl.store(sorted_token_ids_ptr + local_pos, offs, mask=is_local)
+
+
+@triton.jit
+def _fill_expert_block_ids_kernel(
+ expert_ids_ptr,
+ exclusive_offsets_ptr,
+ inclusive_offsets_ptr,
+ BLOCK_SIZE_M: tl.constexpr,
+ BLOCK: tl.constexpr,
+):
+ """Fill expert_ids with expert index for each BLOCK_SIZE_M block.
+
+ Grid: one CTA per expert (parallelised across experts).
+ Inner loop uses vectorised stores of BLOCK elements at a time.
+ """
+ e = tl.program_id(0)
+ start_block = tl.load(exclusive_offsets_ptr + e) // BLOCK_SIZE_M
+ end_block = tl.load(inclusive_offsets_ptr + e) // BLOCK_SIZE_M
+ num_blocks = end_block - start_block
+ for off in tl.range(0, num_blocks, BLOCK):
+ idxs = start_block + off + tl.arange(0, BLOCK)
+ tl.store(expert_ids_ptr + idxs, e, mask=idxs < end_block)
+
+
+def _moe_align_block_size_cuda_graphable(
+ routing_map: torch.Tensor,
+ block_size: int,
+ num_local_experts: int,
+ local_expert_start: int,
+ valid_tokens: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Build indirection tables for the vLLM kernel, fully on-device.
+
+ Replaces the original _moe_align_block_size which used .item() calls
+ and host-side loops. All buffers are allocated at fixed max sizes so
+ the function is safe for CUDA graph capture.
+
+ Args:
+ routing_map: [max_tokens, topk] expert assignments.
+ block_size: BLOCK_SIZE_M for the vLLM kernel.
+ num_local_experts: experts on this rank.
+ local_expert_start: first global expert index on this rank.
+ valid_tokens: scalar int32 CUDA tensor.
+
+ Returns:
+ sorted_token_ids: [max_sorted] int32 indirection table.
+ expert_ids: [max_blocks] int32 expert per block.
+ num_tokens_post_padded: [1] int32 (local expert padded count).
+ """
+ max_tokens, topk = routing_map.shape
+ device = routing_map.device
+
+ max_sorted = max_tokens * topk + block_size * (num_local_experts + 1)
+ max_blocks = _ceil_div(max_sorted, block_size)
+ sentinel = max_tokens * topk
+
+ sorted_token_ids = torch.empty(max_sorted, dtype=torch.int32, device=device)
+ expert_ids = torch.empty(max_blocks, dtype=torch.int32, device=device)
+
+ INIT_BLOCK = 1024
+ init_grid = _ceil_div(max(max_sorted, max_blocks), INIT_BLOCK)
+ _init_sorted_ids_kernel[(init_grid,)](
+ sorted_token_ids, expert_ids, max_sorted, max_blocks, SENTINEL=sentinel, BLOCK=INIT_BLOCK
+ )
+
+ tokens_per_expert = compute_local_tokens_per_expert(
+ routing_map, local_expert_start, num_local_experts, valid_tokens, persistent=True
+ )
+ exclusive_offsets, inclusive_offsets = compute_expert_offsets(
+ tokens_per_expert, alignment=block_size
+ )
+
+ _fill_expert_block_ids_kernel[(num_local_experts,)](
+ expert_ids, exclusive_offsets, inclusive_offsets, BLOCK_SIZE_M=block_size, BLOCK=128
+ )
+
+ max_pairs = max_tokens * topk
+ SCATTER_BLOCK = 256
+ scatter_grid = _ceil_div(max_pairs, SCATTER_BLOCK)
+ _scatter_token_indices_kernel[(scatter_grid,)](
+ routing_map,
+ sorted_token_ids,
+ exclusive_offsets,
+ valid_tokens,
+ topk,
+ local_expert_start,
+ num_local_experts,
+ max_pairs,
+ BLOCK_SIZE=SCATTER_BLOCK,
+ )
+
+ num_tokens_post_padded = inclusive_offsets[-1:]
+ return sorted_token_ids, expert_ids, num_tokens_post_padded
+
+
+# ---------------------------------------------------------------------------
+# Kernel launcher
+# ---------------------------------------------------------------------------
+
+
+def _invoke_fused_moe_kernel(
+ A: torch.Tensor,
+ B: torch.Tensor,
+ C: torch.Tensor,
+ topk_weights: Optional[torch.Tensor],
+ sorted_token_ids: torch.Tensor,
+ expert_ids: torch.Tensor,
+ num_tokens_post_padded: torch.Tensor,
+ mul_routed_weight: bool,
+ top_k: int,
+ config: dict,
+ grid_size: int,
+ fuse_squared_relu: bool = False,
+):
+ """Launch the Triton fused-MoE kernel for one GEMM pass.
+
+ Body matches upstream vLLM `fused_moe_kernel` (1 CTA per (pid_m, pid_n)
+ tile, raw pointer arithmetic with `% N` on the N axis), apart from the
+ optional fused squared-relu activation in fp32.
+
+ `grid_size` is sized host-side from `num_tokens_hint` so launch overhead
+ at decode is small. When the actual padded length exceeds the hinted
+ grid size, each CTA strides over additional tiles via the kernel's outer
+ `tl.range`. The full launch config (tile sizes, warps, stages) is picked
+ host-side by ``_get_default_config`` from M = num_tokens_hint.
+ """
+ M = A.size(0)
+ num_tokens = M * top_k
+
+ _fused_moe_kernel[(grid_size,)](
+ A,
+ B,
+ C,
+ topk_weights,
+ sorted_token_ids,
+ expert_ids,
+ num_tokens_post_padded,
+ B.size(1),
+ B.size(2),
+ num_tokens,
+ A.stride(0),
+ A.stride(1),
+ B.stride(0),
+ B.stride(2),
+ B.stride(1),
+ C.stride(0),
+ C.stride(1),
+ MUL_ROUTED_WEIGHT=mul_routed_weight,
+ FUSE_SQUARED_RELU=fuse_squared_relu,
+ top_k=top_k,
+ BLOCK_SIZE_M=config['BLOCK_SIZE_M'],
+ BLOCK_SIZE_N=config['BLOCK_SIZE_N'],
+ BLOCK_SIZE_K=config['BLOCK_SIZE_K'],
+ GROUP_SIZE_M=config['GROUP_SIZE_M'],
+ num_warps=config['num_warps'],
+ num_stages=config['num_stages'],
+ )
+
+
+# ---------------------------------------------------------------------------
+# Fused topk reduction (replaces torch.sum + copy)
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _moe_sum_kernel(
+ input_ptr,
+ output_ptr,
+ topk_weights_ptr,
+ valid_tokens_ptr,
+ routing_map_ptr,
+ local_expert_start,
+ num_local_experts: tl.constexpr,
+ K,
+ topk: tl.constexpr,
+ BLOCK_M: tl.constexpr,
+ BLOCK_K: tl.constexpr,
+ NUM_K_BLOCKS: tl.constexpr,
+):
+ """Reduce topk dimension with routing weight application.
+
+ input: [max_tokens * topk, K] bf16
+ output: [max_tokens, K] — dtype matches the output buffer (fp32 or bf16)
+
+ For token t < valid_tokens: output[t] = sum of input[t*topk+k] * prob[t*topk+k]
+ over topk slots k where the expert is local. Non-local slots are skipped
+ (their values in `input` are undefined because FC2 only processes
+ local-expert blocks).
+ Rows for t >= valid_tokens are not written; downstream consumers
+ (e.g. reduce-scatter-v) only read the first valid_tokens rows.
+ Routing weight multiplication and accumulation in fp32 for numerical accuracy.
+
+ Persistent grid: launches BLOCK_M CTAs that stride over valid_tokens.
+ CUDA-graph safe (grid is static); the loop bound is loaded device-side.
+ """
+ pid = tl.program_id(0)
+ valid_tokens = tl.load(valid_tokens_ptr)
+
+ for token_id in tl.range(pid, valid_tokens, BLOCK_M):
+ token_id_i64 = token_id.to(tl.int64)
+ base = token_id_i64 * topk * K
+
+ # k_idx outer / topk inner keeps the live accumulator at one BLOCK_K tile.
+ # Swapping (topk outer) would need NUM_K_BLOCKS persistent accumulators
+ # (~NUM_K_BLOCKS * BLOCK_K * 4 B), which spills / cuts occupancy at large K.
+ for k_idx in range(NUM_K_BLOCKS):
+ offs_k = k_idx * BLOCK_K + tl.arange(0, BLOCK_K)
+ k_mask = offs_k < K
+
+ acc = tl.zeros([BLOCK_K], dtype=tl.float32)
+ for t in range(topk):
+ eid = tl.load(routing_map_ptr + token_id * topk + t)
+ lid = eid - local_expert_start
+ if lid >= 0 and lid < num_local_experts:
+ v = tl.load(input_ptr + base + t * K + offs_k, mask=k_mask, other=0.0)
+ w = tl.load(topk_weights_ptr + token_id * topk + t)
+ acc += v.to(tl.float32) * w
+
+ tl.store(output_ptr + token_id_i64 * K + offs_k, acc, mask=k_mask)
+
+
+def _moe_sum(
+ input: torch.Tensor,
+ topk_weights: torch.Tensor,
+ max_tokens: int,
+ topk: int,
+ K: int,
+ valid_tokens: torch.Tensor,
+ routing_map: torch.Tensor,
+ local_expert_start: int,
+ num_local_experts: int,
+ out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+ """Fused topk reduction: [max_tokens*topk, K] bf16 → [max_tokens, K].
+
+ Applies routing weights and reduces over topk in a single kernel.
+ Accumulates in fp32. When `out` is None, allocates and returns an fp32
+ buffer. When `out` is provided (e.g. the RSV symmetric memory tensor),
+ writes directly into it — tl.store handles the cast to the buffer's dtype.
+ Only writes the first valid_tokens rows; rows beyond are left untouched
+ (downstream RSV reads only the valid range). Only accumulates contributions
+ from local experts; non-local topk slots are skipped (their values in
+ `input` are undefined).
+ """
+ if out is None:
+ out = torch.empty(max_tokens, K, dtype=torch.float32, device=input.device)
+ BLOCK_K = min(triton.next_power_of_2(K), 1024)
+ NUM_K_BLOCKS = _ceil_div(K, BLOCK_K)
+ BLOCK_M = _get_num_sms(input.device)
+ _moe_sum_kernel[(BLOCK_M,)](
+ input,
+ out,
+ topk_weights,
+ valid_tokens,
+ routing_map,
+ local_expert_start,
+ num_local_experts,
+ K,
+ topk=topk,
+ BLOCK_M=BLOCK_M,
+ BLOCK_K=BLOCK_K,
+ NUM_K_BLOCKS=NUM_K_BLOCKS,
+ )
+ return out
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def vllm_fused_moe(
+ hidden_states: torch.Tensor,
+ probs: torch.Tensor,
+ fc1_weight: torch.Tensor,
+ fc2_weight: torch.Tensor,
+ activation_type: ActivationType,
+ num_local_experts: int,
+ local_expert_start: int,
+ valid_tokens: torch.Tensor,
+ routing_map: torch.Tensor,
+ out: Optional[torch.Tensor] = None,
+ num_tokens_hint: Optional[int] = None,
+) -> torch.Tensor:
+ """Fused MoE using the vLLM Triton grouped-GEMM kernel (BF16).
+
+ CUDA-graph compatible: indirection tables are built entirely on-device
+ using fixed-size buffers gated by valid_tokens.
+
+ Args:
+ hidden_states: [max_tokens, hidden_size] BF16 input. Only the first
+ valid_tokens rows are valid; the rest are ignored.
+ probs: [max_tokens, topk] fp32 routing probabilities.
+ fc1_weight: [num_local_experts, fc1_out, hidden_size] BF16.
+ fc2_weight: [num_local_experts, hidden_size, fc1_out] BF16.
+ activation_type: ActivationType enum.
+ num_local_experts: experts on this rank.
+ local_expert_start: first global expert index on this rank.
+ valid_tokens: scalar int32 CUDA tensor with number of valid tokens.
+ routing_map: [max_tokens, topk] int expert assignments.
+ out: optional [max_tokens, hidden_size] output buffer (e.g. the RSV
+ symmetric memory tensor). If None, an fp32 buffer is allocated.
+ When provided, tl.store casts to the buffer's dtype automatically.
+ num_tokens_hint: optional host-side int with the expected number of
+ valid tokens (e.g. batch_size * ep_size). Used to select a better
+ BLOCK_SIZE_M instead of using the worst-case buffer size.
+
+ Returns:
+ [max_tokens, hidden_size] output (fp32 when out=None, else out's dtype).
+ tl.store handles the implicit cast when out is a different dtype.
+ """
+ assert (
+ hidden_states.dtype == torch.bfloat16
+ ), f"vllm_fused_moe requires bf16 input, got {hidden_states.dtype}"
+
+ max_tokens = hidden_states.size(0)
+ topk = routing_map.shape[1]
+ effective_tokens = num_tokens_hint if num_tokens_hint is not None else max_tokens
+
+ # Mirror upstream vLLM: pick the full launch config (tile sizes, warps,
+ # stages) host-side from the token-count hint, not from the worst-case
+ # buffer size. Same config is used for both FC1 and FC2 (matches vLLM).
+ config = _get_default_config(M=effective_tokens, E=num_local_experts, top_k=topk)
+
+ sorted_token_ids, expert_ids, num_post_padded = _moe_align_block_size_cuda_graphable(
+ routing_map, config['BLOCK_SIZE_M'], num_local_experts, local_expert_start, valid_tokens
+ )
+ num_valid = max_tokens * topk
+
+ N = fc1_weight.size(1)
+ K = fc1_weight.size(2)
+
+ # Grid sized for the typical-case token count (num_tokens_hint). When the
+ # actual num_tokens_post_padded exceeds this, the kernel's outer tl.range
+ # makes each CTA stride over additional tiles — correct but with reduced
+ # parallelism on rare prefill spikes. EM hint = effective_tokens*topk +
+ # BLOCK_SIZE_M*num_local_experts upper-bounds the per-expert padding.
+ block_m = config['BLOCK_SIZE_M']
+ em_hint = effective_tokens * topk + block_m * num_local_experts
+ num_pid_m_hint = _ceil_div(em_hint, block_m)
+ num_pid_n_fc1 = _ceil_div(N, config['BLOCK_SIZE_N'])
+ num_pid_n_fc2 = _ceil_div(K, config['BLOCK_SIZE_N'])
+ grid_size_fc1 = num_pid_m_hint * num_pid_n_fc1
+ grid_size_fc2 = num_pid_m_hint * num_pid_n_fc2
+
+ topk_weights_flat = probs.reshape(-1).contiguous()
+
+ # FC1 + activation: [max_tokens, K] → [max_tokens*topk, N]
+ assert activation_type == ActivationType.SQUARED_RELU
+ intermediate1 = torch.empty(
+ num_valid, N, dtype=hidden_states.dtype, device=hidden_states.device
+ )
+ _invoke_fused_moe_kernel(
+ hidden_states,
+ fc1_weight,
+ intermediate1,
+ topk_weights_flat,
+ sorted_token_ids,
+ expert_ids,
+ num_post_padded,
+ mul_routed_weight=False,
+ top_k=topk,
+ config=config,
+ grid_size=grid_size_fc1,
+ fuse_squared_relu=True,
+ )
+
+ # FC2: [max_tokens*topk, N] → [max_tokens*topk, K], without routing weights.
+ # Routing weights are applied in the reduction kernel to avoid an extra
+ # bf16 truncation of prob-scaled values before the topk summation.
+ # Only local-expert blocks are processed; non-local positions are left
+ # undefined and skipped by _moe_sum (which checks the routing map).
+ intermediate3 = torch.empty(
+ num_valid, K, dtype=hidden_states.dtype, device=hidden_states.device
+ )
+ _invoke_fused_moe_kernel(
+ intermediate1,
+ fc2_weight,
+ intermediate3,
+ topk_weights_flat,
+ sorted_token_ids,
+ expert_ids,
+ num_post_padded,
+ mul_routed_weight=False,
+ top_k=1,
+ config=config,
+ grid_size=grid_size_fc2,
+ )
+
+ # Reduce over topk: [max_tokens*topk, K] → [max_tokens, K]
+ # Applies routing weights and accumulates in fp32, writes directly to
+ # out (if provided), zeros rows beyond valid_tokens, and skips non-local
+ # expert slots.
+ return _moe_sum(
+ intermediate3,
+ probs,
+ max_tokens,
+ topk,
+ K,
+ valid_tokens,
+ routing_map,
+ local_expert_start,
+ num_local_experts,
+ out=out,
+ )
diff --git a/megatron/core/inference/sampling/__init__.py b/megatron/core/inference/sampling/__init__.py
new file mode 100644
index 00000000000..b2941b33c9e
--- /dev/null
+++ b/megatron/core/inference/sampling/__init__.py
@@ -0,0 +1,7 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+from megatron.core.inference.sampling.base import Sampling
+from megatron.core.inference.sampling.flashinfer_sampling import FlashInferSampling
+from megatron.core.inference.sampling.torch_sampling import TorchSampling
+
+__all__ = ["Sampling", "TorchSampling", "FlashInferSampling"]
diff --git a/megatron/core/inference/sampling/base.py b/megatron/core/inference/sampling/base.py
new file mode 100644
index 00000000000..8aa4c416c27
--- /dev/null
+++ b/megatron/core/inference/sampling/base.py
@@ -0,0 +1,89 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+import torch
+from torch import Tensor
+
+
+class Sampling(ABC):
+ """Abstract base for inference sampling backends.
+
+ Subclasses implement `sample_kernel`. CUDA graphs are added via `CudaGraphManager`.
+ """
+
+ @abstractmethod
+ def sample_kernel(
+ self,
+ logits: Tensor,
+ n: int,
+ context,
+ *,
+ gather_indices: Optional[Tensor] = None,
+ token_to_request_index: Optional[Tensor] = None,
+ eager: bool = False,
+ cache_key: Any = None,
+ ) -> Tensor:
+ """Sample `n` tokens from `logits` and return them.
+
+ Args:
+ logits: Logits tensor of shape `[>=n, vocab_size]`.
+ n: Number of rows to sample.
+ context: The active DynamicInferenceContext.
+ gather_indices: If provided, only sample from `logits[gather_indices[:n], :]`.
+ token_to_request_index: Per-token request mapping; when set, sampling
+ parameters are gathered per-token instead of per-request.
+ eager, cache_key: Consumed by `CudaGraphManager` when it wraps this kernel.
+
+ Returns:
+ Sampled token ids of shape `[n]`. Under CUDA graph replay, this is a static buffer.
+ """
+ ...
+
+ def sample_speculative(
+ self,
+ required_logits: Tensor,
+ num_decode: int,
+ num_prefill: int,
+ num_speculative_tokens: int,
+ context,
+ *,
+ gather_indices: Optional[Tensor] = None,
+ eager: bool = False,
+ cache_key: Any = None,
+ ) -> Tensor:
+ """Sample tokens for the speculative-verify path.
+
+ Decode requests contribute `1 + num_speculative_tokens` rows; prefill requests contribute 1.
+ Builds the per-token request mapping and dispatches to `sample_kernel`.
+ The `sample_kernel` is forced eager so its own `CudaGraphManager` wrapper does not fire.
+
+ When `gather_indices` is supplied, the kernel selects via `logits[gather_indices[:n], :]`.
+ When `gather_indices` is None, `required_logits` is expected to be already pre-gathered to
+ the layout described above (e.g. when `materialize_only_last_token_logits=True` upstream).
+ """
+ # CudaGraphManager consumes these args, if it exists.
+ del eager, cache_key
+
+ n_spec = num_speculative_tokens
+ num_decode_tokens = num_decode * (1 + n_spec)
+ num_tokens = num_decode_tokens + num_prefill
+ device = required_logits.device
+
+ token_to_request_index = torch.cat(
+ [
+ torch.arange(num_decode, device=device).repeat_interleave(
+ 1 + n_spec, output_size=num_decode_tokens
+ ),
+ torch.arange(num_decode, num_decode + num_prefill, device=device),
+ ]
+ )
+ return self.sample_kernel(
+ required_logits,
+ num_tokens,
+ context,
+ gather_indices=gather_indices,
+ token_to_request_index=token_to_request_index,
+ eager=True,
+ )
diff --git a/megatron/core/inference/sampling/flashinfer_sampling.py b/megatron/core/inference/sampling/flashinfer_sampling.py
new file mode 100644
index 00000000000..c89093daeac
--- /dev/null
+++ b/megatron/core/inference/sampling/flashinfer_sampling.py
@@ -0,0 +1,101 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+from typing import Any, Optional
+
+import torch
+from torch import Tensor
+
+try:
+ import flashinfer
+except ImportError:
+ flashinfer = None
+
+from megatron.core.inference.sampling.base import Sampling
+from megatron.core.transformer.cuda_graphs import CudaGraphManager
+
+
+class FlashInferSampling(Sampling):
+ """Fused FlashInfer sampling, with optional CUDA graph capture/replay."""
+
+ def __init__(
+ self, vocab_size: int, rng: torch.Generator, config=None, enable_cuda_graph: bool = False
+ ) -> None:
+ self._vocab_size = vocab_size
+ self._rng = rng
+ if enable_cuda_graph and config is not None and config.cuda_graph_impl == "local":
+ CudaGraphManager(
+ config,
+ self,
+ function_name="sample_kernel",
+ need_backward=False,
+ inline_capture=True,
+ )
+ CudaGraphManager(
+ config,
+ self,
+ function_name="sample_speculative",
+ need_backward=False,
+ inline_capture=True,
+ )
+
+ def sample_kernel(
+ self,
+ logits: Tensor,
+ n: int,
+ context,
+ *,
+ gather_indices: Optional[Tensor] = None,
+ token_to_request_index: Optional[Tensor] = None,
+ eager: bool = False,
+ cache_key: Any = None,
+ ) -> Tensor:
+ """FlashInfer fused top-k / top-p sampling kernel.
+
+ Args:
+ logits: Logits tensor of shape `[>=n, vocab_size]`.
+ n: Number of rows to sample.
+ context: The active DynamicInferenceContext.
+ gather_indices: When set, sample from `logits[gather_indices[:n], :]`.
+ token_to_request_index: When set, sampling parameters are gathered per-token
+ rather than per-request (used by the speculative path).
+ eager, cache_key: Consumed by `CudaGraphManager` when it wraps this kernel.
+
+ Returns:
+ Sampled token ids of shape `[n]`. Under CUDA graph replay, this is a static buffer.
+ """
+ # CudaGraphManager consumes these args, if it exists.
+ del eager, cache_key
+
+ # Read GPU sampling parameters from the per-step gpu_view mirror. The
+ # CPU source-of-truth (`active_request_metadata`) is pinned but resident
+ # on CPU, so reading it here would mix devices with `logits`.
+ gv = context.gpu_view
+ if token_to_request_index is None:
+ temperature = gv.temperature[:n]
+ top_k = gv.top_k[:n]
+ top_p = gv.top_p[:n]
+ else:
+ temperature = gv.temperature[token_to_request_index]
+ top_k = gv.top_k[token_to_request_index]
+ top_p = gv.top_p[token_to_request_index]
+
+ # Clamp temperature to avoid division by 0.
+ temperature = temperature.clamp(min=1e-6)
+ if gather_indices is None:
+ scaled = logits[:n] / temperature.unsqueeze(1)
+ else:
+ scaled = logits[gather_indices[:n], :] / temperature.unsqueeze(1)
+ probs = torch.softmax(scaled, dim=-1)
+
+ # Sentinel values disable filtering:
+ # top_k=vocab_size keeps all tokens, top_p=1.0 keeps the full probability mass.
+ # TODO: Consider changing the disable flags in the `InferenceRequest`.
+ top_k_safe = top_k.masked_fill(top_k == 0, self._vocab_size)
+ top_p_safe = top_p.masked_fill(top_p == 0.0, 1.0)
+ output = torch.empty(n, device=logits.device, dtype=torch.int64)
+ output.copy_(
+ flashinfer.sampling.top_k_top_p_sampling_from_probs(
+ probs, top_k_safe, top_p_safe, generator=self._rng
+ )
+ )
+ return output
diff --git a/megatron/core/inference/sampling/torch_sampling.py b/megatron/core/inference/sampling/torch_sampling.py
new file mode 100644
index 00000000000..79491add5ab
--- /dev/null
+++ b/megatron/core/inference/sampling/torch_sampling.py
@@ -0,0 +1,167 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+from collections import defaultdict
+from typing import Any, List, Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from megatron.core.inference.sampling.base import Sampling
+
+
+class TorchSampling(Sampling):
+ """Sampling via bucketed `torch.multinomial`.
+
+ Groups requests into unique buckets by `(temperature, top_k, top_p)` for separate launches.
+ """
+
+ def __init__(self, rng: torch.Generator, vocab_size: int) -> None:
+ self._rng = rng
+ self._vocab_size = vocab_size
+
+ @staticmethod
+ def sample_from_logits(
+ last_token_logits: Tensor,
+ temperature: float,
+ top_k: int,
+ top_p: float,
+ *,
+ generator: torch.Generator,
+ vocab_size: Optional[int] = None,
+ ) -> Tensor:
+ """Sample tokens from logits with temperature, top-k, and top-p filtering.
+
+ Shared between dynamic batching and static batching.
+
+ Args:
+ last_token_logits: Logits of shape `[batch_size, vocab_size]`.
+ temperature: Temperature scaling factor.
+ top_k: Top-k filtering value (0 = disabled).
+ top_p: Top-p (nucleus) filtering value (0.0 = disabled).
+ generator: RNG used by `torch.multinomial`.
+ vocab_size: When provided, asserts `top_k < vocab_size` and clamps the
+ sampled ids to `[0, vocab_size - 1]`.
+
+ Returns:
+ Sampled token ids of shape `[batch_size]`.
+ """
+ assert isinstance(top_p, float)
+ assert isinstance(top_k, int)
+ assert not (top_k > 0 and top_p > 0.0), "Cannot have top-p and top-k both greater than zero"
+ assert top_p <= 1.0, "top-p should be in (0,1]"
+
+ def modify_logits_for_top_k_filtering(logits, top_k):
+ """Set the logits for none top-k values to -inf."""
+ filter_ = logits < torch.topk(logits, top_k)[0][..., -1, None]
+ logits.masked_fill_(filter_, float("-Inf"))
+
+ def modify_logits_for_top_p_filtering(logits, top_p):
+ """Set the logits for none top-p values to -inf."""
+ sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+ cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
+
+ filter_ = cumulative_probs > top_p
+ # Clone needed: filter_[:, 1:] and filter_[:, :-1] are overlapping views;
+ # without clone, each write would corrupt the next read during the shift.
+ filter_[:, 1:] = filter_[:, :-1].clone()
+ filter_[..., 0] = 0
+
+ filter_ = filter_.scatter(1, sorted_indices, filter_)
+ logits.masked_fill_(filter_, float("-Inf"))
+
+ if top_k == 1:
+ return torch.argmax(last_token_logits, dim=-1)
+
+ # Clone needed: .div_() and masked_fill_() below modify in-place.
+ last_token_logits = last_token_logits.clone()
+ if temperature != 1.0:
+ last_token_logits.div_(temperature)
+ if top_k > 1:
+ assert top_k <= last_token_logits.size(1), "top-k is larger than logit size."
+ if vocab_size:
+ assert top_k < vocab_size, "top-k is larger than vocab size."
+ modify_logits_for_top_k_filtering(last_token_logits, top_k)
+ elif top_p > 0.0:
+ modify_logits_for_top_p_filtering(last_token_logits, top_p)
+
+ probabilities = last_token_logits.softmax(dim=-1)
+ sampled = torch.multinomial(probabilities, num_samples=1, generator=generator).view(-1)
+
+ if vocab_size:
+ sampled = torch.clamp(sampled, min=0, max=(vocab_size - 1))
+
+ return sampled
+
+ def sample_kernel(
+ self,
+ logits: Tensor,
+ n: int,
+ context,
+ *,
+ gather_indices: Optional[Tensor] = None,
+ token_to_request_index: Optional[Tensor] = None,
+ eager: bool = False,
+ cache_key: Any = None,
+ ) -> Tensor:
+ """Bucket active requests by `(temperature, top_k, top_p)` and sample each bucket.
+
+ Args:
+ logits: Logits tensor of shape `[>=n, vocab_size]`.
+ n: Number of rows to sample.
+ context: The active DynamicInferenceContext.
+ gather_indices: When set, sample from `logits[gather_indices[:n], :]`.
+ token_to_request_index: When set, the loop dispatches per-token rather than
+ per-request (used by the speculative path).
+ eager: Accepted for API symmetry; ignored (TorchSampling has no graph wrapper).
+ cache_key: Accepted for API symmetry; ignored.
+
+ Returns:
+ Sampled token ids of shape `[n]`.
+ """
+ # CudaGraphManager consumes these args, if it exists.
+ del eager, cache_key
+
+ # Group active requests into sampling buckets by (temperature, top_k, top_p).
+ active_request_count = context.total_request_count - context.paused_request_count
+ md = context.active_request_metadata
+ device = torch.cuda.current_device()
+
+ bucket_map: dict = defaultdict(list)
+ temp = md["temperature"][:active_request_count].tolist()
+ top_k = md["top_k"][:active_request_count].tolist()
+ top_p = md["top_p"][:active_request_count].tolist()
+ for request_index, (t, k, p) in enumerate(zip(temp, top_k, top_p)):
+ bucket_map[(t, k, p)].append(request_index)
+
+ buckets: List[Tuple] = [(indices, *params) for params, indices in bucket_map.items()]
+ bucket_index_tensors: List[Tensor] = [
+ torch.tensor(indices, device=device, dtype=torch.long) for indices, *_ in buckets
+ ]
+
+ if gather_indices is not None:
+ logits = logits[gather_indices[:n], :]
+
+ output = torch.empty(n, device=logits.device, dtype=torch.int64)
+ token_list = []
+ indices_list = []
+ for idx_tensor, (_, temp, top_k, top_p) in zip(bucket_index_tensors, buckets):
+ if token_to_request_index is None:
+ row_indices = idx_tensor
+ else:
+ row_indices = torch.where(torch.isin(token_to_request_index, idx_tensor))[0]
+ token_list.append(
+ TorchSampling.sample_from_logits(
+ logits[row_indices, :],
+ temp,
+ top_k,
+ top_p,
+ generator=self._rng,
+ vocab_size=self._vocab_size,
+ )
+ )
+ indices_list.append(row_indices)
+
+ sampled_tokens = torch.cat(token_list, dim=0)
+ sampled_indices = torch.cat(indices_list, dim=0)
+ output[sampled_indices] = sampled_tokens
+ return output
diff --git a/megatron/core/inference/symmetric_memory.py b/megatron/core/inference/symmetric_memory.py
index 254d41ce294..a5269989914 100644
--- a/megatron/core/inference/symmetric_memory.py
+++ b/megatron/core/inference/symmetric_memory.py
@@ -39,10 +39,13 @@ class SymmetricMemoryBuffer:
"""
def __init__(self, size_in_mb, process_group):
- if not HAVE_TORCH_SYMM_MEM or not HAVE_TRITON:
- # This should be hit if the user is running an older
- # version of torch, or if they do not have triton
- # installed.
+ self.init_failure_reason: Optional[str] = None
+ if not HAVE_TORCH_SYMM_MEM:
+ self.init_failure_reason = "torch.distributed._symmetric_memory not importable"
+ self.symm_buffer = None
+ self.symm_mem_hdl = None
+ elif not HAVE_TRITON:
+ self.init_failure_reason = "triton not installed"
self.symm_buffer = None
self.symm_mem_hdl = None
else:
@@ -52,8 +55,7 @@ def __init__(self, size_in_mb, process_group):
self.symm_buffer = symm_mem.empty(numel, dtype=torch.uint8, device='cuda')
self.symm_mem_hdl = symm_mem.rendezvous(self.symm_buffer, process_group)
except RuntimeError as e:
- # If symmetric memory initialization fails, set buffer and handle to None
- # This should happen if the process group is not contained within NVlink
+ self.init_failure_reason = f"{type(e).__name__}: {e}"
self.symm_buffer = None
self.symm_mem_hdl = None
@@ -138,7 +140,7 @@ class SymmetricMemoryManager:
"""
_buffers: dict[str, SymmetricMemoryBuffer] = {}
- _default_size_mb: int = 256
+ _default_size_mb: int = 512
@classmethod
def get_buffer(
diff --git a/megatron/core/inference/text_generation_controllers/mtp_utils_pytorch.py b/megatron/core/inference/text_generation_controllers/mtp_utils_pytorch.py
new file mode 100644
index 00000000000..fe5474d0b22
--- /dev/null
+++ b/megatron/core/inference/text_generation_controllers/mtp_utils_pytorch.py
@@ -0,0 +1,255 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+import torch
+
+
+def rewind_kv_cache(
+ accepted_counts,
+ prefill_status,
+ last_kv_block_offset,
+ kv_length_offsets,
+ kv_block_counts,
+ last_kv_block_id,
+ kv_block_ids,
+ num_speculative_tokens,
+ block_size_tokens,
+ num_active_requests=None,
+):
+ """Update the KV cache bookkeeping for speculative decoding.
+
+ After forward pass with speculative tokens, some tokens may be rejected.
+ This function "rewinds" the KV cache bookkeeping to reflect only the accepted tokens.
+
+ When speculative tokens are rejected, we need to:
+ 1. Update kv_length_offsets (total sequence length)
+ 2. Update last_kv_block_offset (position within last block)
+ 3. If rewinding crosses a block boundary:
+ - Reduce kv_block_counts
+ - Update last_kv_block_id to point to the previous block
+ - Clear the entry in kv_block_ids for the released block
+
+ Mutates the input tensors in-place.
+
+ Returns (blocks_to_release, remove_mask).
+ """
+ N = accepted_counts.shape[0]
+ if num_active_requests is None:
+ num_active_requests = N
+
+ # Bulk-extract scalars once via .tolist() instead of per-element .item().
+ # Avoids N round-trips through the Python/C++ boundary inside the loop.
+ accepted_list = accepted_counts.tolist()
+ prefill_list = prefill_status.tolist()
+ offset_list = last_kv_block_offset.tolist()
+ length_list = kv_length_offsets.tolist()
+ block_count_list = kv_block_counts.tolist()
+ last_block_list = last_kv_block_id.tolist()
+ kv_block_ids_list = kv_block_ids.tolist()
+ max_blocks = kv_block_ids.shape[1]
+
+ blocks_to_release = torch.empty_like(last_kv_block_id)
+ remove_mask = torch.empty(N, device=accepted_counts.device, dtype=torch.bool)
+
+ for i in range(N):
+ if i >= num_active_requests:
+ blocks_to_release[i] = 0
+ remove_mask[i] = False
+ continue
+
+ accepted = accepted_list[i]
+ prefill = prefill_list[i]
+ last_offset = offset_list[i]
+ kv_length = length_list[i]
+ block_count = block_count_list[i]
+ last_block = last_block_list[i]
+
+ # Number of tokens to rewind (rejected speculative tokens).
+ # For prefill requests, no speculative tokens were forwarded through the model,
+ # so there is nothing to rewind.
+ num_to_rewind = 0 if prefill == 1 else num_speculative_tokens - accepted
+
+ # Save the original offset BEFORE modifying to correctly detect block boundary crossing.
+ # A request crosses back to a previous block if: original_offset - num_to_rewind < 0
+ diff = last_offset - num_to_rewind
+ remove = diff < 0
+
+ # Update the offsets
+ new_offset = diff % block_size_tokens
+ last_kv_block_offset[i] = new_offset
+ kv_length_offsets[i] = kv_length - num_to_rewind
+
+ # For requests that crossed back to a previous block, we need to:
+ # 1. Reduce the block count by 1
+ # 2. Get the block ID to release (current last_kv_block_id)
+ # 3. Update last_kv_block_id to point to the previous block
+ # 4. Clear the entry in kv_block_ids for the released block
+ # 5. Release the block back to the allocator
+ blocks_to_release[i] = last_block
+
+ # Reduce block counts for requests that crossed back
+ new_block_count = block_count - 1 if remove else block_count
+ kv_block_counts[i] = new_block_count
+
+ # Update last_kv_block_id to point to the previous block (at index new_count - 1)
+ prev_idx = max(new_block_count - 1, 0)
+ prev_block_id = kv_block_ids_list[i][prev_idx]
+ last_kv_block_id[i] = prev_block_id if remove else last_block
+
+ # Clear the released block entry (at index new_count, which was the old last block)
+ scatter_idx = min(new_block_count, max_blocks - 1)
+ if remove:
+ kv_block_ids[i, scatter_idx] = -1
+
+ remove_mask[i] = remove
+
+ return blocks_to_release, remove_mask
+
+
+# pylint: disable=line-too-long
+def verify_speculative_tokens(
+ input_tokens, output_tokens, num_decode_requests, num_prefill_requests, num_speculative_tokens
+):
+ """Verify speculative tokens against input tokens and compute acceptance.
+
+ Creates an accepted tokens mask where:
+ - For prefill requests, the token is always accepted.
+ - For decode requests, the first token (base token) is always accepted, then we compare
+ sampled tokens with input tokens and accept consecutive matches.
+ Then finds the index of the last accepted token per request.
+
+ Example (assume 1, 2, and 0 spec tokens are accepted in the first 3 decode requests):
+ input_tokens_required: [ a5 a6s a7s | b3 b4s b5s | c6 c7s c8s | d2 | e4 ] # Size 11
+ Output tokens [ a6o a7o a8o | b40 b5o b6o | c7o c8o c9o | d3o | e5o ]
+ Output tokens right shift [ d3o a6o a7o | a8o b40 b5o | b6o c7o c8o | c9o | d3o ]
+ Accepted tokens mask [ 1 1 0 | 1 1 1 | 1 0 0 | 1 | 1 ]
+ Last one indices [ 1 | 5 | 6 | 9 | 10 ]
+
+ Returns:
+ tuple: (last_one_indices, accepted_tokens_mask, input_tokens) where
+ last_one_indices contains the index of the last accepted token per request.
+ """
+ if input_tokens.ndim == 2:
+ input_tokens = input_tokens.squeeze(0)
+
+ stride = num_speculative_tokens + 1
+ active_request_count = num_decode_requests + num_prefill_requests
+ decode_len = num_decode_requests * stride
+
+ # Initialize mask with False to prevent boundary bleed
+ accepted_tokens_mask = torch.zeros_like(input_tokens, dtype=torch.bool)
+
+ # Safe decode token verification without cross-batch boundary contamination
+ decode_mask_2d = None
+ if num_decode_requests > 0:
+ decode_inputs = input_tokens[:decode_len].reshape(num_decode_requests, stride)
+ decode_outputs = output_tokens[:decode_len].reshape(num_decode_requests, stride)
+
+ # Shift outputs right by 1 *within* each request to align sampled tokens with input targets
+ decode_outputs_shifted = decode_outputs.roll(1, dims=1)
+ decode_mask_2d = decode_inputs == decode_outputs_shifted
+ # The first token (base token) is always accepted
+ decode_mask_2d[:, 0] = True
+ # Enforce consecutive acceptance: cummin propagates False to the right
+ decode_mask_2d = decode_mask_2d.cummin(dim=1).values
+ accepted_tokens_mask[:decode_len] = decode_mask_2d.flatten()
+
+ # Make all prefill tokens accepted
+ if num_prefill_requests > 0:
+ accepted_tokens_mask[decode_len:] = True
+
+ last_one_indices = torch.full(
+ (active_request_count,), -1, device=input_tokens.device, dtype=torch.long
+ )
+
+ if num_decode_requests > 0:
+ # Summing the consecutive mask gives the count; subtract 1 for the local index
+ local_last_indices = decode_mask_2d.sum(dim=1) - 1
+ row_offsets = torch.arange(num_decode_requests, device=input_tokens.device) * stride
+ last_one_indices[:num_decode_requests] = row_offsets + local_last_indices
+
+ if num_prefill_requests > 0:
+ prefill_valid = torch.nonzero(accepted_tokens_mask[decode_len:]).squeeze(-1) + decode_len
+ last_one_indices[num_decode_requests:] = prefill_valid
+
+ return last_one_indices, accepted_tokens_mask, input_tokens
+
+
+# pylint: disable=line-too-long
+def prepare_next_forward_pass(
+ num_decode_requests,
+ output_tokens,
+ required_logit_indices,
+ last_one_indices,
+ accepted_tokens_mask,
+ input_tokens,
+ sampled_tokens_buf,
+ last_accepted_seq_buf,
+ accepted_tokens_per_request,
+ accepted_token_counts,
+ num_speculative_tokens,
+):
+ """Prepare data for the next forward pass after speculative token verification.
+
+ For each active request:
+ - Store the final sampled tokens for the next forward pass.
+ - Store the last accepted positions in the packed sequence for serial
+ MTP computation after verification.
+
+ For decode requests, extract accepted tokens and counts:
+ input_tokens_required: [ a5 a6s a7s | b3 b4s b5s | c6 c7s c8s | d2 | e4 ]
+ Accepted tokens mask [ 1 1 0 | 1 1 1 | 1 0 0 | 1 | 1 ]
+ Accepted tokens [ [a6s -1] | [b4s b5s] | [-1 -1] ] # Only decode requests (prefill defaults to -1)
+ Accepted token counts [ 1 | 2 | 0 ] # Prefill defaults to 0
+
+ Writes results into the pre-allocated buffers provided by the caller.
+ """
+ active_request_count = last_one_indices.shape[0]
+ stride = num_speculative_tokens + 1
+
+ for pid in range(active_request_count):
+ idx = last_one_indices[pid].item()
+
+ # Store the final sampled tokens for the next forward pass.
+ sampled_tokens_buf[pid] = output_tokens[idx]
+
+ # Store the last accepted positions in the packed sequence for serial
+ # MTP computation after verification.
+ last_accepted_seq_buf[pid] = required_logit_indices[idx]
+
+ # Extract accepted tokens and counts for decode requests.
+ # For prefill it is always set to 1. For decode, the first token is always accepted,
+ # then we compare with input tokens and accept the next tokens if its a match.
+ if pid < num_decode_requests:
+ base = pid * stride
+ # Skip the first token of every decode request (i.e a5, b3, c6)
+ for s in range(num_speculative_tokens):
+ pos = base + 1 + s
+ if accepted_tokens_mask[pos]:
+ accepted_tokens_per_request[pid, s] = input_tokens[pos]
+ else:
+ accepted_tokens_per_request[pid, s] = -1
+
+ count = 0
+ for s in range(num_speculative_tokens):
+ if accepted_tokens_per_request[pid, s].item() != -1:
+ count += 1
+ accepted_token_counts[pid] = count
+
+
+def mamba_state_selective_copy(
+ intermediate_states, current_states, prefill_status, state_idx, accepted_counts, num_layers
+):
+ """Mamba speculative rewind state update.
+
+ For each decode request, copies
+ `intermediate[layer, slot, accepted_count, ...]` →
+ `current[layer, slot, ...]` for every Mamba layer.
+ """
+ N = prefill_status.shape[0]
+ for i in range(N):
+ if prefill_status[i].item() == 1:
+ continue
+ slot = state_idx[i].item()
+ accepted = accepted_counts[i].item()
+ for layer in range(num_layers):
+ current_states[layer, slot] = intermediate_states[layer, slot, accepted]
diff --git a/megatron/core/inference/text_generation_controllers/mtp_utils_triton.py b/megatron/core/inference/text_generation_controllers/mtp_utils_triton.py
new file mode 100644
index 00000000000..37ff55c1e99
--- /dev/null
+++ b/megatron/core/inference/text_generation_controllers/mtp_utils_triton.py
@@ -0,0 +1,456 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+import math
+
+import torch
+
+try:
+ import triton
+ import triton.language as tl
+
+ HAVE_TRITON = True
+except ImportError:
+ from unittest.mock import MagicMock
+
+ from megatron.core.utils import null_decorator
+
+ triton = MagicMock()
+ triton.jit = null_decorator
+ tl = MagicMock()
+ HAVE_TRITON = False
+
+
+# ---------------------------------------------------------------------------
+# Kernel 1: KV-cache rewind for speculative decoding
+# ---------------------------------------------------------------------------
+@triton.jit
+def _rewind_kv_cache_kernel(
+ # Per-request input (read-only)
+ ACCEPTED_COUNTS_PTR,
+ PREFILL_STATUS_PTR,
+ # Per-request state (read-write, updated in-place)
+ LAST_KV_BLOCK_OFFSET_PTR,
+ KV_LENGTH_OFFSETS_PTR,
+ KV_BLOCK_COUNTS_PTR,
+ LAST_KV_BLOCK_ID_PTR,
+ # 2-D table [N, max_blocks] (read-write)
+ KV_BLOCK_IDS_PTR,
+ # Per-request outputs
+ BLOCKS_TO_RELEASE_PTR,
+ REMOVE_MASK_PTR,
+ # Strides / limits
+ kv_block_ids_stride,
+ max_blocks_minus_1,
+ num_active_requests,
+ # Compile-time constants
+ NUM_SPEC_TOKENS: tl.constexpr,
+ BLOCK_SIZE_TOKENS: tl.constexpr,
+):
+ """Rewind KV-cache bookkeeping for one request after speculative verification.
+
+ Grid: may be padded beyond active requests for CUDA-graph compatibility.
+ Each program handles exactly one request. Programs with
+ `pid >= num_active_requests` are padding and produce safe no-op outputs.
+ """
+ pid = tl.program_id(0)
+
+ # Padding programs: write safe defaults and skip all state mutation.
+ if pid >= num_active_requests:
+ tl.store(BLOCKS_TO_RELEASE_PTR + pid, 0)
+ tl.store(REMOVE_MASK_PTR + pid, False)
+ return
+
+ # --- Load per-request scalars ---
+ accepted = tl.load(ACCEPTED_COUNTS_PTR + pid)
+ prefill = tl.load(PREFILL_STATUS_PTR + pid)
+ last_offset = tl.load(LAST_KV_BLOCK_OFFSET_PTR + pid)
+ kv_length = tl.load(KV_LENGTH_OFFSETS_PTR + pid)
+ block_count = tl.load(KV_BLOCK_COUNTS_PTR + pid)
+ last_block_id = tl.load(LAST_KV_BLOCK_ID_PTR + pid)
+
+ # --- Compute rewind (zero for prefill requests) ---
+ num_to_rewind = tl.where(prefill == 1, 0, NUM_SPEC_TOKENS - accepted)
+ diff = last_offset - num_to_rewind
+ remove = diff < 0
+
+ # Python-style modulo: ((diff % M) + M) % M to handle negative diff
+ new_offset = ((diff % BLOCK_SIZE_TOKENS) + BLOCK_SIZE_TOKENS) % BLOCK_SIZE_TOKENS
+ tl.store(LAST_KV_BLOCK_OFFSET_PTR + pid, new_offset)
+ tl.store(KV_LENGTH_OFFSETS_PTR + pid, kv_length - num_to_rewind)
+
+ # Save current last block id (will be released by caller if remove is True)
+ tl.store(BLOCKS_TO_RELEASE_PTR + pid, last_block_id)
+
+ # Decrement block count when a block boundary was crossed
+ new_block_count = tl.where(remove, block_count - 1, block_count)
+ tl.store(KV_BLOCK_COUNTS_PTR + pid, new_block_count)
+
+ # Gather previous block id from the 2-D table
+ kv_row_base = pid.to(tl.int64) * kv_block_ids_stride
+ prev_idx = tl.maximum(new_block_count - 1, 0)
+ prev_block_id = tl.load(KV_BLOCK_IDS_PTR + kv_row_base + prev_idx)
+
+ # Conditionally update last block id
+ tl.store(LAST_KV_BLOCK_ID_PTR + pid, tl.where(remove, prev_block_id, last_block_id))
+
+ # Clear released block entry via scatter
+ scatter_idx = tl.minimum(new_block_count, max_blocks_minus_1)
+ current_val = tl.load(KV_BLOCK_IDS_PTR + kv_row_base + scatter_idx)
+ tl.store(KV_BLOCK_IDS_PTR + kv_row_base + scatter_idx, tl.where(remove, -1, current_val))
+
+ # Output remove mask for the caller (to release blocks outside this kernel)
+ tl.store(REMOVE_MASK_PTR + pid, remove)
+
+
+def rewind_kv_cache(
+ accepted_counts,
+ prefill_status,
+ last_kv_block_offset,
+ kv_length_offsets,
+ kv_block_counts,
+ last_kv_block_id,
+ kv_block_ids,
+ num_speculative_tokens,
+ block_size_tokens,
+ num_active_requests=None,
+):
+ """Launch the KV-cache rewind Triton kernel.
+
+ Args:
+ num_active_requests: Number of real (non-padding) requests. When the
+ grid is padded beyond this count, the kernel skips padding
+ programs so stale data in padding slots cannot corrupt
+ bookkeeping. Defaults to `accepted_counts.shape[0]` (no
+ padding).
+
+ Returns:
+ (blocks_to_release, remove_mask) — same semantics as the original
+ torch.compile'd `_rewind_kv_cache` (KV-cache portion only; Mamba
+ state updates are handled separately by the caller).
+ """
+ N = accepted_counts.shape[0]
+ if num_active_requests is None:
+ num_active_requests = N
+ if N == 0:
+ return (
+ torch.empty(0, device=accepted_counts.device, dtype=last_kv_block_id.dtype),
+ torch.empty(0, device=accepted_counts.device, dtype=torch.bool),
+ )
+
+ blocks_to_release = torch.empty_like(last_kv_block_id)
+ remove_mask = torch.empty(N, device=accepted_counts.device, dtype=torch.bool)
+
+ _rewind_kv_cache_kernel[(N,)](
+ accepted_counts,
+ prefill_status,
+ last_kv_block_offset,
+ kv_length_offsets,
+ kv_block_counts,
+ last_kv_block_id,
+ kv_block_ids,
+ blocks_to_release,
+ remove_mask,
+ kv_block_ids_stride=kv_block_ids.stride(0),
+ max_blocks_minus_1=kv_block_ids.shape[1] - 1,
+ num_active_requests=num_active_requests,
+ NUM_SPEC_TOKENS=num_speculative_tokens,
+ BLOCK_SIZE_TOKENS=block_size_tokens,
+ )
+ return blocks_to_release, remove_mask
+
+
+# ---------------------------------------------------------------------------
+# Kernel 2: Verify speculative tokens
+# ---------------------------------------------------------------------------
+@triton.jit
+def _verify_speculative_tokens_kernel(
+ INPUT_TOKENS_PTR,
+ OUTPUT_TOKENS_PTR,
+ # Outputs
+ ACCEPTED_MASK_PTR,
+ LAST_ONE_INDICES_PTR,
+ # Runtime scalars
+ num_decode_requests,
+ decode_len,
+ # Compile-time constants
+ STRIDE: tl.constexpr, # num_speculative_tokens + 1
+ BLOCK_SIZE: tl.constexpr, # next_power_of_2(STRIDE)
+):
+ """Verify speculative tokens for one request.
+
+ Grid: (active_request_count,)
+ Programs 0..num_decode_requests-1 handle decode requests.
+ Programs num_decode_requests..end handle prefill requests.
+ """
+ pid = tl.program_id(0)
+
+ if pid < num_decode_requests:
+ base = pid * STRIDE
+ offsets = tl.arange(0, BLOCK_SIZE)
+ valid = offsets < STRIDE
+
+ input_toks = tl.load(INPUT_TOKENS_PTR + base + offsets, mask=valid, other=0)
+
+ # Build shifted output: shifted[i] = output[i-1].
+ # Position 0 uses a dummy load (always accepted regardless).
+ safe_shifted = tl.where(offsets > 0, offsets - 1, 0)
+ shifted_output = tl.load(OUTPUT_TOKENS_PTR + base + safe_shifted, mask=valid, other=0)
+
+ # First token is always accepted; rest must match shifted output.
+ match = tl.where(offsets == 0, 1, (input_toks == shifted_output).to(tl.int32))
+ match = tl.where(valid, match, 0)
+
+ # Consecutive acceptance via cumulative-sum trick:
+ # accepted[i] iff cumsum(match)[i] == i + 1
+ cumsum = tl.cumsum(match, axis=0)
+ accepted = (cumsum == (offsets + 1)) & valid
+
+ tl.store(ACCEPTED_MASK_PTR + base + offsets, accepted, mask=valid)
+
+ accepted_count = tl.sum(accepted.to(tl.int32))
+ tl.store(LAST_ONE_INDICES_PTR + pid, (base + accepted_count - 1).to(tl.int64))
+ else:
+ # Prefill request — single token, always accepted
+ prefill_idx = decode_len + (pid - num_decode_requests)
+ tl.store(ACCEPTED_MASK_PTR + prefill_idx, 1)
+ tl.store(LAST_ONE_INDICES_PTR + pid, prefill_idx.to(tl.int64))
+
+
+def verify_speculative_tokens(
+ input_tokens, output_tokens, num_decode_requests, num_prefill_requests, num_speculative_tokens
+):
+ """Launch the speculative-token verification Triton kernel.
+
+ Returns:
+ (last_one_indices, accepted_tokens_mask, input_tokens)
+ matching the original `_verify_speculative_tokens` signature.
+ """
+ if input_tokens.ndim == 2:
+ input_tokens = input_tokens.squeeze(0)
+
+ device = input_tokens.device
+ active_request_count = num_decode_requests + num_prefill_requests
+ stride = num_speculative_tokens + 1
+ decode_len = num_decode_requests * stride
+
+ accepted_tokens_mask = torch.zeros_like(input_tokens, dtype=torch.bool)
+ last_one_indices = torch.full((active_request_count,), -1, device=device, dtype=torch.long)
+
+ if active_request_count > 0:
+ block_size = triton.next_power_of_2(stride)
+ _verify_speculative_tokens_kernel[(active_request_count,)](
+ input_tokens,
+ output_tokens,
+ accepted_tokens_mask,
+ last_one_indices,
+ num_decode_requests=num_decode_requests,
+ decode_len=decode_len,
+ STRIDE=stride,
+ BLOCK_SIZE=block_size,
+ )
+
+ return last_one_indices, accepted_tokens_mask, input_tokens
+
+
+# ---------------------------------------------------------------------------
+# Kernel 3: Prepare speculative tokens for next forward pass
+# ---------------------------------------------------------------------------
+@triton.jit
+def _prepare_next_forward_pass_kernel(
+ OUTPUT_TOKENS_PTR,
+ REQUIRED_LOGIT_INDICES_PTR,
+ LAST_ONE_INDICES_PTR,
+ INPUT_TOKENS_PTR,
+ ACCEPTED_MASK_PTR,
+ # Outputs
+ SAMPLED_TOKENS_OUT_PTR,
+ LAST_ACCEPTED_SEQ_OUT_PTR,
+ ACCEPTED_TOKENS_OUT_PTR,
+ ACCEPTED_COUNTS_OUT_PTR,
+ # Strides
+ accepted_tokens_out_stride,
+ # Runtime scalars
+ num_decode_requests,
+ # Compile-time constants
+ STRIDE: tl.constexpr, # num_speculative_tokens + 1
+ NUM_SPEC_TOKENS: tl.constexpr,
+ SPEC_BLOCK_SIZE: tl.constexpr, # next_power_of_2(NUM_SPEC_TOKENS)
+):
+ """Gather final tokens and extract accepted speculative tokens per request.
+
+ Grid: (active_request_count,)
+ """
+ pid = tl.program_id(0)
+
+ # --- Gather final sampled token and sequence index for every request ---
+ idx = tl.load(LAST_ONE_INDICES_PTR + pid)
+ tl.store(SAMPLED_TOKENS_OUT_PTR + pid, tl.load(OUTPUT_TOKENS_PTR + idx))
+ tl.store(LAST_ACCEPTED_SEQ_OUT_PTR + pid, tl.load(REQUIRED_LOGIT_INDICES_PTR + idx))
+
+ # --- For decode requests: extract accepted tokens and count ---
+ if pid < num_decode_requests:
+ base = pid * STRIDE
+ spec_offsets = tl.arange(0, SPEC_BLOCK_SIZE)
+ spec_valid = spec_offsets < NUM_SPEC_TOKENS
+ token_positions = base + 1 + spec_offsets # skip first (base) token
+
+ tokens = tl.load(INPUT_TOKENS_PTR + token_positions, mask=spec_valid, other=0)
+ mask_val = tl.load(ACCEPTED_MASK_PTR + token_positions, mask=spec_valid, other=0)
+ accepted = mask_val != 0
+
+ result = tl.where(accepted & spec_valid, tokens, -1)
+
+ out_base = pid.to(tl.int64) * accepted_tokens_out_stride
+ tl.store(ACCEPTED_TOKENS_OUT_PTR + out_base + spec_offsets, result, mask=spec_valid)
+
+ count = tl.sum((accepted & spec_valid).to(tl.int64))
+ tl.store(ACCEPTED_COUNTS_OUT_PTR + pid, count)
+
+
+def prepare_next_forward_pass(
+ num_decode_requests,
+ output_tokens,
+ required_logit_indices,
+ last_one_indices,
+ accepted_tokens_mask,
+ input_tokens,
+ sampled_tokens_buf,
+ last_accepted_seq_buf,
+ accepted_tokens_per_request,
+ accepted_token_counts,
+ num_speculative_tokens,
+):
+ """Launch the prepare-next-forward-pass Triton kernel.
+
+ Writes results into the pre-allocated buffers provided by the caller.
+ """
+ active_request_count = last_one_indices.shape[0]
+ if active_request_count == 0:
+ return
+
+ stride = num_speculative_tokens + 1
+ spec_block_size = triton.next_power_of_2(num_speculative_tokens)
+
+ _prepare_next_forward_pass_kernel[(active_request_count,)](
+ output_tokens,
+ required_logit_indices,
+ last_one_indices,
+ input_tokens,
+ accepted_tokens_mask,
+ sampled_tokens_buf,
+ last_accepted_seq_buf,
+ accepted_tokens_per_request,
+ accepted_token_counts,
+ accepted_tokens_out_stride=accepted_tokens_per_request.stride(0),
+ num_decode_requests=num_decode_requests,
+ STRIDE=stride,
+ NUM_SPEC_TOKENS=num_speculative_tokens,
+ SPEC_BLOCK_SIZE=spec_block_size,
+ )
+
+
+# ---------------------------------------------------------------------------
+# Kernel 4: Mamba state selective copy (eliminates temporary allocations)
+# ---------------------------------------------------------------------------
+@triton.jit
+def _mamba_state_selective_copy_kernel(
+ # Source: intermediate states [L, M, S+1, *state_shape]
+ SRC_PTR,
+ # Destination: current states [L, M, *state_shape]
+ DST_PTR,
+ # Per-request index arrays
+ PREFILL_STATUS_PTR, # [N] 0=decode, 1=prefill
+ STATE_IDX_PTR, # [N] maps request → mamba state slot
+ ACCEPTED_PTR, # [N] accepted token index per request
+ # Strides (in elements)
+ src_stride_layer,
+ src_stride_slot,
+ src_stride_spec,
+ dst_stride_layer,
+ dst_stride_slot,
+ # Data size
+ STATE_SIZE,
+ # Compile-time
+ BLOCK_SIZE: tl.constexpr,
+):
+ """Copy intermediate Mamba state to current state for decode requests.
+
+ Grid: (N, L, num_chunks)
+ - dim 0: active request index
+ - dim 1: mamba layer index
+ - dim 2: chunk of the flattened state vector
+
+ No-op for prefill requests.
+ """
+ pid_req = tl.program_id(0)
+ pid_layer = tl.program_id(1)
+ pid_chunk = tl.program_id(2)
+
+ # Skip prefill requests immediately.
+ prefill = tl.load(PREFILL_STATUS_PTR + pid_req)
+ if prefill == 1:
+ return
+
+ state_idx = tl.load(STATE_IDX_PTR + pid_req).to(tl.int64)
+ accepted = tl.load(ACCEPTED_PTR + pid_req).to(tl.int64)
+
+ chunk_start = pid_chunk * BLOCK_SIZE
+ offsets = tl.arange(0, BLOCK_SIZE)
+ elem_offsets = chunk_start + offsets
+ mask = elem_offsets < STATE_SIZE
+
+ src_base = (
+ pid_layer.to(tl.int64) * src_stride_layer
+ + state_idx * src_stride_slot
+ + accepted * src_stride_spec
+ )
+ dst_base = pid_layer.to(tl.int64) * dst_stride_layer + state_idx * dst_stride_slot
+
+ data = tl.load(SRC_PTR + src_base + elem_offsets, mask=mask)
+ tl.store(DST_PTR + dst_base + elem_offsets, data, mask=mask)
+
+
+def mamba_state_selective_copy(
+ intermediate_states, current_states, prefill_status, state_idx, accepted_counts, num_layers
+):
+ """Copy accepted intermediate Mamba states to current states in-place.
+
+ For each decode request, copies
+ `intermediate[layer, slot, accepted_count, ...]` →
+ `current[layer, slot, ...]` for every Mamba layer.
+
+ Args:
+ intermediate_states: `(L, M, S+1, *state_shape)` — intermediate buffer.
+ current_states: `(L, M, *state_shape)` — current state buffer (updated in-place).
+ prefill_status: `(N,)` int tensor — 0 for decode, 1 for prefill.
+ state_idx: `(N,)` int tensor — mamba state slot index per request.
+ accepted_counts: `(N,)` int tensor — accepted token index per request.
+ num_layers: number of Mamba layers (first dim of the state tensors).
+ """
+ N = prefill_status.shape[0]
+ if N == 0:
+ return
+
+ # The state vector to copy per (layer, request) is the product of all
+ # trailing dimensions after the speculative-token axis.
+ # intermediate shape: (L, M, S+1, *state_shape) → state_size = prod(state_shape)
+ state_size = math.prod(intermediate_states.shape[3:])
+
+ BLOCK_SIZE = 1024
+ num_chunks = triton.cdiv(state_size, BLOCK_SIZE)
+ grid = (N, num_layers, num_chunks)
+
+ _mamba_state_selective_copy_kernel[grid](
+ intermediate_states,
+ current_states,
+ prefill_status,
+ state_idx,
+ accepted_counts,
+ src_stride_layer=intermediate_states.stride(0),
+ src_stride_slot=intermediate_states.stride(1),
+ src_stride_spec=intermediate_states.stride(2),
+ dst_stride_layer=current_states.stride(0),
+ dst_stride_slot=current_states.stride(1),
+ STATE_SIZE=state_size,
+ BLOCK_SIZE=BLOCK_SIZE,
+ )
diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
index ba190f799f8..3e788fec0b1 100644
--- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py
+++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py
@@ -4,13 +4,14 @@
import concurrent
import copy
import functools
-import inspect
from collections import defaultdict
from typing import Any, Dict, List, Optional, OrderedDict, Tuple, Union
+import numpy as np
import torch
import torch.nn.functional as F
from torch import Tensor
+from torch.cuda.nvtx import range_pop, range_push
from megatron.core import parallel_state
from megatron.core.inference.async_stream import AsyncStream
@@ -25,14 +26,29 @@
AbstractModelInferenceWrapper,
)
from megatron.core.inference.sampling_params import SamplingParams
-from megatron.core.inference.utils import get_attention_mask, set_decode_expert_padding
+from megatron.core.inference.utils import (
+ get_attention_mask,
+ set_decode_expert_padding,
+ set_moe_metadata_sync,
+)
from megatron.core.models.multimodal.llava_model import LLaVAModel
-from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region
-from megatron.core.transformer.enums import CudaGraphScope
+from megatron.core.tensor_parallel.mappings import (
+ gather_from_sequence_parallel_region,
+ scatter_to_sequence_parallel_region,
+)
from megatron.core.transformer.moe.moe_layer import BaseMoELayer
from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
from megatron.core.transformer.utils import set_model_to_sequence_parallel
-from megatron.core.utils import get_asyncio_loop, get_model_config, get_pg_size, unwrap_model
+from megatron.core.utils import (
+ accepts_parameter,
+ get_asyncio_loop,
+ get_model_config,
+ get_pg_size,
+ nvtx_range_pop,
+ nvtx_range_push,
+ round_up_to_nearest_multiple,
+ unwrap_model,
+)
try:
import transformer_engine as te # pylint: disable=unused-import
@@ -43,6 +59,13 @@
HAVE_TE = False
from megatron.core.inference.batch_dimensions_utils import InferenceBatchDimensions
+from megatron.core.inference.sampling import FlashInferSampling, Sampling, TorchSampling
+from megatron.core.inference.text_generation_controllers.mtp_utils_pytorch import rewind_kv_cache
+from megatron.core.inference.text_generation_controllers.mtp_utils_triton import (
+ mamba_state_selective_copy,
+ prepare_next_forward_pass,
+ verify_speculative_tokens,
+)
# pylint: disable=line-too-long
@@ -84,6 +107,16 @@ def __init__(self, inference_wrapped_model: AbstractModelInferenceWrapper, token
self.num_mtp_heads = self._get_mtp_num_heads()
self.sampling_rng.manual_seed(self.model_config.inference_sampling_seed)
+ if (
+ self.model_config.cuda_graph_impl == "local"
+ and self.model_config.expert_model_parallel_size > 1
+ and self.model_config.transformer_impl != "inference_optimized"
+ ):
+ assert self.model_config.moe_pad_experts_for_cuda_graph_inference, (
+ "--moe-pad-experts-for-cuda-graph-inference must be set when using "
+ "CUDA graphs with expert parallelism"
+ )
+
if self.inference_wrapped_model.inference_context.is_dynamic_batching():
self._init_dynamic_sampling_tensors()
@@ -109,6 +142,11 @@ def _init_dynamic_sampling_tensors(self):
"""Initialize tensors needed for dynamic sampling."""
context = self.inference_wrapped_model.inference_context
max_requests = context.max_requests
+ if context.config.materialize_only_last_token_logits:
+ # Under MTP, each decode request emits (num_speculative_tokens + 1) logit rows
+ max_logits = max_requests * (self.num_speculative_tokens + 1)
+ else:
+ max_logits = context.max_tokens
# Callback to get request IDs that should be marked as finished due to stop words
self._get_stop_word_finished_ids_callback = None
@@ -116,46 +154,79 @@ def _init_dynamic_sampling_tensors(self):
device = torch.cuda.current_device()
logits_dtype = self.inference_wrapped_model.config.params_dtype
- self._sampling_backend = "torch"
- self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
- # Speculative tokens tensor will be allocated later when num_speculative_tokens is set by the engine
- self._accepted_tokens_per_request = None
- # MTP tensor will be allocated later when num_speculative_tokens is set by the engine
- self._sampled_mtp_tokens_cuda = None
- # Last accepted sequence indices for serial MTP computation
- self._last_accepted_seq_indices = None
+ self._sampling_backend = context.config.sampling_backend
+ self._enable_cuda_graph = self.model_config.cuda_graph_impl == "local"
- # Keep track of request metadata.
- self._request_metadata: Dict[str, Tensor] = {}
- for label, dtype, on_gpu in context.request_metadata_types:
- tensor = context.request_metadata[label]
- if not on_gpu:
- # Create pinned tensors for request metadata that lives on CPU.
- # This is metadata which requires D2H copies, such as top_k for torch sampling.
- tensor = torch.empty_like(tensor, device="cpu", pin_memory=True)
- self._request_metadata[label] = tensor
-
- # Used for inefficient torch sampling.
- if self._sampling_backend == "torch":
- self._torch_sampling_buckets: List[Tuple] = []
-
- self._init_mtp_sampling_tensor()
-
- def _init_mtp_sampling_tensor(self):
- """Initialize the MTP sampling tensor after num_speculative_tokens is set."""
- if self.num_speculative_tokens is not None and self.num_speculative_tokens > 0:
- context = self.inference_wrapped_model.inference_context
- max_requests = context.max_requests
- device = torch.cuda.current_device()
- self._sampled_mtp_tokens_cuda = torch.empty(
- [self.num_speculative_tokens, max_requests], dtype=torch.int64, device=device
+ # Initialize bookkeeping tensors.
+ if self._enable_cuda_graph:
+ self._all_logits_cuda = torch.zeros(
+ (1, max_logits, self.vocab_size), dtype=logits_dtype, device=device
)
- self._accepted_tokens_per_request = (
- torch.ones(
- [max_requests, self.num_speculative_tokens], dtype=torch.int64, device=device
- )
- * -1
+ else:
+ self._all_logits_cuda = None
+ # Speculative path:
+ # - `self._sampled_tokens_cuda` is pre-allocated by `_init_mtp_sampling_tensors`.
+ # - The tensor cannot be reused between the Triton kernel and the sampling graph.
+ # Non-speculative path:
+ # - `self._sampled_tokens_cuda` is rebound to the output of `sample_kernel`,
+ # which uses CudaGraphManager syntactic sugar to keep it as a static tensor.
+ self._sampled_tokens_cuda = None
+
+ # Sampling backend: provides the sampling kernel.
+ if self._sampling_backend == "flashinfer":
+ self._sampling: Sampling = FlashInferSampling(
+ self.vocab_size,
+ self.sampling_rng,
+ config=self.model_config,
+ enable_cuda_graph=self._enable_cuda_graph,
)
+ else:
+ self._sampling: Sampling = TorchSampling(self.sampling_rng, self.vocab_size)
+
+ # Cache values that are constant across inference steps.
+ self._unwrapped_model = unwrap_model(self.inference_wrapped_model.model)
+ self._is_last_pp_stage = is_pipeline_last_stage(self.pp_group)
+ self._tp_size = get_pg_size(self.inference_wrapped_model.tp_group)
+ self._sp_enabled = self.model_config.sequence_parallel and self._tp_size > 1
+
+ self._init_mtp_sampling_tensors()
+
+ def _init_mtp_sampling_tensors(self):
+ """Pre-allocate MTP sampling tensors.
+
+ Addresses must be stable across steps for CUDA graph capture.
+ """
+ if not self.num_speculative_tokens:
+ self._sampled_mtp_tokens_cuda = None
+ self._accepted_tokens_per_request = None
+ self._last_accepted_seq_indices = None
+ return
+
+ context = self.inference_wrapped_model.inference_context
+ max_requests = context.max_requests
+ device = torch.cuda.current_device()
+ self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device)
+ self._sampled_mtp_tokens_cuda = torch.empty(
+ [self.num_speculative_tokens, max_requests], dtype=torch.int64, device=device
+ )
+ self._accepted_tokens_per_request = (
+ torch.ones(
+ [max_requests, self.num_speculative_tokens], dtype=torch.int64, device=device
+ )
+ * -1
+ )
+ self._accepted_token_counts_per_request = torch.zeros(
+ max_requests, dtype=torch.int64, device=device
+ )
+ self._last_accepted_seq_indices_buf = torch.empty(
+ max_requests, dtype=torch.int64, device=device
+ )
+ self._last_accepted_seq_indices = None
+ self._num_mtp_depths = min(self.num_speculative_tokens, self.num_mtp_heads)
+ self._mtp_token_ids_buf = torch.empty([1, max_requests], dtype=torch.int64, device=device)
+ self._mtp_position_ids_buf = torch.empty(
+ [1, max_requests], dtype=torch.int64, device=device
+ )
@staticmethod
def tokenize_prompt(tokenizer, prompt: str, add_BOS: bool = False) -> List[int]:
@@ -206,12 +277,7 @@ def detokenize(
while tokens and tokens[-1] == tokenizer.eod:
tokens = tokens[:-1]
- sig_params = inspect.signature(tokenizer.detokenize).parameters.values()
- detok_accepts_skip = any(
- p.name == "skip_special_tokens" or p.kind == inspect.Parameter.VAR_KEYWORD
- for p in sig_params
- )
- if detok_accepts_skip:
+ if accepts_parameter(tokenizer.detokenize, "skip_special_tokens"):
return tokenizer.detokenize(tokens, skip_special_tokens=skip_special_tokens)
else:
return tokenizer.detokenize(tokens)
@@ -269,95 +335,6 @@ def detokenize_generations(
return text, prompts_plus_generations_segments
- def _torch_sampling_func(
- self,
- last_token_logits: torch.Tensor,
- temperature: float,
- top_k: int,
- top_p: float,
- vocab_size: Optional[int] = None,
- ):
- """Samples the logits to generate outputs
-
- Given the logits of the last token, this function samples it
- according to the parameters defined in sampling_params
- and returns the samples. If sampling parameters top_n_logprobs > 0
- at each step it also updates the top_n_logprobs dict.
-
- Args:
- last_token_logits (torch.Tensor): The last token logits. A tensor of
- size [batch_size, vocab_size].
- temperature (float): The temperature to use for sampling.
- top_k (int): The top-k value to use for sampling.
- top_p (float): The top-p value to use for sampling.
- vocab_size (int): Obtained from the tokenizer. Defaults to None.
-
- Returns:
- sampled_logits (torch.Tensor): 1D tensor with [batch_size] elements
- """
- assert isinstance(top_p, float)
- assert isinstance(top_k, int)
- assert not (top_k > 0 and top_p > 0.0), "Cannot have top-p and top-k both greater than zero"
- assert top_p <= 1.0, "top-p should be in (0,1]"
-
- def modify_logits_for_top_k_filtering(logits, top_k):
- """Set the logits for none top-k values to -inf."""
- filter_ = logits < torch.topk(logits, top_k)[0][..., -1, None]
- logits.masked_fill_(filter_, float("-Inf"))
-
- def modify_logits_for_top_p_filtering(logits, top_p):
- """Set the logits for none top-p values to -inf."""
- # First sort and calculate cumulative sum of probabilities.
- sorted_logits, sorted_indices = torch.sort(logits, descending=True)
- cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
-
- # Filteration based on the cumulative sum.
- filter_ = cumulative_probs > top_p
- # This shift by 1 is weird and I cannot justify it. This existed
- # in the original implementation:
- # https://github.com/ari-holtzman/degen/blob/master/gen.py
- # and I guess it is needed so keeping it for now.
- # Clone needed: filter_[:, 1:] and filter_[:, :-1] are overlapping views;
- # without clone, each write would corrupt the next read during the shift.
- filter_[:, 1:] = filter_[:, :-1].clone()
- # Make sure we at least have one token to select from.
- filter_[..., 0] = 0
-
- # Fill in the filtered part
- filter_ = filter_.scatter(1, sorted_indices, filter_)
- logits.masked_fill_(filter_, float("-Inf"))
-
- # Greedy sampling
- if top_k == 1:
- sampled_logits = torch.argmax(last_token_logits, dim=-1)
- else:
- # Clone needed: .div_() and masked_fill_() below modify in-place,
- # which would mutate the caller's tensor without this clone.
- last_token_logits = last_token_logits.clone()
- if temperature != 1.0:
- last_token_logits.div_(temperature)
- if top_k > 1:
- assert top_k <= last_token_logits.size(1), "top-k is larger than logit size."
- if vocab_size:
- assert top_k < vocab_size, "top-k is larger than vocab size."
- modify_logits_for_top_k_filtering(last_token_logits, top_k)
-
- elif top_p > 0.0:
- modify_logits_for_top_p_filtering(last_token_logits, top_p)
-
- # After filtering, we need to recalculate the distribution.
- probabilities = last_token_logits.softmax(dim=-1)
-
- sampled_logits = torch.multinomial(
- probabilities, num_samples=1, generator=self.sampling_rng
- ).view(-1)
-
- # If vocab size is provided, make sure the samples are in in the range [0, vocab-size).
- if vocab_size:
- sampled_logits = torch.clamp(sampled_logits, min=0, max=(vocab_size - 1))
-
- return sampled_logits
-
def sample_from_logits(
self,
last_token_logits: torch.Tensor,
@@ -444,7 +421,14 @@ def sample_from_logits(
top_k = sampling_params.top_k
temperature = sampling_params.temperature
- return self._torch_sampling_func(last_token_logits, temperature, top_k, top_p, vocab_size)
+ return TorchSampling.sample_from_logits(
+ last_token_logits,
+ temperature,
+ top_k,
+ top_p,
+ generator=self.sampling_rng,
+ vocab_size=vocab_size,
+ )
def update_generation_status(
self,
@@ -556,17 +540,37 @@ def _dynamic_step_context_init(
position_ids (Tensor): The active position IDs.
"""
context = self.inference_wrapped_model.inference_context
- active_request_slice = slice(context.paused_request_count, context.total_request_count)
# Remove Float16Module wrapper if it exists
unwrapped_model = unwrap_model(self.inference_wrapped_model.model)
model_config = get_model_config(unwrapped_model)
- # Initialize attention state.
+ # Initialize attention state (100% CPU computation).
+ range_push("initialize_attention_state")
context.initialize_attention_state(
construct_graph_dimensions=construct_graph_dimensions,
is_expert_parallel_dummy_cuda_graph_step=is_dummy_forward,
)
+ range_pop()
+
+ # Single batch CPU-to-GPU transfer of bookkeeping state.
+ range_push("transfer_bookkeeping_to_gpu")
+ context.transfer_bookkeeping_to_gpu()
+ range_pop()
+
+ set_moe_metadata_sync(unwrapped_model)
+
+ # Derive the MTP padded batch size from the existing padded graph dimensions.
+ # For MoE models this is post EP sync. In eager mode MTP uses locally SP-aligned
+ # batch size instead.
+ if context.using_cuda_graph_this_step():
+ self._mtp_resolved_padded_count = context.padded_batch_dimensions.req_count
+ if self._sp_enabled:
+ self._mtp_resolved_padded_count = round_up_to_nearest_multiple(
+ self._mtp_resolved_padded_count, self._tp_size
+ )
+ else:
+ self._mtp_resolved_padded_count = None
# If using symmetric kernels and we are using using nccl
# for prefill turn off symmetric kernels
@@ -597,14 +601,6 @@ def _dynamic_step_context_init(
# Turn off symmetric all reduces for prefill
unwrapped_model.set_symmetric_ar(None)
- # Get request metadata for this step.
- for label, dtype, on_gpu in context.request_metadata_types:
- if not on_gpu:
- # We need a D2H copy from the context to the pinned memory buffer.
- self._request_metadata[label].copy_(
- context.request_metadata[label], non_blocking=True
- )
-
# Get flat tokens, position ids.
# If we are running a dummy forward step we want to use the token count agreed upon
# by all EP ranks rather than the minimum number of tokens.
@@ -615,7 +611,7 @@ def _dynamic_step_context_init(
else:
return context.current_input_and_position_ids()
- def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor) -> Tensor:
+ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor):
"""Forward step the model to get logits for dynamic batching.
This also handles logits-broadcasting for pipeline parallelism.
@@ -625,7 +621,10 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
position_ids (Tensor): The position IDs.
"""
context = self.inference_wrapped_model.inference_context
- active_request_count = context.total_request_count - context.paused_request_count
+ if context.config.materialize_only_last_token_logits:
+ logits_seq_len = context.num_last_token_logits
+ else:
+ logits_seq_len = context.padded_active_token_count
with torch.inference_mode():
logits = self.inference_wrapped_model.run_one_forward_step(
@@ -633,6 +632,9 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
)
# logits shape: [1, seq_len, vocab_size]
+ if not context.config.materialize_only_last_token_logits:
+ assert logits_seq_len == input_ids.shape[1]
+
# Note: When speculative decoding is active (num_speculative_tokens > 0),
# the model skips MTP computation during the forward pass. MTP logits
# will be computed serially after verification to ensure they are
@@ -640,7 +642,7 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
if self.model_is_pipeline_parallel:
if context.config.materialize_only_last_token_logits:
- logits_seq_len = active_request_count
+ logits_seq_len = context.num_last_token_logits
else:
logits_seq_len = input_ids.shape[1]
logits_shape = [1, logits_seq_len, self.vocab_size]
@@ -655,140 +657,77 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor)
pp_group=self.pp_group,
)
- return logits
-
- def _dynamic_step_sample_bookkeeping(self):
- """Perform bookkeeping necessary to sample logits for dynamic batching."""
- context = self.inference_wrapped_model.inference_context
- active_request_slice = slice(context.paused_request_count, context.total_request_count)
-
- if self._sampling_backend == "torch":
- # Bucketize the core sampling parameters.
- # Doing so via list comprehension is orders of magnitude faster than via torch.
- bucket_map = defaultdict(list)
-
- # Shorthands for the dictionary comprehension.
- temp = self._request_metadata["temperature"][active_request_slice].tolist()
- top_k = self._request_metadata["top_k"][active_request_slice].tolist()
- top_p = self._request_metadata["top_p"][active_request_slice].tolist()
-
- for request_index, (t, k, p) in enumerate(zip(temp, top_k, top_p)):
- sampling_params = (t, k, p)
- bucket_map[sampling_params].append(request_index)
-
- # Just unpack the key directly!
- self._torch_sampling_buckets = [
- (indices, *sampling_params) for sampling_params, indices in bucket_map.items()
- ]
+ # Copy logits to contiguous buffer.
+ if self._enable_cuda_graph:
+ self._all_logits_cuda[:, :logits_seq_len, :].copy_(logits[:, :logits_seq_len, :])
+ else:
+ self._all_logits_cuda = logits
- def _rewind_kv_cache(self):
+ def _rewind_kv_cache(self) -> tuple:
"""Update the KV cache bookkeeping for speculative decoding.
After forward pass with speculative tokens, some tokens may be rejected.
- This function "rewinds" the KV cache bookkeeping to reflect only the accepted tokens.
-
- When speculative tokens are rejected, we need to:
- 1. Update request_kv_length_offsets (total sequence length)
- 2. Update request_last_kv_block_offset (position within last block)
- 3. If rewinding crosses a block boundary:
- - Reduce request_kv_block_counts
- - Update request_last_kv_block_id to point to the previous block
- - Clear the entry in request_to_kv_block_ids for the released block
- - Release the block back to the allocator
+ This function "rewinds" the KV cache bookkeeping to reflect only the
+ accepted tokens. The core bookkeeping rewind runs on CPU (mutating the
+ CPU source-of-truth tensors in place); the Mamba hybrid-model state
+ update stays on GPU because it operates on GPU-resident state buffers.
+
+ Returns (blocks_to_release, remove_mask) for the caller to release blocks
+ back to the allocator outside the compiled graph.
"""
context = self.inference_wrapped_model.inference_context
active_request_count = context.total_request_count - context.paused_request_count
active_request_slice = slice(context.paused_request_count, context.total_request_count)
- # Get the accepted token counts for each request
- # Note: _accepted_token_counts is indexed from 0 to active_request_count-1
- accepted_tokens_per_request = self._accepted_token_counts_per_request[:active_request_count]
-
- # Number of tokens to rewind (rejected speculative tokens)
- num_tokens_to_rewind = self.num_speculative_tokens - accepted_tokens_per_request
-
- # For prefill requests, no speculative tokens were forwarded through the model,
- # so there is nothing to rewind.
- request_in_prefill_status = context.request_in_prefill_status_tensor[active_request_slice]
- num_tokens_to_rewind[request_in_prefill_status == 1] = 0
-
- # Save the original offset BEFORE modifying to correctly detect block boundary crossing
- original_offset = context.request_last_kv_block_offset[active_request_slice].clone()
-
- # Check which requests need to rewind to a previous block BEFORE modifying
- # A request crosses back to a previous block if: original_offset - num_tokens_to_rewind < 0
- remove_allocated_blocks_mask = (original_offset - num_tokens_to_rewind) < 0
-
- # Update the offsets
- context.request_last_kv_block_offset[active_request_slice] = (
- original_offset - num_tokens_to_rewind
- ) % context.block_size_tokens
-
- context.request_kv_length_offsets[active_request_slice] = (
- context.request_kv_length_offsets[active_request_slice] - num_tokens_to_rewind
+ # accepted_counts is the only GPU input; D2H a small slice so the
+ # CPU rewind can read its values via .tolist() inside a Python loop.
+ accepted_tokens_per_request_cpu = self._accepted_token_counts_per_request[
+ :active_request_count
+ ].cpu()
+
+ blocks_to_release, remove_mask = rewind_kv_cache(
+ accepted_counts=accepted_tokens_per_request_cpu,
+ prefill_status=context.request_in_prefill_status_tensor[active_request_slice],
+ last_kv_block_offset=context.request_last_kv_block_offset[active_request_slice],
+ kv_length_offsets=context.request_kv_length_offsets[active_request_slice],
+ kv_block_counts=context.request_kv_block_counts[active_request_slice],
+ last_kv_block_id=context.request_last_kv_block_id[active_request_slice],
+ kv_block_ids=context.request_to_kv_block_ids[active_request_slice],
+ num_speculative_tokens=self.num_speculative_tokens,
+ block_size_tokens=context.block_size_tokens,
+ num_active_requests=active_request_count,
)
- # No need to update request_query_lengths (It will be set correctly in the next iteration)
-
- # For requests that crossed back to a previous block, we need to:
- # 1. Reduce the block count by 1
- # 2. Get the block ID to release (current request_last_kv_block_id)
- # 3. Update request_last_kv_block_id to point to the previous block
- # 4. Clear the entry in request_to_kv_block_ids for the released block
- # 5. Release the block back to the allocator
- if remove_allocated_blocks_mask.any():
- # Get indices of requests that need to release a block (relative to active requests)
- requests_needing_release = torch.nonzero(remove_allocated_blocks_mask, as_tuple=True)[0]
- # Convert to absolute indices in the context tensors
- absolute_indices = requests_needing_release + context.paused_request_count
-
- # No clone needed: advanced (fancy) indexing with a tensor already returns
- # a copy, not a view.
- blocks_to_release = context.request_last_kv_block_id[absolute_indices]
-
- # Reduce block counts for requests that crossed back
- context.request_kv_block_counts[absolute_indices] -= 1
-
- # Get the new block counts after decrement
- new_block_counts = context.request_kv_block_counts[absolute_indices]
-
- # Update request_last_kv_block_id to point to the previous block
- # and clear the released block entry in request_to_kv_block_ids
- # Vectorized implementation using advanced indexing:
- # Note: new_block_counts is guaranteed to be > 0 for all requests here, since
- # crossing back to a previous block implies the request had at least 2 blocks.
-
- # Update request_last_kv_block_id to point to the previous block (at index new_count - 1)
- context.request_last_kv_block_id[absolute_indices] = context.request_to_kv_block_ids[
- absolute_indices, new_block_counts - 1
- ]
-
- # Clear the released block entry (at index new_count, which was the old last block)
- context.request_to_kv_block_ids[absolute_indices, new_block_counts] = -1
-
- # Release the blocks back to the allocator
- context.kv_block_allocator.release_memory_blocks(blocks_to_release)
-
- # Mamba speculative rewind state update
+ # Mamba speculative rewind stays on GPU because it mutates GPU-resident
+ # SSM/conv state that the next forward pass reads directly.
if context.is_hybrid_model:
- active_mamba_indices = context.mamba_metadata.request_to_mamba_state_idx[
+ cuda_device = torch.cuda.current_device()
+ # gpu_view.request_in_prefill_status was uploaded by this step's
+ # coalesced H2D and mirrors the active-slice CPU values, so we
+ # don't need to re-upload prefill_status for the Mamba kernels.
+ prefill_status_gpu = context.gpu_view.request_in_prefill_status[:active_request_count]
+ accepted_counts_gpu = self._accepted_token_counts_per_request[:active_request_count]
+ mamba_state_idx = context.mamba_metadata.request_to_mamba_state_idx[
active_request_slice
- ]
- is_decode_mask = context.request_in_prefill_status_tensor[active_request_slice] == 0
- decode_mamba_indices = active_mamba_indices[is_decode_mask]
- accepted_tokens_per_decode_request = accepted_tokens_per_request[is_decode_mask]
-
- if decode_mamba_indices.numel() > 0:
- context.mamba_conv_states[:, decode_mamba_indices] = (
- context.mamba_intermediate_conv_states[
- :, decode_mamba_indices, accepted_tokens_per_decode_request
- ]
- )
- context.mamba_ssm_states[:, decode_mamba_indices] = (
- context.mamba_intermediate_ssm_states[
- :, decode_mamba_indices, accepted_tokens_per_decode_request
- ]
- )
+ ].to(cuda_device, non_blocking=True)
+ mamba_state_selective_copy(
+ intermediate_states=context.mamba_intermediate_conv_states,
+ current_states=context.mamba_conv_states,
+ prefill_status=prefill_status_gpu,
+ state_idx=mamba_state_idx,
+ accepted_counts=accepted_counts_gpu,
+ num_layers=context.num_mamba_layers,
+ )
+ mamba_state_selective_copy(
+ intermediate_states=context.mamba_intermediate_ssm_states,
+ current_states=context.mamba_ssm_states,
+ prefill_status=prefill_status_gpu,
+ state_idx=mamba_state_idx,
+ accepted_counts=accepted_counts_gpu,
+ num_layers=context.num_mamba_layers,
+ )
+
+ return blocks_to_release, remove_mask
def _sample_from_logits_2d(self, logits_2d: Tensor) -> Tensor:
"""Sample tokens from 2D logits using existing sampling parameters.
@@ -799,21 +738,12 @@ def _sample_from_logits_2d(self, logits_2d: Tensor) -> Tensor:
Returns:
Tensor: Sampled tokens of shape [num_requests].
"""
- spec_token_list = []
- indices_list = []
- for request_indices, temp, top_k, top_p in self._torch_sampling_buckets:
- request_indices_tensor = torch.tensor(
- request_indices, device=logits_2d.device, dtype=torch.long
- )
- spec_token_list.append(
- self._torch_sampling_func(logits_2d[request_indices_tensor, :], temp, top_k, top_p)
- )
- indices_list.append(request_indices_tensor)
-
- spec_tokens = torch.empty(logits_2d.shape[0], device=logits_2d.device, dtype=torch.int64)
- for tokens, indices in zip(spec_token_list, indices_list):
- spec_tokens[indices] = tokens
- return spec_tokens
+ return self._sampling.sample_kernel(
+ logits_2d,
+ logits_2d.shape[0],
+ self.inference_wrapped_model.inference_context,
+ eager=True,
+ )
def _compute_serial_mtp_and_sample(self):
"""Compute MTP logits serially after verification and sample speculative tokens.
@@ -822,377 +752,315 @@ def _compute_serial_mtp_and_sample(self):
Each MTP depth receives the correctly sampled token from the previous depth
(or the base token for depth 0) rather than stale speculative tokens from
the previous step.
+
+ When sequence parallelism is active, hidden states are kept in SP format
+ (scattered along the first dimension) between MTP depths to avoid a
+ redundant gather + scatter round-trip per depth.
"""
+ nvtx_range_push("mtp-spec-decoding/serial-mtp-init")
context = self.inference_wrapped_model.inference_context
active_request_count = context.total_request_count - context.paused_request_count
active_slice = slice(context.paused_request_count, context.total_request_count)
- unwrapped_model = unwrap_model(self.inference_wrapped_model.model)
+ unwrapped_model = self._unwrapped_model
# On non-last pipeline stages, the model won't have decoder hidden states.
- has_mtp = is_pipeline_last_stage(self.pp_group) and hasattr(
+ has_mtp = self._is_last_pp_stage and hasattr(
unwrapped_model, '_decoder_hidden_states_cache'
)
if has_mtp:
# Get decoder hidden states at last accepted positions.
hidden_states = unwrapped_model._decoder_hidden_states_cache
+
+ # When SP is active the decoder output is in scattered format
+ # [S/TP, B, H], but _last_accepted_seq_indices are indices into
+ # the full (gathered) sequence.
+ if self._sp_enabled:
+ hidden_states = gather_from_sequence_parallel_region(
+ hidden_states, group=self.inference_wrapped_model.tp_group
+ )
last_accepted_hidden = hidden_states[self._last_accepted_seq_indices, :, :]
# Shape: [active_request_count, 1, hidden_size]
else:
last_accepted_hidden = None
# Compute position IDs for the next tokens.
- # After rewind, request_kv_length_offsets has been adjusted. The actual
- # KV cache length is: adjusted_offset + processed_tokens.
- # The next position to predict starts at that cache length.
- adjusted_offsets = context.request_kv_length_offsets[active_slice]
- processed_tokens = context.request_query_lengths[active_slice]
- base_position = adjusted_offsets + processed_tokens
+ # After rewind, request_kv_length_offsets has been adjusted. Read from
+ # CPU context (post-rewind values), NOT gpu_view (stale pre-rewind snapshot).
+ # The next position to predict is: adjusted_offset + processed_tokens.
+ cuda_device = torch.cuda.current_device()
+ adjusted_offsets = context.request_kv_length_offsets[active_slice].to(
+ cuda_device, non_blocking=True
+ )
+ processed_tokens = context.request_query_lengths[active_slice].to(
+ cuda_device, non_blocking=True
+ )
+ # Cast to int64 to match CUDA graph capture dtype expectations.
+ base_position = (adjusted_offsets + processed_tokens).to(torch.int64)
# Start with the freshly sampled base token.
next_token_ids = self._sampled_tokens_cuda[:active_request_count].clone()
current_hidden = last_accepted_hidden if has_mtp else None
- num_depths = min(self.num_speculative_tokens, self.num_mtp_heads)
- for depth in range(num_depths):
- position_ids = (base_position + depth).unsqueeze(0) # [1, active_request_count]
- token_ids = next_token_ids.unsqueeze(0) # [1, active_request_count]
+ # Compute padding needed to make batch compatible with SP and CUDA graphs.
+ if getattr(self, '_mtp_resolved_padded_count', None) is not None:
+ # CUDA-graph path: use the EP-synced padded count.
+ padded_count = self._mtp_resolved_padded_count
+ assert not self._sp_enabled or padded_count % self._tp_size == 0
+ elif has_mtp:
+ # Eager path: pad only for SP alignment.
+ padded_count = active_request_count
+ if self._sp_enabled:
+ padded_count = round_up_to_nearest_multiple(padded_count, self._tp_size)
+ else:
+ padded_count = active_request_count
+ pad_count = padded_count - active_request_count
+
+ # Pad hidden states and scatter for sequence parallelism.
+ if has_mtp:
+ current_hidden = F.pad(current_hidden, (0, 0, 0, 0, 0, pad_count))
+ if self._sp_enabled:
+ current_hidden = scatter_to_sequence_parallel_region(
+ current_hidden, group=self.inference_wrapped_model.tp_group
+ )
+
+ token_ids_buf = self._mtp_token_ids_buf[:, :padded_count]
+ position_ids_buf = self._mtp_position_ids_buf[:, :padded_count]
+
+ # Zero-fill padding slots so the embedding layer never sees out-of-range IDs.
+ token_ids_buf[0, active_request_count:] = 0
+ position_ids_buf[0, active_request_count:] = 0
+
+ nvtx_range_pop("mtp-spec-decoding/serial-mtp-init")
+ for depth in range(self._num_mtp_depths):
+ nvtx_range_push(f"mtp-spec-decoding/depth-{depth}")
+
+ token_ids_buf[0, :active_request_count] = next_token_ids
+ position_ids_buf[0, :active_request_count] = base_position + depth
mtp_logits_2d = None
if has_mtp:
+ nvtx_range_push(f"mtp-spec-decoding/depth-{depth}/forward")
+ mtp_depth = None if unwrapped_model.mtp.mtp_use_repeated_layer else depth
current_hidden, mtp_logits = unwrapped_model.compute_mtp_single_step(
hidden_states=current_hidden,
- next_token_ids=token_ids,
- position_ids=position_ids,
- depth=depth,
+ next_token_ids=token_ids_buf,
+ position_ids=position_ids_buf,
+ depth=mtp_depth,
+ eager=not context.using_cuda_graph_this_step(),
+ cache_key=(
+ ("mtp", padded_count, mtp_depth)
+ if context.using_cuda_graph_this_step()
+ else None
+ ),
)
+ nvtx_range_pop(f"mtp-spec-decoding/depth-{depth}/forward")
+
+ # Strip padding from logits only. Hidden states stay padded+SP
+ # between depths to avoid redundant gather/scatter round-trips.
+ mtp_logits = mtp_logits[:active_request_count]
+
# mtp_logits: [active_request_count, 1, vocab_size]
mtp_logits_2d = mtp_logits.squeeze(1) # [active_request_count, vocab_size]
# Broadcast MTP logits across pipeline stages.
if self.model_is_pipeline_parallel:
+ nvtx_range_push(f"mtp-spec-decoding/depth-{depth}/pp-broadcast")
mtp_logits_2d = broadcast_from_last_pipeline_stage(
[active_request_count, self.vocab_size],
dtype=self.model_config.params_dtype,
tensor=mtp_logits_2d,
pp_group=self.pp_group,
)
+ nvtx_range_pop(f"mtp-spec-decoding/depth-{depth}/pp-broadcast")
# Sample speculative token using the same sampling parameters.
+ nvtx_range_push(f"mtp-spec-decoding/depth-{depth}/sample")
spec_tokens = self._sample_from_logits_2d(mtp_logits_2d)
self._sampled_mtp_tokens_cuda[depth, :active_request_count] = spec_tokens
+ nvtx_range_pop(f"mtp-spec-decoding/depth-{depth}/sample")
# Use sampled token as input for the next depth.
next_token_ids = spec_tokens
+ nvtx_range_pop(f"mtp-spec-decoding/depth-{depth}")
# Clean up cached hidden states.
if has_mtp:
del unwrapped_model._decoder_hidden_states_cache
- def _get_required_logit_indices(
- self,
- request_in_prefill_status_tensor: Tensor,
- request_query_lengths: Tensor,
- num_decode_requests: int,
- num_prefill_requests: int,
- device: torch.device,
- ) -> Tensor:
- """Get indices into the logits tensor for tokens that need sampling.
-
- For decode requests, all tokens (base + speculative) are needed.
- For prefill requests, only the last token logits are needed.
- Decode requests will always be on the left, followed by prefill requests.
-
- Example with 5 requests (2 spec tokens):
- Assume input ids : [ a5 a6s a7s | b3 b4s b5s | c6 c7s c8s | d1 d2 | e1 e2 e3 e4]
- Request to prefill [ 0 | 0 | 0 | 1 | 1 ]
- Request query lengths [ 3 | 3 | 3 | 2 | 4 ]
- OUTPUT : required_logit_indices [ 0 1 2 | 3 4 5 | 6 7 8 | 10 | 14 ]
-
- Returns:
- Tensor: Indices into the sequence dimension of the logits tensor.
- """
- decode_request_indices = torch.arange(
- num_decode_requests * (self.num_speculative_tokens + 1), device=device
- )
- prefill_request_indices = (
- request_query_lengths.cumsum(dim=0)[request_in_prefill_status_tensor == 1] - 1
- ) # Last token indices for prefill requests
- required_logit_indices = torch.cat([decode_request_indices, prefill_request_indices])
- assert (
- len(required_logit_indices)
- == num_decode_requests * (self.num_speculative_tokens + 1) + num_prefill_requests
- ), (
- f"Expected length of required_logit_indices to be "
- f"num_decode_requests * (self.num_speculative_tokens + 1) + num_prefill_requests, "
- f"but got {len(required_logit_indices)} for num_decode_requests {num_decode_requests} "
- f"and num_prefill_requests {num_prefill_requests}"
- )
- return required_logit_indices
-
- def _sample_speculative_logits(
- self, required_logits: Tensor, request_in_prefill_status_tensor: Tensor
- ) -> tuple:
- """Sample tokens from logits using sampling buckets.
-
- For torch sampling buckets: [request_indices, temp, top_k, top_p]
-
- Example with 5 requests:
- token_to_request_idx : [ 0 0 0 | 1 1 1 | 2 2 2 | 3 | 4 ]
- required_logits : [ a5l a6l a7l | b3l b4l b5l | c6l c7l c8l | d2l | e4l ] # Shape [11, vocab_size]
-
- Sampling buckets: [[[0,2], temp1, top_k1, top_p1], [[1], temp3, top_k3, top_p3], [[3, 4], temp2, top_k2, top_p2]]
-
- Final output tokens : [a5s a6s a7s c6s c7s c8s b3s b4s b5s d2s e4s] # Shape [11]
- (Rearranged from sampling bucket order back to input order using token_order)
-
- Returns:
- tuple: (output_tokens, repeats) where output_tokens has shape [total_required_tokens]
- """
- repeats = torch.where(
- request_in_prefill_status_tensor == 0, 1 + self.num_speculative_tokens, 1
- )
- token_to_request_index = torch.repeat_interleave(
- torch.arange(
- len(request_in_prefill_status_tensor),
- device=request_in_prefill_status_tensor.device,
- ),
- repeats,
- )
-
- output_tokens_jumbled_list = []
- token_order_list = []
-
- for request_indices, temp, top_k, top_p in self._torch_sampling_buckets:
- request_indices_tensor = torch.tensor(
- request_indices, device=token_to_request_index.device
- )
- required_indices = torch.where(
- torch.isin(token_to_request_index, request_indices_tensor)
- )[0]
- output_tokens_jumbled_list.append(
- self._torch_sampling_func(required_logits[required_indices, :], temp, top_k, top_p)
- )
- token_order_list.append(required_indices)
-
- output_tokens_jumbled = torch.cat(output_tokens_jumbled_list, dim=0)
- output_tokens = torch.empty(
- len(output_tokens_jumbled),
- device=output_tokens_jumbled.device,
- dtype=output_tokens_jumbled.dtype,
- )
- token_order = torch.cat(token_order_list, dim=0)
- # Rearrange output tokens from sampling_bucket request order back to input ids order
- output_tokens[token_order] = output_tokens_jumbled
-
- return output_tokens, repeats
-
def _verify_speculative_tokens(
self,
output_tokens: Tensor,
input_tokens_required: Tensor,
- request_in_prefill_status_tensor: Tensor,
- repeats: Tensor,
num_decode_requests: int,
num_prefill_requests: int,
active_request_count: int,
) -> tuple:
- """Verify speculative tokens against input tokens and compute acceptance.
-
- Creates an accepted tokens mask where:
- - For prefill requests, the token is always accepted.
- - For decode requests, the first token (base token) is always accepted, then we compare
- sampled tokens with input tokens and accept consecutive matches.
- Then finds the index of the last accepted token per request.
-
- Example (assume 1, 2, and 0 spec tokens are accepted in the first 3 decode requests):
- input_tokens_required: [ a5 a6s a7s | b3 b4s b5s | c6 c7s c8s | d2 | e4 ] # Size 11
- Output tokens [ a6o a7o a8o | b40 b5o b6o | c7o c8o c9o | d3o | e5o ]
- Output tokens right shift [ d3o a6o a7o | a8o b40 b5o | b6o c7o c8o | c9o | d3o ]
- Accepted tokens mask [ 1 1 0 | 1 1 1 | 1 0 0 | 1 | 1 ]
- Last one indices [ 1 | 5 | 6 | 9 | 10 ]
-
- Returns:
- tuple: (last_one_indices, accepted_tokens_mask, input_tokens_required) where
- last_one_indices contains the index of the last accepted token per request.
- """
- if input_tokens_required.ndim == 2:
- assert (
- input_tokens_required.shape[0] == 1
- ), f"Expected input_tokens_required to have 1 row, but got {input_tokens_required.shape}"
- input_tokens_required = input_tokens_required.squeeze(0)
-
- # Initialize mask with False to prevent boundary bleed
- accepted_tokens_mask = torch.zeros_like(input_tokens_required, dtype=torch.bool)
-
- # Make all prefill tokens accepted
- token_to_prefill_idx = torch.repeat_interleave(request_in_prefill_status_tensor, repeats)
- accepted_tokens_mask[token_to_prefill_idx == 1] = True
-
- # Safe decode token verification without cross-batch boundary contamination
- decode_mask_2d = None
- if num_decode_requests > 0:
- decode_len = num_decode_requests * (self.num_speculative_tokens + 1)
-
- decode_inputs = input_tokens_required[:decode_len].reshape(
- num_decode_requests, self.num_speculative_tokens + 1
- )
- decode_outputs = output_tokens[:decode_len].reshape(
- num_decode_requests, self.num_speculative_tokens + 1
- )
-
- # Shift outputs right by 1 *within* each request to align sampled tokens with input targets
- decode_outputs_shifted = decode_outputs.roll(1, dims=1)
- decode_mask_2d = decode_inputs == decode_outputs_shifted
- # The first token (base token) is always accepted
- decode_mask_2d[:, 0] = True
- # Enforce consecutive acceptance: cummin propagates False to the right
- decode_mask_2d = decode_mask_2d.cummin(dim=1).values
- accepted_tokens_mask[:decode_len] = decode_mask_2d.flatten()
-
- last_one_indices = torch.full(
- (active_request_count,), -1, device=input_tokens_required.device
+ """Verify speculative tokens against input tokens (Triton kernel)."""
+ return verify_speculative_tokens(
+ input_tokens=input_tokens_required,
+ output_tokens=output_tokens,
+ num_decode_requests=num_decode_requests,
+ num_prefill_requests=num_prefill_requests,
+ num_speculative_tokens=self.num_speculative_tokens,
)
- if num_decode_requests > 0:
- # Summing the consecutive mask gives the count; subtract 1 for the local index
- local_last_indices = decode_mask_2d.sum(dim=1) - 1
- row_offsets = torch.arange(num_decode_requests, device=last_one_indices.device) * (
- self.num_speculative_tokens + 1
- )
- last_one_indices[:num_decode_requests] = row_offsets + local_last_indices
-
- if num_prefill_requests > 0:
- decode_len = num_decode_requests * (self.num_speculative_tokens + 1)
- prefill_valid = (
- torch.nonzero(accepted_tokens_mask[decode_len:]).squeeze(-1) + decode_len
- )
- last_one_indices[num_decode_requests:] = prefill_valid
-
- return last_one_indices, accepted_tokens_mask, input_tokens_required
-
- def _dynamic_step_sample_logits_and_verify_tokens(self, logits: Tensor, input_ids: Tensor):
+ def _dynamic_step_sample_logits_and_verify_tokens(self, input_ids: Tensor):
"""
Sample tokens from logits for dynamic batching with speculative tokens and verify the tokens.
"""
context = self.inference_wrapped_model.inference_context
active_request_count = context.total_request_count - context.paused_request_count
- request_in_prefill_status_tensor = context.request_in_prefill_status_tensor[
- context.paused_request_count : context.total_request_count
- ]
- request_query_lengths = context.request_query_lengths[
- context.paused_request_count : context.total_request_count
- ]
-
- num_prefill_requests = request_in_prefill_status_tensor.sum().item()
- num_decode_requests = active_request_count - num_prefill_requests
-
- # Get the logit indices for tokens that need sampling.
- required_logit_indices = self._get_required_logit_indices(
- request_in_prefill_status_tensor,
- request_query_lengths,
- num_decode_requests,
- num_prefill_requests,
- logits.device,
+ # Sampling-side request counts: padded when running a captured graph.
+ # Verify uses the actual counts so the Triton kernels operate on the real workload.
+ use_graph_for_sampling = (
+ self._sampling_backend == "flashinfer"
+ and self._enable_cuda_graph
+ and context.using_cuda_graph_this_step()
)
+ if use_graph_for_sampling:
+ sample_num_decode = context.padded_batch_dimensions.decode_req_count
+ sample_num_prefill = context.padded_batch_dimensions.prefill_req_count
+ else:
+ sample_num_decode = context.num_decode_requests
+ sample_num_prefill = context.num_prefill_requests
+
+ # Logit indices for tokens that need sampling.
+ # Padded under graph capture so the captured `gather_indices` input has a stable shape.
+ # Padded slots resolve to row 0; verify and prepare-next read only the actual prefix,
+ # so the padded-row samples produced by the captured kernel are discarded.
+ nvtx_range_push("mtp-spec-decoding/verify/logit-indices")
+ # Use pre-allocated buffer for CUDA graph compatibility.
+ logits = self._all_logits_cuda
+ # `speculative_required_logit_indices()` already returns padded indices when
+ # running a captured graph (`num_last_token_logits` uses the padded counts and
+ # `pad_active_slices` zero-pads the trailing slots), so the call site does not
+ # need to re-pad here.
+ required_logit_indices = context.speculative_required_logit_indices()
- required_logits = logits.squeeze(0)[
- required_logit_indices, :
- ] # Shape [num_required, vocab_size]
+ if context.config.materialize_only_last_token_logits:
+ # last_token_logits already selected exactly the required positions.
+ sample_logits = logits.squeeze(0)
+ sample_gather_indices = None
+ else:
+ # Push the gather inside the captured kernel:
+ # pass the full per-token logits buffer (constant shape) plus the padded indices.
+ sample_logits = logits.squeeze(0)
+ sample_gather_indices = required_logit_indices
+ nvtx_range_pop("mtp-spec-decoding/verify/logit-indices")
# Sample tokens from logits
- output_tokens, repeats = self._sample_speculative_logits(
- required_logits, request_in_prefill_status_tensor
+ nvtx_range_push("mtp-spec-decoding/verify/sample")
+ output_tokens = self._sampling.sample_speculative(
+ sample_logits,
+ sample_num_decode,
+ sample_num_prefill,
+ self.num_speculative_tokens,
+ context,
+ gather_indices=sample_gather_indices,
+ eager=not use_graph_for_sampling,
+ cache_key=(
+ ("sample_speculative", sample_num_decode, sample_num_prefill)
+ if use_graph_for_sampling
+ else None
+ ),
)
+ nvtx_range_pop("mtp-spec-decoding/verify/sample")
+
+ num_prefill_requests = context.num_prefill_requests
+ num_decode_requests = active_request_count - num_prefill_requests
# Verify speculative tokens against input tokens.
+ nvtx_range_push("mtp-spec-decoding/verify/verify-tokens")
input_tokens_required = input_ids[0, required_logit_indices]
last_one_indices, accepted_tokens_mask, input_tokens_required = (
self._verify_speculative_tokens(
output_tokens,
input_tokens_required,
- request_in_prefill_status_tensor,
- repeats,
num_decode_requests,
num_prefill_requests,
active_request_count,
)
)
+ nvtx_range_pop("mtp-spec-decoding/verify/verify-tokens")
- # Store the final sampled tokens for the next forward pass.
- final_sampled_tokens = output_tokens[last_one_indices]
- self._sampled_tokens_cuda[: len(final_sampled_tokens)] = final_sampled_tokens
-
- # Store the last accepted positions in the packed sequence for serial
- # MTP computation after verification.
- self._last_accepted_seq_indices = required_logit_indices[last_one_indices]
-
- # Extract accepted tokens and counts for decode requests.
- # For prefill it is always set to 1. For decode, the first token is always accepted,
- # then we compare with input tokens and accept the next tokens if its a match.
- #
- # Example (continuing from above):
- # input_tokens_required: [ a5 a6s a7s | b3 b4s b5s | c6 c7s c8s | d2 | e4 ]
- # Accepted tokens mask [ 1 1 0 | 1 1 1 | 1 0 0 | 1 | 1 ]
- # Accepted tokens [ [a6s -1] | [b4s b5s] | [-1 -1] ] # Only decode requests (prefill defaults to -1)
- # Accepted token counts [ 1 | 2 | 0 ] # Prefill defaults to 0
- input_tokens_required[accepted_tokens_mask == 0] = -1 # Mask out non-accepted tokens
- input_tokens_decode_mode = input_tokens_required[
- : num_decode_requests * (self.num_speculative_tokens + 1)
- ]
- input_tokens_reshaped = input_tokens_decode_mode.reshape(
- -1, self.num_speculative_tokens + 1
- ) # shape: [num_decode_requests, num_speculative_tokens + 1]
-
- # Skip the first token of every decode request (i.e a5, b3, c6)
- accepted_tokens = input_tokens_reshaped[:, 1:]
- self._accepted_tokens_per_request[: accepted_tokens.shape[0], :] = accepted_tokens
- self._accepted_token_counts_per_request = (self._accepted_tokens_per_request != -1).sum(
- dim=1
+ nvtx_range_push("mtp-spec-decoding/verify/prepare-next")
+ self._prepare_speculative_tokens_for_next_forward_pass(
+ num_decode_requests,
+ output_tokens,
+ required_logit_indices,
+ last_one_indices,
+ accepted_tokens_mask,
+ input_tokens_required,
)
+ nvtx_range_pop("mtp-spec-decoding/verify/prepare-next")
- def _dynamic_step_sample_logits(self, logits: Tensor):
- """Sample tokens from logits for dynamic batching.
+ def _prepare_speculative_tokens_for_next_forward_pass(
+ self,
+ num_decode_requests: int,
+ output_tokens: torch.Tensor,
+ required_logit_indices: torch.Tensor,
+ last_one_indices: torch.Tensor,
+ accepted_tokens_mask: torch.Tensor,
+ input_tokens_required: torch.Tensor,
+ ):
+ """Prepare accepted speculative tokens for the next forward pass (Triton kernel).
- Args:
- logits (Tensor): The logits from the forward pass.
+ Example:
+ input_tokens_required: [ a5 a6s a7s | b3 b4s b5s | c6 c7s c8s | d2 | e4 ]
+ Accepted tokens mask [ 1 1 0 | 1 1 1 | 1 0 0 | 1 | 1 ]
+ Accepted tokens [ [a6s -1] | [b4s b5s] | [-1 -1] ] (decode only; prefill → -1)
+ Accepted token counts [ 1 | 2 | 0 ] (prefill defaults to 0)
"""
+ active_request_count = last_one_indices.shape[0]
+ prepare_next_forward_pass(
+ num_decode_requests=num_decode_requests,
+ output_tokens=output_tokens,
+ required_logit_indices=required_logit_indices,
+ last_one_indices=last_one_indices,
+ accepted_tokens_mask=accepted_tokens_mask,
+ input_tokens=input_tokens_required,
+ sampled_tokens_buf=self._sampled_tokens_cuda,
+ last_accepted_seq_buf=self._last_accepted_seq_indices_buf,
+ accepted_tokens_per_request=self._accepted_tokens_per_request,
+ accepted_token_counts=self._accepted_token_counts_per_request,
+ num_speculative_tokens=self.num_speculative_tokens,
+ )
+ # Expose the active slice so downstream code sees the right length.
+ self._last_accepted_seq_indices = self._last_accepted_seq_indices_buf[:active_request_count]
+
+ def _dynamic_step_sample_logits(self):
+ """Sample tokens from logits for dynamic batching."""
# TODO(ksanthanam): Evaluate whether it makes more sense to sample on 1 rank
# and then broadcast the sampled tokens rather than broadcasting the raw logits.
- # Last token logits.
context = self.inference_wrapped_model.inference_context
- if context.config.materialize_only_last_token_logits:
- # When materialize_only_last_token_logits is true, last_token_logits is
- # already called in the forward pass of GPT.
- required_token_logits = logits.squeeze(0)
- else:
- # todo : Should do verification here and get approrpiate las token logits
- required_token_logits = context.last_token_logits(logits)
-
- if self._sampling_backend == "torch":
- # Concatenate the outputs once to prevent repeated small writes.
- token_list = []
- indices_list = []
-
- # e.g torch sample buckets will be
- # i.e (for all unique comibnation of t, topk, topk what are the associated
- # requests indices (based on the active slices)
- # [ [req at index 0, req at index 2], t1, topk1, topp1 ]]
- # [ [req at index 1, req at index 3, req at index 4] , t2, topk2, topp2]
- for indices, temp, top_k, top_p in self._torch_sampling_buckets:
- token_list.append(
- self._torch_sampling_func(required_token_logits[indices, :], temp, top_k, top_p)
- )
- indices_list.append(torch.tensor(indices))
-
- # Single write to the output tensor.
- sampled_tokens = torch.cat(token_list, dim=0)
- sampled_indices = torch.cat(indices_list, dim=0)
-
- self._sampled_tokens_cuda[sampled_indices] = sampled_tokens
+ active_request_count = context.total_request_count - context.paused_request_count
+ use_graph = (
+ self._sampling_backend == "flashinfer"
+ and self._enable_cuda_graph
+ and context.using_cuda_graph_this_step()
+ )
+ # Padded count when running a captured graph (cache key buckets); actual otherwise.
+ n = context.padded_active_request_count if use_graph else active_request_count
+ # When `materialize_only_last_token_logits` is true the forward pass already
+ # selected the right rows. Otherwise we point the kernel at the per-request
+ # last-token positions via `gather_indices`; padded slots safely fan in to row 0.
+ gather_indices = (
+ None
+ if context.config.materialize_only_last_token_logits
+ else context.gpu_view.active_request_last_token_idxs
+ )
+ self._sampled_tokens_cuda = self._sampling.sample_kernel(
+ self._all_logits_cuda.squeeze(0),
+ n,
+ context,
+ gather_indices=gather_indices,
+ eager=not use_graph,
+ cache_key=("sample", n) if use_graph else None,
+ )
def _dynamic_step_log_probs_bookkeeping(self) -> Tuple[bool, bool]:
"""Perform bookkeeping necessary to compute log probs for dynamic batching.
@@ -1201,25 +1069,27 @@ def _dynamic_step_log_probs_bookkeeping(self) -> Tuple[bool, bool]:
return_log_probs (bool): Whether to return the sampled log_probs.
"""
context = self.inference_wrapped_model.inference_context
- active_request_slice = slice(context.paused_request_count, context.total_request_count)
-
- return_log_probs = self._request_metadata["return_log_probs"][active_request_slice]
- top_n_log_probs = self._request_metadata["top_n_logprobs"][active_request_slice] > 0
+ active_request_count = context.total_request_count - context.paused_request_count
- return return_log_probs.any(), top_n_log_probs.any()
+ return (
+ (context.active_request_metadata["return_log_probs"][:active_request_count]).any(),
+ (context.active_request_metadata["top_n_logprobs"][:active_request_count] > 0).any(),
+ )
- def _router_record_bookkeeping(self) -> Optional[Dict[int, Tensor]]:
- """Collect and map routing indices per request for MoE router recording.
+ def _router_record_bookkeeping(self) -> Optional[np.ndarray]:
+ """Collect flat routing indices for MoE router recording.
- This method retrieves recorded routing decisions and maps them to individual
- requests using the context's request_ids and query_lengths. Uses the context's
- routing_metadata when available (which handles CUDA graph static buffers automatically).
- Must be called while context attributes are still valid (before request transitions).
+ Retrieves recorded routing decisions via the context's routing_metadata
+ (which handles CUDA graph static buffers), performs the TP all-gather
+ when sequence parallelism is active, strips CUDA padding, and returns
+ a flat CPU numpy array aligned with the context's active-token layout.
+ Must be called while context attributes are still valid (before request
+ transitions).
Returns:
- Optional[Dict[int, Tensor]]: A dictionary mapping request_id to a tensor of
- shape [num_tokens, num_layers, topk]. Returns None if routing replay is
- disabled or no routing data was recorded.
+ Optional[np.ndarray]: Flat routing array of shape
+ [active_token_count, num_layers, topk], or None if routing
+ replay is disabled or no routing data was recorded.
"""
config = self.inference_wrapped_model.model.config
if not config.moe_enable_routing_replay:
@@ -1235,10 +1105,6 @@ def _router_record_bookkeeping(self) -> Optional[Dict[int, Tensor]]:
if stacked_routing is None:
return None
- # Get active request info from context
- active_request_slice = slice(context.paused_request_count, context.total_request_count)
- active_request_ids = context.request_ids[active_request_slice].tolist()
- active_query_lengths = context.request_query_lengths[active_request_slice].tolist()
active_token_count = context.active_token_count
# Get TP group for all-gather if using sequence parallelism
@@ -1249,39 +1115,45 @@ def _router_record_bookkeeping(self) -> Optional[Dict[int, Tensor]]:
# All-gather across TP group if using sequence parallelism (tp_size > 1)
if tp_size > 1 and get_model_config(self.inference_wrapped_model.model).sequence_parallel:
+ # With SP, the model processes padded_active_token_count tokens total,
+ # scattered evenly across TP ranks. Each rank routes
+ # padded_active_token_count // tp_size tokens through MoE layers.
+ #
+ # The CUDA-graph static buffer path in get_routing_indices() may return
+ # a tensor sliced to active_token_count (the global unpadded count),
+ # which can be larger than the per-rank valid count. Truncate to the
+ # true per-rank count before the all-gather so we only gather valid
+ # routing data and reconstruct the full sequence in the correct order.
+ local_token_count = context.padded_active_token_count // tp_size
+
+ stacked_routing = stacked_routing[:local_token_count]
# gather_from_sequence_parallel_region gathers along dim 0
- # [local_token_count, num_layers, topk] -> [global_token_count, num_layers, topk]
+ # [local_token_count, num_layers, topk] -> [padded_token_count, num_layers, topk]
stacked_routing = gather_from_sequence_parallel_region(stacked_routing, group=tp_group)
- # Slice to real tokens (remove CUDA padding)
- stacked_routing = stacked_routing[:active_token_count]
-
- # Split by request along token dimension
- # stacked_routing has shape [active_token_count, num_layers, topk]
- routing_splits = stacked_routing.split(active_query_lengths, dim=0)
-
- # Map to request IDs
- routing_indices_per_request = {}
- for req_id, routing_split in zip(active_request_ids, routing_splits):
- # routing_split has shape [num_tokens_for_request, num_layers, topk]
- routing_indices_per_request[req_id] = routing_split
+ # Slice to real tokens (remove CUDA padding), move to CPU as numpy with target dtype
+ _ri_dtype = np.int16 if (config.num_moe_experts or 0) <= 32768 else np.int32
+ return stacked_routing[:active_token_count].cpu().numpy().astype(_ri_dtype)
- return routing_indices_per_request
-
- def _dynamic_step_calculate_log_probs(self, logits: Tensor) -> Optional[Tensor]:
+ def _dynamic_step_calculate_log_probs(self) -> Optional[Tensor]:
"""Calculate log probs from logits."""
context = self.inference_wrapped_model.inference_context
active_request_count = context.total_request_count - context.paused_request_count
+ # This code cannot be reached when we are using speculative decode.
+ assert self.num_speculative_tokens == 0
+ logits_seq_len = (
+ active_request_count
+ if context.config.materialize_only_last_token_logits
+ else context.padded_active_token_count
+ )
return context.calculate_log_probs(
- logits,
+ self._all_logits_cuda[:, :logits_seq_len, :],
self._sampled_tokens_cuda[:active_request_count],
only_last_token_logits=context.config.materialize_only_last_token_logits,
)
- def _dynamic_step_calculate_log_probs_speculative(
- self, logits: Tensor
- ) -> Tuple[List[List[float]], Tensor]:
+ def _dynamic_step_calculate_log_probs_speculative(self) -> Tuple[List[List[float]], Tensor]:
"""Calculate log probs from logits for speculative decoding.
For decode requests, computes log probs for each accepted speculative token
@@ -1292,9 +1164,6 @@ def _dynamic_step_calculate_log_probs_speculative(
- log_prob(accepted_token[j]) comes from logits at position j
- log_prob(newly_sampled_token) comes from logits at position accepted_count
- Args:
- logits (Tensor): The main model logits [1, seq_len, vocab_size].
-
Returns:
Tuple of (log_probs_list, log_probs_tensor):
log_probs_list: List of lists, one per active request, containing
@@ -1304,18 +1173,23 @@ def _dynamic_step_calculate_log_probs_speculative(
context = self.inference_wrapped_model.inference_context
active_request_count = context.total_request_count - context.paused_request_count
- request_in_prefill_status_tensor = context.request_in_prefill_status_tensor[
- context.paused_request_count : context.total_request_count
- ]
- request_query_lengths = context.request_query_lengths[
- context.paused_request_count : context.total_request_count
+ # Use gpu_view for data consumed by GPU log-probs operations.
+ request_in_prefill_status_tensor = context.gpu_view.request_in_prefill_status[
+ :active_request_count
]
+ request_query_lengths = context.gpu_view.request_query_lengths[:active_request_count]
num_prefill_requests = request_in_prefill_status_tensor.sum().item()
num_decode_requests = active_request_count - num_prefill_requests
+ only_last = context.config.materialize_only_last_token_logits
+ # Use pre-allocated buffer for CUDA graph compatibility.
+ logits = self._all_logits_cuda
logits_squeezed = logits.squeeze(0).float()
- log_probs_tensor = F.log_softmax(logits_squeezed[: context.active_token_count], dim=-1)
+ if only_last:
+ log_probs_tensor = F.log_softmax(logits_squeezed, dim=-1)
+ else:
+ log_probs_tensor = F.log_softmax(logits_squeezed[: context.active_token_count], dim=-1)
log_probs_list_decode = []
@@ -1356,22 +1230,34 @@ def _dynamic_step_calculate_log_probs_speculative(
decode_len = num_decode_requests * (self.num_speculative_tokens + 1)
prefill_log_probs = log_probs_tensor[decode_len:]
- prefill_token_ids = context.token_to_input_ids[
- decode_len : context.active_token_count
- ].roll(-1, 0)
- prefill_query_lengths = request_query_lengths[request_in_prefill_status_tensor == 1]
- new_token_idx = prefill_query_lengths.cumsum(0) - 1
- prefill_new_tokens = self._sampled_tokens_cuda[num_decode_requests:active_request_count]
- prefill_token_ids[new_token_idx] = prefill_new_tokens
-
- prefill_token_count = context.active_token_count - decode_len
- seq_idx = torch.arange(prefill_token_count, device=logits.device)
- selected_log_probs = prefill_log_probs[seq_idx, prefill_token_ids]
-
- prefill_log_probs_split = selected_log_probs.cpu().split(
- prefill_query_lengths.tolist(), dim=0
- )
- log_probs_list_prefill = [lp.tolist() for lp in prefill_log_probs_split]
+ if only_last:
+ # Only last-token logits were materialized per prefill request.
+ prefill_new_tokens = self._sampled_tokens_cuda[
+ num_decode_requests:active_request_count
+ ]
+ selected_log_probs = prefill_log_probs[
+ torch.arange(num_prefill_requests, device=logits.device), prefill_new_tokens
+ ]
+ log_probs_list_prefill = [[lp.item()] for lp in selected_log_probs]
+ else:
+ prefill_token_ids = context.gpu_view.token_to_input_ids[
+ decode_len : context.active_token_count
+ ].roll(-1, 0)
+ prefill_query_lengths = request_query_lengths[request_in_prefill_status_tensor == 1]
+ new_token_idx = prefill_query_lengths.cumsum(0) - 1
+ prefill_new_tokens = self._sampled_tokens_cuda[
+ num_decode_requests:active_request_count
+ ]
+ prefill_token_ids[new_token_idx] = prefill_new_tokens
+
+ prefill_token_count = context.active_token_count - decode_len
+ seq_idx = torch.arange(prefill_token_count, device=logits.device)
+ selected_log_probs = prefill_log_probs[seq_idx, prefill_token_ids]
+
+ prefill_log_probs_split = selected_log_probs.cpu().split(
+ prefill_query_lengths.tolist(), dim=0
+ )
+ log_probs_list_prefill = [lp.tolist() for lp in prefill_log_probs_split]
log_probs_list = log_probs_list_decode + log_probs_list_prefill
@@ -1396,14 +1282,12 @@ def _dynamic_step_calculate_top_n_logprobs_speculative(
"""
context = self.inference_wrapped_model.inference_context
active_request_count = context.total_request_count - context.paused_request_count
- active_request_slice = slice(context.paused_request_count, context.total_request_count)
- request_in_prefill_status_tensor = context.request_in_prefill_status_tensor[
- context.paused_request_count : context.total_request_count
- ]
- request_query_lengths = context.request_query_lengths[
- context.paused_request_count : context.total_request_count
+ # Use gpu_view for data consumed by GPU top-n operations.
+ request_in_prefill_status_tensor = context.gpu_view.request_in_prefill_status[
+ :active_request_count
]
+ request_query_lengths = context.gpu_view.request_query_lengths[:active_request_count]
num_prefill_requests = request_in_prefill_status_tensor.sum().item()
num_decode_requests = active_request_count - num_prefill_requests
@@ -1416,7 +1300,7 @@ def _dynamic_step_calculate_top_n_logprobs_speculative(
num_decode_requests, self.num_speculative_tokens + 1, -1
)
accepted_counts = self._accepted_token_counts_per_request[:num_decode_requests]
- top_n_per_request = self._request_metadata["top_n_logprobs"][active_request_slice][
+ top_n_per_request = context.active_request_metadata["top_n_logprobs"][
:num_decode_requests
]
max_top_n = int(top_n_per_request.max().item())
@@ -1440,47 +1324,72 @@ def _dynamic_step_calculate_top_n_logprobs_speculative(
]
if num_prefill_requests > 0:
+ only_last = context.config.materialize_only_last_token_logits
decode_len = num_decode_requests * (self.num_speculative_tokens + 1)
prefill_log_probs = log_probs_tensor[decode_len:]
- prefill_query_lengths = request_query_lengths[request_in_prefill_status_tensor == 1]
- prefill_log_probs_per_request = prefill_log_probs.split(
- prefill_query_lengths.tolist(), dim=0
- )
- for i in range(num_prefill_requests):
- req_idx = num_decode_requests + i
- top_n = int(
- self._request_metadata["top_n_logprobs"][active_request_slice][req_idx].item()
- )
- if top_n > 0:
- request_lp = prefill_log_probs_per_request[i]
- skip_prompt = bool(
- self._request_metadata["skip_prompt_log_probs"][req_idx].item()
+ # Batch metadata reads: single CPU transfer for all prefill requests.
+ prefill_top_n = context.active_request_metadata["top_n_logprobs"][
+ num_decode_requests:active_request_count
+ ].tolist()
+ max_top_n_prefill = int(max(prefill_top_n)) if prefill_top_n else 0
+
+ if max_top_n_prefill > 0:
+ if only_last:
+ # One logit row per prefill request — single batched topk.
+ topk_results_prefill = torch.topk(
+ prefill_log_probs, k=max_top_n_prefill, dim=-1
)
-
- if skip_prompt and request_lp.size(0) > 1:
- top_n_logits = torch.topk(request_lp[-1], k=top_n)
- top_n_results[req_idx] = [
- (top_n_logits.values.cpu(), top_n_logits.indices.cpu())
- ]
- else:
- top_n_logits = torch.topk(request_lp, k=top_n, dim=-1)
- top_n_values_cpu = top_n_logits.values.cpu()
- top_n_indices_cpu = top_n_logits.indices.cpu()
- top_n_results[req_idx] = [
- (top_n_values_cpu[t], top_n_indices_cpu[t])
- for t in range(request_lp.size(0))
- ]
+ topk_vals_cpu = topk_results_prefill.values.cpu()
+ topk_idxs_cpu = topk_results_prefill.indices.cpu()
+
+ for i in range(num_prefill_requests):
+ top_n = int(prefill_top_n[i])
+ if top_n > 0:
+ req_idx = num_decode_requests + i
+ top_n_results[req_idx] = [
+ (topk_vals_cpu[i, :top_n], topk_idxs_cpu[i, :top_n])
+ ]
+ else:
+ prefill_query_lengths = request_query_lengths[
+ request_in_prefill_status_tensor == 1
+ ]
+ prefill_log_probs_per_request = prefill_log_probs.split(
+ prefill_query_lengths.tolist(), dim=0
+ )
+ prefill_skip_prompt = context.active_request_metadata["skip_prompt_log_probs"][
+ num_decode_requests:active_request_count
+ ].tolist()
+
+ for i in range(num_prefill_requests):
+ top_n = int(prefill_top_n[i])
+ if top_n > 0:
+ req_idx = num_decode_requests + i
+ request_lp = prefill_log_probs_per_request[i]
+ skip_prompt = bool(prefill_skip_prompt[i])
+
+ if skip_prompt and request_lp.size(0) > 1:
+ top_n_logits = torch.topk(request_lp[-1], k=top_n)
+ top_n_results[req_idx] = [
+ (top_n_logits.values.cpu(), top_n_logits.indices.cpu())
+ ]
+ else:
+ top_n_logits = torch.topk(request_lp, k=top_n, dim=-1)
+ top_n_values_cpu = top_n_logits.values.cpu()
+ top_n_indices_cpu = top_n_logits.indices.cpu()
+ top_n_results[req_idx] = [
+ (top_n_values_cpu[t], top_n_indices_cpu[t])
+ for t in range(request_lp.size(0))
+ ]
return top_n_results if top_n_results else None
def _dynamic_step_calculate_top_n_logprobs(
- self, logits: Tensor, log_probs_tensor: Optional[Tensor] = None
+ self, log_probs_tensor: Optional[Tensor] = None
) -> Optional[Dict[int, List[Tuple[Tensor, Tensor]]]]:
"""Calculate top-n log probs from logits for dynamic batching.
Args:
- logits (Tensor): The logits to compute top-n log probs from.
log_probs_tensor (Optional[Tensor]): Pre-computed log probabilities tensor.
If provided, avoids recomputing log_softmax. Should be the tensor
returned by calculate_log_probs.
@@ -1506,9 +1415,7 @@ def _dynamic_step_calculate_top_n_logprobs(
top_n_results = {}
for req_idx in range(active_request_count):
- top_n = int(
- self._request_metadata["top_n_logprobs"][active_request_slice][req_idx].item()
- )
+ top_n = int(context.active_request_metadata["top_n_logprobs"][req_idx].item())
if top_n > 0:
# Get top-n logprobs and indices for this request (single token)
top_n_logits = torch.topk(log_probs[req_idx], k=top_n)
@@ -1530,14 +1437,14 @@ def _dynamic_step_calculate_top_n_logprobs(
top_n_results = {}
for req_idx in range(active_request_count):
- top_n = int(
- self._request_metadata["top_n_logprobs"][active_request_slice][req_idx].item()
- )
+ top_n = int(context.active_request_metadata["top_n_logprobs"][req_idx].item())
if top_n > 0:
request_log_probs = log_probs_per_request[
req_idx
] # [num_tokens_for_request, vocab_size]
- skip_prompt = bool(self._request_metadata["skip_prompt_log_probs"][req_idx].item())
+ skip_prompt = bool(
+ context.active_request_metadata["skip_prompt_log_probs"][req_idx].item()
+ )
# If skip_prompt_log_probs is True, only compute for last token
if skip_prompt and request_log_probs.size(0) > 1:
@@ -1558,42 +1465,23 @@ def _dynamic_step_calculate_top_n_logprobs(
return top_n_results if top_n_results else None
+ @torch.inference_mode()
def dummy_forward(self):
"""Perform a dummy forward pass. This is used in expert model parallelism
on ranks that do not have any real requests. It may run in eager mode."""
context = self.inference_wrapped_model.inference_context
- # if no cuda graphs, directly use dummy forward
- if not context.cuda_graph_batch_dimensions_list:
- self.inference_wrapped_model.dummy_forward()
-
- # Disable MoE padding for MTP computation
- if self.model_config.moe_pad_experts_for_cuda_graph_inference:
- unwrapped_model = unwrap_model(self.inference_wrapped_model.model)
- set_decode_expert_padding(unwrapped_model, False)
-
- self._dummy_serial_mtp_forward()
-
- return
# attempt to use cuda-graph if possible
input_ids, position_ids = self._dynamic_step_context_init(is_dummy_forward=True)
+ self._dynamic_step_forward_logits(input_ids, position_ids)
- # _dynamic_step_context_init tries to find a cuda-graph that is compatible
- # with all EP ranks. It can also return no match, in which case
- # we run in eager mode.
-
- if context.using_cuda_graph_this_step():
- # we found a cuda-graph to run
- self._dynamic_step_forward_logits(input_ids, position_ids)
- else:
- # fallback to eager dummy forward
- self.inference_wrapped_model.dummy_forward()
-
- # Disable MoE padding for MTP computation
+ # Disable MoE padding for MTP computation, unless CUDA graphs
+ # are active (the graphs were captured with padding enabled).
if self.model_config.moe_pad_experts_for_cuda_graph_inference:
- unwrapped_model = unwrap_model(self.inference_wrapped_model.model)
- set_decode_expert_padding(unwrapped_model, False)
+ if not context.using_cuda_graph_this_step():
+ unwrapped_model = unwrap_model(self.inference_wrapped_model.model)
+ set_decode_expert_padding(unwrapped_model, False)
# When speculative decoding is active, the real EP ranks perform serial
# MTP forward passes after the main forward pass. MTP layers may contain
@@ -1625,10 +1513,11 @@ def _dummy_serial_mtp_forward(self):
if self.model_config.expert_model_parallel_size <= 1:
return
- unwrapped_model = unwrap_model(self.inference_wrapped_model.model)
+ unwrapped_model = self._unwrapped_model
- is_last_stage = is_pipeline_last_stage(self.pp_group)
- has_mtp = is_last_stage and hasattr(unwrapped_model, '_decoder_hidden_states_cache')
+ has_mtp = self._is_last_pp_stage and hasattr(
+ unwrapped_model, '_decoder_hidden_states_cache'
+ )
if not has_mtp and not self.model_is_pipeline_parallel:
# No MTP on this rank and no PP broadcast to participate in.
return
@@ -1636,32 +1525,76 @@ def _dummy_serial_mtp_forward(self):
device = torch.cuda.current_device()
dtype = self.model_config.params_dtype
hidden_size = self.model_config.hidden_size
- num_depths = min(self.num_speculative_tokens, self.num_mtp_heads)
+
+ # Use precomputed MTP CUDA graph batch size when available;
+ # otherwise use minimal SP-compatible size.
+ if getattr(self, '_mtp_resolved_padded_count', None) is not None:
+ padded_count = self._mtp_resolved_padded_count
+ assert not self._sp_enabled or padded_count % self._tp_size == 0
+ elif has_mtp:
+ # Eager path: use TP-aligned minimum size for dummy tensors.
+ padded_count = self._tp_size if self._sp_enabled else 1
dummy_hidden = None
if has_mtp:
- # Minimal dummy tensors — just enough to drive the MTP layer forward
+ # Minimal dummy tensors to drive the MTP layer forward
# so that the MoE all-to-all collectives are issued.
- dummy_hidden = torch.zeros((1, 1, hidden_size), device=device, dtype=dtype)
- dummy_token_ids = torch.zeros((1, 1), device=device, dtype=torch.long)
- dummy_position_ids = torch.zeros((1, 1), device=device, dtype=torch.long)
+ dummy_hidden = torch.zeros((padded_count, 1, hidden_size), device=device, dtype=dtype)
+ if self._sp_enabled:
+ dummy_hidden = scatter_to_sequence_parallel_region(
+ dummy_hidden, group=self.inference_wrapped_model.tp_group
+ )
+ dummy_token_ids = torch.zeros((1, padded_count), device=device, dtype=torch.long)
+ dummy_position_ids = torch.zeros((1, padded_count), device=device, dtype=torch.long)
- for depth in range(num_depths):
+ context = self.inference_wrapped_model.inference_context
+
+ for depth in range(self._num_mtp_depths):
+ nvtx_range_push(f"mtp-spec-decoding/dummy-depth-{depth}")
mtp_logits_2d = None
if has_mtp:
+ mtp_depth = None if unwrapped_model.mtp.mtp_use_repeated_layer else depth
dummy_hidden, mtp_logits = unwrapped_model.compute_mtp_single_step(
hidden_states=dummy_hidden,
next_token_ids=dummy_token_ids,
position_ids=dummy_position_ids,
- depth=depth,
+ depth=mtp_depth,
+ eager=not context.using_cuda_graph_this_step(),
+ cache_key=(
+ ("mtp", padded_count, mtp_depth)
+ if context.using_cuda_graph_this_step()
+ else None
+ ),
)
- mtp_logits_2d = mtp_logits.squeeze(1) # [1, vocab_size]
+ mtp_logits_2d = mtp_logits.squeeze(1) # [padded_count, vocab_size]
# Match the PP broadcast that real ranks do in _compute_serial_mtp_and_sample.
if self.model_is_pipeline_parallel:
broadcast_from_last_pipeline_stage(
- [1, self.vocab_size], dtype=dtype, tensor=mtp_logits_2d, pp_group=self.pp_group
+ [padded_count, self.vocab_size],
+ dtype=dtype,
+ tensor=mtp_logits_2d,
+ pp_group=self.pp_group,
)
+ nvtx_range_pop(f"mtp-spec-decoding/dummy-depth-{depth}")
+
+ def _transfer_samples_to_cpu(self, active_request_count: int) -> tuple:
+ """Batch GPU-to-CPU transfer of sampled tokens.
+
+ Called at the boundary between GPU sampling and CPU bookkeeping.
+ After this returns, all sampled data is on CPU and the remainder
+ of the step is 100% CPU.
+
+ Returns:
+ tuple: (sampled_tokens_cpu, sampled_mtp_tokens_cpu) where
+ sampled_mtp_tokens_cpu is None when speculative decoding is off.
+ """
+ sampled_tokens_cpu = self._sampled_tokens_cuda[:active_request_count].cpu()
+ if self.num_speculative_tokens > 0:
+ sampled_mtp_tokens_cpu = self._sampled_mtp_tokens_cuda[:, :active_request_count].cpu()
+ else:
+ sampled_mtp_tokens_cpu = None
+ return sampled_tokens_cpu, sampled_mtp_tokens_cpu
def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]:
"""Update the dynamic inference context after sampling.
@@ -1682,26 +1615,35 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]:
active_request_count = context.total_request_count - context.paused_request_count
active_request_slice = slice(context.paused_request_count, context.total_request_count)
- # Active sequence lengths.
+ # Batch GPU-to-CPU transfer of all sampled tokens.
+ range_push("transfer_samples_to_cpu")
+ sampled_tokens_cpu, sampled_mtp_tokens_cpu = self._transfer_samples_to_cpu(
+ active_request_count
+ )
+ range_pop()
+
+ range_push("active_request_mask")
+ # Everything below is 100% CPU.
active_request_ids = context.request_ids[active_request_slice].long()
active_sequence_lengths = context.get_active_sequence_lengths()
- if self.num_speculative_tokens > 0:
- active_sequence_lengths += (
- self._accepted_token_counts_per_request[:active_request_count] + 1
- )
- else:
- active_sequence_lengths += 1
+ # After the forward pass and KV-cache rewind, get_active_sequence_lengths()
+ # returns kv_offsets + query_lengths which already includes all accepted
+ # speculative tokens (they were part of the query and survived the rewind).
+ # Only the newly sampled base token is not yet in the KV cache, so add 1.
+ active_sequence_lengths += 1
max_sequence_lengths = context.get_max_sequence_lengths()
# Request finished if termination_id or length >= max_sequence_length.
- # Note: termination_id tensor has per-request termination IDs from mixed sampling
+ # Both operands are CPU: sampled_tokens_cpu was D2H'd above, and
+ # active_request_metadata is CPU-pinned.
active_request_mask = (
- self._sampled_tokens_cuda[:active_request_count]
- != self._request_metadata["termination_id"][active_request_slice]
+ sampled_tokens_cpu
+ != context.active_request_metadata["termination_id"][:active_request_count]
).byte() & torch.less(active_sequence_lengths, max_sequence_lengths).byte()
- # Mark requests as finished if they hit stop words (detected in previous step's post_process_requests)
+ # Mark requests as finished if they hit stop words
+ # (detected in previous step's post_process_requests)
if self._get_stop_word_finished_ids_callback is not None:
request_ids_list = active_request_ids.tolist()
stop_word_finished_ids = self._get_stop_word_finished_ids_callback(request_ids_list)
@@ -1715,27 +1657,40 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]:
)
finished_request_ids = context.request_ids[finished_idxs]
+ # Save block IDs for finished requests before update_requests releases them.
+ # Needed for per-block routing reconstruction in the engine.
+ finished_routing_block_ids = {}
+ if context.kv_block_allocator.block_routing and finished_idxs.numel() > 0:
+ for fidx in finished_idxs.tolist():
+ req_id = int(context.request_ids[fidx].item())
+ blocks = context.request_to_kv_block_ids[fidx]
+ valid = blocks[blocks >= 0].tolist()
+ if valid:
+ finished_routing_block_ids[req_id] = valid
+
# Clone needed: update_requests mutates next_tokens in-place via tensor_swap,
- # which would corrupt the reused _sampled_tokens_cuda buffer.
- new_sample_copy = self._sampled_tokens_cuda[:active_request_count].clone()
+ # which would corrupt the reused buffer.
+ new_sample_copy = sampled_tokens_cpu.clone()
+ range_pop()
- # Update requests.
- # _sampled_mtp_tokens_cuda has shape [num_speculative_tokens, max_requests]
- if self.num_speculative_tokens > 0:
- sampled_mtp_tokens_cuda = self._sampled_mtp_tokens_cuda[:, :active_request_count]
- else:
- sampled_mtp_tokens_cuda = None
+ range_push("update_requests")
update_result = context.update_requests(
- active_request_mask, new_sample_copy, sampled_mtp_tokens_cuda
+ active_request_mask, new_sample_copy, sampled_mtp_tokens_cpu
)
+ range_pop()
return {
"active_request_ids": active_request_ids,
"finished_request_ids": finished_request_ids,
+ # Already a CPU tensor (independent of _sampled_tokens_cuda via the
+ # .cpu() in _transfer_samples_to_cpu; update_requests only mutates
+ # the separate new_sample_copy). Returning the CPU copy avoids a
+ # D2H sync when the engine later calls sample.tolist().
+ "sample": sampled_tokens_cpu,
+ "finished_routing_block_ids": finished_routing_block_ids,
**(update_result or {}),
}
- @torch.inference_mode()
async def async_generate_output_tokens_dynamic_batch(
self, skip_bookkeeping: Optional[bool] = False
) -> Optional[Dict]:
@@ -1760,30 +1715,37 @@ async def async_generate_output_tokens_dynamic_batch(
if context.active_token_count == 0 and active_request_count == 0:
return None
- input_ids, position_ids = self._dynamic_step_context_init()
+ with torch.inference_mode():
+ input_ids, position_ids = self._dynamic_step_context_init()
- cuda_graph_request_count = (
- context.padded_active_request_count if context.using_cuda_graph_this_step() else None
- )
+ cuda_graph_request_count = (
+ context.padded_active_request_count
+ if context.using_cuda_graph_this_step()
+ else None
+ )
- # Enable routing recording before forward pass if routing replay is enabled
- config = self.inference_wrapped_model.model.config
- if config.moe_enable_routing_replay:
- RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)
+ # Enable routing recording before forward pass if routing replay is enabled
+ config = self.inference_wrapped_model.model.config
+ if config.moe_enable_routing_replay:
+ RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)
- # Forward pass produces only base logits. When speculative decoding is
- # active, MTP logits are computed serially after verification.
- logits = self._dynamic_step_forward_logits(input_ids, position_ids)
+ # Forward pass produces only base logits. When speculative decoding is
+ # active, MTP logits are computed serially after verification.
+ range_push("forward_pass")
+ self._dynamic_step_forward_logits(input_ids, position_ids)
- # Commit Mamba intermediate states before update_requests, which
- # may swap request indices. The Python lists tracking EOS block IDs
- # and intermediate offsets are not swapped along with tensors, so
- # commit must run while indices are still valid.
- if context.is_hybrid_model and context.mamba_slot_allocator is not None:
- context.mamba_slot_allocator.commit_intermediate_states()
+ # Commit Mamba intermediate states before update_requests, which
+ # may swap request indices. The Python lists tracking EOS block IDs
+ # and intermediate offsets are not swapped along with tensors, so
+ # commit must run while indices are still valid.
+ if context.is_hybrid_model and context.mamba_slot_allocator is not None:
+ context.mamba_slot_allocator.commit_intermediate_states()
- # Collect routing indices per request (must be done before context transitions)
- routing_indices_per_request = self._router_record_bookkeeping()
+ # Collect flat routing indices and scatter them into per-block storage.
+ # Must be done before update_requests while token-to-block mappings are valid.
+ # Reconstruction happens from blocks at request completion.
+ context.kv_block_allocator.store_routing_per_block(self._router_record_bookkeeping())
+ range_pop()
# This is the best place to yield control back to event loop.
# At this point we have enqueued FW pass GPU kernels asynchronously.
@@ -1793,68 +1755,85 @@ async def async_generate_output_tokens_dynamic_batch(
# Todo [Siddharth]: Can we condition the sleep on a cuda event?
# NOTE [TDE]: This will be moved once CPU and GPU methods are separated.
await asyncio.sleep(0)
- return_log_probs, return_top_n_logprobs = self._dynamic_step_log_probs_bookkeeping()
-
- self._dynamic_step_sample_bookkeeping()
-
- if self.num_speculative_tokens > 0:
- # Phase 1: Verify speculative tokens using base logits only.
- self._dynamic_step_sample_logits_and_verify_tokens(logits, input_ids)
- # Phase 2: Rewind KV cache for rejected tokens.
- self._rewind_kv_cache()
-
- # Disable MoE padding for MTP computation
- if self.model_config.moe_pad_experts_for_cuda_graph_inference:
- unwrapped_model = unwrap_model(self.inference_wrapped_model.model)
- set_decode_expert_padding(unwrapped_model, False)
- # Phase 3: Compute MTP serially with correct (verified) inputs.
- self._compute_serial_mtp_and_sample()
- else:
- self._dynamic_step_sample_logits(logits)
+ with torch.inference_mode():
+ range_push("sampling")
+ return_log_probs, return_top_n_logprobs = self._dynamic_step_log_probs_bookkeeping()
- log_probs = None
- top_n_logprobs = None
- if return_log_probs or return_top_n_logprobs:
if self.num_speculative_tokens > 0:
- log_probs, log_probs_tensor = self._dynamic_step_calculate_log_probs_speculative(
- logits
- )
- if return_top_n_logprobs:
- top_n_logprobs = self._dynamic_step_calculate_top_n_logprobs_speculative(
- log_probs_tensor
- )
+ # Phase 1: Verify speculative tokens using base logits only.
+ nvtx_range_push("mtp-spec-decoding/verify")
+ self._dynamic_step_sample_logits_and_verify_tokens(input_ids)
+ nvtx_range_pop("mtp-spec-decoding/verify")
+ # Phase 2: Rewind KV cache for rejected tokens.
+ nvtx_range_push("mtp-spec-decoding/rewind-kv-cache")
+ blocks_to_release, remove_mask = self._rewind_kv_cache()
+ nvtx_range_pop("mtp-spec-decoding/rewind-kv-cache")
+
+ # Disable MoE padding for MTP computation, unless CUDA graphs
+ # are active (the graphs were captured with padding enabled).
+ if self.model_config.moe_pad_experts_for_cuda_graph_inference:
+ if not context.using_cuda_graph_this_step():
+ set_decode_expert_padding(self._unwrapped_model, False)
+
+ # Phase 3: Compute MTP serially with correct (verified) inputs.
+ nvtx_range_push("mtp-spec-decoding/serial-mtp")
+ self._compute_serial_mtp_and_sample()
+ nvtx_range_pop("mtp-spec-decoding/serial-mtp")
+
+ # Phase 4: Release freed blocks. Deferred from Phase 2 so the
+ # data-dependent boolean-mask sync overlaps with MTP GPU work.
+ context.kv_block_allocator.release_memory_blocks(blocks_to_release[remove_mask])
else:
- log_probs, log_probs_tensor = self._dynamic_step_calculate_log_probs(logits)
- if return_top_n_logprobs:
- top_n_logprobs = self._dynamic_step_calculate_top_n_logprobs(
- logits, log_probs_tensor
+ self._dynamic_step_sample_logits()
+
+ log_probs = None
+ top_n_logprobs = None
+ if return_log_probs or return_top_n_logprobs:
+ if self.num_speculative_tokens > 0:
+ log_probs, log_probs_tensor = (
+ self._dynamic_step_calculate_log_probs_speculative()
)
-
- if skip_bookkeeping:
- request_bookkeeping = {}
- else:
- request_bookkeeping = self._dynamic_step_context_bookkeeping()
-
- ret = {
- # Clone needed: _sampled_tokens_cuda is a reused buffer overwritten each step.
- "sample": self._sampled_tokens_cuda[:active_request_count].clone(),
- "accepted_tokens": (
- # Clone needed: .fill_(-1) on line 1480 would corrupt the returned value.
- self._accepted_tokens_per_request.clone()
- if self.num_speculative_tokens > 0
- else None
- ),
- "log_probs": log_probs,
- "top_n_logprobs": top_n_logprobs,
- "routing_indices_per_request": routing_indices_per_request,
- "cuda_graph_request_count": cuda_graph_request_count,
- }
- if self.num_speculative_tokens > 0:
- self._accepted_tokens_per_request.fill_(-1)
- self._accepted_token_counts_per_request.fill_(0)
- ret.update(request_bookkeeping)
- return ret
+ if return_top_n_logprobs:
+ top_n_logprobs = self._dynamic_step_calculate_top_n_logprobs_speculative(
+ log_probs_tensor
+ )
+ else:
+ log_probs, log_probs_tensor = self._dynamic_step_calculate_log_probs()
+ if return_top_n_logprobs:
+ top_n_logprobs = self._dynamic_step_calculate_top_n_logprobs(
+ log_probs_tensor
+ )
+ range_pop()
+
+ if skip_bookkeeping:
+ # _transfer_samples_to_cpu wasn't invoked on this path, so do
+ # a one-shot D2H here to keep "sample" as a CPU tensor for
+ # downstream consumers.
+ request_bookkeeping = {
+ "sample": self._sampled_tokens_cuda[:active_request_count].cpu()
+ }
+ else:
+ # request_bookkeeping supplies "sample" as the already-CPU
+ # tensor produced by _transfer_samples_to_cpu.
+ request_bookkeeping = self._dynamic_step_context_bookkeeping()
+
+ ret = {
+ "accepted_tokens": (
+ # Clone needed: .fill_(-1) on line 1480 would corrupt the returned value.
+ self._accepted_tokens_per_request.clone()
+ if self.num_speculative_tokens > 0
+ else None
+ ),
+ "log_probs": log_probs,
+ "top_n_logprobs": top_n_logprobs,
+ "cuda_graph_request_count": cuda_graph_request_count,
+ }
+ if self.num_speculative_tokens > 0:
+ self._accepted_tokens_per_request.fill_(-1)
+ self._accepted_token_counts_per_request.fill_(0)
+ ret.update(request_bookkeeping)
+ return ret
@torch.inference_mode()
def generate_output_tokens_dynamic_batch(
@@ -1944,10 +1923,7 @@ def generate_all_output_tokens_static_batch(
)
# Check whether CUDA graphs are enabled
- enable_cuda_graph = (
- model_config.cuda_graph_impl == "local"
- and CudaGraphScope.full_iteration not in model_config.cuda_graph_scope
- )
+ enable_cuda_graph = model_config.cuda_graph_impl == "local"
# Pad batch tokens if necessary
batch_size = len(active_requests)
diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py
index 75faefd4b88..460acf39e9b 100644
--- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py
+++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py
@@ -76,16 +76,6 @@ def _get_field(obj, key, default=None):
return getattr(obj, key, default)
-_TRANSFER_TOOL_NAME = "transfer_to_human_agents"
-_TRANSFER_HOLD_MESSAGE = "YOU ARE BEING TRANSFERRED TO A HUMAN AGENT. PLEASE HOLD ON."
-_RESERVATION_UPDATE_TOOLS = {
- "update_reservation_flights",
- "update_reservation_passengers",
- "update_reservation_baggages",
-}
-_RESERVATION_DESTRUCTIVE_TOOLS = {"cancel_reservation", "book_reservation"}
-
-
def _try_parse_jsonish(value):
if not isinstance(value, str):
return value
@@ -199,48 +189,21 @@ def _normalize_tool_calls(tool_calls, tools=None):
"function": {"name": str(fn_name), "arguments": fn_args},
}
)
- return _apply_tool_call_guardrails(normalized)
+ return normalized
-def _apply_tool_call_guardrails(tool_calls):
- """Apply conservative post-parse guardrails to tool call lists.
+def _maybe_filter_parallel_tool_calls(tool_calls, parallel_tool_calls):
+ """Filter to first tool call only when parallel_tool_calls is False.
- If update-style reservation tools are already present in the same response,
- suppress cancel+book style calls to avoid destructive replanning patterns.
+ Matches vLLM's maybe_filter_parallel_tool_calls behavior.
"""
- if not isinstance(tool_calls, list):
+ if parallel_tool_calls:
return tool_calls
-
- call_names = {
- _get_field(_get_field(call, "function", {}), "name")
- for call in tool_calls
- if isinstance(call, dict)
- }
- if call_names & _RESERVATION_UPDATE_TOOLS:
- return [
- call
- for call in tool_calls
- if _get_field(_get_field(call, "function", {}), "name")
- not in _RESERVATION_DESTRUCTIVE_TOOLS
- ]
+ if tool_calls:
+ return tool_calls[:1]
return tool_calls
-def _normalize_assistant_content(message_text, tool_calls):
- """Normalize assistant content for policy-sensitive tool transitions."""
- if not isinstance(message_text, str):
- message_text = "" if message_text is None else str(message_text)
-
- tool_names = {
- _get_field(_get_field(call, "function", {}), "name")
- for call in (tool_calls or [])
- if isinstance(call, dict)
- }
- if _TRANSFER_TOOL_NAME in tool_names:
- return _TRANSFER_HOLD_MESSAGE
- return message_text
-
-
def _coerce_arguments_mapping(arguments):
"""Coerce function.arguments to a mapping for HF/Jinja chat templates.
@@ -397,6 +360,35 @@ def _replace_prefix_tokens(
return previous_turn_token_ids + current_turn_additional_token_ids
+def _coerce_to_token_id_list(result):
+ """Convert the return value of `tokenizer.apply_chat_template` to `list[int]`.
+
+ transformers >= 5.x.x sometimes returns a `BatchEncoding` object instead of a `list[int]`.
+ """
+ # BatchEncoding / dict-like with input_ids
+ if isinstance(result, dict) or hasattr(result, "input_ids"):
+ ids = result["input_ids"]
+ if hasattr(ids, "tolist"):
+ ids = ids.tolist()
+ if ids and isinstance(ids[0], list):
+ ids = ids[0]
+ return list(ids)
+ # Fast-tokenizer Encoding object
+ if hasattr(result, "ids"):
+ ids = result.ids
+ if hasattr(ids, "tolist"):
+ ids = ids.tolist()
+ return list(ids)
+ # Raw tensor / ndarray
+ if hasattr(result, "tolist"):
+ ids = result.tolist()
+ if ids and isinstance(ids[0], list):
+ ids = ids[0]
+ return ids
+ # Plain list
+ return list(result)
+
+
try:
import orjson
@@ -446,7 +438,9 @@ async def chat_completions():
req = await request.get_json()
tools = req.get("tools", None)
- tools_requested = bool(tools)
+ tool_choice = req.get("tool_choice", None)
+ parallel_tool_calls = req.get("parallel_tool_calls", True)
+ tools_requested = bool(tools) and tool_choice != "none"
messages = req.get("messages")
chat_template_kwargs = req.get("chat_template_kwargs", {})
if not isinstance(chat_template_kwargs, dict):
@@ -468,12 +462,14 @@ async def chat_completions():
hasattr(tokenizer, 'apply_chat_template')
and getattr(tokenizer, "chat_template", None) is not None
):
- prompt_tokens = tokenizer.apply_chat_template(
- template_messages,
- tokenize=True,
- add_generation_prompt=True,
- tools=template_tools,
- **chat_template_kwargs,
+ prompt_tokens = _coerce_to_token_id_list(
+ tokenizer.apply_chat_template(
+ template_messages,
+ tokenize=True,
+ add_generation_prompt=True,
+ tools=template_tools,
+ **chat_template_kwargs,
+ )
)
if req.get("prevent_retokenization", True):
@@ -514,12 +510,14 @@ async def chat_completions():
]
# Get the templated tokenization of just the previous generation
- retokenized_previous_turn_token_ids = tokenizer.apply_chat_template(
- messages_to_last_assistant_message,
- tokenize=True,
- add_generation_prompt=False,
- tools=template_tools,
- **chat_template_kwargs,
+ retokenized_previous_turn_token_ids = _coerce_to_token_id_list(
+ tokenizer.apply_chat_template(
+ messages_to_last_assistant_message,
+ tokenize=True,
+ add_generation_prompt=False,
+ tools=template_tools,
+ **chat_template_kwargs,
+ )
)
# Replace the prefix tokens with the tokens from the previous generation.
@@ -640,6 +638,16 @@ async def chat_completions():
error_detail = "; ".join(failed_errors)
status = 400 if has_nontransient_error else 500
logger.error(f"Inference request(s) failed: {error_detail}")
+
+ # NOTE: This exact string is required for compatibility with Nemo-RL, DO NOT MODIFY.
+ if "MaxSequenceLengthOverflowError" in error_detail:
+ error_msg = (
+ f"This model's maximum context length was exceeded. "
+ f"Your messages resulted in {len(prompt_tokens)} tokens. "
+ f"Please reduce the length of the messages. {error_detail}"
+ )
+ return Response(error_msg, status=400)
+
return Response(f"Inference request(s) failed: {error_detail}", status=status)
# --- 5. Format OpenAI Response ---
@@ -699,10 +707,22 @@ async def chat_completions():
)
normalized_tool_calls = metadata.get("tool_calls", [])
- message = {
- "role": "assistant",
- "content": _normalize_assistant_content(message_text, normalized_tool_calls),
- }
+
+ # Apply parallel_tool_calls filtering (matches vLLM behavior)
+ normalized_tool_calls = _maybe_filter_parallel_tool_calls(
+ normalized_tool_calls, parallel_tool_calls
+ )
+
+ # Determine content based on tool_choice (matches vLLM behavior):
+ # - Named tool choice or "required": content is empty string
+ # - Otherwise: content is the parsed message text
+ is_named_tool_choice = isinstance(tool_choice, dict) and "function" in tool_choice
+ if normalized_tool_calls and (is_named_tool_choice or tool_choice == "required"):
+ content = ""
+ else:
+ content = message_text if message_text is not None else ""
+
+ message = {"role": "assistant", "content": content}
if normalized_tool_calls:
message["tool_calls"] = normalized_tool_calls
if "reasoning" in metadata:
@@ -712,14 +732,24 @@ async def chat_completions():
message["prompt_token_ids"] = result["prompt_tokens"]
message["generation_token_ids"] = result["generated_tokens"]
message["generation_log_probs"] = result.get("generated_log_probs", [])
+ message["policy_epoch"] = result["policy_epoch"]
+ message["kv_cache_epoch"] = result["kv_cache_epoch"]
+ message["num_evictions"] = sum(1 for e in result["events"] if e.get("type") == "EVICT")
return_log_probs = sampling_params.return_log_probs
- finish_reason = "tool_calls" if metadata.get("tool_calls", []) else "stop"
+ # Determine finish_reason following vLLM conventions:
+ # - "tool_calls" for auto or required tool choice when tools are called
+ # - "stop" for named tool choice (even when tools are called)
+ # - "length" when max tokens is reached
if (
len(result["generated_tokens"])
>= result["sampling_params"]["num_tokens_to_generate"]
):
finish_reason = "length"
+ elif normalized_tool_calls and not is_named_tool_choice:
+ finish_reason = "tool_calls"
+ else:
+ finish_reason = "stop"
choice_data = {
"index": request_idx,
@@ -733,11 +763,6 @@ async def chat_completions():
"logprobs": {"content": logprobs_content} if return_log_probs else None,
"finish_reason": finish_reason,
}
- choice_data["policy_epoch"] = result["policy_epoch"]
- choice_data["kv_cache_epoch"] = result["kv_cache_epoch"]
- choice_data["num_evictions"] = sum(
- 1 for e in result["events"] if e.get("type") == "EVICT"
- )
if current_app.config['verbose']:
logging.info(_redact_token_id_lists_for_logging(result))
@@ -759,7 +784,7 @@ async def chat_completions():
prompt_token_count = max(prompt_tokens_counts) if prompt_tokens_counts else 0
response = {
- "id": str(uuid.uuid4()),
+ "id": f"chatcmpl-{uuid.uuid4().hex}",
"created": int(time.time()),
"model": "EMPTY",
"object": "chat.completion",
diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/completions.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/completions.py
index d2279b0d07d..6f57a863c1c 100644
--- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/completions.py
+++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/completions.py
@@ -3,6 +3,7 @@
import asyncio
import logging
import time
+import uuid
from megatron.core.inference.inference_request import unwrap_serialized_tensors
from megatron.core.inference.sampling_params import SamplingParams
@@ -92,6 +93,8 @@ async def completions():
if isinstance(stop, str):
stop = [stop]
+ ignore_eos = bool(req.get("ignore_eos", False))
+
sampling_params = SamplingParams(
temperature=temperature,
top_k=top_k,
@@ -101,6 +104,7 @@ async def completions():
skip_prompt_log_probs=skip_prompt_log_probs,
num_tokens_to_generate=int(req.get("max_tokens", 16)),
stop_words=stop,
+ termination_id=-1 if ignore_eos else None,
)
except ValueError as e:
return f"Invalid sampling parameter: {e}", 400
@@ -117,6 +121,7 @@ async def completions():
skip_prompt_log_probs=sampling_params.skip_prompt_log_probs,
num_tokens_to_generate=sampling_params.num_tokens_to_generate,
stop_words=sampling_params.stop_words,
+ termination_id=sampling_params.termination_id,
)
tasks.append(client.add_request(prompt_tokens, per_req_params))
@@ -160,6 +165,8 @@ async def completions():
# --- 5. Format Response (matching old_completions.py) ---
choices = []
+ total_completion_tokens = 0
+ prompt_tokens_counts = []
request_idx = 0
for completed_request in batch_results:
@@ -167,6 +174,17 @@ async def completions():
full_text = result["generated_text"] or ""
text_output = (prompts_as_strings[request_idx] + full_text) if echo else full_text
+ generated_tokens = result.get("generated_tokens") or []
+ prompt_tokens_list = result.get("prompt_tokens") or []
+ total_completion_tokens += len(generated_tokens)
+ prompt_tokens_counts.append(len(prompt_tokens_list))
+
+ finish_reason = "length"
+ sampling_params_result = result.get("sampling_params") or {}
+ num_tokens_requested = sampling_params_result.get("num_tokens_to_generate")
+ if num_tokens_requested is None or len(generated_tokens) < num_tokens_requested:
+ finish_reason = "stop"
+
logprobs_data = None
if sampling_params.return_log_probs:
# Get prompt tokens and logprobs
@@ -230,20 +248,49 @@ async def completions():
"top_logprobs": top_logprobs,
}
- choices.append({"index": request_idx, "text": text_output, "logprobs": logprobs_data})
+ choice_data = {
+ "index": request_idx,
+ "text": text_output,
+ "logprobs": logprobs_data,
+ "finish_reason": finish_reason,
+ "prompt_token_ids": result["prompt_tokens"],
+ "generation_token_ids": result["generated_tokens"],
+ "generation_log_probs": result.get("generated_log_probs", []),
+ }
+ choice_data["policy_epoch"] = result["policy_epoch"]
+ choice_data["kv_cache_epoch"] = result["kv_cache_epoch"]
+ choice_data["num_evictions"] = sum(
+ 1 for e in result["events"] if e.get("type") == "EVICT"
+ )
+
if result["routing_indices"] is not None:
- choices[-1]["moe_topk_indices"] = result["routing_indices"]
+ choice_data["moe_topk_indices"] = result["routing_indices"]
prompt_length = (
len(result["prompt_tokens"]) if result["prompt_tokens"] is not None else 0
)
if prompt_length:
- choices[-1]["prompt_moe_topk_indices"] = result["routing_indices"][
+ choice_data["prompt_moe_topk_indices"] = result["routing_indices"][
:prompt_length
]
+ choices.append(choice_data)
request_idx += 1
- return jsonify({"choices": choices})
+ prompt_token_count = max(prompt_tokens_counts) if prompt_tokens_counts else 0
+ return jsonify(
+ {
+ "id": str(uuid.uuid4()),
+ "object": "text_completion", # as per the openAI spec
+ "created": int(time.time()),
+ "model": "EMPTY",
+ "choices": choices,
+ "usage": {
+ "prompt_tokens": prompt_token_count,
+ "completion_tokens": total_completion_tokens,
+ "total_tokens": prompt_token_count + total_completion_tokens,
+ },
+ }
+ )
except ImportError as e:
logger.warning(f"Could not import quart: {e}")
diff --git a/megatron/core/inference/text_generation_server/run_mcore_engine.py b/megatron/core/inference/text_generation_server/run_mcore_engine.py
index e278fcde3ee..3ba25687cd1 100644
--- a/megatron/core/inference/text_generation_server/run_mcore_engine.py
+++ b/megatron/core/inference/text_generation_server/run_mcore_engine.py
@@ -1,12 +1,11 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
-import inspect
-
from megatron.core import mpu
from megatron.core.inference.communication_utils import broadcast_float_list
from megatron.core.inference.inference_request import InferenceRequest
from megatron.core.inference.sampling_params import SamplingParams
from megatron.core.inference.text_generation_server.tokenization import tokenize_prompts
+from megatron.core.utils import accepts_parameter
def run_mcore_engine(
@@ -60,18 +59,11 @@ def run_mcore_engine(
for p, l in zip(context_tokens_tensor, context_length_tensor):
tokenized_prompts.append(p[:l].cpu().numpy().tolist())
- # detect if detokenize supports skip_special_tokens or **kwargs
- sig_params = inspect.signature(tokenizer.detokenize).parameters.values()
- accepts_skip = any(
- p.name == "skip_special_tokens" or p.kind == inspect.Parameter.VAR_KEYWORD
- for p in sig_params
- )
-
# Detokenize prompts into strings to pass through the engine
detokenized_prompts = [
(
tokenizer.detokenize(p, skip_special_tokens=True)
- if accepts_skip
+ if accepts_parameter(tokenizer.detokenize, "skip_special_tokens")
else tokenizer.detokenize(p)
)
for p in tokenized_prompts
@@ -89,10 +81,11 @@ def run_mcore_engine(
result = engine.generate(inference_requests=requests)
- # Only post-process on first stage.
- if mpu.is_pipeline_first_stage():
+ # Only post-process on the server rank (first stage with prompts)
+ if mpu.is_pipeline_first_stage() and prompts is not None:
response_dict = {
- "text": [x.prompt + x.generated_text for x in result],
+ # Send original prompts, not x.prompt, to circumvent tokenization artifacts
+ "text": [p + x.generated_text for p, x in zip(prompts, result)],
"tokens": [x.prompt_tokens + x.generated_tokens.tolist() for x in result],
}
if sampling_params.return_log_probs:
diff --git a/megatron/core/inference/utils.py b/megatron/core/inference/utils.py
index 0914b81f005..f20debe2589 100644
--- a/megatron/core/inference/utils.py
+++ b/megatron/core/inference/utils.py
@@ -1,6 +1,7 @@
# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
import asyncio
+import contextlib
import logging
import multiprocessing
import sys
@@ -8,7 +9,6 @@
import torch
-from megatron.core.transformer.moe.moe_layer import MoELayer
from megatron.core.utils import get_model_config
try:
@@ -17,6 +17,42 @@
FLASHINFER_JIT_CACHE_VERSION = None
+class InferenceMode:
+ """Process-wide flag indicating whether an inference engine is currently using the model.
+
+ Modules that need to distinguish between inference and non-inference (e.g. training,
+ RL logprobs) paths should read `InferenceMode.is_active()` rather than relying on
+ `self.training`, `torch.is_grad_enabled()`, or `inference_context is not None`.
+ """
+
+ _is_active: bool = False
+
+ @classmethod
+ def is_active(cls) -> bool:
+ """Return True while an inference engine is currently using the model."""
+ return cls._is_active
+
+ @classmethod
+ def set_active(cls) -> None:
+ """Mark the inference engine as active. Idempotent."""
+ cls._is_active = True
+
+ @classmethod
+ def unset_active(cls) -> None:
+ """Mark the inference engine as inactive. Idempotent."""
+ cls._is_active = False
+
+ @classmethod
+ @contextlib.contextmanager
+ def active(cls):
+ """Context manager: set the flag for the duration of the `with` block."""
+ cls.set_active()
+ try:
+ yield
+ finally:
+ cls.unset_active()
+
+
def device_memory_summary() -> str:
"""One-line GPU memory summary for torch_memory_saver logging."""
dev = torch.cuda.current_device()
@@ -73,12 +109,15 @@ def get_attention_mask(seq_length: int) -> torch.Tensor:
# Initialize cache for sequence parallel modules
moe_layer_cache = None
+_moe_metadata_sync_initialized = False
def _init_moe_expert_cache(model):
"""
Initialize the cache of MoE layers once
"""
+ from megatron.core.transformer.moe.moe_layer import MoELayer
+
global moe_layer_cache
if moe_layer_cache is not None:
return # already initialized
@@ -100,6 +139,25 @@ def walk(module):
walk(model)
+def set_moe_metadata_sync(model) -> None:
+ """Set _runs_metadata_sync on inference dispatchers.
+
+ Exactly one dispatcher per model — the first MoE layer — fires update_metadata
+ each step. All subsequent layers skip it to avoid redundant collective calls.
+ Must be called once after the model is built and put into eval mode.
+ """
+ global moe_layer_cache, _moe_metadata_sync_initialized
+ if _moe_metadata_sync_initialized:
+ return
+ if moe_layer_cache is None:
+ _init_moe_expert_cache(model)
+ for i, moe_layer in enumerate(moe_layer_cache):
+ dispatcher = getattr(moe_layer, '_inference_token_dispatcher', None)
+ if dispatcher is not None:
+ dispatcher._runs_metadata_sync = i == 0
+ _moe_metadata_sync_initialized = True
+
+
def set_decode_expert_padding(model, set_to: bool = False, capacity_factor: int = None):
"""
Toggle MoE drop-and-pad for decode.
@@ -201,34 +259,6 @@ def check_flashinfer_jit_cache_installed(log_version: bool = False):
)
-def set_inference_cuda_graphed_iteration_for_ep_inference(model):
- """Enable CUDA graph compatibility for expert parallel inference.
-
- Sets a flag in all MoELayers indicating the current iteration is being
- captured/executed in a CUDA graph. This allows the dispatcher to adjust
- its behavior for CUDA graph compatibility.
- """
- global moe_layer_cache
- if moe_layer_cache is None:
- _init_moe_expert_cache(model)
-
- for moe_layer in moe_layer_cache:
- moe_layer.set_inference_cuda_graphed_iteration()
-
-
-def unset_inference_cuda_graphed_iteration_for_ep_inference(model):
- """Disable CUDA graph compatibility for expert parallel inference.
-
- Clears the flag in all MoELayers, restoring standard dispatcher behavior.
- """
- global moe_layer_cache
- if moe_layer_cache is None:
- _init_moe_expert_cache(model)
-
- for moe_layer in moe_layer_cache:
- moe_layer.unset_inference_cuda_graphed_iteration()
-
-
def tensor_swap(x, src_idxs, dst_idxs):
"""
Swap x[src_idxs] and x[dst_idxs]
diff --git a/megatron/core/model_parallel_config.py b/megatron/core/model_parallel_config.py
index d5cd5397d56..dabe0d0aced 100644
--- a/megatron/core/model_parallel_config.py
+++ b/megatron/core/model_parallel_config.py
@@ -261,6 +261,15 @@ class ModelParallelConfig:
delay_wgrad_compute: bool = False
"""Delay the weight gradient computation to improve batch-level communication overlapping"""
+ overlap_dispatch_backward_with_experts_wgrad: bool = False
+ """Delay the weight gradient computation for TE Grouped GEMM MoE experts.
+ When enabled with FSDP, the expert weight gradients are computed on a separate
+ CUDA stream after the data gradients finish, allowing overlap of wgrad compute
+ with EP A2A communication. The FSDP gradient reduce-scatter for
+ expert parameters is deferred until the delayed wgrad computation completes.
+ This requires transformer_engine with GroupedLinear support (TE >= 2.3.0).
+ """
+
ep_overlap_early_attn_memory_release: bool = False
"""Enable early memory release of attention activations during EP overlap.
EP overlap can increase peak memory usage when the overlapped forward module allocates
diff --git a/megatron/core/models/T5/t5_spec.py b/megatron/core/models/T5/t5_spec.py
index 9f465df5c21..0b273b8f9e7 100644
--- a/megatron/core/models/T5/t5_spec.py
+++ b/megatron/core/models/T5/t5_spec.py
@@ -1,4 +1,6 @@
# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+from functools import partial
+
from megatron.core.extensions.transformer_engine import HAVE_TE
from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
@@ -63,14 +65,14 @@ def encoder_model_with_transformer_engine_default_spec() -> ModuleSpec:
submodules=SelfAttentionSubmodules(
linear_qkv=not_none(TELayerNormColumnParallelLinear),
core_attention=not_none(TEDotProductAttention),
- linear_proj=TERowParallelLinear,
+ linear_proj=not_none(TERowParallelLinear),
q_layernorm=IdentityOp,
k_layernorm=IdentityOp,
),
),
self_attn_bda=get_bias_dropout_add,
- mlp=ModuleSpec(
- module=MLP,
+ mlp=partial(
+ MLP.as_mlp_submodule,
submodules=MLPSubmodules(
linear_fc1=not_none(TELayerNormColumnParallelLinear),
linear_fc2=not_none(TERowParallelLinear),
@@ -93,7 +95,7 @@ def decoder_model_with_transformer_engine_default_spec() -> ModuleSpec:
submodules=SelfAttentionSubmodules(
linear_qkv=not_none(TELayerNormColumnParallelLinear),
core_attention=not_none(TEDotProductAttention),
- linear_proj=TERowParallelLinear,
+ linear_proj=not_none(TERowParallelLinear),
q_layernorm=IdentityOp,
k_layernorm=IdentityOp,
),
@@ -107,12 +109,12 @@ def decoder_model_with_transformer_engine_default_spec() -> ModuleSpec:
linear_q=not_none(TEColumnParallelLinear),
linear_kv=not_none(TEColumnParallelLinear),
core_attention=not_none(TEDotProductAttention),
- linear_proj=TERowParallelLinear,
+ linear_proj=not_none(TERowParallelLinear),
),
),
cross_attn_bda=get_bias_dropout_add,
- mlp=ModuleSpec(
- module=MLP,
+ mlp=partial(
+ MLP.as_mlp_submodule,
submodules=MLPSubmodules(
linear_fc1=not_none(TELayerNormColumnParallelLinear),
linear_fc2=not_none(TERowParallelLinear),
@@ -143,8 +145,8 @@ def encoder_model_with_local_spec() -> ModuleSpec:
),
self_attn_bda=get_bias_dropout_add,
pre_mlp_layernorm=LNImpl,
- mlp=ModuleSpec(
- module=MLP,
+ mlp=partial(
+ MLP.as_mlp_submodule,
submodules=MLPSubmodules(
linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear
),
@@ -190,8 +192,8 @@ def decoder_model_with_local_spec() -> ModuleSpec:
),
cross_attn_bda=get_bias_dropout_add,
pre_mlp_layernorm=LNImpl,
- mlp=ModuleSpec(
- module=MLP,
+ mlp=partial(
+ MLP.as_mlp_submodule,
submodules=MLPSubmodules(
linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear
),
diff --git a/megatron/core/models/backends.py b/megatron/core/models/backends.py
index b019d527342..a270161ddd6 100644
--- a/megatron/core/models/backends.py
+++ b/megatron/core/models/backends.py
@@ -103,7 +103,7 @@ def column_parallel_linear(self) -> type:
"""Which column parallel linear module the backend uses"""
return ColumnParallelLinear
- def row_parallel_linear(self) -> type:
+ def row_parallel_linear(self) -> type[RowParallelLinear]:
"""Which row parallel linear module the backend uses"""
return RowParallelLinear
@@ -157,8 +157,8 @@ def column_parallel_linear(self) -> type:
"""Which column parallel linear module TE backend uses"""
return InferenceColumnParallelLinear
- def row_parallel_linear(self) -> type:
- """Which row parallel linear module TE backend uses"""
+ def row_parallel_linear(self) -> type[InferenceRowParallelLinear]:
+ """Which row parallel linear module Inference backend uses"""
return InferenceRowParallelLinear
def fuse_layernorm_and_linear(self) -> bool:
diff --git a/megatron/core/models/bert/bert_layer_specs.py b/megatron/core/models/bert/bert_layer_specs.py
index 53cc0f4280d..dc0099fa66e 100644
--- a/megatron/core/models/bert/bert_layer_specs.py
+++ b/megatron/core/models/bert/bert_layer_specs.py
@@ -1,5 +1,6 @@
# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
import warnings
+from functools import partial
from megatron.core.extensions.transformer_engine import HAVE_TE
from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
@@ -66,8 +67,8 @@ def get_bert_layer_with_transformer_engine_submodules() -> TransformerLayerSubmo
),
),
self_attn_bda=get_bias_dropout_add,
- mlp=ModuleSpec(
- module=MLP,
+ mlp=partial(
+ MLP.as_mlp_submodule,
submodules=MLPSubmodules(
linear_fc1=not_none(TELayerNormColumnParallelLinear),
linear_fc2=not_none(TERowParallelLinear),
@@ -117,8 +118,8 @@ def __getattr__(name):
),
self_attn_bda=get_bias_dropout_add,
pre_mlp_layernorm=LNImpl,
- mlp=ModuleSpec(
- module=MLP,
+ mlp=partial(
+ MLP.as_mlp_submodule,
submodules=MLPSubmodules(linear_fc1=ColumnParallelLinear, linear_fc2=RowParallelLinear),
),
mlp_bda=get_bias_dropout_add,
diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
index 3c6b7c4ab8d..17c73a33ae8 100644
--- a/megatron/core/models/common/language_module/language_module.py
+++ b/megatron/core/models/common/language_module/language_module.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
import logging
import os
from typing import Optional, Tuple
@@ -8,6 +8,7 @@
from megatron.core import parallel_state, tensor_parallel
from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.transformer.cuda_graphs import CudaGraphManager
try:
from megatron.core.extensions.transformer_engine import te_parallel_cross_entropy
@@ -21,7 +22,7 @@
is_vp_last_stage,
)
from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.core.transformer.enums import AttnBackend, CudaGraphScope
+from megatron.core.transformer.enums import AttnBackend
from megatron.core.transformer.module import MegatronModule
from megatron.core.transformer.multi_token_prediction import tie_word_embeddings_state_dict
from megatron.core.transformer.transformer_config import TransformerConfig
@@ -63,6 +64,20 @@ def __init__(
self.vp_stage = None
self.vp_size = self.config.virtual_pipeline_model_parallel_size
+ def _setup_mtp_cuda_graphs(self):
+ """Wrap `compute_mtp_single_step` with a CudaGraphManager.
+
+ Must be called by subclasses after `self.mtp` is created.
+ """
+ if self.config.cuda_graph_impl == "local":
+ self._mtp_cudagraph_manager = CudaGraphManager(
+ self.config,
+ base_module=self,
+ function_name="compute_mtp_single_step",
+ need_backward=False,
+ inline_capture=True,
+ )
+
def _is_in_embd_group(self):
if self.embd_group is None:
return False
@@ -144,8 +159,8 @@ def compute_language_model_loss(self, labels: Tensor, logits: Tensor) -> Tensor:
labels = torch.as_strided(labels, labels.size(), (labels.size()[1], 1))
# Use is_cg_capturable=True for full iteration CUDA graphs to avoid torch.equal checks
is_cg_capturable = (
- hasattr(self.config, 'cuda_graph_scope')
- and CudaGraphScope.full_iteration in self.config.cuda_graph_scope
+ hasattr(self.config, 'cuda_graph_impl')
+ and self.config.cuda_graph_impl == "full_iteration"
)
if is_cg_capturable and not is_te_min_version("2.7.0"):
from megatron.core.utils import get_te_version
@@ -154,7 +169,7 @@ def compute_language_model_loss(self, labels: Tensor, logits: Tensor) -> Tensor:
raise AssertionError(
f"CUDA graph compatible cross entropy requires TransformerEngine >= 2.7.0, "
f"but found version {current_version}. Please upgrade TransformerEngine "
- f"or set cuda_graph_scope to a value other than 'full_iteration'."
+ f"or set cuda_graph_impl to a value other than 'full_iteration'."
)
loss = te_parallel_cross_entropy(
@@ -187,7 +202,12 @@ def setup_embeddings_and_output_layer(self) -> None:
# Mark embedding and output layer for decoupled_lr and other features.
# This is the original Megatron attribute used by decoupled_lr, Muon, FSDP, etc.
- if self.pre_process and hasattr(self, 'embedding'):
+ # Include MTP-stage embedding too: it is a duplicated copy of the pre_process
+ # embedding (kept in sync via cross-stage all-reduce). Without this tag, the
+ # LayerWise distributed optimizer routes it to its Muon-managed buffer and
+ # `_emit_bucket(shared_embedding=True)` replicates the (vocab x hidden) tensor
+ # across all dp_size shards, blowing up the chunk's buffer by ~8x.
+ if (self.pre_process or getattr(self, 'mtp_process', False)) and hasattr(self, 'embedding'):
self.embedding.word_embeddings.weight.is_embedding_or_output_parameter = True
if (
self.post_process
@@ -323,6 +343,55 @@ def shared_embedding_or_output_weight(self) -> Tensor:
return self.output_layer.weight
return None
+ @torch.inference_mode()
+ def compute_mtp_single_step(
+ self,
+ hidden_states: Tensor,
+ next_token_ids: Tensor,
+ position_ids: Tensor,
+ depth: Optional[int] = None,
+ eager: bool = False,
+ cache_key=None,
+ ) -> tuple:
+ """Compute a single MTP depth for speculative decoding.
+
+ This is called after speculative token verification to compute MTP
+ predictions conditioned on verified tokens only.
+
+ Args:
+ hidden_states (Tensor): Hidden states at last accepted positions.
+ next_token_ids (Tensor): Correct next token IDs [1, N].
+ position_ids (Tensor): Position IDs for the next tokens [1, N].
+ depth (int, optional): MTP depth index. Only needed when `mtp_use_repeated_layer` is
+ False (each depth uses a distinct layer). Omit for repeated-layer models so that a
+ single CUDA graph can serve all depths.
+ eager, cache_key: The `CudaGraphManager` works by monkey-patching this argument onto the
+ function signature. Explictly including them removes the need for a monkey-patch,
+ and makes it straightforward to call the same method with and without eager mode.
+ These arguments are consumed by `CudaGraphManager`, if it exists.
+
+ Returns:
+ tuple: (new_hidden_states, logits [N, 1, vocab_size]).
+ """
+ # CudaGraphManager consumes these args, if it exists
+ del eager, cache_key
+ layer_idx = 0 if depth is None else depth
+ mtp_hidden = self.mtp.layers[layer_idx].forward_single_position(
+ hidden_states=hidden_states,
+ next_token_ids=next_token_ids,
+ position_ids=position_ids,
+ embedding=self.embedding,
+ )
+
+ output_weight = None
+ if self.share_embeddings_and_output_weights:
+ output_weight = self.shared_embedding_or_output_weight()
+
+ logits, _ = self.output_layer(mtp_hidden, weight=output_weight, runtime_gather_output=True)
+ logits = self._scale_logits(logits)
+
+ return mtp_hidden, logits
+
def sharded_state_dict(
self,
prefix: str = '',
diff --git a/megatron/core/models/common/model_chunk_schedule_plan.py b/megatron/core/models/common/model_chunk_schedule_plan.py
index 2b9d72d5f35..8358e05a612 100644
--- a/megatron/core/models/common/model_chunk_schedule_plan.py
+++ b/megatron/core/models/common/model_chunk_schedule_plan.py
@@ -1,7 +1,7 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
from contextlib import nullcontext
-from typing import Optional
+from typing import Any, Callable, Optional
import torch
from torch import Tensor
@@ -14,6 +14,7 @@
get_comm_stream,
get_comp_stream,
)
+from megatron.core.utils import nvtx_range_pop, nvtx_range_push
class ModelChunkState:
@@ -172,6 +173,46 @@ def create_node(stream, module, name):
else:
self.mtp_post_process = NoopScheduleNode()
+ def set_fsdp_reshard_hooks(self, post_forward_hook, post_backward_hook):
+ """Wire FSDP parameter release callbacks for the fine-grained overlap schedule.
+
+ The EP overlap schedule bypasses the normal FSDP forward/backward hooks
+ (registered on the FSDP unit module) because it calls sub-modules directly
+ instead of going through TransformerLayer.forward(). This method attaches
+ explicit release hooks to individual schedule nodes so that all-gathered
+ parameters are freed at the right time.
+
+ Args:
+ post_forward_hook: Callable(module) that releases forward-pass params
+ (bwd=False). Typically ``fsdp_wrapper.post_forward_release_module``.
+ post_backward_hook: Callable(module) that releases backward-pass params
+ (bwd=True). Typically ``fsdp_wrapper.post_backward_release_module``.
+ """
+ from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionLayer
+ from megatron.core.transformer.transformer_layer import TransformerLayer
+
+ assert isinstance(self.layer, (TransformerLayer, MultiTokenPredictionLayer)), (
+ f"Megatron FSDP with EP Overlap only supports TransformerLayer, "
+ f"but got {type(self.layer).__name__}."
+ )
+
+ if isinstance(self.layer, TransformerLayer):
+ hook_module = self.layer
+ else:
+ hook_module = self.layer.mtp_model_layer
+
+ # After the last backward op (attn), release backward-pass params.
+ self.attn.set_post_backward_hook(lambda: post_backward_hook(hook_module))
+
+ # Determine the last node in forward order.
+ if isinstance(self.moe_combine, NoopScheduleNode):
+ last_fwd_node = self.mlp
+ else:
+ last_fwd_node = self.moe_combine
+
+ # After the last forward op, release forward-pass params.
+ last_fwd_node.set_post_forward_hook(lambda: post_forward_hook(hook_module))
+
def get_fp8_context(self):
"""
Get the fp8 context for the transformer layer.
@@ -240,11 +281,14 @@ def run(f_layer, b_layer, f_input=None, b_grad=None, is_last_layer_in_bwd=False)
if f_layer is not None:
with f_layer.get_fp8_context():
f_input = f_layer.moe_combine.forward(f_input)
- f_input = f_layer.mtp_post_process.forward(f_input)
if b_layer is not None and not b_layer.config.ep_overlap_early_attn_memory_release:
b_grad = b_layer.attn.backward(b_grad)
+ if f_layer is not None:
+ with f_layer.get_fp8_context():
+ f_input = f_layer.mtp_post_process.forward(f_input)
+
# Delay the last attn_dw in backward pass (attn_dw of the first layer)
# for overlapping with the p2p comm
if b_layer is not None and not is_last_layer_in_bwd:
@@ -281,6 +325,9 @@ def __init__(
runtime_gather_output: Optional[bool] = None,
loss_mask: Optional[Tensor] = None,
padding_mask=None,
+ *,
+ output_processor: Optional[Callable[..., Tensor]] = None,
+ output_processor_context: Optional[Any] = None,
):
"""Initialize the schedule plan of all Transformer layers' sub-modules.
@@ -298,6 +345,10 @@ def __init__(
extra_block_kwargs: Additional keyword arguments for blocks.
runtime_gather_output: Whether to gather output at runtime.
loss_mask (torch.Tensor): Used to mask out some portions of the loss
+ output_processor (Callable): Custom postprocess hook to run instead of the
+ default logits/loss path.
+ output_processor_context (Any): User-defined context object forwarded to
+ `output_processor`.
Returns:
The model chunk schedule plan.
@@ -323,6 +374,8 @@ def __init__(
self._model_chunk_state.padding_mask = padding_mask
self._model_chunk_state.extra_block_kwargs = extra_block_kwargs
self._model_chunk_state.runtime_gather_output = runtime_gather_output
+ self._model_chunk_state.output_processor = output_processor
+ self._model_chunk_state.output_processor_context = output_processor_context
self._model_chunk_state.model = model
self._model_chunk_state.context = None
self._model_chunk_state.context_mask = None
@@ -473,7 +526,8 @@ def run(
for i in range(overlapped_layers):
f_layer = f_schedule_plan.get_layer(i)
b_layer = b_schedule_plan.pop_layer()
- torch.cuda.nvtx.range_push(f"layer_{i}f-layer_{b_schedule_plan.num_layers()}b")
+ nvtx_msg = f"layer_{i}f-layer_{b_schedule_plan.num_layers()}b"
+ nvtx_range_push(nvtx_msg)
f_input, b_grad = TransformerLayerSchedulePlan.run(
f_layer,
b_layer,
@@ -483,25 +537,27 @@ def run(
)
if i < b_num_layers - 1:
b_layer.release_state()
- torch.cuda.nvtx.range_pop()
+ nvtx_range_pop(nvtx_msg)
# backward pass for the remaining layers
for i in range(overlapped_layers, b_num_layers):
b_layer = b_schedule_plan.pop_layer()
- torch.cuda.nvtx.range_push(f"layer_{b_schedule_plan.num_layers()}b")
+ nvtx_msg = f"layer_{b_schedule_plan.num_layers()}b"
+ nvtx_range_push(nvtx_msg)
_, b_grad = TransformerLayerSchedulePlan.run(
None, b_layer, b_grad=b_grad, is_last_layer_in_bwd=(i == b_num_layers - 1)
)
if i < b_num_layers - 1:
b_layer.release_state()
- torch.cuda.nvtx.range_pop()
+ nvtx_range_pop(nvtx_msg)
# forward pass for the remaining layers
for i in range(overlapped_layers, f_num_layers):
f_layer = f_schedule_plan.get_layer(i)
- torch.cuda.nvtx.range_push(f"layer_{i}f")
+ nvtx_msg = f"layer_{i}f"
+ nvtx_range_push(nvtx_msg)
f_input, _ = TransformerLayerSchedulePlan.run(f_layer, None, f_input=f_input)
- torch.cuda.nvtx.range_pop()
+ nvtx_range_pop(nvtx_msg)
if f_schedule_plan is not None and post_forward is not None:
# post_forward()/send_forward_recv_forward() is running in the communication stream,
diff --git a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py
index 6608073136c..8f6b1a1a3f8 100644
--- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py
+++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py
@@ -24,10 +24,12 @@
)
from megatron.core.transformer.transformer_config import TransformerConfig
from megatron.core.transformer.transformer_layer import (
+ MlpBuilder,
TransformerLayer,
TransformerLayerSubmodules,
get_transformer_layer_offset,
)
+from megatron.core.typed_torch import not_none
try:
import transformer_engine as te # type: ignore[import-untyped] # pylint: disable=unused-import
@@ -123,6 +125,7 @@ def get_dsa_module_spec_for_backend(
q_layernorm=IdentityOp,
kv_layernorm=IdentityOp,
),
+ metainfo={"fuse_input_layernorm": False},
)
return attention
@@ -138,6 +141,8 @@ def get_experimental_attention_variant_module_spec(
if config.experimental_attention_variant == "gated_delta_net":
return get_gated_delta_net_module_spec(config=config, backend=backend)
+ elif config.experimental_attention_variant == "dsa":
+ return get_dsa_module_spec_for_backend(config=config, backend=backend)
else:
raise ValueError(
f"Invalid experimental attention variant: {config.experimental_attention_variant}"
@@ -213,14 +218,18 @@ def get_transformer_block_with_experimental_attention_variant_spec(
moe_layer_pattern = [0] * config.num_layers
if 1 in moe_layer_pattern:
- moe_layer_spec = _get_moe_module_spec(config=config, backend=backend)
+ moe_layer_spec, fuse_layernorm_pre_moe = _get_moe_module_spec(
+ config=config, backend=backend
+ )
else:
- moe_layer_spec = None
+ moe_layer_spec, fuse_layernorm_pre_moe = None, False
if 0 in moe_layer_pattern:
- dense_mlp_layer_spec = _get_dense_mlp_module_spec(config=config, backend=backend)
+ dense_mlp_layer_spec, fuse_layernorm_pre_dense = _get_dense_mlp_module_spec(
+ config=config, backend=backend
+ )
else:
- dense_mlp_layer_spec = None
+ dense_mlp_layer_spec, fuse_layernorm_pre_dense = None, False
# Get GPT decoder block layer specs
rms_norm = config.normalization == "RMSNorm"
@@ -232,6 +241,11 @@ def get_transformer_block_with_experimental_attention_variant_spec(
else standard_attention_spec
)
mlp = moe_layer_spec if moe_layer_pattern[layer_number] == 1 else dense_mlp_layer_spec
+ fuse_pre_mlp_layernorm = (
+ fuse_layernorm_pre_moe
+ if moe_layer_pattern[layer_number] == 1
+ else fuse_layernorm_pre_dense
+ )
input_layernorm = (
IdentityOp
if attention.metainfo["fuse_input_layernorm"]
@@ -239,7 +253,7 @@ def get_transformer_block_with_experimental_attention_variant_spec(
)
pre_mlp_layernorm = (
IdentityOp
- if mlp.metainfo["fuse_pre_mlp_layernorm"]
+ if fuse_pre_mlp_layernorm
else backend.layer_norm(rms_norm=rms_norm, for_qk=False)
)
@@ -251,7 +265,7 @@ def get_transformer_block_with_experimental_attention_variant_spec(
self_attention=attention,
self_attn_bda=get_bias_dropout_add,
pre_mlp_layernorm=pre_mlp_layernorm,
- mlp=mlp,
+ mlp=not_none(mlp),
mlp_bda=get_bias_dropout_add,
),
)
@@ -410,41 +424,50 @@ def _get_self_attention_module_spec(
def _get_dense_mlp_module_spec(
config: TransformerConfig, backend: BackendSpecProvider = None
-) -> ModuleSpec:
+) -> tuple[MlpBuilder, bool]:
"""Get dense MLP module spec.
For hybrid models that mix dense MLP and experimental attention architectures.
- Warning: This function may be deprecated in the future."""
+ Warning: This function may be deprecated in the future.
+
+ Returns:
+ A tuple of (MLP module spec, whether to fuse pre-MLP layernorm)
+ """
if backend is None:
backend = _get_backend_spec_provider(config=config)
from megatron.core.models.gpt.gpt_layer_specs import get_mlp_module_spec_for_backend
- mlp_spec = get_mlp_module_spec_for_backend(backend=backend, num_experts=None)
- mlp_spec.metainfo["fuse_pre_mlp_layernorm"] = backend.fuse_layernorm_and_linear()
-
- return mlp_spec
+ return (
+ get_mlp_module_spec_for_backend(backend=backend, num_experts=None),
+ backend.fuse_layernorm_and_linear(),
+ )
def _get_moe_module_spec(
config: TransformerConfig, backend: BackendSpecProvider = None
-) -> ModuleSpec:
+) -> tuple[MlpBuilder, bool]:
"""Get MoE module spec.
For hybrid models that mix MoE and experimental attention architectures.
- Warning: This function may be deprecated in the future."""
+ Warning: This function may be deprecated in the future.
+
+ Returns:
+ A tuple of (MoE module spec, whether to fuse pre-MoE layernorm)
+ """
if backend is None:
backend = _get_backend_spec_provider(config=config)
from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec_for_backend
- moe_spec = get_moe_module_spec_for_backend(
- backend=backend,
- num_experts=config.num_moe_experts,
- moe_grouped_gemm=config.moe_grouped_gemm,
- use_te_activation_func=config.use_te_activation_func,
+ return (
+ get_moe_module_spec_for_backend(
+ backend=backend,
+ num_experts=config.num_moe_experts,
+ moe_grouped_gemm=config.moe_grouped_gemm,
+ use_te_activation_func=config.use_te_activation_func,
+ ),
+ False,
)
- moe_spec.metainfo["fuse_pre_mlp_layernorm"] = False
- return moe_spec
diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py
index 93f3748de4d..4b50dfe359f 100644
--- a/megatron/core/models/gpt/fine_grained_callables.py
+++ b/megatron/core/models/gpt/fine_grained_callables.py
@@ -14,7 +14,7 @@
FineGrainedActivationOffloadingInterface as off_interface,
)
from megatron.core.pipeline_parallel.utils import ScheduleNode, make_viewless
-from megatron.core.transformer.enums import CudaGraphScope
+from megatron.core.transformer.enums import CudaGraphModule
from megatron.core.transformer.module import GraphableMegatronModule, float16_to_fp32
from megatron.core.transformer.moe.moe_layer import MoELayer
from megatron.core.transformer.multi_token_prediction import (
@@ -23,7 +23,7 @@
)
from megatron.core.transformer.transformer_layer import TransformerLayer, make_viewless_tensor
from megatron.core.typed_torch import apply_module, copy_signature
-from megatron.core.utils import internal_api
+from megatron.core.utils import internal_api, nvtx_range_pop, nvtx_range_push
def weak_method(method):
@@ -99,7 +99,7 @@ def should_free_input(name, is_moe, config, num_local_experts):
# If moe_preprocess is in cuda graph scope, tokens and probs are fixed size tensors,
# so they cannot be freed.
"moe_dispatch": not (enable_deepep or enable_hybridep or enable_mori)
- and (CudaGraphScope.moe_preprocess not in config.cuda_graph_scope),
+ and (CudaGraphModule.moe_preprocess not in config.cuda_graph_modules),
}
return free_input_nodes.get(name, False)
@@ -234,6 +234,8 @@ def forward_impl(self, hidden_states):
sequence_len_offset=self.chunk_state.sequence_len_offset,
runtime_gather_output=self.chunk_state.runtime_gather_output,
extra_block_kwargs=self.chunk_state.extra_block_kwargs,
+ output_processor=self.chunk_state.output_processor,
+ output_processor_context=self.chunk_state.output_processor_context,
)
# For now, 1f1b only supports fp16 module
@@ -271,7 +273,7 @@ def __init__(
bwd_dw_callables (list): List of weight gradient functions for the layer.
extra_args (dict): Extra arguments for the node: is_moe, config.
"""
- # determine whether to free input memory
+ # Determine whether to free input memory
config = extra_args.get("config", None)
assert config is not None, "model config must be passed to TransformerLayerNode."
is_moe = extra_args.get("is_moe", False)
@@ -279,6 +281,9 @@ def __init__(
free_input = should_free_input(name, is_moe, config, num_local_experts)
self.delay_wgrad_compute = extra_args.get("delay_wgrad_compute", False)
+ self.is_layer_first_node = None
+ self.is_layer_last_node = None
+
super().__init__(
weak_method(self.forward_impl),
stream,
@@ -293,6 +298,7 @@ def __init__(
self.detached = tuple()
self.before_detached = tuple()
self.is_mtp = extra_args.get("is_mtp", False)
+ self.post_wgrad_grad_acc_hooks = None
# Create flags to indicate first and last layer
self.is_first_layer = extra_args.get("is_first_layer", False)
@@ -322,16 +328,24 @@ def backward_impl(self, outputs, output_grad):
detached_grad = tuple([e.grad for e in self.detached])
grads = output_grad + detached_grad
self.default_backward_func(outputs + self.before_detached, grads)
- # release the output grad memory after backward finishes,
- # except when delay_wgrad_comptue is enabled, the grad should be
- # kept until all modules' backward_dw has been invoked.
- if self.delay_wgrad_compute:
- self.output_grads = grads
- self.delay_grads_release = len(self.bwd_dw_callables) > 0
# return grads for record stream
return grads
+ def forward(self, *inputs):
+ """Execute forward pass and corresponding hooks."""
+ output = super().forward(*inputs)
+ if self.is_layer_last_node:
+ self._post_forward_hook()
+ return output
+
+ def backward(self, *output_grad):
+ """Execute backward pass and corresponding hooks."""
+ grads = super().backward(*output_grad)
+ if not self.delay_wgrad_compute and self.is_layer_first_node:
+ self._post_backward_hook()
+ return grads
+
def backward_dw(self):
"""Computes the weight gradients for the transformer layer node."""
if not self.delay_wgrad_compute:
@@ -339,20 +353,51 @@ def backward_dw(self):
if isinstance(self.stream, Callable):
self.stream = self.stream()
with torch.cuda.stream(self.stream):
- torch.cuda.nvtx.range_push(f"{self.name} wgrad")
+ nvtx_msg = f"{self.name} wgrad"
+ nvtx_range_push(nvtx_msg)
for module in self.bwd_dw_callables:
module.backward_dw()
- torch.cuda.nvtx.range_pop()
-
- # the output grad memory is last used in wgrad compute, should be safe to release.
- assert self.delay_grads_release, "output grad memory should be valid before wgrad."
- if self.manual_release_grads:
- for tensor in self.output_grads:
- tensor.untyped_storage().resize_(0)
- self.output_grads = None
+ nvtx_range_pop(nvtx_msg)
+ # Collecting gradient acc hooks if there is `post_wgrad_grad_acc_hook`
+ # attribute attached to param, o.w. the wgrad hook wouldn't be fired.
+ if self.post_wgrad_grad_acc_hooks is None:
+ self.post_wgrad_grad_acc_hooks = []
+ for module in self.bwd_dw_callables:
+ for param in module.parameters():
+ # Collect hook only if the gradient is generated in current
+ # TransformerLayerNode, because the grad_acc hook needs
+ # to be executed right after `backward_dw` finishes.
+ # For example: Shared expert's hook should be collected in
+ # `attn` Node, even if the param belongs to `mlp` Node.
+ if (
+ getattr(param, "post_wgrad_grad_acc_hook", False)
+ and param.requires_grad
+ and param.grad is not None
+ ):
+ self.post_wgrad_grad_acc_hooks.append(param.post_wgrad_grad_acc_hook)
+
+ # Execute gradient accumulation hooks after wgrad compute.
+ if self.post_wgrad_grad_acc_hooks:
+ with torch.cuda.stream(self.stream):
+ for hook in self.post_wgrad_grad_acc_hooks:
+ hook()
+
+ # Execute TransformerLayer backward hook.
+ if self.is_layer_first_node:
+ self._post_backward_hook()
self.bwd_dw_callables = None
+ def set_post_forward_hook(self, hook):
+ """Register post_forward_hook at TransformerLayer level."""
+ self.is_layer_last_node = True
+ self._post_forward_hook = hook
+
+ def set_post_backward_hook(self, hook):
+ """Register post_backward_hook at TransformerLayer level."""
+ self.is_layer_first_node = True
+ self._post_backward_hook = hook
+
def __del__(self):
# Release reference as early as possible, this helps avoid memory leak.
self.before_detached = None
@@ -385,22 +430,25 @@ def __init__(self, layer):
self.layer = layer
self.graphed_backward_dw_callable = None
self.attn_dw_callable = layer.self_attention.backward_dw
+ self.submodules = [layer.self_attention]
if layer.is_moe_layer:
self.shared_expert_dw_callable = partial(
layer.mlp.backward_dw, routed_experts=False, shared_experts=True
)
+ if layer.mlp.use_shared_expert:
+ self.submodules.append(layer.mlp.shared_experts)
else:
self.shared_expert_dw_callable = None
- self.cuda_graph_scope = layer.config.cuda_graph_scope
+ self.cuda_graph_modules = layer.config.cuda_graph_modules
def backward_dw(self):
"""Execute weight gradients, skipping CUDA graphed components during replay."""
is_replay = hasattr(self.layer, 'cuda_graphs') and self.layer.cuda_graphs
if self.shared_expert_dw_callable is not None and (
- not is_replay or CudaGraphScope.moe_router not in self.cuda_graph_scope
+ not is_replay or CudaGraphModule.moe_router not in self.cuda_graph_modules
):
self.shared_expert_dw_callable()
- if not is_replay or CudaGraphScope.attn not in self.cuda_graph_scope:
+ if not is_replay or CudaGraphModule.attn not in self.cuda_graph_modules:
self.attn_dw_callable()
if is_replay and self.graphed_backward_dw_callable is not None:
self.graphed_backward_dw_callable()
@@ -410,6 +458,17 @@ def set_graphed_backward_dw_callable(self, graphed_backward_dw_callable):
"""Store the CUDA graphed backward weight gradient callable."""
self.graphed_backward_dw_callable = graphed_backward_dw_callable
+ def parameters(self):
+ """Returns an iterator over module parameters.
+
+ This method mimics the behavior of torch.nn.Module.parameters() by yielding
+ all parameters from the submodules managed by this wrapper. It is used to
+ collect parameters that require gradient computation during the backward pass.
+ """
+ for module in self.submodules:
+ for param in module.parameters():
+ yield param
+
def build_transformer_layer_callables(layer: TransformerLayer):
"""Create callables for transformer layer nodes.
@@ -504,6 +563,18 @@ def forward_func(
hidden_states
)
+ # When using fused residual norm (e.g. TEFusedResidualRMSNorm),
+ # the layernorm returns (normalized_output, residual). Unpack
+ # and use the fused residual for the downstream BDA connection.
+ if isinstance(pre_mlp_layernorm_output, tuple):
+ if len(pre_mlp_layernorm_output) != 2:
+ raise ValueError(
+ f"When the output of pre_mlp_layernorm is a tuple, it is "
+ f"expected to have 2 elements (output, residual), but "
+ f"got {len(pre_mlp_layernorm_output)}"
+ )
+ pre_mlp_layernorm_output, hidden_states = pre_mlp_layernorm_output
+
shared_expert_output = layer.mlp.shared_experts_compute(pre_mlp_layernorm_output)
padding_mask = node.chunk_state.padding_mask
if padding_mask is not None:
@@ -669,14 +740,17 @@ def submodule_mtp_attn_forward(node, hidden_states):
node.chunk_state.mtp_hidden_states = list(torch.chunk(hidden_states, 1 + offset, dim=0))
hidden_states = node.chunk_state.mtp_hidden_states[offset]
- input_ids, position_ids, decoder_input, hidden_states = layer._get_embeddings(
+ input_ids, position_ids, padding_mask, decoder_input, hidden_states = layer._get_embeddings(
input_ids=node.chunk_state.input_ids,
position_ids=node.chunk_state.position_ids,
embedding=node.chunk_state.model.embedding,
hidden_states=hidden_states,
+ packed_seq_params=node.chunk_state.packed_seq_params,
+ padding_mask=node.chunk_state.padding_mask,
)
node.chunk_state.input_ids = input_ids
node.chunk_state.position_ids = position_ids
+ node.chunk_state.padding_mask = padding_mask
# MTP Layer Preprocess
# norm, linear projection and transformer
diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py
index 0d2ca5fa6a7..c09545b6db1 100755
--- a/megatron/core/models/gpt/gpt_layer_specs.py
+++ b/megatron/core/models/gpt/gpt_layer_specs.py
@@ -1,5 +1,6 @@
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
import warnings
+from functools import partial
from typing import Optional, Union
from megatron.core.extensions.transformer_engine import HAVE_TE
@@ -34,18 +35,23 @@
)
from megatron.core.transformer.transformer_config import TransformerConfig
from megatron.core.transformer.transformer_layer import (
+ MlpBuilder,
TransformerLayer,
TransformerLayerSubmodules,
get_transformer_layer_offset,
)
-from megatron.core.typed_torch import copy_signature
+from megatron.core.typed_torch import copy_signature, not_none
from megatron.core.utils import is_te_min_version
if HAVE_TE:
- from megatron.core.extensions.transformer_engine import TEFusedMLP, TENorm
+ from megatron.core.extensions.transformer_engine import (
+ TEFusedMLP,
+ TEFusedMLPWithGroupedLinear,
+ TENorm,
+ )
from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
else:
- TEFusedMLP, TENorm, TESpecProvider = None, None, None
+ TEFusedMLPWithGroupedLinear, TEFusedMLP, TENorm, TESpecProvider = None, None, None, None
try:
from megatron.core.extensions.kitchen import HAVE_KITCHEN, KitchenSpecProvider
@@ -183,6 +189,7 @@ def get_gpt_layer_with_transformer_engine_submodules(
use_kitchen_attention: bool = False,
kitchen_attention_backend: str = "sdpa",
mla_down_proj_fusion: bool = False,
+ use_grouped_gemm_for_dense_mlp: bool = False,
) -> TransformerLayerSubmodules:
"""Use these submodules to use lower-level Transformer Engine modules (required for fp8
training).
@@ -231,6 +238,7 @@ def get_gpt_layer_with_transformer_engine_submodules(
moe_grouped_gemm=moe_grouped_gemm,
use_te_op_fuser=use_te_op_fuser,
use_te_activation_func=use_te_activation_func,
+ use_grouped_gemm_for_dense_mlp=use_grouped_gemm_for_dense_mlp,
)
if multi_latent_attention:
@@ -485,7 +493,7 @@ def get_mlp_module_spec(
moe_grouped_gemm: Optional[bool] = False,
fp8: Optional[str] = None, # pylint: disable=unused-argument
use_te_op_fuser: Optional[bool] = False,
-) -> ModuleSpec:
+) -> MlpBuilder:
"""Helper function to get module spec for MLP/MoE"""
if fp8 is not None:
warnings.warn(
@@ -516,7 +524,8 @@ def get_mlp_module_spec_for_backend(
moe_grouped_gemm: Optional[bool] = False,
use_te_op_fuser: Optional[bool] = False,
use_te_activation_func: bool = False,
-) -> ModuleSpec:
+ use_grouped_gemm_for_dense_mlp: bool = False,
+) -> MlpBuilder:
"""Helper function to get module spec for MLP/MoE"""
linear_fc2 = backend.row_parallel_linear()
@@ -524,14 +533,19 @@ def get_mlp_module_spec_for_backend(
if num_experts is None:
# Dense MLP w/ or w/o TE modules.
- module = TEFusedMLP if use_te_op_fuser else MLP
+ if use_grouped_gemm_for_dense_mlp and use_te_op_fuser:
+ module = not_none(TEFusedMLPWithGroupedLinear).as_mlp_submodule
+ elif use_te_op_fuser:
+ module = not_none(TEFusedMLP).as_mlp_submodule
+ else:
+ module = MLP.as_mlp_submodule
if backend.fuse_layernorm_and_linear():
linear_fc1 = backend.column_parallel_layer_norm_linear()
assert linear_fc1 is not None
else:
linear_fc1 = backend.column_parallel_linear()
- return ModuleSpec(
- module=module,
+ return partial(
+ module,
submodules=MLPSubmodules(
linear_fc1=linear_fc1, linear_fc2=linear_fc2, activation_func=activation_func
),
@@ -754,15 +768,22 @@ def get_gpt_mtp_block_spec_for_backend(
mtp_model_layer_spec=transformer_layer_spec, backend=backend
)
mtp_num_layers = config.mtp_num_layers if config.mtp_num_layers else 0
- mtp_layer_specs = [mtp_layer_spec] * mtp_num_layers
+ if config.mtp_use_repeated_layer:
+ mtp_layer_specs = [mtp_layer_spec]
+ else:
+ mtp_layer_specs = [mtp_layer_spec] * mtp_num_layers
+
+ if not config.mtp_use_repeated_layer:
+ offset = get_mtp_layer_offset(config, vp_stage=vp_stage)
+ # Split the MTP layer specs to only include the layers that are built in this
+ # pipeline stage.
+ mtp_layer_specs = mtp_layer_specs[offset : offset + num_layers_to_build]
+ if len(mtp_layer_specs) > 0:
+ assert (
+ len(mtp_layer_specs) == config.mtp_num_layers
+ ), f"All MTP layers must reside in the same pipeline stage"
- offset = get_mtp_layer_offset(config, vp_stage=vp_stage)
- # split the mtp layer specs to only include the layers that are built in this pipeline stage.
- mtp_layer_specs = mtp_layer_specs[offset : offset + num_layers_to_build]
if len(mtp_layer_specs) > 0:
- assert (
- len(mtp_layer_specs) == config.mtp_num_layers
- ), f"currently all of the mtp layers must stage in the same pipeline stage."
mtp_block_spec = MultiTokenPredictionBlockSubmodules(layer_specs=mtp_layer_specs)
else:
mtp_block_spec = None
diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
index d63b2c1ddfa..99853939f4c 100644
--- a/megatron/core/models/gpt/gpt_model.py
+++ b/megatron/core/models/gpt/gpt_model.py
@@ -1,7 +1,7 @@
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
from collections import OrderedDict
-from typing import Dict, Literal, Optional
+from typing import Any, Callable, Dict, Literal, Optional
import torch
from torch import Tensor
@@ -9,7 +9,10 @@
from megatron.core import tensor_parallel
from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.extensions.transformer_engine import TELMHeadColumnParallelLinear
+from megatron.core.fp8_utils import is_mxfp8_output_proj_active
from megatron.core.inference.contexts import BaseInferenceContext
+from megatron.core.inference.utils import InferenceMode
from megatron.core.models.common.embeddings import YarnRotaryEmbedding
from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
from megatron.core.models.common.embeddings.rotary_pos_embedding import (
@@ -24,7 +27,8 @@
from megatron.core.process_groups_config import ProcessGroupCollection
from megatron.core.quantization.utils import get_quant_config_or_none
from megatron.core.tensor_parallel import gather_from_sequence_parallel_region
-from megatron.core.transformer.enums import CudaGraphScope, ModelType
+from megatron.core.transformer.enums import ModelType
+from megatron.core.transformer.moe.paged_stash import paged_stash_init_chunk_handler
from megatron.core.transformer.multi_token_prediction import (
MultiTokenPredictionBlock,
mtp_on_this_rank,
@@ -143,7 +147,10 @@ def __init__(
self.rotary_scaling = rope_scaling
self.mtp_block_spec = mtp_block_spec
self.mtp_process = mtp_block_spec is not None and mtp_on_this_rank(
- self.config, ignore_virtual=False, vp_stage=vp_stage
+ layout=self.config.pipeline_model_parallel_layout,
+ mtp_num_layers=self.config.mtp_num_layers,
+ ignore_virtual=False,
+ vp_stage=vp_stage,
)
if self.pre_process or self.mtp_process:
@@ -223,6 +230,8 @@ def __init__(
pg_collection=self.pg_collection,
)
+ self._setup_mtp_cuda_graphs()
+
# Output
if self.post_process:
@@ -241,7 +250,12 @@ def __init__(
self.embedding_activation_buffer = None
self.grad_output_buffer = None
- self.output_layer = tensor_parallel.ColumnParallelLinear(
+ output_layer_cls = (
+ TELMHeadColumnParallelLinear
+ if is_mxfp8_output_proj_active(config)
+ else tensor_parallel.ColumnParallelLinear
+ )
+ self.output_layer = output_layer_cls(
config.hidden_size,
self.vocab_size,
config=config,
@@ -306,7 +320,7 @@ def _preprocess(
# If decoder_input is provided (not None), then input_ids and position_ids are ignored.
# Otherwise, apply embedding layer on input_ids and position_ids to get decoder_input.
- in_inference_mode = inference_context is not None and not self.training
+ in_inference_mode = InferenceMode.is_active()
# Decoder embedding.
if decoder_input is not None:
@@ -343,7 +357,11 @@ def _preprocess(
hasattr(inference_context, 'use_flashinfer_fused_rope')
and inference_context.use_flashinfer_fused_rope
)
- if in_inference_mode and (self.config.flash_decode or use_flash_infer_fused_rope):
+ if (
+ in_inference_mode
+ and inference_context is not None
+ and (self.config.flash_decode or use_flash_infer_fused_rope)
+ ):
assert (
not self.config.flash_decode
) or inference_context.is_static_batching(), (
@@ -375,7 +393,7 @@ def _preprocess(
cp_group=packed_seq_params.cp_group if packed_seq_params is not None else None,
)
elif self.position_embedding_type == 'yarn':
- if self.training or not self.config.flash_decode:
+ if not InferenceMode.is_active() or not self.config.flash_decode:
rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
inference_context, self.decoder, decoder_input, self.config, packed_seq_params
)
@@ -391,7 +409,7 @@ def _preprocess(
"YarnRotaryEmbedding yet."
)
elif self.position_embedding_type == 'mrope' and not self.config.multi_latent_attention:
- if self.training or not self.config.flash_decode:
+ if not InferenceMode.is_active() or not self.config.flash_decode:
rotary_pos_emb = self.rotary_pos_emb(
position_ids,
self.mrope_section,
@@ -406,13 +424,8 @@ def _preprocess(
if (
in_inference_mode
- and (
- (
- self.config.cuda_graph_impl == "local"
- and CudaGraphScope.full_iteration not in self.config.cuda_graph_scope
- )
- or self.config.flash_decode
- )
+ and inference_context is not None
+ and (self.config.cuda_graph_impl == "local" or self.config.flash_decode)
and inference_context.is_static_batching()
):
current_batch_size = input_ids.shape[0]
@@ -427,8 +440,10 @@ def _preprocess(
if in_inference_mode:
# Clear the outputs for padding tokens when using dynamic batching with
# quantization scales to avoid corrupting amax calculations
- if inference_context.is_dynamic_batching() and is_using_quantization_scales(
- self.config
+ if (
+ inference_context is not None
+ and inference_context.is_dynamic_batching()
+ and is_using_quantization_scales(self.config)
):
decoder_input[inference_context.padding_slice] = 0.0
@@ -462,6 +477,7 @@ def preprocess_for_fine_grained_offloading(self):
vp_size=self.config.virtual_pipeline_model_parallel_size,
vp_stage=self.vp_stage,
min_offloaded_tensor_size=self.config.min_offloaded_tensor_size,
+ max_inflight_offloads=self.config.fine_grained_offloading_max_inflight_offloads,
)
if self.disable_param_offloading:
for param in self.decoder.parameters():
@@ -474,6 +490,12 @@ def preprocess_for_fine_grained_offloading(self):
off_interface.mark_not_offloadable(param)
self.disable_param_offloading = False
+ def preprocess_for_paged_stash(self):
+ """Preprocess for paged stash."""
+ return paged_stash_init_chunk_handler(
+ vp_size=self.config.virtual_pipeline_model_parallel_size, vp_stage=self.vp_stage
+ )
+
def forward(
self,
input_ids: Tensor,
@@ -489,7 +511,8 @@ def forward(
inference_params: Optional[BaseInferenceContext] = None,
loss_mask: Optional[Tensor] = None,
padding_mask: Optional[Tensor] = None,
- is_spec_decode: Optional[bool] = None,
+ output_processor: Optional[Callable[..., Tensor]] = None,
+ output_processor_context: Optional[Any] = None,
) -> Tensor:
"""Forward function of the GPT Model This function passes the input tensors
through the embedding layer, and then the decoder and finally into the post
@@ -503,13 +526,17 @@ def forward(
padding_mask (Tensor, optional): Padding mask for MoE routing.
Shape [bsz, seq_length]. True = padding (exclude), False = valid (include).
Only used for MoE layers to exclude padding tokens from routing computations.
- is_spec_decode (bool, optional): Explicitly override whether speculative
- decoding is active. When ``None`` (default) the flag is inferred from
- ``inference_context.num_speculative_tokens``.
+ output_processor (Callable, optional): Custom postprocess hook that receives
+ decoder hidden states and output-layer helpers, then returns the model output.
+ output_processor_context (Any, optional): User-defined context object forwarded to
+ `output_processor`.
"""
if self.config.fine_grained_activation_offloading:
self.preprocess_for_fine_grained_offloading()
+ if self.config.moe_paged_stash:
+ self.preprocess_for_paged_stash()
+
inference_context = deprecate_inference_params(inference_context, inference_params)
preproc_output = self._preprocess(
@@ -559,13 +586,15 @@ def forward(
loss_mask=loss_mask,
decoder_input=decoder_input,
attention_mask=attention_mask,
+ padding_mask=padding_mask,
inference_params=inference_params,
packed_seq_params=packed_seq_params,
sequence_len_offset=sequence_len_offset,
runtime_gather_output=runtime_gather_output,
extra_block_kwargs=extra_block_kwargs,
inference_context=inference_context,
- is_spec_decode=is_spec_decode,
+ output_processor=output_processor,
+ output_processor_context=output_processor_context,
)
def _postprocess(
@@ -581,32 +610,34 @@ def _postprocess(
loss_mask=None,
decoder_input=None,
attention_mask=None,
+ padding_mask=None,
inference_params=None,
packed_seq_params=None,
sequence_len_offset=None,
runtime_gather_output=None,
extra_block_kwargs=None,
inference_context=None,
- is_spec_decode=None,
+ output_processor=None,
+ output_processor_context=None,
):
"""Postprocesses decoder hidden states to generate logits or compute loss.
Applies Multi-Token Prediction if enabled, generates output logits through
the output layer, and computes language model loss when labels are provided.
"""
- in_inference_mode = inference_context is not None and not self.training
+ in_inference_mode = InferenceMode.is_active()
if in_inference_mode:
assert runtime_gather_output, "Inference must always gather TP logits"
# Check if speculative decoding is active. When it is, MTP must be
# computed *after* verification so that it is conditioned on verified
# tokens rather than stale speculative tokens from the previous step.
- if is_spec_decode is None:
- is_spec_decode = (
- in_inference_mode
- and inference_context.is_dynamic_batching()
- and inference_context.num_speculative_tokens > 0
- )
+ is_spec_decode = (
+ in_inference_mode
+ and inference_context is not None
+ and inference_context.is_dynamic_batching()
+ and inference_context.num_speculative_tokens > 0
+ )
# logits and loss
output_weight = None
@@ -624,6 +655,7 @@ def _postprocess(
rotary_pos_sin=rotary_pos_sin,
packed_seq_params=packed_seq_params,
sequence_len_offset=sequence_len_offset,
+ padding_mask=padding_mask,
embedding=self.embedding,
**(extra_block_kwargs or {}),
)
@@ -655,7 +687,31 @@ def _postprocess(
)
sequence_parallel_override = False
- if in_inference_mode and inference_context.config.materialize_only_last_token_logits:
+ if output_processor is not None:
+ return output_processor(
+ hidden_states=hidden_states,
+ output_layer=self.output_layer,
+ output_weight=output_weight,
+ labels=labels,
+ loss_mask=loss_mask,
+ input_ids=input_ids,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ decoder_input=decoder_input,
+ inference_context=inference_context,
+ packed_seq_params=packed_seq_params,
+ runtime_gather_output=runtime_gather_output,
+ context=output_processor_context,
+ compute_language_model_loss=self.compute_language_model_loss,
+ scale_logits=self._scale_logits,
+ config=self.config,
+ )
+
+ if (
+ in_inference_mode
+ and inference_context is not None
+ and inference_context.config.materialize_only_last_token_logits
+ ):
if inference_context.is_static_batching():
hidden_states = hidden_states[-1:, :, :]
else:
@@ -710,49 +766,6 @@ def _postprocess(
return loss
- @torch.inference_mode()
- def compute_mtp_single_step(
- self,
- hidden_states: Tensor,
- next_token_ids: Tensor,
- position_ids: Tensor,
- depth: int,
- runtime_gather_output: bool = True,
- ) -> tuple:
- """Compute a single MTP depth for speculative decoding.
-
- This is called after speculative token verification to compute MTP
- predictions conditioned on verified tokens only.
-
- Args:
- hidden_states (Tensor): Hidden states at last accepted positions [N, 1, H].
- next_token_ids (Tensor): Correct next token IDs [1, N].
- position_ids (Tensor): Position IDs for the next tokens [1, N].
- depth (int): MTP depth index (0-indexed).
- runtime_gather_output (bool): Whether to gather output across TP.
-
- Returns:
- tuple: (new_hidden_states [N, 1, H], logits [N, 1, vocab_size]).
- """
- layer_idx = 0 if self.mtp.mtp_use_repeated_layer else depth
- mtp_hidden = self.mtp.layers[layer_idx].forward_single_position(
- hidden_states=hidden_states,
- next_token_ids=next_token_ids,
- position_ids=position_ids,
- embedding=self.embedding,
- )
-
- output_weight = None
- if self.share_embeddings_and_output_weights:
- output_weight = self.shared_embedding_or_output_weight()
-
- logits, _ = self.output_layer(
- mtp_hidden, weight=output_weight, runtime_gather_output=runtime_gather_output
- )
- logits = self._scale_logits(logits)
-
- return mtp_hidden, logits
-
def build_schedule_plan(
self,
input_ids: Tensor,
@@ -767,6 +780,9 @@ def build_schedule_plan(
inference_params: Optional[BaseInferenceContext] = None,
loss_mask: Optional[Tensor] = None,
padding_mask: Optional[Tensor] = None,
+ *,
+ output_processor: Optional[Callable[..., Tensor]] = None,
+ output_processor_context: Optional[Any] = None,
):
"""Builds a computation schedule plan for the model.
@@ -793,6 +809,10 @@ def build_schedule_plan(
Parameters for inference. Defaults to None.
loss_mask (Optional[Tensor], optional): Loss mask. Defaults to None.
padding_mask (Optional[Tensor], optional): Padding mask. Defaults to None.
+ output_processor (Callable, optional): Custom postprocess hook to run in the
+ schedule-plan postprocess node instead of the default logits/loss path.
+ output_processor_context (Any, optional): User-defined context object forwarded to
+ `output_processor`.
Returns:
TransformerModelChunkSchedulePlan: The model chunk schedule plan.
@@ -800,6 +820,8 @@ def build_schedule_plan(
if self.config.fine_grained_activation_offloading:
self.preprocess_for_fine_grained_offloading()
+ if self.config.moe_paged_stash:
+ self.preprocess_for_paged_stash()
from ..common.model_chunk_schedule_plan import TransformerModelChunkSchedulePlan
@@ -815,6 +837,8 @@ def build_schedule_plan(
runtime_gather_output,
loss_mask,
padding_mask,
+ output_processor=output_processor,
+ output_processor_context=output_processor_context,
)
def sharded_state_dict(
diff --git a/megatron/core/models/gpt/heterogeneous/heterogeneous_layer_specs.py b/megatron/core/models/gpt/heterogeneous/heterogeneous_layer_specs.py
index f4385429422..2c2b26f2290 100644
--- a/megatron/core/models/gpt/heterogeneous/heterogeneous_layer_specs.py
+++ b/megatron/core/models/gpt/heterogeneous/heterogeneous_layer_specs.py
@@ -1,6 +1,7 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
import warnings
+from functools import partial
from typing import Optional
from megatron.core.extensions.transformer_engine import HAVE_TE
@@ -118,7 +119,7 @@ def _get_heterogenous_attention_spec(
not_none(TELayerNormColumnParallelLinear) if use_te else ColumnParallelLinear
),
core_attention=not_none(TEDotProductAttention) if use_te else DotProductAttention,
- linear_proj=TERowParallelLinear if use_te else RowParallelLinear,
+ linear_proj=not_none(TERowParallelLinear) if use_te else RowParallelLinear,
q_layernorm=ln,
k_layernorm=ln,
),
@@ -128,17 +129,19 @@ def _get_heterogenous_attention_spec(
def _get_heterogenous_mlp_spec(mlp_config: MLPConfig, use_te: bool):
if mlp_config.no_op:
- mlp = ModuleSpec(module=IdentityOp)
+ return IdentityOp
elif mlp_config.replace_with_linear:
- mlp = ModuleSpec(
- module=(
- TELayerNormColumnParallelLinearGathered if use_te else ColumnParallelLinearGathered
+ return partial(
+ (
+ not_none(TELayerNormColumnParallelLinearGathered)
+ if use_te
+ else ColumnParallelLinearGathered
),
- params={"tp_comm_buffer_name": "linear_mlp"},
+ tp_comm_buffer_name="linear_mlp",
)
else:
- mlp = ModuleSpec(
- module=MLP,
+ return partial(
+ MLP.as_mlp_submodule,
submodules=MLPSubmodules(
linear_fc1=(
not_none(TELayerNormColumnParallelLinear) if use_te else ColumnParallelLinear
@@ -146,7 +149,6 @@ def _get_heterogenous_mlp_spec(mlp_config: MLPConfig, use_te: bool):
linear_fc2=not_none(TERowParallelLinear) if use_te else RowParallelLinear,
),
)
- return mlp
def _get_sharded_state_dict_keys_map(block_config: TransformerBlockConfig, use_te: bool):
diff --git a/megatron/core/models/gpt/moe_module_specs.py b/megatron/core/models/gpt/moe_module_specs.py
index 53bca85f502..e9a86ff3bad 100755
--- a/megatron/core/models/gpt/moe_module_specs.py
+++ b/megatron/core/models/gpt/moe_module_specs.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved.
from functools import partial
from typing import Optional
@@ -13,17 +13,17 @@
from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules
from megatron.core.transformer.moe.router import InferenceTopKRouter
from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
-from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_layer import MlpBuilder
def get_moe_module_spec(
use_te: Optional[bool] = True,
num_experts: Optional[int] = None,
moe_grouped_gemm: Optional[bool] = False,
-) -> ModuleSpec:
+) -> MlpBuilder:
"""Helper function to get module spec for MoE.
- Called by mamba_layer_specs.py for standard (non-inference) MoE specs.
+ Called by hybrid_layer_specs.py for standard (non-inference) MoE specs.
The GPT layer specs call get_moe_module_spec_for_backend directly.
Args:
@@ -46,7 +46,7 @@ def get_moe_module_spec_for_backend(
num_experts: Optional[int] = None,
moe_grouped_gemm: Optional[bool] = False,
use_te_activation_func: bool = False,
-) -> ModuleSpec:
+) -> MlpBuilder:
"""Helper function to get module spec for MoE"""
assert num_experts is not None
@@ -63,22 +63,19 @@ def get_moe_module_spec_for_backend(
shared_experts = partial(SharedExpertMLP, submodules=mlp)
# MoE module spec
- moe_module_spec = ModuleSpec(
- module=MoELayer,
- submodules=MoESubmodules(experts=experts, shared_experts=shared_experts),
- metainfo={"fuse_pre_mlp_layernorm": False},
+ return partial(
+ MoELayer, submodules=MoESubmodules(experts=experts, shared_experts=shared_experts)
)
- return moe_module_spec
-def get_inference_optimized_moe_spec() -> ModuleSpec:
+def get_inference_optimized_moe_spec() -> MlpBuilder:
"""MoE module spec for inference-optimized transformer impl.
Uses InferenceSpecProvider to select inference-optimized modules:
InferenceTopKRouter, InferenceGroupedMLP. MoELayer detects inference mode
via config.transformer_impl and sets up the inference dispatcher internally.
- Called by mamba_layer_specs.py and gpt_layer_specs.py.
+ Called by hybrid_layer_specs.py and gpt_layer_specs.py.
"""
backend = InferenceSpecProvider()
activation_func = backend.activation_func()
@@ -93,10 +90,9 @@ def get_inference_optimized_moe_spec() -> ModuleSpec:
),
)
- return ModuleSpec(
- module=MoELayer,
+ return partial(
+ MoELayer,
submodules=MoESubmodules(
router=InferenceTopKRouter, experts=experts, shared_experts=shared_experts
),
- metainfo={"fuse_pre_mlp_layernorm": False},
)
diff --git a/megatron/core/models/huggingface/fastconformer_model.py b/megatron/core/models/huggingface/fastconformer_model.py
new file mode 100644
index 00000000000..25265871240
--- /dev/null
+++ b/megatron/core/models/huggingface/fastconformer_model.py
@@ -0,0 +1,102 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+import torch
+
+from megatron.core.models.huggingface import HuggingFaceModule
+
+# NeMo model loading is slow, so cache the (preprocessor, encoder) tuple per
+# `sound_model_type`. Keying by model id avoids returning a stale cached encoder
+# when the same process constructs more than one Parakeet variant.
+_NEMO_SOUND_MODEL_CACHE: dict[str, tuple] = {}
+
+
+def get_nemo_sound_model(sound_model_type):
+ """Load (and cache) a NeMo ASR encoder + preprocessor for the given ``nemo://`` model id."""
+ if sound_model_type not in _NEMO_SOUND_MODEL_CACHE:
+ import nemo.collections.asr as nemo_asr
+
+ asr_model = nemo_asr.models.ASRModel.from_pretrained(
+ model_name=sound_model_type.split("nemo://")[1]
+ )
+ # Avoid hangs from an unnecessary max-seq-len NCCL sync in some edge cases.
+ asr_model.encoder.sync_max_audio_length = False
+ for layer in asr_model.encoder.layers:
+ layer.self_attn.use_pytorch_sdpa = True
+ _NEMO_SOUND_MODEL_CACHE[sound_model_type] = (asr_model.preprocessor, asr_model.encoder)
+ return _NEMO_SOUND_MODEL_CACHE[sound_model_type]
+
+
+class ParakeetHuggingFaceModel(HuggingFaceModule):
+ """Wrapper for Parakeet sound encoders.
+
+ Supports two backends, selected by ``config.sound_model_type`` prefix:
+
+ - ``nemo://`` loads a NeMo ASR encoder + preprocessor.
+ - ``hf://`` loads the upstream Hugging Face FastConformer model
+ via ``transformers.AutoModel`` / ``AutoFeatureExtractor``.
+ """
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.use_nemo = config.sound_model_type.startswith("nemo://")
+ if self.use_nemo:
+ self.feature_extractor, self.model = get_nemo_sound_model(config.sound_model_type)
+
+ for module in self.model.modules():
+ if module.__class__.__name__.lower() == "dropout":
+ module.p = config.hidden_dropout
+
+ if config.recompute_granularity is not None:
+ from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+ checkpoint_wrapper,
+ )
+
+ self.model = checkpoint_wrapper(self.model)
+ elif config.sound_model_type.startswith("hf://"):
+ from transformers import AutoFeatureExtractor, AutoModel
+
+ sound_model_type = config.sound_model_type.split("hf://")[1]
+ self.feature_extractor = AutoFeatureExtractor.from_pretrained(sound_model_type)
+ self.model = AutoModel.from_pretrained(sound_model_type)
+
+ if config.recompute_granularity is not None:
+ self.model.gradient_checkpointing_enable()
+ else:
+ raise ValueError(f"Unknown sound model type: {config.sound_model_type}")
+
+ def _model_dtype(self) -> torch.dtype:
+ """Return the dtype of the encoder's first parameter (defaults to bf16)."""
+ for param in self.model.parameters():
+ return param.dtype
+ return torch.bfloat16
+
+ def _sampling_rate(self) -> int:
+ """Return the sampling rate the feature extractor expects (default 16 kHz)."""
+ return int(getattr(self.feature_extractor, "sampling_rate", 16000))
+
+ def forward(self, *args, **kwargs):
+ """Forward pass returning (hidden_states, lengths).
+
+ Args:
+ args[0]: Sound clips tensor.
+ args[1]: Sound length tensor (used by NeMo backend; ignored for HF).
+ """
+ if self.use_nemo:
+ features = self.feature_extractor(input_signal=args[0], length=args[1])
+ y = self.model(audio_signal=features[0], length=features[1])
+ # NeMo encoder returns [B, H, T]; LLaVA expects [B, T, H].
+ return y[0].permute(0, 2, 1), y[1]
+ else:
+ # HF feature extractor expects audio as the first arg only,
+ # not (audio, length) as in NeMo.
+ sound_clips = args[0]
+ features = self.feature_extractor(
+ sound_clips,
+ **kwargs,
+ return_tensors="pt",
+ sampling_rate=self._sampling_rate(),
+ return_attention_mask=True,
+ )
+ y = self.model(features.input_features.to(self._model_dtype()), features.attention_mask)
+ lengths = features.attention_mask.sum(dim=-1).to(y.last_hidden_state.device)
+ return y.last_hidden_state, lengths
diff --git a/megatron/core/models/huggingface/module.py b/megatron/core/models/huggingface/module.py
index 5c78fc96708..2d874c7513b 100644
--- a/megatron/core/models/huggingface/module.py
+++ b/megatron/core/models/huggingface/module.py
@@ -68,6 +68,27 @@ def get_hf_model_type(model_path):
"please install it with `pip install transformers`"
)
+ # Parakeet is a special case: its model id may be `nemo://...`, which
+ # AutoConfig cannot resolve, so detect it from the prefix. Require the
+ # `nemo://` or `hf://` scheme so unrelated local paths that happen to
+ # contain "parakeet" (e.g. a user directory) don't get misrouted.
+ lowered = model_path.lower()
+ if lowered.startswith(("nemo://", "hf://")):
+ model_id = lowered.split("://", 1)[1]
+ # Match a path segment whose name begins with "parakeet" (e.g.
+ # `nvidia/parakeet-tdt-0.6b-v2`). Substring-anywhere matches like
+ # `myparakeet-clone` are intentionally rejected.
+ if any(seg.startswith("parakeet") for seg in model_id.split("/")):
+ return "parakeet"
+ # Any other `nemo://` model can't be resolved by AutoConfig below;
+ # raise a clear error rather than letting `split("hf://")[1]` raise
+ # an IndexError with no context.
+ if lowered.startswith("nemo://"):
+ raise NotImplementedError(
+ f"nemo:// scheme is currently only supported for parakeet models, "
+ f"got {model_path}"
+ )
+
hf_config = AutoConfig.from_pretrained(model_path.split("hf://")[1])
model_type = hf_config.architectures[0].lower()
@@ -91,6 +112,10 @@ def build_hf_model(config, model_path):
from megatron.core.models.huggingface.clip_model import SiglipHuggingFaceModel
model = SiglipHuggingFaceModel(config)
+ elif "parakeet" in model_type:
+ from megatron.core.models.huggingface.fastconformer_model import ParakeetHuggingFaceModel
+
+ model = ParakeetHuggingFaceModel(config)
else:
raise NotImplementedError(f"unsupported huggingface model {config.hf_config}")
diff --git a/megatron/core/models/hybrid/__init__.py b/megatron/core/models/hybrid/__init__.py
new file mode 100644
index 00000000000..d8a0a817ee3
--- /dev/null
+++ b/megatron/core/models/hybrid/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/models/hybrid/hybrid_block.py b/megatron/core/models/hybrid/hybrid_block.py
new file mode 100644
index 00000000000..99745d98d3d
--- /dev/null
+++ b/megatron/core/models/hybrid/hybrid_block.py
@@ -0,0 +1,425 @@
+# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2024, Tri Dao, Albert Gu.
+
+# Some of this code was adopted from https://github.com/state-spaces/mamba/
+# This source code is licensed under the Apache license found in the
+# LICENSE file in the root directory of this source tree.
+
+from contextlib import nullcontext
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+from torch import Tensor, nn
+
+from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+from megatron.core.dist_checkpointing.utils import replace_prefix_for_sharding
+from megatron.core.enums import Fp8Recipe
+from megatron.core.extensions.transformer_engine import TENorm
+from megatron.core.fp4_utils import get_fp4_context
+from megatron.core.fp8_utils import get_fp8_context
+from megatron.core.inference.contexts import BaseInferenceContext
+from megatron.core.inference.utils import InferenceMode
+from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols as LayerSymbols
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.recompute import checkpointed_forward
+from megatron.core.transformer import TransformerConfig
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.transformer.transformer_layer import TransformerLayer
+from megatron.core.transformer.utils import sharded_state_dict_default
+from megatron.core.utils import WrappedTensor, deprecate_inference_params, make_viewless_tensor
+
+
+@dataclass
+class HybridStackSubmodules:
+ """
+ A class for the module specs for the HybridStack.
+ """
+
+ mamba_layer: Union[ModuleSpec, type] = IdentityOp
+ gdn_layer: Union[ModuleSpec, type] = IdentityOp
+ attention_layer: Union[ModuleSpec, type] = IdentityOp
+ dsa_layer: Union[ModuleSpec, type] = IdentityOp
+ mlp_layer: Union[ModuleSpec, type] = IdentityOp
+ moe_layer: Union[ModuleSpec, type] = IdentityOp
+ mtp_block_spec: Optional[ModuleSpec] = None
+
+
+class HybridStack(MegatronModule):
+ """
+ Constructor for the HybridStack class.
+
+ Args:
+ config (TransformerConfig): the model configuration
+ submodules (HybridStackSubmodules): the submodules for the stack
+ pre_process (bool, optional): whether to include an embedding layer.
+ Defaults to True.
+ layer_type_list (list, optional): pre-computed list of layer type symbols for
+ this pipeline segment. When provided (by HybridModel), pipeline stage
+ selection has already been done via '|' separators in the pattern.
+ pp_layer_offset (int, optional): the global layer offset for this pipeline
+ segment. Defaults to 0.
+ post_layer_norm (bool, optional): whether to include a final layer norm.
+ Defaults to True.
+ post_process (bool, optional): whether to include an output layer.
+ Defaults to True.
+ device (optional): the device to use. Defaults to None.
+ dtype (optional): the data type to use. Defaults to None.
+ pg_collection (ProcessGroupCollection): the required model communication
+ process groups to use.
+ is_mtp_layer (bool, optional): whether this is an MTP layer. Defaults to False.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ submodules: HybridStackSubmodules,
+ pre_process: bool = True,
+ layer_type_list: Optional[list[str]] = None,
+ pp_layer_offset: int = 0,
+ post_layer_norm: bool = True,
+ post_process: bool = True,
+ device=None,
+ dtype=None,
+ pg_collection: ProcessGroupCollection = None,
+ is_mtp_layer: bool = False,
+ name: str | None = None,
+ ) -> None:
+ """
+ Args:
+ name (str | None): module instance name passed top-down from its paranet module
+ """
+ super().__init__(config=config)
+ self.pre_process = pre_process
+ self.post_layer_norm = post_layer_norm
+ self.post_process = post_process
+ self.is_mtp_layer = is_mtp_layer
+
+ assert pg_collection is not None, "pg_collection must be provided for HybridStack"
+
+ self.pp_group = pg_collection.pp
+ self.tp_group = pg_collection.tp
+
+ # Required for pipeline parallel schedules
+ self.input_tensor = None
+ self.pg_collection = pg_collection
+
+ assert layer_type_list is not None, (
+ "layer_type_list must be provided. It should be pre-computed from "
+ "--hybrid-layer-pattern by HybridModel."
+ )
+ self.layer_type_list = layer_type_list
+
+ # Build layers from the pre-selected segment
+ self.layers = nn.ModuleList()
+ for i, layer_type in enumerate(self.layer_type_list):
+ layer_number = i + 1 + pp_layer_offset
+ if self.config.fp8:
+ quant_init_context = get_fp8_context(self.config, i + pp_layer_offset, is_init=True)
+ elif self.config.fp4:
+ quant_init_context = get_fp4_context(self.config, i + pp_layer_offset, is_init=True)
+ else:
+ quant_init_context = nullcontext()
+ with quant_init_context:
+ if layer_type == LayerSymbols.MAMBA:
+ layer = build_module(
+ submodules.mamba_layer,
+ config=self.config,
+ layer_number=layer_number,
+ pp_layer_offset=pp_layer_offset,
+ pg_collection=pg_collection,
+ name=(name + f".layers.{i}") if name is not None else None,
+ )
+ elif layer_type == LayerSymbols.ATTENTION:
+ layer = build_module(
+ submodules.attention_layer,
+ config=self.config,
+ layer_number=layer_number,
+ pg_collection=pg_collection,
+ is_mtp_layer=is_mtp_layer,
+ add_layer_offset=False,
+ pp_layer_offset=pp_layer_offset,
+ name=(name + f".layers.{i}") if name is not None else None,
+ )
+ elif layer_type == LayerSymbols.DS_ATTENTION:
+ layer = build_module(
+ submodules.dsa_layer,
+ config=self.config,
+ layer_number=layer_number,
+ pg_collection=pg_collection,
+ is_mtp_layer=is_mtp_layer,
+ add_layer_offset=False,
+ pp_layer_offset=pp_layer_offset,
+ name=(name + f".layers.{i}") if name is not None else None,
+ )
+ elif layer_type == LayerSymbols.MLP:
+ layer = build_module(
+ submodules.mlp_layer,
+ config=self.config,
+ layer_number=layer_number,
+ pg_collection=pg_collection,
+ add_layer_offset=False,
+ name=(name + f".layers.{i}") if name is not None else None,
+ )
+ elif layer_type == LayerSymbols.MOE:
+ layer = build_module(
+ submodules.moe_layer,
+ config=self.config,
+ layer_number=layer_number,
+ pg_collection=pg_collection,
+ add_layer_offset=False,
+ name=(name + f".layers.{i}") if name is not None else None,
+ )
+ elif layer_type == LayerSymbols.GDN:
+ layer = build_module(
+ submodules.gdn_layer,
+ config=self.config,
+ layer_number=layer_number,
+ pg_collection=pg_collection,
+ # Set to False as we do not want to change offset.
+ add_layer_offset=False,
+ name=(name + f".layers.{i}") if name is not None else None,
+ )
+ else:
+ raise ValueError("unexpected layer_type")
+ self.layers.append(layer)
+
+ # Required for activation recomputation
+ self.num_layers_per_pipeline_rank = len(self.layers)
+
+ if self.post_process and self.post_layer_norm:
+ # Final layer norm before output.
+ self.final_norm = TENorm(
+ config=self.config,
+ hidden_size=self.config.hidden_size,
+ eps=self.config.layernorm_epsilon,
+ )
+
+ def set_input_tensor(self, input_tensor: Tensor):
+ """Set input tensor to be used instead of forward()'s input.
+
+ When doing pipeline parallelism the input from the previous
+ stage comes from communication, not from the input, so the
+ model's forward_step_func won't have it. This function is thus
+ used by internal code to bypass the input provided by the
+ forward_step_func"""
+ self.input_tensor = input_tensor
+
+ def mamba_state_shapes_per_request(self) -> Optional[Tuple[Tuple[int], Tuple[int]]]:
+ """
+ Returns the Mamba conv and ssm states shapes per input sequence
+ if this block contains Mamba layers (this may not be the case with PP > 1).
+ """
+ for layer_type, layer in zip(self.layer_type_list, self.layers):
+ if layer_type == LayerSymbols.MAMBA:
+ return layer.mamba_state_shapes_per_request()
+ return None
+
+ def forward(
+ self,
+ hidden_states: Union[Tensor, WrappedTensor],
+ attention_mask: Tensor,
+ inference_context: Optional[BaseInferenceContext] = None,
+ rotary_pos_emb: Optional[Tensor] = None,
+ *,
+ inference_params: Optional[BaseInferenceContext] = None,
+ packed_seq_params: Optional[PackedSeqParams] = None,
+ padding_mask=None,
+ ):
+ """
+ Forward function of the HybridStack class.
+
+ It either returns the Loss values if labels are given or the
+ final hidden units
+
+ Args:
+ hidden_states (Union[Tensor, WrappedTensor]): the input tensor.
+ Can be passed as a WrappedTensor during inference to avoid an obsolete
+ reference in the calling function.
+ attention_mask (Tensor): the attention mask.
+ inference_context (BaseInferenceContext): the inference parameters.
+ rotary_pos_emb (Tensor, optional): the rotary positional embeddings.
+ Defaults to None.
+ Returns:
+ Tensor: the output tensor.
+ """
+
+ inference_context = deprecate_inference_params(inference_context, inference_params)
+
+ if not self.pre_process:
+ # See set_input_tensor()
+ hidden_states = self.input_tensor
+
+ # Delete the obsolete reference to the initial input tensor if necessary
+ if isinstance(hidden_states, WrappedTensor):
+ hidden_states = hidden_states.unwrap()
+
+ if inference_context and inference_context.is_static_batching():
+ # NOTE(bnorick): match BaseInferenceContext attributes for
+ # mamba_ssm.utils.generation.BaseInferenceContext,
+ # this hack supports eval
+ inference_context.max_seqlen = inference_context.max_sequence_length
+ inference_context.seqlen_offset = inference_context.sequence_len_offset
+
+ if (
+ (self.config.cuda_graph_impl == "local" or self.config.flash_decode)
+ and inference_context
+ and inference_context.is_static_batching()
+ and InferenceMode.is_active()
+ ):
+ current_batch_size = hidden_states.shape[1]
+ sequence_len_offset = torch.tensor(
+ [inference_context.sequence_len_offset] * current_batch_size,
+ dtype=torch.int32,
+ device='cuda',
+ )
+ else:
+ sequence_len_offset = None
+
+ # If fp8_recipe is delayed, wrap the entire pass with get_fp8_context(),
+ # otherwise do nothing extra at the outer level
+ # if we are using other fp8 recipes, then the context manager enter&exit are free
+ # we can wrap fp8_context within the for loop over layers, so that we can fine-grained
+ # control which layer will be fp8 or bf16
+ use_outer_fp8_context = self.config.fp8 and self.config.fp8_recipe == Fp8Recipe.delayed
+ use_inner_fp8_context = self.config.fp8 and self.config.fp8_recipe != Fp8Recipe.delayed
+ use_fp4_context = self.config.fp4 is not None
+ outer_fp8_context = get_fp8_context(self.config) if use_outer_fp8_context else nullcontext()
+
+ if use_inner_fp8_context:
+
+ def get_inner_quant_context(config, layer_number):
+ return get_fp8_context(config, layer_number)
+
+ elif use_fp4_context:
+
+ def get_inner_quant_context(config, layer_number):
+ return get_fp4_context(config, layer_number)
+
+ else:
+
+ def get_inner_quant_context(config, layer_number):
+ return nullcontext()
+
+ with outer_fp8_context:
+ if self.config.recompute_granularity == 'full' and self.training:
+ hidden_states = checkpointed_forward(
+ self,
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ context=None,
+ context_mask=None,
+ rotary_pos_emb=rotary_pos_emb,
+ attention_bias=None,
+ packed_seq_params=packed_seq_params,
+ padding_mask=padding_mask,
+ use_inner_quantization_context=(use_inner_fp8_context or use_fp4_context),
+ )
+ else:
+ for layer in self.layers:
+ # Layers have 1-indexed layer numbers attribute.
+ inner_quant_context = get_inner_quant_context(
+ self.config, layer.layer_number - 1
+ )
+ with inner_quant_context:
+ if isinstance(layer, TransformerLayer):
+ hidden_states, _ = layer(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ inference_context=inference_context,
+ rotary_pos_emb=rotary_pos_emb,
+ sequence_len_offset=sequence_len_offset,
+ packed_seq_params=packed_seq_params,
+ padding_mask=padding_mask,
+ )
+ else: # MambaLayer, Expert, or MLP
+ hidden_states = layer(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ inference_context=inference_context,
+ packed_seq_params=packed_seq_params,
+ )
+
+ # The attention layer (currently a simplified transformer layer)
+ # outputs a tuple of (hidden_states, context). Context is intended
+ # for cross-attention, and is not needed in our model.
+ if isinstance(hidden_states, tuple):
+ hidden_states = hidden_states[0]
+
+ # Final layer norm.
+ if self.post_process and self.post_layer_norm:
+ hidden_states = self.final_norm(hidden_states)
+
+ # Ensure that the tensor passed between pipeline parallel stages is
+ # viewless. See related notes in TransformerBlock and TransformerLayer
+ hidden_states = make_viewless_tensor(
+ inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True
+ )
+
+ return hidden_states
+
+ def sharded_state_dict(
+ self,
+ prefix: str = '',
+ sharded_offsets: Optional[tuple] = None,
+ metadata: Optional[dict] = None,
+ ) -> ShardedStateDict:
+ """
+ Returns a sharded state dictionary for the current object.
+
+ This function constructs a sharded state dictionary by iterating over the layers
+ in the current object, computing the sharded state dictionary for each layer,
+ and combining the results into a single dictionary.
+
+ Parameters:
+ prefix (str): The prefix to use for the state dictionary keys.
+ sharded_offsets (tuple): The sharded offsets to use for the state dictionary.
+ metadata (dict): Additional metadata to use when computing the sharded state dictionary.
+
+ Returns:
+ dict: The sharded state dictionary for the current object.
+ """
+
+ sharded_state_dict = {}
+ layer_prefix = f'{prefix}layers.'
+
+ for local_layer_idx, layer in enumerate(self.layers):
+
+ global_layer_offset = layer.layer_number - 1 # self.layer_number starts at 1
+ state_dict_prefix = (
+ f'{layer_prefix}{local_layer_idx}.' # module list index in HybridStack
+ )
+
+ sharded_prefix = f'{layer_prefix}{global_layer_offset}.'
+ sharded_pp_offset = []
+
+ layer_sharded_state_dict = layer.sharded_state_dict(
+ state_dict_prefix, sharded_pp_offset, metadata
+ )
+
+ replace_prefix_for_sharding(layer_sharded_state_dict, state_dict_prefix, sharded_prefix)
+
+ sharded_state_dict.update(layer_sharded_state_dict)
+
+ # Add modules other than self.layers
+ for name, module in self.named_children():
+ if not module is self.layers:
+ sharded_state_dict.update(
+ sharded_state_dict_default(
+ module,
+ f'{prefix}{name}.',
+ sharded_offsets,
+ metadata,
+ tp_group=self.tp_group,
+ )
+ )
+
+ return sharded_state_dict
+
+
+# Backward-compatible aliases
+MambaStackSubmodules = HybridStackSubmodules
+MambaStack = HybridStack
diff --git a/megatron/core/models/hybrid/hybrid_layer_allocation.py b/megatron/core/models/hybrid/hybrid_layer_allocation.py
new file mode 100644
index 00000000000..f1ba94ef7fa
--- /dev/null
+++ b/megatron/core/models/hybrid/hybrid_layer_allocation.py
@@ -0,0 +1,498 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+from megatron.core.utils import log_on_each_pipeline_stage, log_single_rank
+
+logger = logging.getLogger(__name__)
+
+
+class Symbols:
+ """Symbols for different layer types and pattern separators."""
+
+ MAMBA = "M"
+ GDN = 'G'
+ ATTENTION = "*"
+ DS_ATTENTION = "D"
+ MLP = "-"
+ MOE = 'E'
+ PIPE = '|'
+ MTP_SEPARATOR = "/"
+ VALID_LAYERS = {MAMBA, GDN, ATTENTION, DS_ATTENTION, MLP, MOE}
+
+ @classmethod
+ def name_sorted_valid_layer_symbols(cls) -> list[str]:
+ """Return the valid layer symbols sorted lexicographically by their public attribute
+ name.
+ """
+ valid_layer_attrs = []
+ for name, value in vars(cls).items():
+ if not name.startswith('_') and value in cls.VALID_LAYERS:
+ valid_layer_attrs.append((name, value))
+ valid_layer_attrs.sort()
+ return [value for (_, value) in valid_layer_attrs]
+
+
+@dataclass
+class ParsedHybridPattern:
+ """Result of parsing a unified hybrid pattern string.
+
+ A unified pattern encodes both the main decoder pattern and the MTP pattern
+ in a single string using "/" as a separator. The main pattern may also
+ contain "|" pipe symbols to define pipeline stage boundaries for flexible
+ virtual pipeline parallelism (fVPP).
+
+ Format: "///..."
+
+ Examples:
+ - "M*M*" -> main="M*M*", mtp=None, depths=0 (no MTP)
+ - "M*M*/MM/MM" -> main="M*M*", mtp="MM", depths=2
+ - "MMMM/*M/*M/*M" -> main="MMMM", mtp="*M", depths=3
+ - "M-M-|M-M*-/MM/MM" -> main="M-M-|M-M*-" (2 PP stages), mtp="MM", depths=2
+
+ The "/" symbol introduces MTP patterns. Each repeated pattern after the main
+ decoder represents one MTP prediction depth.
+
+ The "|" symbol in the main pattern defines pipeline stage boundaries.
+
+ Attributes:
+ main_pattern: The main decoder layer pattern (e.g., "M*M*" or "M-M-|M-M*-")
+ mtp_pattern: The MTP layer pattern per depth (e.g., "MM"), or None if no MTP
+ mtp_num_depths: Number of MTP prediction depths (0 if no MTP)
+ """
+
+ main_pattern: Optional[str]
+ mtp_pattern: Optional[str]
+ mtp_num_depths: int
+
+
+def pattern_from_ratios(
+ num_layers: int, attention_ratio: float = 0.0, mlp_ratio: float = 0.0
+) -> str:
+ """Convert deprecated ratio arguments to a layer pattern string.
+
+ Generates an evenly-spaced hybrid layer pattern from target attention and MLP
+ ratios. This exists for backward compatibility with code that uses the deprecated
+ hybrid_attention_ratio and hybrid_mlp_ratio parameters.
+
+ Args:
+ num_layers: Total number of layers.
+ attention_ratio: Target ratio of attention layers to total layers.
+ mlp_ratio: Target ratio of MLP layers to total layers.
+
+ Returns:
+ A layer pattern string (e.g., "MMM*MMM*MM").
+ """
+ assert num_layers > 0
+ assert 0.0 <= attention_ratio <= 1.0
+ assert 0.0 <= mlp_ratio <= 1.0
+ assert attention_ratio + mlp_ratio <= 1.0
+
+ # Allocate attention layers (evenly spaced, starting and ending with mamba)
+ attention_count = round(num_layers * attention_ratio)
+ mamba_count = num_layers - attention_count
+ sections = attention_count + 1
+ section_len = mamba_count / sections
+
+ layer_types = [Symbols.MAMBA] * num_layers
+ x = section_len
+ for i in range(num_layers):
+ if x < 0.5:
+ layer_types[i] = Symbols.ATTENTION
+ x += section_len
+ else:
+ x -= 1
+
+ # Allocate MLP layers (evenly distributed, not replacing attention)
+ mlp_count = round(num_layers * mlp_ratio)
+ if mlp_count > 0:
+ mamba_count -= mlp_count
+ ratio = mamba_count / mlp_count
+ x = ratio
+ for i in range(num_layers):
+ if layer_types[i] == Symbols.MAMBA:
+ if x < 0.5:
+ layer_types[i] = Symbols.MLP
+ x += ratio
+ else:
+ x -= 1
+
+ return ''.join(layer_types)
+
+
+def get_hybrid_total_layer_count(pattern: str) -> int:
+ """Returns the total number of main decoder layers in a hybrid layer pattern.
+
+ Extracts the main pattern (before the first MTP separator '/'), strips
+ pipeline stage separators '|', and returns the character count.
+
+ Args:
+ pattern: Full hybrid layer pattern, possibly including MTP and pipe separators.
+
+ Returns:
+ Total number of layers in the main decoder pattern.
+ """
+ main_pattern = pattern.split(Symbols.MTP_SEPARATOR)[0]
+ _validate_pattern(main_pattern, "main", allow_pipe=True)
+ return len(main_pattern.replace(Symbols.PIPE, ''))
+
+
+def get_hybrid_total_pipeline_segment_count(pattern: str) -> int:
+ """Returns the number of pipeline segments in a hybrid layer pattern.
+
+ Extracts the main pattern (before the first MTP separator '/') and counts
+ the number of segments delimited by '|'.
+
+ Args:
+ pattern: Full hybrid layer pattern, possibly including MTP and pipe separators.
+
+ Returns:
+ Number of pipeline segments (pipe count + 1).
+ """
+ main_pattern = pattern.split(Symbols.MTP_SEPARATOR)[0]
+ return main_pattern.count(Symbols.PIPE) + 1
+
+
+def get_hybrid_layer_counts(pattern: str) -> Dict[str, int]:
+ """Count layers by type across the full hybrid pattern (main + MTP).
+
+ Parses the pattern to extract main and MTP components, then counts
+ each layer type. Main pattern '|' separators are skipped. MTP layers
+ are counted once per MTP depth.
+
+ Args:
+ pattern: Full hybrid layer pattern string.
+
+ Returns:
+ Dictionary mapping layer symbol to count. Keys are all valid layer symbols
+ (Symbols.VALID_LAYERS).
+
+ Examples:
+ >>> get_hybrid_layer_counts("M*M*")
+ {'*': 2, 'G': 0, 'D': 0, 'M': 2, '-': 0, 'E': 0}
+
+ >>> get_hybrid_layer_counts("M-M-|M-M*-/MM/MM")
+ {'*': 1, 'G': 0, 'D': 0, 'M': 8, '-': 4, 'E': 0}
+ """
+ parsed = parse_hybrid_pattern(pattern)
+ counts = {symbol: 0 for symbol in Symbols.name_sorted_valid_layer_symbols()}
+
+ # Count main decoder layers (skip '|' pipe separators)
+ if parsed.main_pattern:
+ for char in parsed.main_pattern:
+ if char in counts:
+ counts[char] += 1
+
+ # Count MTP layers (pattern repeated mtp_num_depths times)
+ if parsed.mtp_pattern and parsed.mtp_num_depths > 0:
+ for char in parsed.mtp_pattern:
+ if char in counts:
+ counts[char] += parsed.mtp_num_depths
+
+ return counts
+
+
+def parse_hybrid_pattern(pattern: Optional[str]) -> ParsedHybridPattern:
+ """Parse a unified hybrid pattern string into main and MTP components.
+
+ The pattern uses "/" as a separator between the main decoder pattern and
+ MTP patterns. Each MTP pattern after the separator represents one prediction
+ depth. The main pattern may contain "|" pipe symbols for pipeline stage
+ boundaries.
+
+ Format: "///..."
+
+ Args:
+ pattern: Unified pattern string, e.g., "M*M*/MM/MM" or just "M*M*"
+
+ Returns:
+ ParsedHybridPattern with main_pattern, mtp_pattern, and mtp_num_depths
+
+ Raises:
+ ValueError: If MTP patterns are inconsistent (all must be identical)
+ ValueError: If pattern contains invalid layer symbols
+
+ Examples:
+ >>> parse_hybrid_pattern("M*M*")
+ ParsedHybridPattern(main_pattern="M*M*", mtp_pattern=None, mtp_num_depths=0)
+
+ >>> parse_hybrid_pattern("M*M*/MM/MM")
+ ParsedHybridPattern(main_pattern="M*M*", mtp_pattern="MM", mtp_num_depths=2)
+
+ >>> parse_hybrid_pattern("MMMM/*M/*M/*M")
+ ParsedHybridPattern(main_pattern="MMMM", mtp_pattern="*M", mtp_num_depths=3)
+
+ >>> parse_hybrid_pattern("M-M-|M-M*-/MM/MM")
+ ParsedHybridPattern(main_pattern="M-M-|M-M*-", mtp_pattern="MM", mtp_num_depths=2)
+ """
+ if pattern is None:
+ return ParsedHybridPattern(main_pattern=None, mtp_pattern=None, mtp_num_depths=0)
+
+ parts = pattern.split(Symbols.MTP_SEPARATOR)
+
+ if len(parts) == 1:
+ # No MTP separator found - pattern is main decoder only
+ main_pattern = parts[0]
+ _validate_pattern(main_pattern, "main", allow_pipe=True)
+ return ParsedHybridPattern(main_pattern=main_pattern, mtp_pattern=None, mtp_num_depths=0)
+
+ # First part is main decoder pattern
+ main_pattern = parts[0]
+ if main_pattern:
+ _validate_pattern(main_pattern, "main", allow_pipe=True)
+
+ # Remaining parts are MTP patterns (one per depth)
+ mtp_parts = parts[1:]
+
+ if not mtp_parts or all(p == "" for p in mtp_parts):
+ # No MTP patterns after separator
+ return ParsedHybridPattern(
+ main_pattern=main_pattern if main_pattern else None, mtp_pattern=None, mtp_num_depths=0
+ )
+
+ # Validate all MTP patterns are identical
+ mtp_pattern = mtp_parts[0]
+ for i, part in enumerate(mtp_parts[1:], start=2):
+ if part != mtp_pattern:
+ raise ValueError(
+ f"All MTP patterns must be identical. "
+ f"Pattern 1 is '{mtp_pattern}', but pattern {i} is '{part}'. "
+ f"Full pattern: '{pattern}'"
+ )
+
+ _validate_pattern(mtp_pattern, "MTP", allow_pipe=False)
+
+ return ParsedHybridPattern(
+ main_pattern=main_pattern if main_pattern else None,
+ mtp_pattern=mtp_pattern,
+ mtp_num_depths=len(mtp_parts),
+ )
+
+
+def _validate_pattern(pattern: str, pattern_name: str, allow_pipe: bool = False) -> None:
+ """Validate that a pattern contains only valid layer symbols.
+
+ Args:
+ pattern: Layer pattern string to validate
+ pattern_name: Name of pattern for error messages (e.g., "main" or "MTP")
+ allow_pipe: Whether to allow the pipe '|' separator (for main patterns)
+
+ Raises:
+ ValueError: If pattern contains invalid symbols
+ """
+ valid_chars = Symbols.VALID_LAYERS | {Symbols.PIPE} if allow_pipe else Symbols.VALID_LAYERS
+ for char in pattern:
+ if char not in valid_chars:
+ raise ValueError(
+ f"In {pattern_name} pattern, '{char}' is not a valid layer symbol. "
+ f"Valid symbols are: {valid_chars}"
+ )
+
+ # Disallow Attention + MLA/DSA hybridity.
+ if Symbols.ATTENTION in pattern and Symbols.DS_ATTENTION in pattern:
+ raise ValueError("Not supported to have both Attention and MLA/DSA in one model")
+
+
+def validate_segment_layers(segment: str) -> List[str]:
+ """Validate and convert a single pipeline segment pattern to a layer type list.
+
+ This is used after the main pattern has been split by '|' into segments.
+ Each segment should contain only valid layer symbols (no '|').
+
+ Args:
+ segment: A single pipeline segment pattern string (e.g., "M-M*-")
+
+ Returns:
+ List of layer type characters.
+
+ Raises:
+ ValueError: If segment contains invalid layer symbols.
+ """
+ layer_type_list = list(segment)
+ for layer_char in layer_type_list:
+ if layer_char not in Symbols.VALID_LAYERS:
+ raise ValueError(
+ f"In hybrid layer pattern segment, '{layer_char}' is not "
+ f"one of {Symbols.VALID_LAYERS}"
+ )
+
+ # Disallow Attention + MLA/DSA hybridity.
+ if Symbols.ATTENTION in segment and Symbols.DS_ATTENTION in segment:
+ raise ValueError("Not supported to have both Attention and MLA/DSA in one model")
+
+ return layer_type_list
+
+
+def select_pipeline_segment(
+ main_pattern: str,
+ pp_group: Optional[torch.distributed.ProcessGroup],
+ vp_stage: Optional[int],
+ first_stage_layers: Optional[int] = None,
+ last_stage_layers: Optional[int] = None,
+) -> Tuple[List[str], int]:
+ """Select and validate the pipeline segment for the given PP rank and VP stage.
+
+ When the main pattern contains '|' pipe separators, splits by '|' into
+ pipeline segments and selects the segment for the current PP rank / VP stage.
+
+ When the pattern has no pipes but pp_size > 1, falls back to runtime layer
+ slicing (for backwards compatibility), supporting both even and uneven PP splits
+ via first_stage_layers / last_stage_layers.
+
+ Args:
+ main_pattern: Main decoder pattern (may contain '|' separators).
+ Empty string is allowed (produces one empty segment).
+ pp_group: Pipeline parallel process group, or None if not using PP.
+ vp_stage: Virtual pipeline stage, or None if not using VPP.
+ first_stage_layers: Number of layers on the first pipeline stage for
+ uneven PP. Only valid when the pattern has no pipe separators.
+ last_stage_layers: Number of layers on the last pipeline stage for
+ uneven PP. Only valid when the pattern has no pipe separators.
+
+ Returns:
+ Tuple of (layer_type_list, layer_offset) where layer_type_list is
+ the list of layer type characters for this segment, and layer_offset
+ is the sum of layer counts from all preceding segments.
+
+ Raises:
+ ValueError: If the segment contains invalid layer symbols, if
+ first/last_stage_layers are used with pipe separators, if VPP is
+ requested without pipe separators, or if layer counts are not
+ evenly divisible across pipeline stages.
+ """
+ segments = main_pattern.split(Symbols.PIPE) if main_pattern else ['']
+
+ pp_rank = torch.distributed.get_rank(pp_group) if pp_group is not None else 0
+ pp_size = torch.distributed.get_world_size(pp_group) if pp_group is not None else 1
+
+ if len(segments) > 1 and (first_stage_layers is not None or last_stage_layers is not None):
+ raise ValueError(
+ "Cannot specify num_layers_in_first_pipeline_stage or "
+ "num_layers_in_last_pipeline_stage when hybrid_layer_pattern "
+ "contains pipe ('|') separators. The pipeline layout is already "
+ "explicitly defined by the pipe separators."
+ )
+
+ if len(segments) == 1 and pp_size > 1:
+ if vp_stage is not None:
+ raise ValueError(
+ "Virtual pipeline parallelism (vp_stage != None) is not supported "
+ "when hybrid_layer_pattern has no pipe ('|') separators. "
+ "Add '|' separators to define explicit pipeline/virtual-pipeline "
+ "stage boundaries."
+ )
+ log_single_rank(
+ logger,
+ logging.WARNING,
+ "DEPRECATION: Using hybrid_layer_pattern without pipe ('|') separators "
+ "with pipeline_model_parallel_size > 1 is deprecated. Please add '|' "
+ "separators to explicitly define pipeline stage boundaries. "
+ "Example: 'M*M*M*M*' with pp_size=2 should become 'M*M*|M*M*'.",
+ )
+ full_pattern = segments[0]
+ layer_type_list = validate_segment_layers(full_pattern)
+ num_layers = len(layer_type_list)
+
+ if first_stage_layers is not None or last_stage_layers is not None:
+ first = first_stage_layers or 0
+ last = last_stage_layers or 0
+ middle_num_layers = num_layers - first - last
+ middle_stages = pp_size - sum(
+ 1 for x in (first_stage_layers, last_stage_layers) if x is not None
+ )
+ if middle_stages > 0:
+ if middle_num_layers % middle_stages != 0:
+ raise ValueError(
+ f"Middle layers ({middle_num_layers}) must be evenly divisible "
+ f"by middle pipeline stages ({middle_stages})."
+ )
+ layers_per_middle = middle_num_layers // middle_stages
+ else:
+ layers_per_middle = 0
+
+ is_first = first_stage_layers is not None and pp_rank == 0
+ is_last = last_stage_layers is not None and pp_rank == pp_size - 1
+
+ if is_first:
+ offset = 0
+ count = first
+ elif is_last:
+ offset = num_layers - last
+ count = last
+ else:
+ middle_rank = pp_rank if first_stage_layers is None else pp_rank - 1
+ offset = middle_rank * layers_per_middle + first
+ count = layers_per_middle
+ else:
+ if num_layers % pp_size != 0:
+ raise ValueError(
+ f"Number of layers ({num_layers}) must be evenly divisible "
+ f"by pipeline-model-parallel-size ({pp_size}) when no pipe "
+ f"separators are specified in the pattern."
+ )
+ layers_per_rank = num_layers // pp_size
+ offset = pp_rank * layers_per_rank
+ count = layers_per_rank
+
+ selected = layer_type_list[offset : offset + count]
+ log_on_each_pipeline_stage(
+ logger,
+ logging.INFO,
+ f"HybridModel: pp_rank={pp_rank}/{pp_size}, vp_stage={vp_stage}, "
+ f"layers='{''.join(selected)}' ({len(selected)} layers), "
+ f"layer_offset={offset} (auto-split)",
+ )
+ return selected, offset
+
+ # Pipe-based segment selection
+ if len(segments) > 1 and len(segments) % pp_size != 0:
+ raise ValueError(
+ f"The number of pipe-delimited segments ({len(segments)}) in "
+ f"hybrid_layer_pattern must be evenly divisible by "
+ f"pipeline_model_parallel_size ({pp_size})."
+ )
+
+ vp_rel = vp_stage if vp_stage is not None else 0
+ segment_index = vp_rel * pp_size + pp_rank
+
+ if segment_index >= len(segments):
+ raise ValueError(
+ f"Pipeline segment index {segment_index} (pp_rank={pp_rank}, "
+ f"vp_stage={vp_rel}) is out of range for {len(segments)} segments. "
+ f"The pattern does not define enough pipe-delimited segments for "
+ f"the current PP/VPP configuration."
+ )
+
+ layer_offset = sum(len(segments[i]) for i in range(segment_index))
+ my_segment = segments[segment_index]
+
+ layer_type_list = validate_segment_layers(my_segment)
+
+ log_on_each_pipeline_stage(
+ logger,
+ logging.INFO,
+ f"HybridModel: pp_rank={pp_rank}/{pp_size}, vp_stage={vp_rel}, "
+ f"segment_index={segment_index}/{len(segments)}, "
+ f"layers='{my_segment}' ({len(layer_type_list)} layers), "
+ f"layer_offset={layer_offset}",
+ )
+
+ return layer_type_list, layer_offset
+
+
+def get_layer_maps_from_layer_type_list(layer_type_list: list[str]) -> dict[str, dict[int, int]]:
+ """
+ Returns maps from global layer index to the corresponding layer index
+ for each valid layer type (those in Symbols.VALID_LAYERS) given a layer type list.
+ """
+ layer_types = [symbol for symbol in Symbols.name_sorted_valid_layer_symbols()]
+ layer_maps = {layer_type: {} for layer_type in layer_types}
+ for global_layer_idx, layer_type in enumerate(layer_type_list):
+ layer_map = layer_maps[layer_type]
+ local_layer_idx = len(layer_map)
+ layer_map[global_layer_idx] = local_layer_idx
+ return layer_maps
diff --git a/megatron/core/models/hybrid/hybrid_layer_specs.py b/megatron/core/models/hybrid/hybrid_layer_specs.py
new file mode 100755
index 00000000000..5b968f720c0
--- /dev/null
+++ b/megatron/core/models/hybrid/hybrid_layer_specs.py
@@ -0,0 +1,309 @@
+# Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved.
+from functools import partial
+
+from megatron.core.extensions.transformer_engine import (
+ TEColumnParallelLinear,
+ TEDotProductAttention,
+ TELayerNormColumnParallelLinear,
+ TELinear,
+ TENorm,
+ TERowParallelLinear,
+)
+from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
+from megatron.core.models.gpt.moe_module_specs import (
+ get_inference_optimized_moe_spec,
+ get_moe_module_spec,
+)
+from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules
+from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules
+from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules
+from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules
+from megatron.core.ssm.mlp_layer import MLPLayer
+from megatron.core.tensor_parallel import (
+ InferenceColumnParallelLinear,
+ InferenceLayerNormColumnParallelLinear,
+ InferenceRowParallelLinear,
+)
+from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.experimental_attention_variant.dsa import (
+ DSAIndexer,
+ DSAIndexerSubmodules,
+ DSAttention,
+ DSAttentionSubmodules,
+)
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.mlp import MLP, MLPSubmodules
+from megatron.core.transformer.multi_latent_attention import (
+ MLASelfAttention,
+ MLASelfAttentionSubmodules,
+)
+from megatron.core.transformer.multi_token_prediction import (
+ MultiTokenPredictionBlock,
+ MultiTokenPredictionBlockSubmodules,
+ MultiTokenPredictionLayer,
+ MultiTokenPredictionLayerSubmodules,
+)
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_layer import (
+ MoETransformerLayer,
+ TransformerLayer,
+ TransformerLayerSubmodules,
+)
+
+# This should be private and should not be used outside of this file.
+moe = get_moe_module_spec(
+ use_te=True,
+ num_experts=8, # Can be any positive integer (must not be None).
+ moe_grouped_gemm=True,
+)
+
+# Inference-optimized MoE spec
+moe_inference = get_inference_optimized_moe_spec()
+
+
+# MTP block spec - provides norms and projection only.
+# Inner layers are built by MultiTokenPredictionLayer using nested HybridStack
+_hybrid_mtp_block_spec = ModuleSpec(
+ module=MultiTokenPredictionBlock,
+ submodules=MultiTokenPredictionBlockSubmodules(
+ layer_specs=[
+ ModuleSpec(
+ module=MultiTokenPredictionLayer,
+ submodules=MultiTokenPredictionLayerSubmodules(
+ enorm=TENorm,
+ hnorm=TENorm,
+ eh_proj=TEColumnParallelLinear,
+ mtp_model_layer=None, # Built via pattern + hybrid_submodules
+ layer_norm=TENorm,
+ ),
+ )
+ ]
+ ),
+)
+
+
+hybrid_stack_spec = ModuleSpec(
+ module=HybridStack,
+ submodules=HybridStackSubmodules(
+ mamba_layer=ModuleSpec(
+ module=MambaLayer,
+ submodules=MambaLayerSubmodules(
+ mixer=ModuleSpec(
+ module=MambaMixer,
+ submodules=MambaMixerSubmodules(
+ in_proj=TELayerNormColumnParallelLinear, out_proj=TERowParallelLinear
+ ),
+ ),
+ mamba_bda=get_bias_dropout_add,
+ ),
+ ),
+ gdn_layer=ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ self_attention=ModuleSpec(
+ module=GatedDeltaNet,
+ submodules=GatedDeltaNetSubmodules(
+ in_proj=TELayerNormColumnParallelLinear,
+ out_norm=TENorm,
+ out_proj=TERowParallelLinear,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ ),
+ ),
+ # Started with spec from gpt_layer_specs.py (with MLP removed)
+ # Using the TE spec because we had problems getting the non-TE spec
+ # working
+ attention_layer=ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.causal},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=TELayerNormColumnParallelLinear,
+ core_attention=TEDotProductAttention,
+ linear_proj=TERowParallelLinear,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ ),
+ ),
+ dsa_layer=ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ input_layernorm=TENorm,
+ self_attention=ModuleSpec(
+ module=MLASelfAttention,
+ params={"attn_mask_type": AttnMaskType.causal},
+ submodules=MLASelfAttentionSubmodules(
+ linear_q_proj=TEColumnParallelLinear,
+ linear_q_down_proj=TELinear,
+ linear_q_up_proj=TEColumnParallelLinear,
+ linear_kv_down_proj=TELinear,
+ linear_kv_up_proj=TEColumnParallelLinear,
+ core_attention=ModuleSpec(
+ module=DSAttention,
+ submodules=DSAttentionSubmodules(
+ indexer=ModuleSpec(
+ module=DSAIndexer,
+ submodules=DSAIndexerSubmodules(
+ linear_wq_b=TELinear,
+ linear_wk=TELinear,
+ k_norm=TENorm,
+ linear_weights_proj=TELinear,
+ ),
+ )
+ ),
+ ),
+ linear_proj=TERowParallelLinear,
+ q_layernorm=IdentityOp,
+ kv_layernorm=IdentityOp,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ ),
+ ),
+ # Started with spec from gpt_layer_specs.py
+ # Using the TE spec because we had problems getting the non-TE spec
+ # working
+ mlp_layer=ModuleSpec(
+ module=MLPLayer,
+ submodules=TransformerLayerSubmodules(
+ mlp=partial(
+ MLP.as_mlp_submodule,
+ submodules=MLPSubmodules(
+ linear_fc1=TELayerNormColumnParallelLinear, linear_fc2=TERowParallelLinear
+ ),
+ ),
+ mlp_bda=get_bias_dropout_add,
+ ),
+ ),
+ moe_layer=ModuleSpec(
+ module=MoETransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ pre_mlp_layernorm=TENorm, mlp=moe, mlp_bda=get_bias_dropout_add
+ ),
+ ),
+ mtp_block_spec=_hybrid_mtp_block_spec,
+ ),
+)
+
+
+hybrid_inference_stack_spec = ModuleSpec(
+ module=HybridStack,
+ submodules=HybridStackSubmodules(
+ mamba_layer=ModuleSpec(
+ module=MambaLayer,
+ submodules=MambaLayerSubmodules(
+ mixer=ModuleSpec(
+ module=MambaMixer,
+ submodules=MambaMixerSubmodules(
+ in_proj=InferenceLayerNormColumnParallelLinear,
+ out_proj=InferenceRowParallelLinear,
+ ),
+ ),
+ mamba_bda=get_bias_dropout_add,
+ ),
+ ),
+ # Started with spec from gpt_layer_specs.py (with MLP removed)
+ # Using the TE spec because we had problems getting the non-TE spec
+ # working
+ attention_layer=ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ self_attention=ModuleSpec(
+ module=SelfAttention,
+ params={"attn_mask_type": AttnMaskType.causal},
+ submodules=SelfAttentionSubmodules(
+ linear_qkv=InferenceLayerNormColumnParallelLinear,
+ core_attention=TEDotProductAttention,
+ linear_proj=InferenceRowParallelLinear,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ ),
+ ),
+ dsa_layer=ModuleSpec(
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ input_layernorm=TENorm,
+ self_attention=ModuleSpec(
+ module=MLASelfAttention,
+ params={"attn_mask_type": AttnMaskType.causal},
+ submodules=MLASelfAttentionSubmodules(
+ linear_q_proj=TEColumnParallelLinear,
+ linear_q_down_proj=TELinear,
+ linear_q_up_proj=TEColumnParallelLinear,
+ linear_kv_down_proj=TELinear,
+ linear_kv_up_proj=TEColumnParallelLinear,
+ core_attention=ModuleSpec(
+ module=DSAttention,
+ submodules=DSAttentionSubmodules(
+ indexer=ModuleSpec(
+ module=DSAIndexer,
+ submodules=DSAIndexerSubmodules(
+ linear_wq_b=TELinear,
+ linear_wk=TELinear,
+ k_norm=TENorm,
+ linear_weights_proj=TELinear,
+ ),
+ )
+ ),
+ ),
+ linear_proj=InferenceRowParallelLinear,
+ q_layernorm=IdentityOp,
+ kv_layernorm=IdentityOp,
+ ),
+ ),
+ self_attn_bda=get_bias_dropout_add,
+ ),
+ ),
+ # Started with spec from gpt_layer_specs.py
+ # Using the TE spec because we had problems getting the non-TE spec
+ # working
+ mlp_layer=ModuleSpec(
+ module=MLPLayer,
+ submodules=TransformerLayerSubmodules(
+ mlp=partial(
+ MLP.as_mlp_submodule,
+ submodules=MLPSubmodules(
+ linear_fc1=InferenceLayerNormColumnParallelLinear,
+ linear_fc2=InferenceRowParallelLinear,
+ ),
+ ),
+ mlp_bda=get_bias_dropout_add,
+ ),
+ ),
+ moe_layer=ModuleSpec(
+ # Use inference-optimized MoE layer for end-to-end CUDA graph support
+ module=TransformerLayer,
+ submodules=TransformerLayerSubmodules(
+ pre_mlp_layernorm=TENorm, mlp=moe_inference, mlp_bda=get_bias_dropout_add
+ ),
+ ),
+ mtp_block_spec=ModuleSpec(
+ module=MultiTokenPredictionBlock,
+ submodules=MultiTokenPredictionBlockSubmodules(
+ layer_specs=[
+ ModuleSpec(
+ module=MultiTokenPredictionLayer,
+ submodules=MultiTokenPredictionLayerSubmodules(
+ enorm=TENorm,
+ hnorm=TENorm,
+ eh_proj=InferenceColumnParallelLinear,
+ mtp_model_layer=None, # Built via pattern + hybrid_submodules
+ layer_norm=TENorm,
+ ),
+ )
+ ]
+ ),
+ ),
+ ),
+)
+
+
+# Backward-compatible aliases
+mamba_stack_spec = hybrid_stack_spec
+mamba_inference_stack_spec = hybrid_inference_stack_spec
diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py
new file mode 100644
index 00000000000..511b24673b0
--- /dev/null
+++ b/megatron/core/models/hybrid/hybrid_model.py
@@ -0,0 +1,592 @@
+# Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+from typing import Literal, Optional
+
+from torch import Tensor
+
+from megatron.core import tensor_parallel
+from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
+from megatron.core.inference.contexts import BaseInferenceContext
+from megatron.core.inference.utils import InferenceMode
+from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
+from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
+from megatron.core.models.common.embeddings.yarn_rotary_pos_embedding import YarnRotaryEmbedding
+from megatron.core.models.common.language_module.language_module import LanguageModule
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.pipeline_parallel.fine_grained_activation_offload import (
+ FineGrainedActivationOffloadingInterface as off_interface,
+)
+from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.quantization.utils import get_quant_config_or_none
+from megatron.core.tensor_parallel import gather_from_sequence_parallel_region
+from megatron.core.transformer import TransformerConfig
+from megatron.core.transformer.enums import InferenceCudaGraphScope, ModelType
+from megatron.core.transformer.module import GraphableMegatronModule
+from megatron.core.transformer.moe.paged_stash import paged_stash_init_chunk_handler
+from megatron.core.transformer.multi_token_prediction import (
+ MultiTokenPredictionBlock,
+ mtp_on_this_rank,
+ process_mtp_loss,
+)
+from megatron.core.transformer.spec_utils import ModuleSpec, build_module
+from megatron.core.utils import (
+ WrappedTensor,
+ deprecate_inference_params,
+ is_using_quantization_scales,
+ log_single_rank,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class HybridModel(LanguageModule, GraphableMegatronModule):
+ """Hybrid language model.
+
+ Args:
+ config (TransformerConfig): Model config
+ hybrid_stack_spec (ModuleSpec): Specifies the modules to use for the various layer types
+ vocab_size (int): Vocabulary size
+ max_sequence_length (int): maximum size of sequence.
+ This is used for positional embedding
+ hybrid_layer_pattern (str): Unified hybrid layer pattern with optional MTP and
+ pipeline stage boundaries.
+ Format: "///..."
+ The main pattern may contain "|" to define pipeline stage boundaries.
+ Examples:
+ - "M*M*" -> main decoder only, no MTP
+ - "M*M*/MM/MM" -> main="M*M*", mtp="MM", 2 depths
+ - "M-M-|M-M*-|M-M-|M-M*-" -> 4 pipeline segments
+ hybrid_attention_ratio (float, optional): Deprecated. Use hybrid_layer_pattern instead.
+ If set to a value > 0.0 and hybrid_layer_pattern is None, a pattern will be
+ generated from the ratio with a deprecation warning.
+ hybrid_mlp_ratio (float, optional): Deprecated. Use hybrid_layer_pattern instead.
+ If set to a value > 0.0 and hybrid_layer_pattern is None, a pattern will be
+ generated from the ratio with a deprecation warning.
+ hybrid_override_pattern (str, optional): Deprecated. Use hybrid_layer_pattern instead.
+ If set and hybrid_layer_pattern is None, the value is copied to hybrid_layer_pattern
+ with a deprecation warning.
+ pre_process (bool, optional): Include embedding layer
+ (used with pipeline parallelism). Defaults to True.
+ post_process (bool, optional): Include an output layer (used with pipeline parallelism).
+ Defaults to True.
+ fp16_lm_cross_entropy (bool, optional): Defaults to False.
+ parallel_output (bool, optional): Do not gather the outputs, keep them split across tensor
+ parallel ranks. Defaults to True.
+ share_embeddings_and_output_weights (bool, optional): When True, input embeddings and
+ output logit weights are shared. Defaults to False.
+ position_embedding_type (Literal[learned_absolute,rope,yarn,none], optional): Position
+ embedding type. Defaults to 'none'.
+ rotary_percent (float, optional): Percent of rotary dimension to use for rotary position
+ embeddings. Ignored unless position_embedding_type is 'rope'. Defaults to 1.0.
+ rotary_base (int, optional): Base period for rotary position embeddings. Ignored unless
+ position_embedding_type is 'rope'. Defaults to 10000.
+ seq_len_interpolation_factor (Optional[float], optional): scale of linearly
+ interpolating RoPE for longer sequences. The value must be a float larger than 1.0.
+ Defaults to None.
+ pg_collection (ProcessGroupCollection, optional): Model communication process groups.
+ vp_stage (Optional[int], optional): Virtual pipeline stage index. Defaults to None.
+ """
+
+ def __init__(
+ self,
+ config: TransformerConfig,
+ hybrid_stack_spec: ModuleSpec,
+ vocab_size: int,
+ max_sequence_length: int,
+ hybrid_layer_pattern: Optional[str] = None,
+ hybrid_attention_ratio: Optional[float] = None,
+ hybrid_mlp_ratio: Optional[float] = None,
+ hybrid_override_pattern: Optional[str] = None,
+ pre_process: bool = True,
+ post_process: bool = True,
+ fp16_lm_cross_entropy: bool = False,
+ parallel_output: bool = True,
+ share_embeddings_and_output_weights: bool = False,
+ # Mamba with no attention has no need for position embeddings, so none is default
+ position_embedding_type: Literal['learned_absolute', 'rope', 'yarn', 'none'] = 'none',
+ rotary_percent: float = 1.0,
+ rotary_base: int = 10000,
+ scatter_embedding_sequence_parallel: bool = True,
+ seq_len_interpolation_factor: Optional[float] = None,
+ pg_collection: Optional[ProcessGroupCollection] = None,
+ vp_stage: Optional[int] = None,
+ ) -> None:
+ super().__init__(config=config, pg_collection=pg_collection)
+
+ if has_config_logger_enabled(config):
+ log_config_to_disk(config, locals(), prefix=type(self).__name__)
+
+ if self.config.use_mup and not getattr(HybridModel, "mup_warning_printed", False):
+ log_single_rank(
+ logger,
+ logging.WARNING,
+ "MuP for HybridModel is experimental and not fully validated yet.",
+ )
+ HybridModel.mup_warning_printed = True
+
+ self.hybrid_stack_spec: ModuleSpec = hybrid_stack_spec
+ self.vocab_size = vocab_size
+ self.max_sequence_length = max_sequence_length
+ self.hybrid_layer_pattern = hybrid_layer_pattern
+ self.pre_process = pre_process
+ self.post_process = post_process
+ self.fp16_lm_cross_entropy = fp16_lm_cross_entropy
+ self.parallel_output = parallel_output
+ self.share_embeddings_and_output_weights = share_embeddings_and_output_weights
+ self.position_embedding_type = position_embedding_type
+ self.vp_stage = vp_stage
+ self.disable_param_offloading = True
+
+ # Backward compatibility for deprecated hybrid parameters
+ if hybrid_override_pattern is not None:
+ if self.hybrid_layer_pattern is None:
+ log_single_rank(
+ logger,
+ logging.WARNING,
+ "hybrid_override_pattern has been deprecated. "
+ "Use hybrid_layer_pattern instead.",
+ )
+ self.hybrid_layer_pattern = hybrid_override_pattern
+ else:
+ raise ValueError(
+ "hybrid_override_pattern and hybrid_layer_pattern cannot both be set. "
+ "hybrid_override_pattern has been deprecated; use hybrid_layer_pattern instead."
+ )
+ if (hybrid_attention_ratio is not None and hybrid_attention_ratio > 0.0) or (
+ hybrid_mlp_ratio is not None and hybrid_mlp_ratio > 0.0
+ ):
+ if hybrid_layer_pattern is not None:
+ raise ValueError(
+ "hybrid_layer_pattern cannot be used together with "
+ "hybrid_attention_ratio or hybrid_mlp_ratio. "
+ "These ratios have been deprecated; use hybrid_layer_pattern alone."
+ )
+ log_single_rank(
+ logger,
+ logging.WARNING,
+ "hybrid_attention_ratio and hybrid_mlp_ratio have been deprecated. "
+ "Use hybrid_layer_pattern instead.",
+ )
+ if self.hybrid_layer_pattern is None:
+ from megatron.core.models.hybrid.hybrid_layer_allocation import pattern_from_ratios
+
+ attn_ratio = hybrid_attention_ratio if hybrid_attention_ratio else 0.0
+ mlp_ratio = hybrid_mlp_ratio if hybrid_mlp_ratio else 0.0
+ self.hybrid_layer_pattern = pattern_from_ratios(
+ config.num_layers, attn_ratio, mlp_ratio
+ )
+
+ # Parse unified pattern to extract main and MTP components, and
+ # determine the pipeline segment for this model instance.
+ from megatron.core.models.hybrid.hybrid_layer_allocation import (
+ parse_hybrid_pattern,
+ select_pipeline_segment,
+ )
+
+ parsed = parse_hybrid_pattern(self.hybrid_layer_pattern)
+ self.mtp_pattern = parsed.mtp_pattern
+ self.mtp_num_depths = parsed.mtp_num_depths
+
+ layer_type_list, layer_offset = select_pipeline_segment(
+ parsed.main_pattern or '',
+ self.pg_collection.pp,
+ vp_stage,
+ first_stage_layers=self.config.num_layers_in_first_pipeline_stage,
+ last_stage_layers=self.config.num_layers_in_last_pipeline_stage,
+ )
+
+ # Determine if MTP is needed (based on pattern parsing)
+ self.mtp_process = (
+ self.mtp_pattern is not None
+ and self.mtp_num_depths > 0
+ # The following forces MTP to be on the final pipeline stage. It might be more optimal
+ # to split the hybrid layer pattern into pipeline stages before parsing the pattern for
+ # the current pipeline stage. This could also enable MTP standalone (MTP in a pipeline
+ # stage separate from loss) to be supported in the hybrid model.
+ and mtp_on_this_rank(
+ layout=self.config.pipeline_model_parallel_layout,
+ mtp_num_layers=self.config.mtp_num_layers,
+ ignore_virtual=False,
+ vp_stage=self.vp_stage,
+ )
+ )
+
+ # megatron core pipelining currently depends on model type
+ # TODO: remove this dependency ?
+ self.model_type = ModelType.encoder_or_decoder
+
+ if self.pre_process or self.mtp_process:
+ self.embedding = LanguageModelEmbedding(
+ config=self.config,
+ vocab_size=self.vocab_size,
+ max_sequence_length=self.max_sequence_length,
+ position_embedding_type=position_embedding_type,
+ scatter_to_sequence_parallel=scatter_embedding_sequence_parallel,
+ tp_group=self.pg_collection.tp,
+ )
+
+ # MLA (also used by DeepSeek Sparse Attention) uses its own decoupled RoPE, therefore we do
+ # not build standard RoPE here when using MLA.
+ if self.position_embedding_type == 'rope' and not self.config.multi_latent_attention:
+ self.rotary_pos_emb = RotaryEmbedding(
+ kv_channels=self.config.kv_channels,
+ rotary_percent=rotary_percent,
+ seq_len_interpolation_factor=seq_len_interpolation_factor,
+ rotary_base=rotary_base,
+ use_cpu_initialization=self.config.use_cpu_initialization,
+ cp_group=self.pg_collection.cp,
+ )
+ elif self.position_embedding_type == 'yarn':
+ self.rotary_pos_emb = YarnRotaryEmbedding(
+ kv_channels=self.config.kv_channels,
+ rotary_percent=rotary_percent,
+ seq_len_interpolation_factor=seq_len_interpolation_factor,
+ rotary_base=rotary_base,
+ scaling_factor=getattr(self.config, "yarn_rotary_scaling_factor"),
+ original_max_position_embeddings=getattr(
+ self.config, "yarn_original_max_position_embeddings"
+ ),
+ beta_fast=getattr(self.config, "yarn_beta_fast"),
+ beta_slow=getattr(self.config, "yarn_beta_slow"),
+ mscale=getattr(self.config, "yarn_mscale"),
+ mscale_all_dim=getattr(self.config, "yarn_mscale_all_dim"),
+ correction_range_round_to_int=getattr(
+ self.config, "yarn_correction_range_round_to_int"
+ ),
+ use_cpu_initialization=self.config.use_cpu_initialization,
+ cp_group=self.pg_collection.cp,
+ )
+ self.decoder = build_module(
+ hybrid_stack_spec,
+ self.config,
+ pre_process=self.pre_process,
+ layer_type_list=layer_type_list,
+ pp_layer_offset=layer_offset,
+ post_process=self.post_process,
+ dtype=config.params_dtype,
+ pg_collection=self.pg_collection,
+ name="decoder",
+ )
+
+ # MTP block - uses mtp_block_spec from hybrid_stack_spec.submodules
+ if self.mtp_process:
+ hybrid_submodules = hybrid_stack_spec.submodules
+ mtp_block_spec = hybrid_submodules.mtp_block_spec
+ assert mtp_block_spec is not None, (
+ "MTP pattern specified but mtp_block_spec is None in hybrid_stack_spec.submodules. "
+ "Ensure hybrid_stack_spec includes mtp_block_spec for MTP support."
+ )
+
+ self.mtp = MultiTokenPredictionBlock(
+ config=self.config,
+ spec=mtp_block_spec,
+ pg_collection=self.pg_collection,
+ vp_stage=self.vp_stage,
+ mtp_layer_pattern=self.mtp_pattern,
+ mtp_num_depths=self.mtp_num_depths,
+ hybrid_submodules=hybrid_submodules,
+ name="mtp",
+ )
+ self._setup_mtp_cuda_graphs()
+
+ # Output
+ if post_process or self.mtp_process:
+ self.output_layer = tensor_parallel.ColumnParallelLinear(
+ config.hidden_size,
+ self.vocab_size,
+ config=config,
+ init_method=(
+ config.embedding_init_method
+ if config.use_mup and not self.share_embeddings_and_output_weights
+ else config.init_method
+ ),
+ bias=False,
+ skip_bias_add=False,
+ gather_output=not self.parallel_output,
+ skip_weight_param_allocation=self.pre_process
+ and self.share_embeddings_and_output_weights,
+ tp_group=self.pg_collection.tp,
+ )
+
+ if self.pre_process or self.post_process or self.mtp_process:
+ self.setup_embeddings_and_output_layer()
+
+ for name, module in self.named_modules():
+ if hasattr(module, 'finish_init'):
+ quant_config = get_quant_config_or_none(name, self.config.quant_recipe)
+ module.finish_init(quant_config)
+
+ def set_input_tensor(self, input_tensor: Tensor) -> None:
+ """Sets input tensor to the model.
+
+ See megatron.model.transformer.set_input_tensor()
+
+ Args:
+ input_tensor (Tensor): Sets the input tensor for the model.
+ """
+ # This is usually handled in schedules.py but some inference code still
+ # gives us non-lists or None
+ if not isinstance(input_tensor, list):
+ input_tensor = [input_tensor]
+
+ assert len(input_tensor) == 1, 'input_tensor should only be length 1 for gpt/bert'
+ self.decoder.set_input_tensor(input_tensor[0])
+
+ def preprocess_for_fine_grained_offloading(self):
+ """Preprocess for fine-grained activation offloading."""
+ off_interface.init_chunk_handler(
+ vp_size=self.config.virtual_pipeline_model_parallel_size,
+ vp_stage=self.vp_stage,
+ min_offloaded_tensor_size=self.config.min_offloaded_tensor_size,
+ max_inflight_offloads=self.config.fine_grained_offloading_max_inflight_offloads,
+ )
+ if self.disable_param_offloading:
+ for param in self.decoder.parameters():
+ off_interface.mark_not_offloadable(param)
+ if self.mtp_process:
+ for param in self.mtp.parameters():
+ off_interface.mark_not_offloadable(param)
+ if self.post_process:
+ for param in self.output_layer.parameters():
+ off_interface.mark_not_offloadable(param)
+ self.disable_param_offloading = False
+
+ def preprocess_for_paged_stash(self):
+ """Preprocess for paged stash."""
+ return paged_stash_init_chunk_handler(
+ vp_size=self.config.virtual_pipeline_model_parallel_size, vp_stage=self.vp_stage
+ )
+
+ def _should_call_local_cudagraph(self, *args, **kwargs):
+ """
+ Check if we should call the local cudagraph path.
+ """
+ if (
+ InferenceMode.is_active()
+ and hasattr(self, 'cudagraph_manager')
+ and (
+ kwargs.get('inference_context') is not None
+ or kwargs.get('inference_params') is not None
+ )
+ and self.config.inference_cuda_graph_scope == InferenceCudaGraphScope.block
+ ):
+ if kwargs['inference_context'].is_static_batching():
+ using_cuda_graph = kwargs['inference_context'].is_decode_only()
+ else:
+ using_cuda_graph = kwargs['inference_context'].using_cuda_graph_this_step()
+
+ if using_cuda_graph:
+ return True
+ return False
+
+ def __call__(self, *args, **kwargs):
+ if self._should_call_local_cudagraph(*args, **kwargs):
+ return super().__call__(*args, **kwargs)[0]
+ return super().__call__(*args, **kwargs)
+
+ def create_mcore_cudagraph_manager(self, config):
+ """
+ Create the cudagraph manager for the full iteration inference scope
+ """
+ if config.inference_cuda_graph_scope == InferenceCudaGraphScope.block:
+ from megatron.core.transformer.cuda_graphs import CudaGraphManager
+
+ self.cudagraph_manager = CudaGraphManager(config)
+
+ def forward(
+ self,
+ input_ids: Tensor,
+ position_ids: Tensor,
+ attention_mask: Tensor,
+ decoder_input: Tensor = None,
+ labels: Tensor = None,
+ inference_context: BaseInferenceContext = None,
+ runtime_gather_output: Optional[bool] = None,
+ *,
+ inference_params: Optional[BaseInferenceContext] = None,
+ loss_mask: Optional[Tensor] = None,
+ packed_seq_params: Optional[PackedSeqParams] = None,
+ padding_mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ """Forward function of the Hybrid model. This function passes the input tensors
+ through the embedding layer, and then the decoder and finally into the post
+ processing layer (optional).
+
+ It either returns the Loss values if labels are given or the final hidden units
+ """
+ # If decoder_input is provided (not None), then input_ids and position_ids are ignored.
+ # Otherwise, apply embedding layer on input_ids and position_ids to get decoder_input.
+
+ if self.config.fine_grained_activation_offloading:
+ self.preprocess_for_fine_grained_offloading()
+
+ if self.config.moe_paged_stash:
+ self.preprocess_for_paged_stash()
+
+ inference_context = deprecate_inference_params(inference_context, inference_params)
+
+ in_inference_mode = InferenceMode.is_active()
+
+ if in_inference_mode:
+ assert runtime_gather_output, "Inference must always gather TP logits"
+
+ # Decoder embedding.
+ if decoder_input is not None:
+ pass
+ elif self.pre_process:
+ decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
+
+ # Clear the outputs for padding tokens when using dynamic batching with
+ # quantization scales to avoid corrupting amax calculations
+ if (
+ in_inference_mode
+ and inference_context is not None
+ and inference_context.is_dynamic_batching()
+ and is_using_quantization_scales(self.config)
+ ):
+ decoder_input[inference_context.padding_slice] = 0.0
+ else:
+ # intermediate stage of pipeline
+ # decoder will get hidden_states from encoder.input_tensor
+ decoder_input = None
+
+ rotary_pos_emb = None
+ if self.position_embedding_type == 'rope' and not self.config.multi_latent_attention:
+ rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
+ inference_context, self.decoder, decoder_input, self.config, packed_seq_params
+ )
+ rotary_pos_emb = self.rotary_pos_emb(
+ rotary_seq_len,
+ packed_seq=packed_seq_params is not None and packed_seq_params.qkv_format == 'thd',
+ )
+ elif self.position_embedding_type == 'yarn':
+ rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
+ inference_context, self.decoder, decoder_input, self.config, packed_seq_params
+ )
+ # YarnRotaryEmbedding.forward returns (emb, mscale); discard mscale here
+ rotary_pos_emb, _ = self.rotary_pos_emb(
+ rotary_seq_len,
+ packed_seq=packed_seq_params is not None and packed_seq_params.qkv_format == 'thd',
+ )
+
+ # Wrap decoder_input to allow the decoder (HybridStack) to delete the
+ # reference held by this caller function, enabling early garbage collection
+ # for inference.
+ if in_inference_mode:
+ decoder_input = WrappedTensor(decoder_input)
+
+ # The following assert will currently fail when running inference.
+ # Commented out for now.
+ # TODO (duncan/rwaleffe): (1) confirm that the externally-generated
+ # attention mask is not needed and is ignored by the model in
+ # inference mode, (2) reduce the size of the externally-generated
+ # attention mask to prevent CPU OOM (as we did for training), (3)
+ # force the attention mask passed to the model in inference mode to
+ # be None, so this assert will succeed.
+ # assert attention_mask is None, "The attention mask is ignored and should be set to None"
+
+ # Run decoder.
+ hidden_states = self.decoder(
+ hidden_states=decoder_input,
+ attention_mask=attention_mask,
+ inference_context=inference_context,
+ rotary_pos_emb=rotary_pos_emb,
+ packed_seq_params=packed_seq_params,
+ padding_mask=padding_mask,
+ )
+
+ output_weight = None
+ if self.share_embeddings_and_output_weights:
+ output_weight = self.shared_embedding_or_output_weight()
+
+ # Check if speculative decoding is active. When it is, MTP must be
+ # computed *after* verification so that it is conditioned on verified
+ # tokens rather than stale speculative tokens from the previous step.
+ is_spec_decode = (
+ in_inference_mode
+ and inference_context is not None
+ and inference_context.is_dynamic_batching()
+ and inference_context.num_speculative_tokens > 0
+ )
+
+ mtp_forward_ran = self.mtp_process and not (in_inference_mode or is_spec_decode)
+ if mtp_forward_ran:
+ hidden_states = self.mtp(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ inference_params=inference_params,
+ rotary_pos_emb=rotary_pos_emb,
+ packed_seq_params=packed_seq_params,
+ embedding=self.embedding,
+ )
+
+ if not self.post_process:
+ return hidden_states
+
+ if self.config.mtp_num_layers is not None and self.mtp_process:
+ assert self.config.mtp_num_layers > 0
+ if in_inference_mode or is_spec_decode:
+ self._decoder_hidden_states_cache = hidden_states
+ else:
+ hidden_states = process_mtp_loss(
+ hidden_states=hidden_states,
+ labels=labels,
+ loss_mask=loss_mask,
+ output_layer=self.output_layer,
+ output_weight=output_weight,
+ runtime_gather_output=runtime_gather_output,
+ is_training=self.training,
+ compute_language_model_loss=self.compute_language_model_loss,
+ config=self.config,
+ cp_group=self.pg_collection.cp,
+ packed_seq_params=packed_seq_params,
+ scale_logits_fn=self._scale_logits if self.config.use_mup else None,
+ )
+ sequence_parallel_override = False
+ if (
+ in_inference_mode
+ and inference_context is not None
+ and inference_context.config.materialize_only_last_token_logits
+ ):
+ if inference_context.is_static_batching():
+ hidden_states = hidden_states[-1:, :, :]
+ else:
+ if self.output_layer.sequence_parallel:
+ # Perform the sequence parallel gather here instead of after the output layer
+ # because we need to slice the last token logits from the full view of the
+ # packed logits across all requests.
+ hidden_states = gather_from_sequence_parallel_region(
+ hidden_states, group=self.pg_collection.tp
+ )
+ self.output_layer.sequence_parallel = False
+ sequence_parallel_override = True
+
+ # Reshape [S, B, H] (with B=1) to [1, S, H] for logit extraction,
+ # then back to [S', B, H] for the output layer.
+ reshaped = hidden_states.squeeze(1).unsqueeze(0)
+ hidden_states = inference_context.last_token_logits(reshaped).unsqueeze(1)
+
+ logits, _ = self.output_layer(
+ hidden_states, weight=output_weight, runtime_gather_output=runtime_gather_output
+ )
+ logits = self._scale_logits(logits)
+
+ # Restore sequence parallel execution to the output layer if necessary.
+ if sequence_parallel_override:
+ assert (
+ in_inference_mode
+ and inference_context.is_dynamic_batching()
+ and inference_context.config.materialize_only_last_token_logits
+ )
+ self.output_layer.sequence_parallel = True
+
+ if labels is None:
+ # [s b h] => [b s h]
+ return logits.transpose(0, 1).contiguous()
+
+ loss = self.compute_language_model_loss(labels, logits)
+
+ return loss
diff --git a/megatron/core/models/mamba/__init__.py b/megatron/core/models/mamba/__init__.py
index 5aaf8524018..4de391a62c9 100644
--- a/megatron/core/models/mamba/__init__.py
+++ b/megatron/core/models/mamba/__init__.py
@@ -1,2 +1,6 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
-from .mamba_model import MambaModel
+# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.
+
+# Backward-compatible re-exports. The canonical location is now
+# megatron.core.models.hybrid.
+from megatron.core.models.hybrid.hybrid_model import HybridModel
+from megatron.core.models.mamba.mamba_model import MambaModel
diff --git a/megatron/core/models/mamba/mamba_layer_specs.py b/megatron/core/models/mamba/mamba_layer_specs.py
old mode 100755
new mode 100644
index d2a85d004ef..5fb9e49a0dd
--- a/megatron/core/models/mamba/mamba_layer_specs.py
+++ b/megatron/core/models/mamba/mamba_layer_specs.py
@@ -1,204 +1,5 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved.
-from megatron.core.extensions.transformer_engine import (
- TEColumnParallelLinear,
- TEDotProductAttention,
- TELayerNormColumnParallelLinear,
- TENorm,
- TERowParallelLinear,
-)
-from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
-from megatron.core.models.gpt.moe_module_specs import (
- get_inference_optimized_moe_spec,
- get_moe_module_spec,
-)
-from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules
-from megatron.core.ssm.mamba_block import MambaStack, MambaStackSubmodules
-from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules
-from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules
-from megatron.core.ssm.mlp_layer import MLPLayer
-from megatron.core.tensor_parallel import (
- InferenceLayerNormColumnParallelLinear,
- InferenceRowParallelLinear,
-)
-from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
-from megatron.core.transformer.enums import AttnMaskType
-from megatron.core.transformer.mlp import MLP, MLPSubmodules
-from megatron.core.transformer.multi_token_prediction import (
- MultiTokenPredictionBlock,
- MultiTokenPredictionBlockSubmodules,
- MultiTokenPredictionLayer,
- MultiTokenPredictionLayerSubmodules,
-)
-from megatron.core.transformer.spec_utils import ModuleSpec
-from megatron.core.transformer.transformer_layer import (
- MoETransformerLayer,
- TransformerLayer,
- TransformerLayerSubmodules,
-)
-
-# This should be private and should not be used outside of this file.
-moe = get_moe_module_spec(
- use_te=True,
- num_experts=8, # Can be any positive integer (must not be None).
- moe_grouped_gemm=True,
-)
-
-# Inference-optimized MoE spec
-moe_inference = get_inference_optimized_moe_spec()
-
-
-# MTP block spec for Mamba - provides norms and projection only.
-# Inner layers are built by MultiTokenPredictionLayer using nested MambaStack
-_mamba_mtp_block_spec = ModuleSpec(
- module=MultiTokenPredictionBlock,
- submodules=MultiTokenPredictionBlockSubmodules(
- layer_specs=[
- ModuleSpec(
- module=MultiTokenPredictionLayer,
- submodules=MultiTokenPredictionLayerSubmodules(
- enorm=TENorm,
- hnorm=TENorm,
- eh_proj=TEColumnParallelLinear,
- mtp_model_layer=None, # Built via pattern + mamba_submodules
- layer_norm=TENorm,
- ),
- )
- ]
- ),
-)
-
-
-mamba_stack_spec = ModuleSpec(
- module=MambaStack,
- submodules=MambaStackSubmodules(
- mamba_layer=ModuleSpec(
- module=MambaLayer,
- submodules=MambaLayerSubmodules(
- mixer=ModuleSpec(
- module=MambaMixer,
- submodules=MambaMixerSubmodules(
- in_proj=TELayerNormColumnParallelLinear, out_proj=TERowParallelLinear
- ),
- ),
- mamba_bda=get_bias_dropout_add,
- ),
- ),
- gdn_layer=ModuleSpec(
- module=TransformerLayer,
- submodules=TransformerLayerSubmodules(
- self_attention=ModuleSpec(
- module=GatedDeltaNet,
- submodules=GatedDeltaNetSubmodules(
- in_proj=TELayerNormColumnParallelLinear,
- out_norm=TENorm,
- out_proj=TERowParallelLinear,
- ),
- ),
- self_attn_bda=get_bias_dropout_add,
- ),
- ),
- # Started with spec from gpt_layer_specs.py (with MLP removed)
- # Using the TE spec because we had problems getting the non-TE spec
- # working
- attention_layer=ModuleSpec(
- module=TransformerLayer,
- submodules=TransformerLayerSubmodules(
- self_attention=ModuleSpec(
- module=SelfAttention,
- params={"attn_mask_type": AttnMaskType.causal},
- submodules=SelfAttentionSubmodules(
- linear_qkv=TELayerNormColumnParallelLinear,
- core_attention=TEDotProductAttention,
- linear_proj=TERowParallelLinear,
- ),
- ),
- self_attn_bda=get_bias_dropout_add,
- ),
- ),
- # Started with spec from gpt_layer_specs.py
- # Using the TE spec because we had problems getting the non-TE spec
- # working
- mlp_layer=ModuleSpec(
- module=MLPLayer,
- submodules=TransformerLayerSubmodules(
- mlp=ModuleSpec(
- module=MLP,
- submodules=MLPSubmodules(
- linear_fc1=TELayerNormColumnParallelLinear, linear_fc2=TERowParallelLinear
- ),
- ),
- mlp_bda=get_bias_dropout_add,
- ),
- ),
- moe_layer=ModuleSpec(
- module=MoETransformerLayer,
- submodules=TransformerLayerSubmodules(
- pre_mlp_layernorm=TENorm, mlp=moe, mlp_bda=get_bias_dropout_add
- ),
- ),
- mtp_block_spec=_mamba_mtp_block_spec,
- ),
-)
-
-
-mamba_inference_stack_spec = ModuleSpec(
- module=MambaStack,
- submodules=MambaStackSubmodules(
- mamba_layer=ModuleSpec(
- module=MambaLayer,
- submodules=MambaLayerSubmodules(
- mixer=ModuleSpec(
- module=MambaMixer,
- submodules=MambaMixerSubmodules(
- in_proj=InferenceLayerNormColumnParallelLinear,
- out_proj=InferenceRowParallelLinear,
- ),
- ),
- mamba_bda=get_bias_dropout_add,
- ),
- ),
- # Started with spec from gpt_layer_specs.py (with MLP removed)
- # Using the TE spec because we had problems getting the non-TE spec
- # working
- attention_layer=ModuleSpec(
- module=TransformerLayer,
- submodules=TransformerLayerSubmodules(
- self_attention=ModuleSpec(
- module=SelfAttention,
- params={"attn_mask_type": AttnMaskType.causal},
- submodules=SelfAttentionSubmodules(
- linear_qkv=InferenceLayerNormColumnParallelLinear,
- core_attention=TEDotProductAttention,
- linear_proj=InferenceRowParallelLinear,
- ),
- ),
- self_attn_bda=get_bias_dropout_add,
- ),
- ),
- # Started with spec from gpt_layer_specs.py
- # Using the TE spec because we had problems getting the non-TE spec
- # working
- mlp_layer=ModuleSpec(
- module=MLPLayer,
- submodules=TransformerLayerSubmodules(
- mlp=ModuleSpec(
- module=MLP,
- submodules=MLPSubmodules(
- linear_fc1=InferenceLayerNormColumnParallelLinear,
- linear_fc2=InferenceRowParallelLinear,
- ),
- ),
- mlp_bda=get_bias_dropout_add,
- ),
- ),
- moe_layer=ModuleSpec(
- # Use inference-optimized MoE layer for end-to-end CUDA graph support
- module=TransformerLayer,
- submodules=TransformerLayerSubmodules(
- pre_mlp_layernorm=TENorm, mlp=moe_inference, mlp_bda=get_bias_dropout_add
- ),
- ),
- mtp_block_spec=_mamba_mtp_block_spec,
- ),
-)
+# Backward-compatible re-export. The canonical location is now
+# megatron.core.models.hybrid.hybrid_layer_specs.
+from megatron.core.models.hybrid.hybrid_layer_specs import * # noqa: F401,F403
diff --git a/megatron/core/models/mamba/mamba_model.py b/megatron/core/models/mamba/mamba_model.py
index e295c3d6b01..13964286daf 100644
--- a/megatron/core/models/mamba/mamba_model.py
+++ b/megatron/core/models/mamba/mamba_model.py
@@ -1,519 +1,26 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved.
import logging
-from typing import Literal, Optional
-import torch
-from torch import Tensor
-
-from megatron.core import tensor_parallel
-from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk
-from megatron.core.inference.contexts import BaseInferenceContext
-from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
-from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
-from megatron.core.models.common.language_module.language_module import LanguageModule
-from megatron.core.packed_seq_params import PackedSeqParams
-from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.core.quantization.utils import get_quant_config_or_none
-from megatron.core.tensor_parallel import gather_from_sequence_parallel_region
-from megatron.core.transformer import TransformerConfig
-from megatron.core.transformer.enums import ModelType
-from megatron.core.transformer.multi_token_prediction import (
- MultiTokenPredictionBlock,
- mtp_on_this_rank,
- process_mtp_loss,
-)
-from megatron.core.transformer.spec_utils import ModuleSpec, build_module
-from megatron.core.utils import (
- WrappedTensor,
- deprecate_inference_params,
- is_using_quantization_scales,
- log_single_rank,
-)
+from megatron.core.models.hybrid.hybrid_model import * # noqa: F401,F403 # pylint: disable=unused-import
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.utils import log_single_rank
logger = logging.getLogger(__name__)
-class MambaModel(LanguageModule):
- """Mamba language model.
-
- Args:
- config (TransformerConfig): Model config
- mamba_stack_spec (ModuleSpec): Specifies the modules to use for the various layer types
- vocab_size (int): Vocabulary size
- max_sequence_length (int): maximum size of sequence.
- This is used for positional embedding
- hybrid_layer_pattern (str): Unified hybrid layer pattern with optional MTP and
- pipeline stage boundaries.
- Format: "///..."
- The main pattern may contain "|" to define pipeline stage boundaries.
- Examples:
- - "M*M*" -> main decoder only, no MTP
- - "M*M*/MM/MM" -> main="M*M*", mtp="MM", 2 depths
- - "M-M-|M-M*-|M-M-|M-M*-" -> 4 pipeline segments
- hybrid_attention_ratio (float, optional): Deprecated. Use hybrid_layer_pattern instead.
- If set to a value > 0.0 and hybrid_layer_pattern is None, a pattern will be
- generated from the ratio with a deprecation warning.
- hybrid_mlp_ratio (float, optional): Deprecated. Use hybrid_layer_pattern instead.
- If set to a value > 0.0 and hybrid_layer_pattern is None, a pattern will be
- generated from the ratio with a deprecation warning.
- hybrid_override_pattern (str, optional): Deprecated. Use hybrid_layer_pattern instead.
- If set and hybrid_layer_pattern is None, the value is copied to hybrid_layer_pattern
- with a deprecation warning.
- pre_process (bool, optional): Include embedding layer
- (used with pipeline parallelism). Defaults to True.
- post_process (bool, optional): Include an output layer (used with pipeline parallelism).
- Defaults to True.
- fp16_lm_cross_entropy (bool, optional): Defaults to False.
- parallel_output (bool, optional): Do not gather the outputs, keep them split across tensor
- parallel ranks. Defaults to True.
- share_embeddings_and_output_weights (bool, optional): When True, input embeddings and
- output logit weights are shared. Defaults to False.
- position_embedding_type (Literal[learned_absolute,rope,none], optional): Position
- embedding type. Defaults to 'none'.
- rotary_percent (float, optional): Percent of rotary dimension to use for rotary position
- embeddings. Ignored unless position_embedding_type is 'rope'. Defaults to 1.0.
- rotary_base (int, optional): Base period for rotary position embeddings. Ignored unless
- position_embedding_type is 'rope'. Defaults to 10000.
- seq_len_interpolation_factor (Optional[float], optional): scale of linearly
- interpolating RoPE for longer sequences. The value must be a float larger than 1.0.
- Defaults to None.
- pg_collection (ProcessGroupCollection, optional): Model communication process groups.
- vp_stage (Optional[int], optional): Virtual pipeline stage index. Defaults to None.
- """
-
- def __init__(
- self,
- config: TransformerConfig,
- mamba_stack_spec: ModuleSpec,
- vocab_size: int,
- max_sequence_length: int,
- hybrid_layer_pattern: Optional[str] = None,
- hybrid_attention_ratio: Optional[float] = None,
- hybrid_mlp_ratio: Optional[float] = None,
- hybrid_override_pattern: Optional[str] = None,
- pre_process: bool = True,
- post_process: bool = True,
- fp16_lm_cross_entropy: bool = False,
- parallel_output: bool = True,
- share_embeddings_and_output_weights: bool = False,
- # Mamba with no attention has no need for position embeddings, so none is default
- position_embedding_type: Literal['learned_absolute', 'rope', 'none'] = 'none',
- rotary_percent: float = 1.0,
- rotary_base: int = 10000,
- scatter_embedding_sequence_parallel: bool = True,
- seq_len_interpolation_factor: Optional[float] = None,
- pg_collection: Optional[ProcessGroupCollection] = None,
- vp_stage: Optional[int] = None,
- ) -> None:
- super().__init__(config=config, pg_collection=pg_collection)
-
- if has_config_logger_enabled(config):
- log_config_to_disk(config, locals(), prefix=type(self).__name__)
-
- if self.config.use_mup and not getattr(MambaModel, "mup_warning_printed", False):
- log_single_rank(
- logger,
- logging.WARNING,
- "MuP for MambaModel is experimental and not fully validated yet.",
- )
- MambaModel.mup_warning_printed = True
-
- self.mamba_stack_spec: ModuleSpec = mamba_stack_spec
- self.vocab_size = vocab_size
- self.max_sequence_length = max_sequence_length
- self.hybrid_layer_pattern = hybrid_layer_pattern
- self.pre_process = pre_process
- self.post_process = post_process
- self.fp16_lm_cross_entropy = fp16_lm_cross_entropy
- self.parallel_output = parallel_output
- self.share_embeddings_and_output_weights = share_embeddings_and_output_weights
- self.position_embedding_type = position_embedding_type
- self.vp_stage = vp_stage
-
- # Backward compatibility for deprecated hybrid parameters
- if hybrid_override_pattern is not None:
- if self.hybrid_layer_pattern is None:
- log_single_rank(
- logger,
- logging.WARNING,
- "hybrid_override_pattern has been deprecated. "
- "Use hybrid_layer_pattern instead.",
- )
- self.hybrid_layer_pattern = hybrid_override_pattern
- else:
- raise ValueError(
- "hybrid_override_pattern and hybrid_layer_pattern cannot both be set. "
- "hybrid_override_pattern has been deprecated; use hybrid_layer_pattern instead."
- )
- if (hybrid_attention_ratio is not None and hybrid_attention_ratio > 0.0) or (
- hybrid_mlp_ratio is not None and hybrid_mlp_ratio > 0.0
- ):
- if hybrid_layer_pattern is not None:
- raise ValueError(
- "hybrid_layer_pattern cannot be used together with "
- "hybrid_attention_ratio or hybrid_mlp_ratio. "
- "These ratios have been deprecated; use hybrid_layer_pattern alone."
- )
- log_single_rank(
- logger,
- logging.WARNING,
- "hybrid_attention_ratio and hybrid_mlp_ratio have been deprecated. "
- "Use hybrid_layer_pattern instead.",
- )
- if self.hybrid_layer_pattern is None:
- from megatron.core.ssm.mamba_hybrid_layer_allocation import pattern_from_ratios
-
- attn_ratio = hybrid_attention_ratio if hybrid_attention_ratio else 0.0
- mlp_ratio = hybrid_mlp_ratio if hybrid_mlp_ratio else 0.0
- self.hybrid_layer_pattern = pattern_from_ratios(
- config.num_layers, attn_ratio, mlp_ratio
- )
-
- # Parse unified pattern to extract main and MTP components, and
- # determine the pipeline segment for this model instance.
- from megatron.core.ssm.mamba_hybrid_layer_allocation import (
- parse_hybrid_pattern,
- select_pipeline_segment,
- )
-
- parsed = parse_hybrid_pattern(self.hybrid_layer_pattern)
- self.mtp_pattern = parsed.mtp_pattern
- self.mtp_num_depths = parsed.mtp_num_depths
-
- layer_type_list, layer_offset = select_pipeline_segment(
- parsed.main_pattern or '',
- self.pg_collection.pp,
- vp_stage,
- first_stage_layers=self.config.num_layers_in_first_pipeline_stage,
- last_stage_layers=self.config.num_layers_in_last_pipeline_stage,
- )
+class MambaModel(HybridModel):
+ """Backward-compatible wrapper that accepts the deprecated mamba_stack_spec kwarg."""
- # Determine if MTP is needed (based on pattern parsing)
- self.mtp_process = (
- self.mtp_pattern is not None
- and self.mtp_num_depths > 0
- # The following forces MTP to be on the final pipeline stage. It might be more optimal
- # to split the hybrid layer pattern into pipeline stages before parsing the pattern for
- # the current pipeline stage. This could also enable MTP standalone (MTP in a pipeline
- # stage separate from loss) to be supported in the hybrid model.
- and mtp_on_this_rank(self.config, ignore_virtual=False, vp_stage=self.vp_stage)
+ def __init__(self, *args, mamba_stack_spec: ModuleSpec = None, **kwargs):
+ log_single_rank(
+ logger, logging.WARNING, "MambaModel has been deprecated. Use HybridModel instead."
)
-
- # megatron core pipelining currently depends on model type
- # TODO: remove this dependency ?
- self.model_type = ModelType.encoder_or_decoder
-
- if self.pre_process or self.mtp_process:
- self.embedding = LanguageModelEmbedding(
- config=self.config,
- vocab_size=self.vocab_size,
- max_sequence_length=self.max_sequence_length,
- position_embedding_type=position_embedding_type,
- scatter_to_sequence_parallel=scatter_embedding_sequence_parallel,
- tp_group=self.pg_collection.tp,
- )
-
- if self.position_embedding_type == 'rope':
- self.rotary_pos_emb = RotaryEmbedding(
- kv_channels=self.config.kv_channels,
- rotary_percent=rotary_percent,
- seq_len_interpolation_factor=seq_len_interpolation_factor,
- rotary_base=rotary_base,
- use_cpu_initialization=self.config.use_cpu_initialization,
- cp_group=self.pg_collection.cp,
- )
-
- self.decoder = build_module(
- mamba_stack_spec,
- self.config,
- pre_process=self.pre_process,
- layer_type_list=layer_type_list,
- pp_layer_offset=layer_offset,
- post_process=self.post_process,
- dtype=config.params_dtype,
- pg_collection=self.pg_collection,
- )
-
- # MTP block - uses mtp_block_spec from mamba_stack_spec.submodules
- if self.mtp_process:
- mamba_submodules = mamba_stack_spec.submodules
- mtp_block_spec = mamba_submodules.mtp_block_spec
- assert mtp_block_spec is not None, (
- "MTP pattern specified but mtp_block_spec is None in mamba_stack_spec.submodules. "
- "Ensure mamba_stack_spec includes mtp_block_spec for MTP support."
- )
-
- self.mtp = MultiTokenPredictionBlock(
- config=self.config,
- spec=mtp_block_spec,
- pg_collection=self.pg_collection,
- vp_stage=self.vp_stage,
- mtp_layer_pattern=self.mtp_pattern,
- mtp_num_depths=self.mtp_num_depths,
- mamba_submodules=mamba_submodules,
- )
-
- # Output
- if post_process or self.mtp_process:
- self.output_layer = tensor_parallel.ColumnParallelLinear(
- config.hidden_size,
- self.vocab_size,
- config=config,
- init_method=(
- config.embedding_init_method
- if config.use_mup and not self.share_embeddings_and_output_weights
- else config.init_method
- ),
- bias=False,
- skip_bias_add=False,
- gather_output=not self.parallel_output,
- skip_weight_param_allocation=self.pre_process
- and self.share_embeddings_and_output_weights,
- tp_group=self.pg_collection.tp,
- )
-
- if self.pre_process or self.post_process or self.mtp_process:
- self.setup_embeddings_and_output_layer()
-
- for name, module in self.named_modules():
- if hasattr(module, 'finish_init'):
- quant_config = get_quant_config_or_none(name, self.config.quant_recipe)
- module.finish_init(quant_config)
-
- def set_input_tensor(self, input_tensor: Tensor) -> None:
- """Sets input tensor to the model.
-
- See megatron.model.transformer.set_input_tensor()
-
- Args:
- input_tensor (Tensor): Sets the input tensor for the model.
- """
- # This is usually handled in schedules.py but some inference code still
- # gives us non-lists or None
- if not isinstance(input_tensor, list):
- input_tensor = [input_tensor]
-
- assert len(input_tensor) == 1, 'input_tensor should only be length 1 for gpt/bert'
- self.decoder.set_input_tensor(input_tensor[0])
-
- def forward(
- self,
- input_ids: Tensor,
- position_ids: Tensor,
- attention_mask: Tensor,
- decoder_input: Tensor = None,
- labels: Tensor = None,
- inference_context: BaseInferenceContext = None,
- runtime_gather_output: Optional[bool] = None,
- *,
- inference_params: Optional[BaseInferenceContext] = None,
- loss_mask: Optional[Tensor] = None,
- packed_seq_params: Optional[PackedSeqParams] = None,
- padding_mask: Optional[Tensor] = None,
- is_spec_decode: Optional[bool] = None,
- ) -> Tensor:
- """Forward function of the Mamba model. This function passes the input tensors
- through the embedding layer, and then the decoder and finally into the post
- processing layer (optional).
-
- It either returns the Loss values if labels are given or the final hidden units
- """
- # If decoder_input is provided (not None), then input_ids and position_ids are ignored.
- # Otherwise, apply embedding layer on input_ids and position_ids to get decoder_input.
-
- inference_context = deprecate_inference_params(inference_context, inference_params)
-
- in_inference_mode = inference_context is not None and not self.training
-
- if in_inference_mode:
- assert runtime_gather_output, "Inference must always gather TP logits"
-
- # Decoder embedding.
- if decoder_input is not None:
- pass
- elif self.pre_process:
- decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids)
-
- # Clear the outputs for padding tokens when using dynamic batching with
- # quantization scales to avoid corrupting amax calculations
- if (
- in_inference_mode
- and inference_context.is_dynamic_batching()
- and is_using_quantization_scales(self.config)
- ):
- decoder_input[inference_context.padding_slice] = 0.0
- else:
- # intermediate stage of pipeline
- # decoder will get hidden_states from encoder.input_tensor
- decoder_input = None
-
- rotary_pos_emb = None
- if self.position_embedding_type == 'rope':
- rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(
- inference_context, self.decoder, decoder_input, self.config, packed_seq_params
- )
- rotary_pos_emb = self.rotary_pos_emb(
- rotary_seq_len,
- packed_seq=packed_seq_params is not None and packed_seq_params.qkv_format == 'thd',
- )
-
- # Wrap decoder_input to allow the decoder (MambaBlock) to delete the
- # reference held by this caller function, enabling early garbage collection
- # for inference.
- if in_inference_mode:
- decoder_input = WrappedTensor(decoder_input)
-
- # The following assert will currently fail when running inference.
- # Commented out for now.
- # TODO (duncan/rwaleffe): (1) confirm that the externally-generated
- # attention mask is not needed and is ignored by the model in
- # inference mode, (2) reduce the size of the externally-generated
- # attention mask to prevent CPU OOM (as we did for training), (3)
- # force the attention mask passed to the model in inference mode to
- # be None, so this assert will succeed.
- # assert attention_mask is None, "The attention mask is ignored and should be set to None"
-
- # Run decoder.
- hidden_states = self.decoder(
- hidden_states=decoder_input,
- attention_mask=attention_mask,
- inference_context=inference_context,
- rotary_pos_emb=rotary_pos_emb,
- packed_seq_params=packed_seq_params,
- padding_mask=padding_mask,
- )
-
- output_weight = None
- if self.share_embeddings_and_output_weights:
- output_weight = self.shared_embedding_or_output_weight()
-
- # Check if speculative decoding is active. When it is, MTP must be
- # computed *after* verification so that it is conditioned on verified
- # tokens rather than stale speculative tokens from the previous step.
- if is_spec_decode is None:
- is_spec_decode = (
- in_inference_mode
- and inference_context.is_dynamic_batching()
- and inference_context.num_speculative_tokens > 0
- )
-
- mtp_forward_ran = self.mtp_process and not (in_inference_mode or is_spec_decode)
- if mtp_forward_ran:
- hidden_states = self.mtp(
- input_ids=input_ids,
- position_ids=position_ids,
- hidden_states=hidden_states,
- attention_mask=attention_mask,
- inference_params=inference_params,
- rotary_pos_emb=rotary_pos_emb,
- packed_seq_params=packed_seq_params,
- embedding=self.embedding,
- )
-
- if not self.post_process:
- return hidden_states
-
- if self.config.mtp_num_layers is not None and self.mtp_process:
- assert self.config.mtp_num_layers > 0
- if in_inference_mode or is_spec_decode:
- self._decoder_hidden_states_cache = hidden_states
- else:
- hidden_states = process_mtp_loss(
- hidden_states=hidden_states,
- labels=labels,
- loss_mask=loss_mask,
- output_layer=self.output_layer,
- output_weight=output_weight,
- runtime_gather_output=runtime_gather_output,
- is_training=self.training,
- compute_language_model_loss=self.compute_language_model_loss,
- config=self.config,
- cp_group=self.pg_collection.cp,
- packed_seq_params=packed_seq_params,
- scale_logits_fn=self._scale_logits if self.config.use_mup else None,
+ if mamba_stack_spec is not None:
+ if 'hybrid_stack_spec' in kwargs or (args and len(args) >= 2):
+ raise ValueError(
+ "Cannot specify both hybrid_stack_spec and mamba_stack_spec. "
+ "mamba_stack_spec has been deprecated; use hybrid_stack_spec instead."
)
- sequence_parallel_override = False
- if in_inference_mode and inference_context.config.materialize_only_last_token_logits:
- if inference_context.is_static_batching():
- hidden_states = hidden_states[-1:, :, :]
- else:
- if self.output_layer.sequence_parallel:
- # Perform the sequence parallel gather here instead of after the output layer
- # because we need to slice the last token logits from the full view of the
- # packed logits across all requests.
- hidden_states = gather_from_sequence_parallel_region(
- hidden_states, group=self.pg_collection.tp
- )
- self.output_layer.sequence_parallel = False
- sequence_parallel_override = True
-
- # Reshape [S, B, H] (with B=1) to [1, S, H] for logit extraction,
- # then back to [S', B, H] for the output layer.
- reshaped = hidden_states.squeeze(1).unsqueeze(0)
- hidden_states = inference_context.last_token_logits(reshaped).unsqueeze(1)
-
- logits, _ = self.output_layer(
- hidden_states, weight=output_weight, runtime_gather_output=runtime_gather_output
- )
- logits = self._scale_logits(logits)
-
- # Restore sequence parallel execution to the output layer if necessary.
- if sequence_parallel_override:
- assert (
- in_inference_mode
- and inference_context.is_dynamic_batching()
- and inference_context.config.materialize_only_last_token_logits
- )
- self.output_layer.sequence_parallel = True
-
- if labels is None:
- # [s b h] => [b s h]
- return logits.transpose(0, 1).contiguous()
-
- loss = self.compute_language_model_loss(labels, logits)
-
- return loss
-
- @torch.inference_mode()
- def compute_mtp_single_step(
- self,
- hidden_states: Tensor,
- next_token_ids: Tensor,
- position_ids: Tensor,
- depth: int,
- runtime_gather_output: bool = True,
- ) -> tuple:
- """Compute a single MTP depth for speculative decoding.
-
- This is called after speculative token verification to compute MTP
- predictions conditioned on verified tokens only.
-
- Args:
- hidden_states (Tensor): Hidden states at last accepted positions [N, 1, H].
- next_token_ids (Tensor): Correct next token IDs [1, N].
- position_ids (Tensor): Position IDs for the next tokens [1, N].
- depth (int): MTP depth index (0-indexed).
- runtime_gather_output (bool): Whether to gather output across TP.
-
- Returns:
- tuple: (new_hidden_states [N, 1, H], logits [N, 1, vocab_size]).
- """
- layer_idx = 0 if self.mtp.mtp_use_repeated_layer else depth
- mtp_hidden = self.mtp.layers[layer_idx].forward_single_position(
- hidden_states=hidden_states,
- next_token_ids=next_token_ids,
- position_ids=position_ids,
- embedding=self.embedding,
- )
-
- output_weight = None
- if self.share_embeddings_and_output_weights:
- output_weight = self.shared_embedding_or_output_weight()
-
- logits, _ = self.output_layer(
- mtp_hidden, weight=output_weight, runtime_gather_output=runtime_gather_output
- )
- logits = self._scale_logits(logits)
-
- return mtp_hidden, logits
+ kwargs['hybrid_stack_spec'] = mamba_stack_spec
+ super().__init__(*args, **kwargs)
diff --git a/megatron/core/models/mimo/comm/__init__.py b/megatron/core/models/mimo/comm/__init__.py
new file mode 100644
index 00000000000..26496bfed70
--- /dev/null
+++ b/megatron/core/models/mimo/comm/__init__.py
@@ -0,0 +1 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
diff --git a/megatron/core/models/mimo/comm/colocated_communicator.py b/megatron/core/models/mimo/comm/colocated_communicator.py
new file mode 100644
index 00000000000..4c43dcdf3cd
--- /dev/null
+++ b/megatron/core/models/mimo/comm/colocated_communicator.py
@@ -0,0 +1,325 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+import logging
+from dataclasses import dataclass
+from enum import Enum
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.distributed as dist
+
+from megatron.core.hyper_comm_grid import HyperCommGrid
+
+
+@dataclass
+class SliceInfo:
+ """Batch dimension slice information for a rank's data partition."""
+
+ start: int
+ size: int
+
+
+class BridgeDirection(str, Enum):
+ """Which side of the bridge scales up, if any.
+
+ ``FAN_IN`` — src has more DP replicas than dest; forward all-gathers
+ src outputs along the batch dim, backward narrows the sibling dest
+ gradient down to this src rank's slot.
+
+ ``FAN_OUT`` — dest has more DP replicas; forward narrows, backward
+ all-gathers across the sibling dest DP ranks (the adjoint of narrow
+ is not zero-pad-and-scatter because every dest rank consumes a
+ different slice of the same src activation).
+
+ ``EQUAL`` — matching DP; the bridge is a pure passthrough.
+ """
+
+ FAN_IN = "fan_in"
+ FAN_OUT = "fan_out"
+ EQUAL = "equal"
+
+
+class ColocatedBridgeCommunicator:
+ """Bridges tensors between colocated modules with different TP/DP layouts.
+
+ Default ``dim_mapping`` assumes 3D ``(b, s, h)``. Callers bridging
+ ``MimoModel``'s pre-flattened ``(s*b, h)`` encoder output should pass
+ ``dim_mapping={'b': 0, 'h': 1}``; this relies on a uniform token count per
+ sample so dim 0 divides evenly by the DP scale.
+
+ Precondition: the input must be TP-replicated across the src TP group —
+ i.e. all TP ranks inside a src DP replica hold the same tensor on the
+ batch dim. The bridge never gathers along TP; violating this silently
+ produces wrong results.
+ """
+
+ def __init__(
+ self,
+ src_grid: HyperCommGrid,
+ dest_grid: HyperCommGrid,
+ src_module_name: str = "src",
+ dest_module_name: str = "dest",
+ dim_mapping: Optional[Dict[str, int]] = None,
+ ):
+ self.src_grid = src_grid
+ self.dest_grid = dest_grid
+ self.src_module_name = src_module_name
+ self.dest_module_name = dest_module_name
+ self.dim_mapping = dim_mapping or {'b': 0, 's': 1, 'h': 2}
+ self.current_rank = dist.get_rank()
+
+ self._validate_grids()
+ self._extract_parallelism_info()
+ self._build_rank_mappings()
+
+ # At most one direction is active; fan-in and fan-out are mutually
+ # exclusive (one of ``src_dp / dest_dp`` is >1, the other is 1).
+ # Equal DP uses no collective at all. Unify behind a single
+ # ``gather_pg`` + ``direction`` + ``scale`` rather than a fan-in
+ # and fan-out pair of attributes.
+ self.gather_pg: Optional[dist.ProcessGroup] = None
+ self.gather_group_ranks: List[List[int]] = []
+
+ if self.src_dp_size > self.dest_dp_size:
+ self.direction = BridgeDirection.FAN_IN
+ self.scale = self.src_dp_size // self.dest_dp_size
+ self.gather_group_ranks = self._build_gather_groups(
+ iter_size=self.dest_dp_size,
+ sibling_tp_size=self.src_tp_size,
+ scale=self.scale,
+ rank_to_pos=self.rank_to_src_pos,
+ )
+ self.gather_pg, _ = dist.new_subgroups_by_enumeration(
+ self.gather_group_ranks, backend='nccl'
+ )
+ elif self.dest_dp_size > self.src_dp_size:
+ self.direction = BridgeDirection.FAN_OUT
+ self.scale = self.dest_dp_size // self.src_dp_size
+ self.gather_group_ranks = self._build_gather_groups(
+ iter_size=self.src_dp_size,
+ sibling_tp_size=self.dest_tp_size,
+ scale=self.scale,
+ rank_to_pos=self.rank_to_dest_pos,
+ )
+ self.gather_pg, _ = dist.new_subgroups_by_enumeration(
+ self.gather_group_ranks, backend='nccl'
+ )
+ else:
+ self.direction = BridgeDirection.EQUAL
+ self.scale = 1
+
+ logging.info(
+ f"[Rank {self.current_rank}] ColocatedBridgeCommunicator: "
+ f"{src_module_name}({self.src_tp_size}TP/{self.src_dp_size}DP) -> "
+ f"{dest_module_name}({self.dest_tp_size}TP/{self.dest_dp_size}DP), "
+ f"direction={self.direction.value}, scale={self.scale}"
+ )
+
+ def _validate_grids(self):
+ if self.src_grid.size != self.dest_grid.size:
+ raise ValueError(
+ f"Grids must span same number of ranks: "
+ f"src={self.src_grid.size}, dest={self.dest_grid.size}"
+ )
+
+ if self.src_grid.rank_offset != self.dest_grid.rank_offset:
+ raise ValueError(
+ f"Grids must have same rank offset: "
+ f"src={self.src_grid.rank_offset}, dest={self.dest_grid.rank_offset}"
+ )
+
+ # Per-grid dim checks: tp/dp required; pp and cp (if present) must be 1.
+ # CP>1 also corrupts dp_idx when iterating get_rank_enum(['tp']) groups.
+ for name, grid in [("src", self.src_grid), ("dest", self.dest_grid)]:
+ for required in ('tp', 'dp'):
+ if required not in grid.dim_names:
+ raise ValueError(
+ f"{name} grid must have '{required}' dimension, "
+ f"got dim_names={grid.dim_names}"
+ )
+ for singleton in ('pp', 'cp'):
+ if singleton in grid.dim_names:
+ size = grid.shape[grid.dim_names.index(singleton)]
+ if size != 1:
+ raise ValueError(
+ f"{name} {singleton.upper()} must be 1 for "
+ f"ColocatedBridgeCommunicator, got {size}"
+ )
+
+ src_dp = self.src_grid.shape[self.src_grid.dim_names.index('dp')]
+ dest_dp = self.dest_grid.shape[self.dest_grid.dim_names.index('dp')]
+ if src_dp % dest_dp != 0 and dest_dp % src_dp != 0:
+ raise ValueError(
+ f"DP sizes must be evenly divisible: src_dp={src_dp}, dest_dp={dest_dp}"
+ )
+
+ def _extract_parallelism_info(self):
+ self.src_tp_size = self.src_grid.shape[self.src_grid.dim_names.index('tp')]
+ self.src_dp_size = self.src_grid.shape[self.src_grid.dim_names.index('dp')]
+ self.dest_tp_size = self.dest_grid.shape[self.dest_grid.dim_names.index('tp')]
+ self.dest_dp_size = self.dest_grid.shape[self.dest_grid.dim_names.index('dp')]
+
+ def _build_rank_mappings(self):
+ self.rank_to_src_pos: Dict[int, Tuple[int, int]] = {}
+ self.rank_to_dest_pos: Dict[int, Tuple[int, int]] = {}
+
+ src_tp_groups = self.src_grid.get_rank_enum(['tp'])
+ for dp_idx, tp_group in enumerate(src_tp_groups):
+ for tp_idx, rank in enumerate(tp_group):
+ self.rank_to_src_pos[rank] = (dp_idx, tp_idx)
+
+ dest_tp_groups = self.dest_grid.get_rank_enum(['tp'])
+ for dp_idx, tp_group in enumerate(dest_tp_groups):
+ for tp_idx, rank in enumerate(tp_group):
+ self.rank_to_dest_pos[rank] = (dp_idx, tp_idx)
+
+ @staticmethod
+ def _build_gather_groups(
+ iter_size: int, sibling_tp_size: int, scale: int, rank_to_pos: Dict[int, Tuple[int, int]]
+ ) -> List[List[int]]:
+ """Build ``iter_size * sibling_tp_size`` gather groups of ``scale`` ranks.
+
+ For each slot on the "iterating" side and each TP shard on the
+ sibling side, collect the ``scale`` sibling ranks whose DP indices
+ map into that slot. Append order equals group-local-rank order,
+ which ``all_gather_into_tensor`` uses to concatenate outputs — do
+ not sort.
+ """
+ groups: List[List[int]] = []
+ for iter_idx in range(iter_size):
+ sibling_dp_indices = range(iter_idx * scale, (iter_idx + 1) * scale)
+ for sibling_tp_idx in range(sibling_tp_size):
+ group_ranks = []
+ for sibling_dp_idx in sibling_dp_indices:
+ for rank, (dp, tp) in rank_to_pos.items():
+ if dp == sibling_dp_idx and tp == sibling_tp_idx:
+ group_ranks.append(rank)
+ break
+ groups.append(group_ranks)
+ return groups
+
+ def is_fan_in(self) -> bool:
+ """True if src DP > dest DP (forward all-gathers)."""
+ return self.direction is BridgeDirection.FAN_IN
+
+ def is_fan_out(self) -> bool:
+ """True if src DP < dest DP (forward narrows)."""
+ return self.direction is BridgeDirection.FAN_OUT
+
+ def get_slice_info(self, batch_size: int) -> SliceInfo:
+ """Compute this rank's slice of ``batch_size`` on the narrowing side.
+
+ For FAN_OUT this is the forward narrow; for FAN_IN it is the
+ backward narrow against the post-gather batch. EQUAL returns the
+ identity slice.
+
+ Raises ``ValueError`` if ``batch_size`` is not divisible by ``scale``.
+ """
+ if self.direction is BridgeDirection.EQUAL:
+ return SliceInfo(start=0, size=batch_size)
+ self._check_divisible(batch_size)
+ if self.direction is BridgeDirection.FAN_OUT:
+ dp_idx = self.rank_to_dest_pos[self.current_rank][0]
+ else: # FAN_IN
+ dp_idx = self.rank_to_src_pos[self.current_rank][0]
+ slot = dp_idx % self.scale
+ slice_size = batch_size // self.scale
+ return SliceInfo(start=slot * slice_size, size=slice_size)
+
+ def _check_divisible(self, batch_size: int) -> None:
+ if batch_size % self.scale != 0:
+ raise ValueError(
+ f"ColocatedBridgeCommunicator: batch dim size {batch_size} is "
+ f"not divisible by {self.direction.value} scale={self.scale}."
+ )
+
+ def communicate(self, tensor: torch.Tensor) -> torch.Tensor:
+ """Transform ``tensor`` from src TP/DP layout to dest TP/DP layout.
+
+ Raises ``ValueError`` when FAN_OUT and the batch dim is not
+ divisible by ``scale``; FAN_IN only slices on the backward pass
+ and re-checks via ``get_slice_info`` there.
+ """
+ if self.direction is BridgeDirection.FAN_OUT:
+ self._check_divisible(tensor.shape[self.dim_mapping['b']])
+ return _ColocatedCommunicate.apply(tensor, self)
+
+ def destroy(self) -> None:
+ """Release the NCCL subgroup created by this communicator.
+
+ NCCL caps concurrent communicators; long-lived or repeated
+ construction leaks PGs without this call.
+ """
+ if self.gather_pg is not None:
+ dist.destroy_process_group(self.gather_pg)
+ self.gather_pg = None
+
+
+class _ColocatedCommunicate(torch.autograd.Function):
+ """Autograd function for colocated communication with correct backward pass."""
+
+ @staticmethod
+ def forward(ctx, tensor: torch.Tensor, comm: ColocatedBridgeCommunicator) -> torch.Tensor:
+ """Reshape the batch dim across the bridge: narrow on fan-out, all-gather on fan-in."""
+ ctx.comm = comm
+ ctx.batch_dim = comm.dim_mapping['b']
+
+ if comm.direction is BridgeDirection.FAN_OUT:
+ # Narrow this rank's slice out of the full src batch.
+ slice_info = comm.get_slice_info(tensor.shape[ctx.batch_dim])
+ return tensor.narrow(ctx.batch_dim, slice_info.start, slice_info.size).contiguous()
+
+ if comm.direction is BridgeDirection.FAN_IN:
+ # All-gather sibling src outputs into a single full-batch tensor.
+ return _all_gather_along_batch_dim(tensor, comm.gather_pg, ctx.batch_dim)
+
+ # EQUAL: pure passthrough.
+ return tensor.contiguous()
+
+ @staticmethod
+ def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]:
+ """Adjoint of forward: narrow for fan-in, all-gather for fan-out.
+
+ Fan-out's forward is ``narrow``, whose naive adjoint is zero-pad.
+ That would leave each src rank with only its own dest rank's
+ slice of the gradient, missing the contributions from every
+ other dest rank that consumed a different slice of the same src
+ activation. Instead we all-gather across the fan-out sibling
+ group, reconstructing the full src-batch gradient (symmetric
+ with the fan-in forward's all-gather).
+ """
+ comm = ctx.comm
+ batch_dim = ctx.batch_dim
+
+ if comm.direction is BridgeDirection.FAN_OUT:
+ return _all_gather_along_batch_dim(grad_output, comm.gather_pg, batch_dim), None
+
+ if comm.direction is BridgeDirection.FAN_IN:
+ slice_info = comm.get_slice_info(grad_output.shape[batch_dim])
+ return (
+ grad_output.narrow(batch_dim, slice_info.start, slice_info.size).contiguous(),
+ None,
+ )
+
+ return grad_output.contiguous(), None
+
+
+def _all_gather_along_batch_dim(
+ tensor: torch.Tensor, group: dist.ProcessGroup, batch_dim: int
+) -> torch.Tensor:
+ """All-gather ``tensor`` along an arbitrary batch dim into a single tensor.
+
+ ``all_gather_into_tensor`` concatenates along dim 0, so when the
+ batch dim is not 0 we move it, gather, then restore.
+ """
+ world_size = dist.get_world_size(group)
+ src = tensor.contiguous()
+ if batch_dim != 0:
+ src = src.movedim(batch_dim, 0).contiguous()
+ out_shape = list(src.shape)
+ out_shape[0] *= world_size
+ out = torch.empty(out_shape, dtype=tensor.dtype, device=tensor.device)
+ dist.all_gather_into_tensor(out, src, group=group)
+ if batch_dim != 0:
+ out = out.movedim(0, batch_dim).contiguous()
+ return out
diff --git a/megatron/core/models/mimo/config/base_configs.py b/megatron/core/models/mimo/config/base_configs.py
index a92484a5a48..0eda09465e0 100644
--- a/megatron/core/models/mimo/config/base_configs.py
+++ b/megatron/core/models/mimo/config/base_configs.py
@@ -23,9 +23,11 @@ class MimoModelConfig:
in the input_ids to insert the modality embeddings at the correct positions.
module_to_grid_map (Optional[Dict[str, HyperCommGrid]]):
Dictionary mapping module keys (e.g., "vision", "language") to their
- corresponding HyperCommGrid configurations for non-colocated pipeline
- parallelism. The language model must use the key MIMO_LANGUAGE_MODULE_KEY.
- When None, all modules are assumed to be colocated on the same ranks.
+ corresponding HyperCommGrid configurations. The language model must use
+ the key MIMO_LANGUAGE_MODULE_KEY.
+ When grids span the same ranks → colocated (same or different TP/DP).
+ When grids span disjoint ranks → non-colocated (pipeline parallel).
+ When None → colocated with legacy global parallel_state.
kv_format (str):
Key-value format for attention: "sbhd" (seq-batch-head-dim) or "thd" (total-head-dim).
Default is "sbhd".
@@ -43,3 +45,18 @@ class MimoModelConfig:
special_token_ids: Dict[str, int] = field(default_factory=dict)
module_to_grid_map: Optional[Dict[str, HyperCommGrid]] = None
kv_format: str = "sbhd"
+
+ def __post_init__(self):
+ if not self.module_to_grid_map:
+ return
+ # Local import avoids circular imports at dataclass-module import time.
+ from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY
+
+ expected_keys = set(self.modality_submodules_spec.keys()) | {MIMO_LANGUAGE_MODULE_KEY}
+ grid_keys = set(self.module_to_grid_map.keys())
+ if grid_keys != expected_keys:
+ raise ValueError(
+ f"module_to_grid_map keys must match modality module names + "
+ f"'{MIMO_LANGUAGE_MODULE_KEY}'. Missing: {expected_keys - grid_keys}, "
+ f"Extra: {grid_keys - expected_keys}"
+ )
diff --git a/megatron/core/models/mimo/config/role.py b/megatron/core/models/mimo/config/role.py
index 77c2512e8e6..411791f1e5c 100644
--- a/megatron/core/models/mimo/config/role.py
+++ b/megatron/core/models/mimo/config/role.py
@@ -5,7 +5,7 @@
import logging
from dataclasses import dataclass, field
from enum import Enum
-from typing import Dict, List
+from typing import Dict, List, Optional
import torch.distributed as dist
@@ -24,22 +24,17 @@ class ModuleLayout(Enum):
Determines how modules are distributed across ranks and which
forward path is used.
- UNIFIED: No module_to_grid_map. All modules share same ranks and
- parallelism. Uses the unified forward path (_forward_all_modules).
+ COLOCATED: All modules share the same ranks. Covers both legacy
+ (no grid map, global parallel_state) and heterogeneous TP/DP
+ (grid map with overlapping ranks). Uses _forward_all_modules.
NON_COLOCATED: module_to_grid_map is set with non-overlapping rank
ranges. Each rank runs EITHER encoder(s) OR the language model.
Uses role-based dispatch with separate forward paths.
-
- COLOCATED: (future) module_to_grid_map is set with overlapping rank
- ranges. Encoder(s) and language model share ranks but have
- different parallelism configs. Uses role-based dispatch but
- allows both module types on the same rank.
"""
- UNIFIED = "unified"
- NON_COLOCATED = "non_colocated"
COLOCATED = "colocated"
+ NON_COLOCATED = "non_colocated"
@dataclass
@@ -70,50 +65,50 @@ class RankRole:
"""
modules: Dict[str, ModuleStageInfo] = field(default_factory=dict)
- mode: ModuleLayout = ModuleLayout.UNIFIED
+ mode: ModuleLayout = ModuleLayout.COLOCATED
+
+ @classmethod
+ def build(
+ cls,
+ modality_module_names: List[str],
+ module_to_grid_map: Optional[Dict[str, 'HyperCommGrid']] = None,
+ ) -> 'RankRole':
+ """Build a RankRole, dispatching by whether grids share ranks.
+
+ No grid map or all grids span the same ranks → COLOCATED.
+ Grids differ → NON_COLOCATED with PP-stage info per module.
+ """
+ if module_to_grid_map is None or cls._all_grids_colocated(module_to_grid_map):
+ return cls._colocated(modality_module_names)
+ return cls._from_grid_map(module_to_grid_map)
+
+ @staticmethod
+ def _all_grids_colocated(module_to_grid_map: Dict[str, 'HyperCommGrid']) -> bool:
+ grids = list(module_to_grid_map.values())
+ first = grids[0]
+ return all(g.rank_offset == first.rank_offset and g.size == first.size for g in grids[1:])
@classmethod
- def unified(cls, module_names: List[str]) -> 'RankRole':
- """Create a role for the unified case: every module, first+last stage."""
+ def _colocated(cls, modality_module_names: List[str]) -> 'RankRole':
+ """Colocated layout: every module on every rank, PP=1."""
+ all_module_names = list(modality_module_names) + [MIMO_LANGUAGE_MODULE_KEY]
return cls(
modules={
name: ModuleStageInfo(is_first_stage=True, is_last_stage=True)
- for name in module_names
+ for name in all_module_names
},
- mode=ModuleLayout.UNIFIED,
+ mode=ModuleLayout.COLOCATED,
)
@classmethod
- def from_grid_map(
- cls, module_to_grid_map: Dict[str, HyperCommGrid], modality_module_names: List[str]
- ) -> 'RankRole':
- """Create a role from a module-to-grid mapping for non-colocated PP.
-
- Determines which modules the current rank participates in and its
- pipeline stage position within each module.
+ def _from_grid_map(cls, module_to_grid_map: Dict[str, HyperCommGrid]) -> 'RankRole':
+ """Non-colocated role for this rank from a module-to-grid mapping.
- Args:
- module_to_grid_map: Dict mapping module names to HyperCommGrid objects.
- Must contain keys matching modality_module_names + MIMO_LANGUAGE_MODULE_KEY.
- modality_module_names: List of modality module names (e.g., ["images", "audio"]).
-
- Returns:
- RankRole for the current rank.
+ Grid map keys are validated by ``MimoModelConfig.__post_init__``.
Raises:
- ValueError: If grid map keys don't match expected module names.
RuntimeError: If current rank is not in any module grid.
"""
- # Validate keys
- expected_keys = set(modality_module_names) | {MIMO_LANGUAGE_MODULE_KEY}
- grid_keys = set(module_to_grid_map.keys())
- if grid_keys != expected_keys:
- raise ValueError(
- f"module_to_grid_map keys must match modality module names + "
- f"'{MIMO_LANGUAGE_MODULE_KEY}'. Missing: {expected_keys - grid_keys}, "
- f"Extra: {grid_keys - expected_keys}"
- )
-
current_rank = dist.get_rank()
modules = {}
@@ -131,7 +126,7 @@ def from_grid_map(
is_first = pp_rank == 0
is_last = pp_rank == pp_size - 1
logger.info(
- f"[RankRole.from_grid_map] Rank {current_rank}: module={module_name}, "
+ f"[RankRole._from_grid_map] Rank {current_rank}: module={module_name}, "
f"pp_rank={pp_rank}/{pp_size}, is_first_stage={is_first}, is_last_stage={is_last}"
)
modules[module_name] = ModuleStageInfo(is_first_stage=is_first, is_last_stage=is_last)
diff --git a/megatron/core/models/mimo/model/base.py b/megatron/core/models/mimo/model/base.py
index b1c12f521c3..372c20b4e8e 100644
--- a/megatron/core/models/mimo/model/base.py
+++ b/megatron/core/models/mimo/model/base.py
@@ -7,6 +7,7 @@
import torch
from megatron.core.distributed import DistributedDataParallel
+from megatron.core.models.mimo.comm.colocated_communicator import ColocatedBridgeCommunicator
from megatron.core.models.mimo.config import MimoModelConfig
from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY, ModuleLayout, RankRole
from megatron.core.models.mimo.partition.utils import PartitionAdapter, PartitionConfig
@@ -59,10 +60,12 @@ def __init__(self, mimo_config: MimoModelConfig, cp_group=None, tp_group=None) -
self.mimo_config = mimo_config
modality_names = list(mimo_config.modality_submodules_spec.keys())
- if mimo_config.module_to_grid_map:
- self.role = RankRole.from_grid_map(mimo_config.module_to_grid_map, modality_names)
- else:
- self.role = RankRole.unified(modality_names + [MIMO_LANGUAGE_MODULE_KEY])
+ self.colocated_comms = {}
+ self.role = RankRole.build(modality_names, mimo_config.module_to_grid_map)
+ if self.role.mode is ModuleLayout.COLOCATED and mimo_config.module_to_grid_map:
+ # Per-encoder bridge needed iff modules share ranks but may differ
+ # in TP/DP within those ranks.
+ self._build_colocated_communicators()
# Use special token IDs from the config
self.special_token_ids = (
@@ -295,9 +298,14 @@ def get_text_embeddings(
batch_idx, seq_idx = text_mask.nonzero(as_tuple=True)
input_ids_text = input_ids[batch_idx, seq_idx].unsqueeze(0)
- position_ids_text = (
- position_ids[batch_idx, seq_idx].unsqueeze(0) if position_ids is not None else None
- )
+ if position_ids is None:
+ position_ids_text = None
+ elif position_ids.dim() == 3:
+ # Multimodal RoPE can carry [rope_dim, batch, seq] ids. Text
+ # embedding lookup only needs a single absolute position channel.
+ position_ids_text = position_ids[0, batch_idx, seq_idx].unsqueeze(0)
+ else:
+ position_ids_text = position_ids[batch_idx, seq_idx].unsqueeze(0)
text_embeddings = (
unwrap_model(self.language_model)
@@ -358,7 +366,7 @@ def forward(
# Get any tensors passed via set_input_tensor
input_tensors = getattr(self, 'input_tensors', None)
- if self.role.mode == ModuleLayout.UNIFIED:
+ if self.role.mode == ModuleLayout.COLOCATED:
return self._forward_all_modules(
input_ids,
position_ids,
@@ -371,7 +379,7 @@ def forward(
if self.role.mode == ModuleLayout.NON_COLOCATED:
if self.role.has_modality_modules:
- return self._forward_encoders(modality_inputs, input_tensors), loss_mask
+ return self._forward_encoders(input_ids, modality_inputs, input_tensors), loss_mask
if self.role.has_language_module:
return (
@@ -387,6 +395,7 @@ def forward(
def _forward_encoders(
self,
+ input_ids: Optional[torch.Tensor],
modality_inputs: Optional[Dict[str, Dict[str, Any]]],
input_tensors: Optional[Dict[str, torch.Tensor]],
) -> Dict[str, torch.Tensor]:
@@ -406,16 +415,88 @@ def _forward_encoders(
continue
submodule = self.modality_submodules[encoder_name]
- output = submodule.forward(
- encoder_inputs=modality_inputs.get(encoder_name) if modality_inputs else None,
- hidden_states=input_tensors.get(encoder_name) if input_tensors else None,
- )
+ encoder_inputs = modality_inputs.get(encoder_name) if modality_inputs else None
+ hidden_states = input_tensors.get(encoder_name) if input_tensors else None
+ output = submodule.forward(encoder_inputs=encoder_inputs, hidden_states=hidden_states)
+ if output is None and encoder_inputs is None and hidden_states is None:
+ if self._has_encoder_tokens(input_ids, encoder_name):
+ raise RuntimeError(
+ f"{encoder_name} inputs are missing, but matching special tokens exist"
+ )
+ output = self._empty_encoder_output(encoder_name)
if output is not None:
+ self._attach_modality_split_sizes(output, input_ids, encoder_name)
outputs[encoder_name] = output
return outputs
+ def _attach_modality_split_sizes(
+ self, output: torch.Tensor, input_ids: Optional[torch.Tensor], encoder_name: str
+ ) -> None:
+ """Annotate flat modality outputs with per-sample split sizes for bridge fan-out.
+
+ Only attaches when per-sample token counts are non-uniform. Uniform counts
+ give equal splits, which the bridge's ``torch.tensor_split`` fallback
+ already produces, so the metadata would be a no-op.
+
+ TODO(mimo): non-uniform per-sample counts in fan-in (encoder DP > LM DP)
+ are not supported. Multiple encoder ranks contribute slices to a single
+ LM peer, and the receiver-side ``torch.cat`` path in BridgeCommunicator
+ has no metadata channel today, so per-sample boundaries are lost on the
+ LM rank. Lift this by routing per-sample sizes through the bridge
+ alongside the activations and adding a sample-aligned concat path.
+ """
+ token_id = self.special_token_ids.get(encoder_name)
+ if token_id is None or input_ids is None or output.ndim != 2 or input_ids.size(0) <= 1:
+ return
+
+ split_sizes = (input_ids == token_id).sum(dim=1).to(torch.long).tolist()
+ if sum(split_sizes) != output.size(0):
+ return
+ if len(set(split_sizes)) <= 1:
+ # Uniform counts — tensor_split fallback gives the same result.
+ return
+
+ if self.role.mode is ModuleLayout.NON_COLOCATED:
+ grid_map = self.mimo_config.module_to_grid_map
+ encoder_grid = grid_map[encoder_name]
+ language_grid = grid_map[MIMO_LANGUAGE_MODULE_KEY]
+ encoder_dp = encoder_grid.shape[encoder_grid.dim_names.index("dp")]
+ language_dp = language_grid.shape[language_grid.dim_names.index("dp")]
+ assert encoder_dp <= language_dp, (
+ f"Bridge fan-out split metadata with non-uniform per-sample sizes "
+ f"requires encoder DP <= LM DP (got encoder='{encoder_name}' "
+ f"DP={encoder_dp}, LM DP={language_dp}). Fan-in with variable "
+ f"modality token counts is not supported yet — see TODO in "
+ f"_attach_modality_split_sizes."
+ )
+
+ output._mimo_bridge_split_sizes = split_sizes
+
+ def _has_encoder_tokens(self, input_ids: Optional[torch.Tensor], encoder_name: str) -> bool:
+ """Return whether the batch contains tokens for an encoder module."""
+ if input_ids is None or encoder_name not in self.special_token_ids:
+ return False
+ return bool((input_ids == self.special_token_ids[encoder_name]).any().item())
+
+ def _empty_encoder_output(self, encoder_name: str) -> torch.Tensor:
+ """Return the bridge payload for text-only non-colocated batches."""
+ language_config = self.mimo_config.language_model_spec.params['config']
+ hidden_size = getattr(language_config, 'hidden_size', None)
+ if hidden_size is None:
+ raise ValueError(
+ "Language model config must define hidden_size for empty modality output"
+ )
+
+ output_dtype = getattr(language_config, 'params_dtype', None) or torch.float32
+ return torch.empty(
+ (0, hidden_size),
+ device=torch.cuda.current_device(),
+ dtype=output_dtype,
+ requires_grad=True,
+ )
+
def _forward_language_module(
self,
input_ids: torch.Tensor,
@@ -461,8 +542,11 @@ def _forward_language_module(
)
lm_output = self.language_model(
+ # decoder_input replaces the embedding lookup, so input_ids is
+ # unused here; position_ids is still consumed by mRoPE in models
+ # such as Qwen3-VL.
input_ids=None,
- position_ids=None,
+ position_ids=position_ids,
decoder_input=combined_embeddings,
labels=labels,
attention_mask=attention_mask,
@@ -478,8 +562,10 @@ def _forward_language_module(
underlying_lm.set_input_tensor(hidden_states)
lm_output = self.language_model(
+ # Hidden states arrive via set_input_tensor; position_ids is
+ # still consumed by mRoPE on non-first PP stages.
input_ids=None,
- position_ids=None,
+ position_ids=position_ids,
decoder_input=None,
labels=labels,
attention_mask=attention_mask,
@@ -491,6 +577,47 @@ def _forward_language_module(
return lm_output
+ def _build_colocated_communicators(self):
+ grid_map = self.mimo_config.module_to_grid_map
+ if any(
+ 'tp' not in grid.dim_names or 'dp' not in grid.dim_names for grid in grid_map.values()
+ ):
+ logger.info(
+ "Skipping colocated communicator setup because module_to_grid_map "
+ "does not define TP/DP topology for every module."
+ )
+ return
+
+ lang_key = MIMO_LANGUAGE_MODULE_KEY
+ lang_grid = grid_map[lang_key]
+ for mod_name in self.mimo_config.modality_submodules_spec:
+ if mod_name == lang_key:
+ continue
+ self.colocated_comms[(mod_name, lang_key)] = ColocatedBridgeCommunicator(
+ src_grid=grid_map[mod_name],
+ dest_grid=lang_grid,
+ src_module_name=mod_name,
+ dest_module_name=lang_key,
+ dim_mapping={'b': 0, 'h': 1},
+ )
+
+ def destroy(self) -> None:
+ """Release process groups owned by this MimoModel."""
+ for comm in self.colocated_comms.values():
+ comm.destroy()
+ self.colocated_comms.clear()
+
+ def _apply_colocated_comms(self, modality_embeddings):
+ """Transform encoder embeddings from encoder TP/DP to LLM TP/DP layout."""
+ lang_key = MIMO_LANGUAGE_MODULE_KEY
+ for modality_name in list(modality_embeddings.keys()):
+ comm = self.colocated_comms.get((modality_name, lang_key))
+ if comm is not None:
+ modality_embeddings[modality_name] = comm.communicate(
+ modality_embeddings[modality_name]
+ )
+ return modality_embeddings
+
def _forward_all_modules(
self,
input_ids: torch.Tensor,
@@ -533,6 +660,10 @@ def _forward_all_modules(
f"Generated embeddings for {modality_name} with shape {embeddings.shape}"
)
+ # Apply colocated communication if configured (no-op when colocated_comms is empty)
+ if self.colocated_comms:
+ modality_embeddings = self._apply_colocated_comms(modality_embeddings)
+
# Get text embeddings
text_embeddings = self.get_text_embeddings(input_ids, position_ids, self.special_token_ids)
logger.debug(f"Generated text embeddings with shape {text_embeddings.shape}")
@@ -570,8 +701,11 @@ def _forward_all_modules(
# 5. Forward pass through language model
lm_output = self.language_model(
+ # decoder_input replaces the embedding lookup, so input_ids is
+ # unused here; position_ids is still consumed by mRoPE in models
+ # such as Qwen3-VL.
input_ids=None,
- position_ids=None,
+ position_ids=position_ids,
decoder_input=combined_embeddings,
labels=labels,
attention_mask=None,
diff --git a/megatron/core/models/mimo/optimizer.py b/megatron/core/models/mimo/optimizer.py
index 1a79c1f91ff..6d23998490d 100644
--- a/megatron/core/models/mimo/optimizer.py
+++ b/megatron/core/models/mimo/optimizer.py
@@ -153,6 +153,7 @@ def load_state_dict(self, state_dict: Dict):
for sub_sd, inner_opt in _iter_optimizer_sub_dicts(module_sd, info.optimizer):
_restore_param_groups(sub_sd, inner_opt, name)
+ _restore_param_state_sharding_type(sub_sd)
_restore_grad_scaler(sub_sd)
info.optimizer.load_state_dict(module_sd)
@@ -175,6 +176,7 @@ def sharded_state_dict(self, model_sharded_state_dict, is_loading: bool = False,
):
suffix = f'.{idx}' if idx > 0 else ''
_extract_param_groups(sub_sd, name, suffix, replica_id)
+ _extract_param_state_sharding_type(sub_sd, name, suffix, replica_id)
_extract_grad_scaler(sub_sd, name, suffix, replica_id)
sharded_state[name] = module_sd
@@ -218,6 +220,8 @@ def _extract_param_groups(sub_sd, module_name, suffix, replica_id):
replica_id=replica_id,
)
del opt_sub['param_groups']
+ if not opt_sub:
+ del sub_sd['optimizer']
def _extract_grad_scaler(sub_sd, module_name, suffix, replica_id):
@@ -232,6 +236,18 @@ def _extract_grad_scaler(sub_sd, module_name, suffix, replica_id):
)
+def _extract_param_state_sharding_type(sub_sd, module_name, suffix, replica_id):
+ """Save: extract param_state_sharding_type into a ShardedObject."""
+ if 'param_state_sharding_type' in sub_sd:
+ sub_sd[f'_mimo_param_state_sharding_type{suffix}'] = ShardedObject(
+ f'optimizer.mimo.{module_name}{suffix}.param_state_sharding_type',
+ sub_sd.pop('param_state_sharding_type'),
+ (1,),
+ (0,),
+ replica_id=replica_id,
+ )
+
+
def _restore_param_groups(sub_sd, inner_optimizer, module_name):
"""Load: restore param_groups with current param IDs from the inner optimizer."""
# Find the _mimo_param_groups key (may have a suffix for chained optimizers)
@@ -253,7 +269,21 @@ def _restore_param_groups(sub_sd, inner_optimizer, module_name):
)
for loaded_g, current_g in zip(loaded_pg, current_pg):
loaded_g['params'] = current_g['params']
- sub_sd['optimizer']['param_groups'] = loaded_pg
+ # `sub_sd['optimizer']` may be absent on load: when the per-module state_dict
+ # produced by DistributedOptimizer.state_dict() only contains `param_groups`
+ # under the 'optimizer' key, `_extract_param_groups` removes it at save time
+ # and the resulting empty dict can be dropped during dist_checkpointing
+ # common-state save/load. Use setdefault so the restored param_groups land
+ # in the right place regardless.
+ sub_sd.setdefault('optimizer', {})['param_groups'] = loaded_pg
+
+
+def _restore_param_state_sharding_type(sub_sd):
+ """Load: restore param_state_sharding_type from ShardedObject key."""
+ for k in list(sub_sd.keys()):
+ if k.startswith('_mimo_param_state_sharding_type'):
+ sub_sd['param_state_sharding_type'] = sub_sd.pop(k)
+ break
def _restore_grad_scaler(sub_sd):
@@ -267,17 +297,21 @@ def _restore_grad_scaler(sub_sd):
def _get_replica_id(pg_collection: Optional[ProcessGroupCollection]) -> tuple:
"""Build replica_id tuple for ShardedObject deduplication.
- Includes pp_rank so only one PP stage writes the metadata,
- and dp_rank so only dp_rank=0 writes (others are replicas).
+ Returns (tp_rank, pp_rank, dp_rank) so only (0, 0, 0) within each
+ module's parallelism group is the main replica; all other ranks
+ in the same module are non-main replicas of the same object.
"""
assert pg_collection is not None, "pg_collection required for checkpoint replica_id"
+ assert (
+ hasattr(pg_collection, 'tp') and pg_collection.tp is not None
+ ), "pg_collection.tp must be set for checkpoint deduplication"
assert (
hasattr(pg_collection, 'pp') and pg_collection.pp is not None
), "pg_collection.pp must be set for checkpoint deduplication"
assert (
hasattr(pg_collection, 'dp') and pg_collection.dp is not None
), "pg_collection.dp must be set for checkpoint deduplication"
- return (0, pg_collection.pp.rank(), pg_collection.dp.rank())
+ return (pg_collection.tp.rank(), pg_collection.pp.rank(), pg_collection.dp.rank())
def _get_pg_collection_for_optimizer(grid) -> ProcessGroupCollection:
diff --git a/megatron/core/models/mimo/partition/utils.py b/megatron/core/models/mimo/partition/utils.py
index 0b43e5548ff..592a6253b4a 100644
--- a/megatron/core/models/mimo/partition/utils.py
+++ b/megatron/core/models/mimo/partition/utils.py
@@ -235,7 +235,7 @@ def _apply_context_parallel(
batch["attention_mask"] = attention_mask
if packed_seq_params is None or getattr(packed_seq_params, 'qkv_format', 'sbhd') == 'sbhd':
- batch = get_batch_on_this_cp_rank(batch)
+ batch = get_batch_on_this_cp_rank(batch, is_hybrid_cp=False, cp_group=self.cfg.cp_group)
else:
assert _HAVE_TEX and is_te_min_version("1.10.0"), (
"Please update Transformer Engine to >= 1.10 "
diff --git a/megatron/core/models/mimo/submodules/base.py b/megatron/core/models/mimo/submodules/base.py
index 3b54fd737f2..f05ecc6b15c 100644
--- a/megatron/core/models/mimo/submodules/base.py
+++ b/megatron/core/models/mimo/submodules/base.py
@@ -234,6 +234,15 @@ def encode(self, encoders_data_batch: Dict) -> List[torch.Tensor]:
encoder_inputs = encoders_data_batch[name]
encoder_outputs = encoder(**encoder_inputs)
+ # Some encoders return (embeddings, aux_state). MIMO consumes the
+ # primary embedding tensor here; model-specific aux handling should
+ # live in a modality-specific submodule.
+ if (
+ isinstance(encoder_outputs, tuple)
+ and encoder_outputs
+ and torch.is_tensor(encoder_outputs[0])
+ ):
+ encoder_outputs = encoder_outputs[0]
logger.debug(f"Encoder '{name}' output shape: {encoder_outputs.shape}")
if encoder_outputs.ndim == 3:
diff --git a/megatron/core/models/multimodal/context_parallel.py b/megatron/core/models/multimodal/context_parallel.py
index 6a3cb8bdf48..ceee2d0af7d 100644
--- a/megatron/core/models/multimodal/context_parallel.py
+++ b/megatron/core/models/multimodal/context_parallel.py
@@ -1,9 +1,16 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
"""Multimodal Sequence Parallel (SP) and Context Parallel (CP) functionality."""
+import math
+
import torch
from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.parallel_state import (
+ get_context_parallel_group,
+ get_context_parallel_rank,
+ get_context_parallel_world_size,
+)
def get_padding(
@@ -109,3 +116,421 @@ def get_packed_seq_params(tokens, img_seq_len, padding_needed, cp_size, use_pack
)
return packed_seq_params
+
+
+def split_to_context_parallel_ranks(global_t, pad_value=0):
+ """Split the tensor global_t into context parallel world size parts.
+
+ Args:
+ global_t: [batch, ...]
+ pad_value: Value to pad the last rank with.
+
+ Returns:
+ local_t: [samples_per_rank, ...]. samples_per_rank is the # of samples per CP rank.
+ global_pad: Total padding to have equal samples_per_rank across context parallel ranks.
+ """
+ cp_size = get_context_parallel_world_size()
+ cp_rank = get_context_parallel_rank()
+
+ samples_per_rank = (global_t.shape[0] + cp_size - 1) // cp_size
+ local_t = global_t[cp_rank * samples_per_rank : (cp_rank + 1) * samples_per_rank]
+ global_pad = samples_per_rank * cp_size - global_t.shape[0]
+
+ if local_t.shape[0] < samples_per_rank:
+ local_pad = samples_per_rank - local_t.shape[0]
+ zeros = torch.full(
+ (local_pad, *local_t.shape[1:]), pad_value, device=local_t.device, dtype=local_t.dtype
+ )
+ local_t = torch.cat([local_t, zeros], dim=0)
+
+ return local_t, global_pad
+
+
+def _gather_along_second_dim(local_t):
+ group = get_context_parallel_group()
+ cp_size = get_context_parallel_world_size()
+ if cp_size == 1:
+ return local_t
+
+ tensor_list = [
+ torch.empty(local_t.shape, device=local_t.device, dtype=local_t.dtype)
+ for _ in range(cp_size)
+ ]
+ torch.distributed.all_gather(tensor_list, local_t, group=group)
+ return torch.cat(tensor_list, dim=1)
+
+
+def _reduce_scatter_along_second_dim(global_t):
+ cp_size = get_context_parallel_world_size()
+ if cp_size == 1:
+ return global_t
+
+ assert global_t.shape[1] % cp_size == 0
+ samples_per_rank = global_t.shape[1] // cp_size
+
+ tensor_list = [
+ global_t[:, cp_rank * samples_per_rank : (cp_rank + 1) * samples_per_rank]
+ for cp_rank in range(cp_size)
+ ]
+
+ local_t = torch.zeros(
+ global_t.shape[0],
+ samples_per_rank,
+ *global_t.shape[2:],
+ device=global_t.device,
+ dtype=global_t.dtype,
+ )
+
+ torch.distributed.reduce_scatter(local_t, tensor_list, group=get_context_parallel_group())
+ return local_t
+
+
+class GatherFromContextParallelRanks(torch.autograd.Function):
+ """Gather the input from context parallel ranks."""
+
+ @staticmethod
+ def symbolic(graph, input_):
+ """Symbolic forward used during ``torch.jit`` tracing."""
+ return _gather_along_second_dim(input_)
+
+ @staticmethod
+ def forward(ctx, input_):
+ """All-gather ``input_`` along its second dimension across CP ranks."""
+ return _gather_along_second_dim(input_)
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ """Reduce-scatter the gradient along the second dimension."""
+ return _reduce_scatter_along_second_dim(grad_output)
+
+
+def gather_from_context_parallel_ranks(local_t, global_pad):
+ """Gather ``local_t`` across CP ranks, removing ``global_pad`` trailing pad tokens."""
+ global_t = GatherFromContextParallelRanks.apply(local_t)
+ if global_pad > 0:
+ global_t = global_t[:, :-global_pad]
+ return global_t
+
+
+def gather_from_context_parallel_ranks_dynamic_res(local_t, num_padded_imgs=0):
+ """Gather dynamic-resolution tensors (variable seq per rank) from CP ranks."""
+ cp_size = get_context_parallel_world_size()
+ shape = torch.as_tensor(local_t.shape, device=local_t.device)
+ shapes = [torch.empty_like(shape) for _ in range(cp_size)]
+
+ torch.distributed.all_gather(shapes, shape, group=get_context_parallel_group())
+
+ inputs = [local_t] * cp_size
+ outputs = [torch.empty(*s, dtype=local_t.dtype, device=local_t.device) for s in shapes]
+ torch.distributed.nn.functional.all_to_all(outputs, inputs, group=get_context_parallel_group())
+
+ if num_padded_imgs > 0:
+ outputs = outputs[:-num_padded_imgs]
+
+ return torch.cat(outputs, dim=0)
+
+
+def _compute_tubelet_aware_split_points(num_frames, temporal_patch_size, cp_size, total_frames):
+ """Compute frame-space split points that respect tubelet boundaries within videos.
+
+ Returns ``cp_size + 1`` split points in **frame** indices (not tubelet indices),
+ since callers slice per-frame ``cu_seqlens`` and ``imgs_sizes`` with these bounds.
+ Splits land on either media boundaries or tubelet boundaries inside a media so
+ that no rank receives a partial tubelet.
+ """
+ T = temporal_patch_size
+ target_per_rank = total_frames / cp_size
+
+ media_boundaries = [0]
+ for nf in num_frames:
+ media_boundaries.append(media_boundaries[-1] + nf)
+ boundary_set = set(media_boundaries)
+
+ split_points = [0]
+ for rank in range(1, cp_size):
+ target_split = int(rank * target_per_rank)
+
+ # If the target lands exactly on a media boundary, split there cleanly
+ # without forcing a cut into the next media.
+ if target_split in boundary_set:
+ split_point = target_split
+ else:
+ media_idx = 0
+ for i, boundary in enumerate(media_boundaries[1:], 1):
+ if boundary > target_split:
+ media_idx = i - 1
+ break
+ else:
+ media_idx = len(num_frames) - 1
+
+ media_start = media_boundaries[media_idx]
+ media_end = media_boundaries[media_idx + 1]
+ nf = num_frames[media_idx]
+ num_tubelets = math.ceil(nf / T)
+
+ if num_tubelets <= 1:
+ if target_split - media_start < media_end - target_split:
+ split_point = media_start
+ else:
+ split_point = media_end
+ else:
+ offset_in_media = target_split - media_start
+ tubelet_idx = round(offset_in_media / T)
+ tubelet_idx = max(1, min(tubelet_idx, num_tubelets - 1))
+ split_point = media_start + tubelet_idx * T
+ split_point = min(split_point, media_end)
+
+ split_point = max(split_point, split_points[-1])
+ split_points.append(split_point)
+
+ split_points.append(total_frames)
+ return split_points
+
+
+def _split_num_frames(num_frames, lb, ub):
+ """Return per-media frame counts clipped to the frame range ``[lb, ub)``.
+
+ ``lb`` and ``ub`` are frame indices (the same coordinate system used by
+ :func:`_compute_tubelet_aware_split_points` and the per-frame ``seqlens``
+ array in :func:`split_to_context_parallel_ranks_dynamic_res`). The returned
+ list has one entry per media that contributes at least one frame to the
+ range, with the value being the number of frames of that media in the
+ range.
+ """
+ new_num_frames = []
+ frame_idx = 0
+ for nf in num_frames:
+ media_start = frame_idx
+ media_end = frame_idx + nf
+ overlap_start = max(media_start, lb)
+ overlap_end = min(media_end, ub)
+ if overlap_start < overlap_end:
+ new_num_frames.append(overlap_end - overlap_start)
+ frame_idx = media_end
+ return new_num_frames
+
+
+def split_to_context_parallel_ranks_dynamic_res(
+ global_t,
+ global_imgs_sizes,
+ global_packed_seq_params,
+ *,
+ patch_dim,
+ fp8_enabled=False,
+ fp8_recipe=None,
+ num_frames=None,
+ temporal_patch_size=1,
+):
+ """Split patched vision input across CP ranks.
+
+ ``global_packed_seq_params`` provides per-image seqlens; the split respects them
+ so each rank owns an integer number of images. When ``temporal_patch_size > 1``,
+ splits also respect tubelet boundaries and ``num_frames`` is required.
+
+ Args:
+ global_t: ``[1, total_patches, C * patch_dim * patch_dim]`` patched tokens
+ (pre-embedder). The last dim must equal ``3 * patch_dim * patch_dim``.
+ global_imgs_sizes: ``[num_imgs, 2]`` per-image (H, W) in pixels.
+ global_packed_seq_params: ``PackedSeqParams`` with per-image ``cu_seqlens_q``.
+ patch_dim: Patch size of the vision backbone (e.g. 14 for SigLIP, 16 for
+ many ViTs). Required because dummy padding tensors are sized in patch
+ units and the default would silently mismatch some backbones.
+ fp8_enabled: If True, pad each rank's local sequence to the FP8 multiple
+ (16 by default; 32 for ``mxfp8``).
+ fp8_recipe: Forwarded to :func:`get_padding` so the FP8 padding multiple
+ matches the active recipe.
+ num_frames: Per-media frame count, required when ``temporal_patch_size > 1``.
+ temporal_patch_size: Tubelet size for temporal compression.
+
+ Returns:
+ (local_t, local_imgs_sizes, local_packed_seq_params, has_padding,
+ num_padded_ranks, local_num_frames)
+ """
+ cp_size = get_context_parallel_world_size()
+ cp_rank = get_context_parallel_rank()
+
+ use_tubelet_aware_split = temporal_patch_size > 1
+ if use_tubelet_aware_split:
+ assert num_frames is not None, (
+ f"num_frames must be provided when using temporal compression "
+ f"(temporal_patch_size={temporal_patch_size})"
+ )
+ num_frames_list = num_frames.tolist() if hasattr(num_frames, "tolist") else list(num_frames)
+
+ cu_seqlens = global_packed_seq_params.cu_seqlens_q
+
+ num_imgs = len(global_imgs_sizes)
+ if use_tubelet_aware_split:
+ T = temporal_patch_size
+ total_tubelets = sum(math.ceil(nf / T) for nf in num_frames_list)
+ num_padded_imgs = max(0, cp_size - total_tubelets)
+ else:
+ num_padded_imgs = max(0, cp_size - num_imgs)
+
+ # This function operates on pre-embedder patches, so the hidden dim is
+ # exactly ``3 * patch_dim * patch_dim``. Both the dummy padding image and
+ # the FP8 right-pad tensor below assume this layout.
+ expected_hidden = 3 * patch_dim * patch_dim
+ assert int(global_t.shape[2]) == expected_hidden, (
+ f"split_to_context_parallel_ranks_dynamic_res expects pre-embedder patches "
+ f"with hidden dim 3*patch_dim*patch_dim={expected_hidden}, got "
+ f"{int(global_t.shape[2])} (patch_dim={patch_dim})."
+ )
+
+ dummy_img_size = torch.tensor(
+ [[patch_dim, patch_dim]], device=global_imgs_sizes.device, dtype=global_imgs_sizes.dtype
+ )
+ hidden_dim = expected_hidden
+ dummy_seqlen = 1
+ dummy_img = torch.zeros(
+ [1, dummy_seqlen, hidden_dim], device=global_t.device, dtype=global_t.dtype
+ )
+
+ def _add_dummies(n, global_t, global_imgs_sizes, cu_seqlens, num_frames_list):
+ seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
+ for _ in range(n):
+ global_imgs_sizes = torch.cat([global_imgs_sizes, dummy_img_size], dim=0)
+ global_t = torch.cat([global_t, dummy_img], dim=1)
+ seqlens = torch.cat(
+ [seqlens, torch.tensor([dummy_seqlen], device=seqlens.device, dtype=seqlens.dtype)]
+ )
+ if use_tubelet_aware_split:
+ num_frames_list = num_frames_list + [1] * n
+ cu_seqlens = torch.cat(
+ [
+ torch.tensor([0], device=cu_seqlens.device, dtype=cu_seqlens.dtype),
+ torch.cumsum(seqlens, dim=0),
+ ]
+ )
+ return global_t, global_imgs_sizes, cu_seqlens, num_frames_list
+
+ if num_padded_imgs > 0:
+ global_t, global_imgs_sizes, cu_seqlens, num_frames_list = _add_dummies(
+ num_padded_imgs,
+ global_t,
+ global_imgs_sizes,
+ cu_seqlens,
+ num_frames_list if use_tubelet_aware_split else None,
+ )
+
+ seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
+ total_frames = len(global_imgs_sizes)
+ num_padded_ranks = num_padded_imgs
+
+ if use_tubelet_aware_split:
+ for _retry in range(cp_size):
+ total_frames = len(global_imgs_sizes)
+ split_points = _compute_tubelet_aware_split_points(
+ num_frames_list, temporal_patch_size, cp_size, total_frames
+ )
+ num_empty = sum(1 for k in range(cp_size) if split_points[k] == split_points[k + 1])
+ if num_empty == 0:
+ break
+ global_t, global_imgs_sizes, cu_seqlens, num_frames_list = _add_dummies(
+ num_empty, global_t, global_imgs_sizes, cu_seqlens, num_frames_list
+ )
+ num_padded_imgs += num_empty
+ seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
+
+ original_total_frames = total_frames - num_padded_imgs
+ if num_padded_imgs > 0 and original_total_frames not in split_points:
+ for k in range(cp_size):
+ if split_points[k] < original_total_frames < split_points[k + 1]:
+ split_points[k + 1] = original_total_frames
+ break
+
+ num_padded_ranks = 0
+ if num_padded_imgs > 0:
+ for i in range(cp_size - 1, -1, -1):
+ if split_points[i] >= original_total_frames:
+ num_padded_ranks += 1
+ else:
+ break
+
+ lb = split_points[cp_rank]
+ ub = split_points[cp_rank + 1]
+ local_num_frames = _split_num_frames(num_frames_list, lb, ub)
+ else:
+ seq_per_rank = total_frames // cp_size
+ lb = cp_rank * seq_per_rank
+ # The last rank absorbs the remainder so the union of [lb, ub) ranges
+ # exactly covers the [0, total_frames) image set.
+ ub = (cp_rank + 1) * seq_per_rank if cp_rank < cp_size - 1 else total_frames
+ local_num_frames = None
+
+ seqlens_local = torch.cat([torch.tensor([0], device=seqlens.device), seqlens[lb:ub]])
+ cu_seqlens_local = torch.cumsum(seqlens_local, dim=0).to(torch.int32)
+
+ final_seqlen = cu_seqlens_local[-1]
+
+ pad_img = None
+ if fp8_enabled:
+ padding_needed = get_padding(
+ final_seqlen, 1, 1, False, fp8_enabled=True, fp8_recipe=fp8_recipe
+ )
+ if padding_needed > 0:
+ pad_img = torch.zeros(
+ [1, padding_needed, patch_dim * patch_dim * 3],
+ device=global_t.device,
+ dtype=global_t.dtype,
+ )
+ cu_seqlens_local = torch.cat(
+ [
+ cu_seqlens_local,
+ torch.tensor(
+ [final_seqlen + padding_needed],
+ device=cu_seqlens_local.device,
+ dtype=cu_seqlens_local.dtype,
+ ),
+ ]
+ )
+
+ has_padding = pad_img is not None
+
+ local_packed_seq_params = PackedSeqParams(
+ qkv_format="thd",
+ cu_seqlens_q=cu_seqlens_local,
+ cu_seqlens_kv=cu_seqlens_local,
+ cu_seqlens_q_padded=None,
+ cu_seqlens_kv_padded=None,
+ )
+ max_seqlen_local = max(seqlens_local).to(torch.int32)
+ local_packed_seq_params.max_seqlen_q = max_seqlen_local
+ local_packed_seq_params.max_seqlen_kv = max_seqlen_local
+
+ local_imgs_sizes = global_imgs_sizes[lb:ub]
+ if has_padding:
+ local_imgs_sizes = torch.cat(
+ [
+ local_imgs_sizes,
+ torch.tensor(
+ [[patch_dim, patch_dim * padding_needed]],
+ device=local_imgs_sizes.device,
+ dtype=local_imgs_sizes.dtype,
+ ),
+ ]
+ )
+
+ offset = torch.cumsum(seqlens[:lb], dim=0)[-1] if lb > 0 else 0
+
+ if not has_padding:
+ local_t = global_t[:, offset + cu_seqlens_local[0] : offset + cu_seqlens_local[-1]]
+ else:
+ local_t = torch.cat(
+ [global_t[:, offset + cu_seqlens_local[0] : offset + cu_seqlens_local[-2]], pad_img],
+ dim=1,
+ )
+
+ if local_num_frames is not None:
+ local_num_frames = torch.tensor(
+ local_num_frames, dtype=torch.int32, device=global_imgs_sizes.device
+ )
+
+ return (
+ local_t,
+ local_imgs_sizes,
+ local_packed_seq_params,
+ has_padding,
+ num_padded_ranks,
+ local_num_frames,
+ )
diff --git a/megatron/core/models/multimodal/llava_model.py b/megatron/core/models/multimodal/llava_model.py
index 86ce04521a7..a5b5e7cce58 100644
--- a/megatron/core/models/multimodal/llava_model.py
+++ b/megatron/core/models/multimodal/llava_model.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.
import logging
from collections import namedtuple
from functools import partial
@@ -11,7 +11,11 @@
from megatron.core.extensions.transformer_engine import HAVE_TE
from megatron.core.inference.contexts import BaseInferenceContext
from megatron.core.models.gpt import GPTModel
-from megatron.core.models.mamba import MambaModel
+from megatron.core.models.hybrid.hybrid_model import HybridModel
+from megatron.core.models.multimodal.context_parallel import (
+ gather_from_context_parallel_ranks_dynamic_res,
+ split_to_context_parallel_ranks_dynamic_res,
+)
from megatron.core.models.vision.clip_vit_model import CLIPViTModel, get_num_image_embeddings
from megatron.core.models.vision.multimodal_projector import MultimodalProjector
from megatron.core.models.vision.radio import RADIOViTModel
@@ -43,8 +47,10 @@
IGNORE_INDEX = -100 # ID for labels that should be ignored.
# Image token index can be tokenizer dependent so the default value does not work in all cases.
DEFAULT_IMAGE_TOKEN_INDEX = -200
+DEFAULT_SOUND_TOKEN_INDEX = -300
IMAGE_TOKEN = ""
VIDEO_TOKEN = "