Support PackedSeqParams in TE CUDA graphs

megatron/core/packed_seq_params.py

@@ -1,10 +1,28 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 from dataclasses import dataclass
+from typing import Mapping, MutableMapping
 
 import torch
 import torch.distributed as dist
 from torch import Tensor
 
+CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX = "_packed_seq_params_"
+
+PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS = (
+    "cu_seqlens_q",
+    "cu_seqlens_kv",
+    "cu_seqlens_q_padded",
+    "cu_seqlens_kv_padded",
+)
+
+PACKED_SEQ_PARAMS_CUDA_GRAPH_STATIC_FIELDS = (
+    "qkv_format",
+    "max_seqlen_q",
+    "max_seqlen_kv",
+    "local_cp_size",
+    "cp_group",
+)
+
 
 @dataclass
 class PackedSeqParams:
@@ -65,3 +83,90 @@ def __post_init__(self):
                 .to(torch.int32)
                 .unsqueeze(0)  # Add a batch dimension
             )
+
+
+def _cuda_graph_packed_seq_params_key(field_name: str, prefix: str) -> str:
+    return f"{prefix}{field_name}"
+
+
+def split_packed_seq_params_for_cuda_graph(
+    packed_seq_params: PackedSeqParams | None, prefix: str = CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX
+) -> tuple[dict[str, Tensor | None], dict[str, object]]:
+    """Split ``PackedSeqParams`` into graph Tensor inputs and static metadata.
+
+    Transformer Engine CUDA graph inputs must be tensors or ``None``. ``PackedSeqParams`` mixes
+    dynamic Tensor fields, such as cumulative sequence lengths, with static metadata, such as THD
+    format and max sequence lengths. This helper keeps only the fields TE attention consumes;
+    Mamba-only fields such as ``total_tokens`` and ``seq_idx`` stay outside this graph boundary.
+    """
+    if packed_seq_params is None:
+        return {}, {}
+
+    tensor_kwargs = {}
+    for field_name in PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS:
+        value = getattr(packed_seq_params, field_name)
+        if value is not None and not isinstance(value, Tensor):
+            raise TypeError(
+                f"PackedSeqParams.{field_name} must be a Tensor or None for CUDA graphs, "
+                f"got {type(value).__name__}."
+            )
+        if value is not None:
+            tensor_kwargs[_cuda_graph_packed_seq_params_key(field_name, prefix)] = value
+
+    static_metadata = {}
+    for field_name in PACKED_SEQ_PARAMS_CUDA_GRAPH_STATIC_FIELDS:
+        value = getattr(packed_seq_params, field_name)
+        if isinstance(value, Tensor):
+            raise TypeError(
+                f"PackedSeqParams.{field_name} is static CUDA graph metadata and must not be "
+                "a Tensor."
+            )
+        static_metadata[field_name] = value
+
+    return tensor_kwargs, static_metadata
+
+
+def has_packed_seq_params_cuda_graph_kwargs(
+    kwargs: Mapping[str, object], prefix: str = CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX
+) -> bool:
+    """Return whether ``kwargs`` contains flattened ``PackedSeqParams`` Tensor fields."""
+    return any(
+        _cuda_graph_packed_seq_params_key(field_name, prefix) in kwargs
+        for field_name in PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS
+    )
+
+
+def build_packed_seq_params_from_cuda_graph_kwargs(
+    kwargs: MutableMapping[str, object],
+    static_metadata: Mapping[str, object] | None,
+    prefix: str = CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX,
+    remove_from_kwargs: bool = True,
+) -> PackedSeqParams | None:
+    """Rebuild ``PackedSeqParams`` from flattened CUDA graph kwargs.
+
+    Args:
+        kwargs: Graph kwargs that may contain flattened packed-sequence Tensor fields.
+        static_metadata: Non-Tensor metadata produced by
+            :func:`split_packed_seq_params_for_cuda_graph`.
+        prefix: Prefix used for flattened Tensor fields.
+        remove_from_kwargs: Whether to pop consumed flattened fields from ``kwargs``.
+    """
+    packed_seq_params_kwargs = dict(static_metadata or {})
+    found_tensor_field = False
+    for field_name in PACKED_SEQ_PARAMS_CUDA_GRAPH_TENSOR_FIELDS:
+        key = _cuda_graph_packed_seq_params_key(field_name, prefix)
+        if key not in kwargs:
+            continue
+        found_tensor_field = True
+        value = kwargs.pop(key) if remove_from_kwargs else kwargs[key]
+        if value is not None and not isinstance(value, Tensor):
+            raise TypeError(
+                f"Flattened PackedSeqParams field {key} must be a Tensor or None, "
+                f"got {type(value).__name__}."
+            )
+        packed_seq_params_kwargs[field_name] = value
+
+    if not packed_seq_params_kwargs and not found_tensor_field:
+        return None
+
+    return PackedSeqParams(**packed_seq_params_kwargs)

megatron/core/transformer/cuda_graphs.py

@@ -21,6 +21,7 @@
 from torch.utils._pytree import tree_map as tree_map_pyt
 
 from megatron.core.num_microbatches_calculator import get_num_microbatches
+from megatron.core.packed_seq_params import split_packed_seq_params_for_cuda_graph
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel.random import (
     CudaRNGStatesTracker,
@@ -1668,6 +1669,31 @@ def _layer_is_graphable(layer, config):
     return False
 
 
+def _add_packed_seq_params_to_te_cuda_graph_sample_kwargs(
+    layer, sample_kwargs, sample_packed_seq_params
+):
+    """Add flattened ``PackedSeqParams`` Tensor inputs to TE graph sample kwargs."""
+    if sample_packed_seq_params is None:
+        return
+
+    tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(
+        sample_packed_seq_params
+    )
+    duplicate_keys = set(sample_kwargs) & set(tensor_kwargs)
+    assert not duplicate_keys, (
+        "PackedSeqParams CUDA graph Tensor kwargs overlap with existing sample kwargs: "
+        f"{', '.join(sorted(duplicate_keys))}."
+    )
+    assert hasattr(layer, '_set_te_cuda_graph_packed_seq_params_static_metadata'), (
+        "Transformer layers using TE CUDA graph packed sequence samples must support "
+        "PackedSeqParams static metadata."
+    )
+    layer._set_te_cuda_graph_packed_seq_params_static_metadata(
+        static_metadata, tensor_kwargs.keys()
+    )
+    sample_kwargs.update(tensor_kwargs)
+
+
 class TECudaGraphHelper:
     """
     Helper class to capture CUDA Graphs using TE make_graphed_callables().
@@ -1678,7 +1704,14 @@ class TECudaGraphHelper:
     """
 
     def __init__(
-        self, model, config, seq_length, micro_batch_size, optimizers=[], pg_collection=None
+        self,
+        model,
+        config,
+        seq_length,
+        micro_batch_size,
+        optimizers=[],
+        pg_collection=None,
+        sample_packed_seq_params=None,
     ):
         assert HAVE_TE_GRAPHS, "CUDA Graphs are not supported without TE."
         assert (
@@ -1692,11 +1725,16 @@ def __init__(
             "CUDA Graph with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True."
         )
         self.model = model
+        assert sample_packed_seq_params is None or is_te_min_version("1.10.0"), (
+            "TE CUDA graph packed_seq_params support requires Transformer Engine >= 1.10.0 "
+            "because packed-sequence Tensor fields are passed as keyword arguments."
+        )
         self.config = config
         self.seq_length = seq_length
         self.micro_batch_size = micro_batch_size
         self.optimizers = optimizers
         self.pg_collection = pg_collection
+        self.sample_packed_seq_params = sample_packed_seq_params
         if self.pg_collection is None:
             self.pg_collection = ProcessGroupCollection.use_mpu_process_groups()
         self.tp_group = self.pg_collection.tp
@@ -1925,6 +1963,9 @@ def get_rotary_pos_emb(transformer_module, transformer_input):
                     rotary_pos_emb = get_rotary_pos_emb(chunk_of_the_layer, hidden_states)
                     if rotary_pos_emb is not None:
                         static_inputs["rotary_pos_emb"] = rotary_pos_emb
+                    _add_packed_seq_params_to_te_cuda_graph_sample_kwargs(
+                        layer, static_inputs, self.sample_packed_seq_params
+                    )
                 _sample_kwargs = static_inputs
             elif contains_self_attn:
                 _sample_args = (

megatron/core/transformer/transformer_layer.py

@@ -16,7 +16,13 @@
 from megatron.core.dist_checkpointing.mapping import ShardedStateDict
 from megatron.core.dist_checkpointing.utils import apply_prefix_mapping
 from megatron.core.inference.utils import InferenceMode
-from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.packed_seq_params import (
+    CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX,
+    PackedSeqParams,
+    build_packed_seq_params_from_cuda_graph_kwargs,
+    has_packed_seq_params_cuda_graph_kwargs,
+    split_packed_seq_params_for_cuda_graph,
+)
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.cuda_graphs import is_graph_capturing, is_graph_warmup, make_weakref
 from megatron.core.transformer.enums import CudaGraphModule, InferenceCudaGraphScope, LayerType
@@ -1135,6 +1141,110 @@ def _get_submodules_under_cudagraphs(self):
                 submodules += [self.mlp.shared_experts]
         return submodules
 
+    def _set_te_cuda_graph_packed_seq_params_static_metadata(
+        self, static_metadata, tensor_kwarg_names=None
+    ):
+        """Store non-Tensor ``PackedSeqParams`` metadata used during TE graph capture."""
+        self._te_cuda_graph_packed_seq_params_static_metadata = dict(static_metadata)
+        self._te_cuda_graph_packed_seq_params_tensor_kwarg_names = (
+            None if tensor_kwarg_names is None else tuple(sorted(tensor_kwarg_names))
+        )
+
+    def _get_te_cuda_graph_packed_seq_params_static_metadata(self):
+        """Return the static ``PackedSeqParams`` metadata used for this TE graph."""
+        return getattr(self, '_te_cuda_graph_packed_seq_params_static_metadata', None)
+
+    def _validate_te_cuda_graph_packed_seq_params_static_metadata(self, static_metadata):
+        """Validate that replay uses the same static packed-sequence contract as capture."""
+        expected_static_metadata = self._get_te_cuda_graph_packed_seq_params_static_metadata()
+        assert expected_static_metadata is not None, (
+            "TE CUDA graph replay received packed_seq_params, but the graph was captured without "
+            "packed-sequence sample inputs. Recapture the graph with matching PackedSeqParams "
+            "static metadata."
+        )
+
+        mismatched_fields = []
+        for field_name in sorted(set(expected_static_metadata) | set(static_metadata)):
+            expected_value = expected_static_metadata.get(field_name)
+            actual_value = static_metadata.get(field_name)
+            if expected_value is actual_value:
+                continue
+            if expected_value != actual_value:
+                mismatched_fields.append(field_name)
+
+        assert not mismatched_fields, (
+            "TE CUDA graph replay received PackedSeqParams with static metadata that differs "
+            "from capture. Recapture the graph for changed fields: "
+            f"{', '.join(mismatched_fields)}."
+        )
+
+    def _get_te_cuda_graph_packed_seq_params_tensor_kwarg_names(self):
+        """Return flattened ``PackedSeqParams`` Tensor kwargs used for this TE graph."""
+        return getattr(self, '_te_cuda_graph_packed_seq_params_tensor_kwarg_names', None)
+
+    def _validate_te_cuda_graph_packed_seq_params_tensor_kwargs(self, tensor_kwargs):
+        """Validate replay uses the same flattened Tensor field set as capture."""
+        expected_names = self._get_te_cuda_graph_packed_seq_params_tensor_kwarg_names()
+        if expected_names is None:
+            return
+
+        expected_names = set(expected_names)
+        actual_names = set(tensor_kwargs)
+        missing_names = sorted(expected_names - actual_names)
+        extra_names = sorted(actual_names - expected_names)
+        assert not missing_names and not extra_names, (
+            "TE CUDA graph replay received PackedSeqParams with Tensor fields that differ "
+            "from capture. Recapture the graph for missing fields "
+            f"{missing_names} and extra fields {extra_names}."
+        )
+
+    def _rebuild_te_cuda_graph_packed_seq_params(self, kwargs):
+        """Rebuild ``PackedSeqParams`` from flattened TE graph capture kwargs."""
+        if not has_packed_seq_params_cuda_graph_kwargs(kwargs):
+            return
+
+        assert kwargs.get('packed_seq_params') is None, (
+            "PackedSeqParams must be passed either as flattened TE CUDA graph kwargs or as "
+            "packed_seq_params, but not both."
+        )
+        static_metadata = self._get_te_cuda_graph_packed_seq_params_static_metadata()
+        assert static_metadata is not None, (
+            "Flattened PackedSeqParams Tensor fields require static metadata captured on the "
+            "TransformerLayer."
+        )
+        tensor_kwargs = {
+            key: value
+            for key, value in kwargs.items()
+            if key.startswith(CUDA_GRAPH_PACKED_SEQ_PARAMS_PREFIX)
+        }
+        self._validate_te_cuda_graph_packed_seq_params_tensor_kwargs(tensor_kwargs)
+
+        kwargs['packed_seq_params'] = build_packed_seq_params_from_cuda_graph_kwargs(
+            kwargs, static_metadata
+        )
+
+    def _flatten_te_cuda_graph_packed_seq_params(self, kwargs):
+        """Flatten replay-time ``PackedSeqParams`` into Tensor kwargs for TE graphs."""
+        packed_seq_params = kwargs.pop('packed_seq_params', None)
+        expected_static_metadata = self._get_te_cuda_graph_packed_seq_params_static_metadata()
+        if packed_seq_params is None:
+            assert expected_static_metadata is None, (
+                "TE CUDA graph was captured with packed_seq_params, so replay must also pass "
+                "packed_seq_params with matching static metadata."
+            )
+            return
+
+        tensor_kwargs, static_metadata = split_packed_seq_params_for_cuda_graph(packed_seq_params)
+        self._validate_te_cuda_graph_packed_seq_params_static_metadata(static_metadata)
+        self._validate_te_cuda_graph_packed_seq_params_tensor_kwargs(tensor_kwargs)
+
+        duplicate_keys = set(kwargs) & set(tensor_kwargs)
+        assert not duplicate_keys, (
+            "PackedSeqParams CUDA graph Tensor kwargs overlap with existing replay kwargs: "
+            f"{', '.join(sorted(duplicate_keys))}."
+        )
+        kwargs.update(tensor_kwargs)
+
     def _te_cuda_graph_capture(self, *args, **kwargs):
         """
         CUDA Graph capture for this layer using TE interface.
@@ -1155,6 +1265,8 @@ def _te_cuda_graph_capture(self, *args, **kwargs):
                 hidden_states = kwargs.pop("hidden_states")
                 hidden_states = self.off_interface.backward_record(hidden_states)
                 kwargs["hidden_states"] = hidden_states
+        self._rebuild_te_cuda_graph_packed_seq_params(kwargs)
+
         context = None
         if (
             not self.config.cuda_graph_modules
@@ -1197,7 +1309,8 @@ def _te_cuda_graph_replay(self, *args, **kwargs):
         CUDA graph replay for this layer and microbatch `self.current_microbatch` using TE
         interface. TransformerEngine versions>=1.10 allow keyword arguments with CUDA graph.
         However, CUDA graph accepts only Tensor inputs.
-        Hence, `inference_context` and `packed_seq_params` are excluded from input list.
+        Hence, `inference_context` is excluded from input list. `packed_seq_params` is split
+        into Tensor graph inputs and static metadata when attention is in the graph scope.
         """
         context = None
         if (
@@ -1207,12 +1320,13 @@ def _te_cuda_graph_replay(self, *args, **kwargs):
             hidden_states, context = self._forward_attention(*args, **kwargs)
             args = (hidden_states,)
             kwargs = {}
+        else:
+            self._flatten_te_cuda_graph_packed_seq_params(kwargs)
 
-        assert (kwargs.get('inference_context') is None) and (
-            kwargs.get('packed_seq_params') is None
-        ), (
+        assert kwargs.get('inference_context') is None, (
             "CUDA graph accepts only Tensor inputs. "
-            "inference_context and packed_seq_params are excluded from input list. "
+            "inference_context is excluded from input list; packed_seq_params must be "
+            "flattened into Tensor kwargs with matching static metadata. "
             "For inference cuda graph, please use cuda_graph_impl=local instead."
         )