[graph_trainer] Add CUDA graph kernel annotation pass

Adds `insert_kernel_annotations_pass` that labels CUDA graph kernels with
their originating nn.Module path in profiler traces.

How it works:
1. `annotate_module_fqns(model)` wraps each submodule's forward with
   `torch.fx.traceback.annotate_fn({"module_fqn": fqn})`, setting
   `node.meta["custom"]["module_fqn"]` on FX nodes during tracing.
2. `insert_kernel_annotations_pass` reads the `module_fqn` metadata and
   inserts `mark_kernels` enter/exit calls at module boundaries.  Bundled
   with the cudagraph pass.
3. `CUDAGraphWrapper` captures annotations during graph recording when the
   pass is active.
4. `Profiler.Config.trace_post_processors` (new field) runs
   `cudagraph_annotate_trace_post_processor()` on each exported trace so
   profiler traces automatically carry `module_fqn` fields on graphed
   kernel events.

Requires PyTorch with `torch.cuda._graph_annotations` and the
`enable_annotations` kwarg on `torch.cuda.graph()`.  Falls back gracefully
when unavailable (pass is a no-op).

Authored with Claude.

Author: yushangdi

.github/workflows/integration_test_8gpu_graph_trainer.yaml

@@ -69,6 +69,10 @@ jobs:
         # Run the graph passes unit tests
         torchrun --nproc-per-node=8 -m pytest torchtitan/experiments/graph_trainer/tests/test_passes.py::TestReassignToPgPass -v
         pytest torchtitan/experiments/graph_trainer/tests/test_passes.py::TestApplySACPass -v
+        pytest torchtitan/experiments/graph_trainer/tests/test_passes.py::TestAnnotateModuleFqns -v
+
+        # Run kernel annotation profiler tests (skips if cuda-compat < 13.1)
+        pytest torchtitan/experiments/graph_trainer/tests/test_profiler.py -v
 
         # Run the make_fx tracer unit tests
         pytest torchtitan/experiments/graph_trainer/tests/test_trace_module.py -v -k "TestTraceModule or TestTraceDTensor or TestMetadataPropagation"

torchtitan/experiments/graph_trainer/.claude/CLAUDE.md

@@ -208,3 +208,19 @@ This verifies that the aot_fx_trace path produces bitwise identical losses
 and gradients across runs, and matches eager numerics exactly. Any change
 that breaks this test must be investigated and fixed before proceeding with
 other tests.
+
+### CUDA Graph Kernel Annotations
+
+The `insert_kernel_annotations_pass` labels CUDA graph kernels with their
+originating `nn.Module` path in profiler traces. It runs automatically in the
+`aot_fx_trace` path (bundled with the cudagraph pass). The post-processor
+is attached via ``Profiler.Config.trace_post_processors`` (see
+``cudagraph_annotate_trace_post_processor``) so exported traces are
+annotated automatically — no manual post-processing is needed.
+
+Requirements: `cuda-python` package and CUDA toolkit/driver >= 13.1
+(or `cuda-compat >= 13.1` on `LD_LIBRARY_PATH`). The pass is a no-op when
+these are unavailable.
+
+To view annotated traces, open the exported JSON in https://ui.perfetto.dev.
+Kernel events will have `module_fqn` fields like `layers.0.attention.wq`.

torchtitan/experiments/graph_trainer/common_utils.py

@@ -25,7 +25,10 @@ def annotate_module_fqns(model: nn.Module) -> None:
     """Annotate all modules' forward with their fully-qualified names.
 
     Every named submodule (excluding the root) gets its forward method wrapped
-    with annotate_fn so that FX nodes carry the module_fqn metadata.
+    with ``annotate_fn`` so that FX nodes carry ``module_fqn`` in
+    ``node.meta["custom"]``.
+
+    Call once after model construction, before tracing/compilation.
     """
     for fqn, submodule in model.named_modules():
         if fqn:  # skip root module

torchtitan/experiments/graph_trainer/configs.py

@@ -62,6 +62,7 @@ def to_graph_trainer_config(
     from the graph_trainer model_registry. The compile field is removed and
     left as the GraphTrainer.Config default; callers should explicitly set it.
     """
+    from .cudagraph import cudagraph_annotate_trace_post_processor
     from .trainer import GraphTrainer
 
     d = {f.name: getattr(base_config, f.name) for f in fields(base_config)}
@@ -89,4 +90,11 @@ def to_graph_trainer_config(
     if ac is not None and ac.mode != "none":
         d["activation_checkpoint"] = ActivationCheckpointConfig(mode="selective")
 
+    # Merge CUDA graph kernel annotations into profiler traces when profiling
+    # is active.  No-op otherwise (and no-op when requirements aren't met).
+    # It's also a no-op if there is CUDA graph is not enabled.
+    profiler = d.get("profiler")
+    if profiler is not None:
+        profiler.trace_post_processor = cudagraph_annotate_trace_post_processor()
+
     return GraphTrainer.Config(**d)

torchtitan/experiments/graph_trainer/cudagraph.py

@@ -11,7 +11,7 @@
 during compilation.
 """
 
-import logging
+import json
 import warnings
 from collections.abc import Callable, Sequence
 from typing import Any
@@ -21,7 +21,8 @@
 from torch._library.opaque_object import is_opaque_value
 from torch.utils._ordered_set import OrderedSet
 
-logger = logging.getLogger(__name__)
+from torchtitan.config.function import Function
+from torchtitan.tools.logging import logger
 
 
 class _CUDAGraphManager:
@@ -31,6 +32,10 @@ def __init__(self) -> None:
         self._initialized = False
         self._cudagraph_wrappers: list["CUDAGraphWrapper"] = []
         self._teardown_called = False
+        # toolsId (graph_id << 32 | node_id) -> list of annotation dicts
+        # (e.g. [{"module_fqn": "layers.0.attention.wq"}]).
+        self.all_annotations: dict[int, list] = {}
+        self.enable_annotations: bool = False
 
     def maybe_initialize(self) -> None:
         if self._initialized:
@@ -100,6 +105,54 @@ def cudagraph_teardown() -> None:
     _cg_manager.teardown()
 
 
+def get_cudagraph_annotations() -> dict[int, list]:
+    """Return all kernel annotations accumulated across CUDA graph captures."""
+    return _cg_manager.all_annotations
+
+
+def enable_cudagraph_annotations() -> None:
+    """Enable kernel annotation capture on subsequent CUDA graph recordings."""
+    _cg_manager.enable_annotations = True
+
+
+def cudagraph_annotate_trace_post_processor() -> Function.Config:
+    """Return a ``Function.Config`` that merges captured CUDA graph kernel
+    annotations into a profiler trace file.
+
+    Attach this to ``Profiler.Config.trace_post_processor`` so that exported
+    profiler traces automatically carry ``module_fqn`` fields on graphed kernel
+    events.
+    """
+    return Function.Config(fn=_cudagraph_annotate_trace_file)
+
+
+def _cudagraph_annotate_trace_file(trace_path: str) -> None:
+    """Post-process a profiler trace with CUDA graph kernel annotations."""
+    annotations = _cg_manager.all_annotations
+    if not annotations:
+        return
+
+    try:
+        from torch.cuda._annotate_cuda_graph_trace import (  # pyrefly: ignore[missing-import]
+            annotate_trace,
+        )
+    except ImportError:
+        logger.warning(
+            "torch.cuda._annotate_cuda_graph_trace not available. "
+            "Upgrade PyTorch to enable trace CUDA graph kernel annotation."
+        )
+        return
+
+    with open(trace_path) as f:
+        trace = json.load(f)
+
+    count = annotate_trace(trace, annotations)
+    if count > 0:
+        with open(trace_path, "w") as f:
+            json.dump(trace, f)
+        logger.info(f"Annotated {count} CUDAGraph kernel event(s) in profiler trace")
+
+
 class CUDAGraphWrapper:
     """Wraps a callable with cudagraph. It warms up the callable, records cudagraph,
     and replays cudagraph during runtime. It also handles static input tensors, which
@@ -219,10 +272,16 @@ def __call__(self, *args):
                 self._cudagraph,
                 pool=_cg_manager.graph_pool,
                 stream=_cg_manager.stream,
+                enable_annotations=_cg_manager.enable_annotations,
             ):
                 # `output` is managed by pytorch's cudagraph pool
                 self._output = self._runnable(*args)
 
+            if _cg_manager.enable_annotations:
+                from torch.cuda._graph_annotations import get_kernel_annotations
+
+                _cg_manager.all_annotations.update(get_kernel_annotations())
+
         if self._should_check_address:
             self._check_static_inputs_address()
 

torchtitan/experiments/graph_trainer/deepseek_v3/parallelize.py

@@ -43,7 +43,8 @@ def annotate_deepseekv3(model: GraphTrainerDeepSeekV3Model) -> None:
     - AC region annotation: Tags each transformer block's forward with a unique
       ac_region_id so that apply_sac_pass can assign per-block ac_graph_id
       boundaries for the min-cut partitioner.
-
+    - Module FQN annotation: Tags each submodule's forward with its
+      fully-qualified name for downstream passes.
     """
     from torchtitan.models.common.moe import MoE
     from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher

torchtitan/experiments/graph_trainer/llama3/parallelize.py

@@ -34,8 +34,7 @@ def annotate_llama(model: GraphTrainerLlama3Model) -> None:
     """Attach annotations to FX graph nodes with ``torch.fx.traceback.annotate_fn``
 
     - Module FQN annotation: Tags each submodule's forward with its
-      fully-qualified name so that the transformer_block_bucketing pass
-      can identify which graph nodes belong to which module.
+      fully-qualified name for downstream passes.
     - AC region annotation: Tags each transformer block's forward with a unique
       ac_region_id so that apply_sac_pass can assign per-block ac_graph_id
       boundaries for the min-cut partitioner.

torchtitan/experiments/graph_trainer/passes.py

@@ -61,6 +61,8 @@ def _is_backward_node(node: torch.fx.Node) -> bool:
 
 def compile_time_passes(
     traced_result: "TracedResult",
+    *,
+    cudagraph_compatible: bool = False,
 ) -> list[Callable]:
     """Cleanup, FlexAttention annotation, and regional_inductor passes.
 
@@ -76,7 +78,7 @@ def compile_time_passes(
     # )
     from torchtitan.models.common.attention import FlexAttention
 
-    return [
+    passes: list[Callable] = [
         remove_detach_pass,
         remove_identity_view_pass,
         remove_identity_slice_pass,
@@ -95,17 +97,23 @@ def compile_time_passes(
             flex_compile_config=FlexAttention.inductor_configs,
         ),
         regional_inductor_pass,
-        # TODO: Switch to upstream PyTorch implementation when
-        # https://github.com/pytorch/pytorch/pull/178246 lands.
-        # custom_codegen_pass saves the FX graph to disk for:
-        # 1. Debugging: inspect the generated graph code directly
-        # 2. Profiling provenance: dual-path codegen with _RecordFunctionFast
-        #    gives fine-grained operator-level attribution in profiler traces
-        # 3. User-editable codegen: users can directly modify the generated
-        #    program on disk for fine-grain scheduling optimizations, with
-        #    hot-reload picking up changes at runtime
-        custom_codegen_pass,
     ]
+    if cudagraph_compatible:
+        # Must run before custom_codegen_pass (last in pre_passes)
+        # which replaces the GraphModule's forward().
+        # Also must run before cudagraph_pass.
+        passes.append(insert_kernel_annotations_pass)
+    # TODO: Switch to upstream PyTorch implementation when
+    # https://github.com/pytorch/pytorch/pull/178246 lands.
+    # custom_codegen_pass saves the FX graph to disk for:
+    # 1. Debugging: inspect the generated graph code directly
+    # 2. Profiling provenance: dual-path codegen with _RecordFunctionFast
+    #    gives fine-grained operator-level attribution in profiler traces
+    # 3. User-editable codegen: users can directly modify the generated
+    #    program on disk for fine-grain scheduling optimizations, with
+    #    hot-reload picking up changes at runtime
+    passes.append(custom_codegen_pass)
+    return passes
 
 
 def construct_default_graph_passes(
@@ -124,11 +132,14 @@ def construct_default_graph_passes(
     from torchtitan.experiments.graph_trainer.cudagraph import is_cudagraph_compatible
 
     passes: list[Callable] = []
+    cudagraph_compatible = is_cudagraph_compatible(traced_result.gm)
+
     if not precompiled:
-        passes.extend(compile_time_passes(traced_result))
+        passes.extend(
+            compile_time_passes(traced_result, cudagraph_compatible=cudagraph_compatible)
+        )
 
-    # cudagraph should be the last pass.
-    if is_cudagraph_compatible(traced_result.gm):
+    if cudagraph_compatible:
         static_input_indices = list(range(traced_result.num_static_inputs))
         passes.append(
             functools.partial(
@@ -380,6 +391,96 @@ def _get_fake_mode_from_gm(gm: torch.fx.GraphModule):
     return gm
 
 
+def insert_kernel_annotations_pass(
+    gm: torch.fx.GraphModule,
+    example_inputs: tuple | None = None,
+) -> torch.fx.GraphModule:
+    """Insert mark_kernels() calls at module boundaries in the FX graph.
+
+    Reads ``node.meta["custom"]["module_fqn"]`` (set via
+    ``annotate_module_fqns``) and inserts enter/exit calls so that
+    CUDA graph capture records the annotations.
+
+    Requires ``cuda-python`` package and CUDA toolkit/driver >= 13.1
+    (or cuda-compat >= 13.1).  Returns the graph unchanged when unavailable.
+
+    Also enables annotation capture on :class:`CUDAGraphWrapper` so that
+    ``enable_annotations=True`` is passed to ``torch.cuda.graph()``.
+
+    Alternative approaches:
+
+    1. **fx.Interpreter**: During cudagraph capture, run the graph via an
+       ``fx.Interpreter`` subclass that reads ``module_fqn`` metadata and
+       calls ``mark_kernels`` enter/exit around each node — avoids mutating
+       the graph.
+    2. **Custom CodeGen**: Use a custom ``torch.fx.graph.CodeGen`` to emit
+       enter/exit lines (or ``with`` blocks) directly in the generated
+       Python code.
+
+    The current graph-pass approach is the least invasive.
+    """
+    from torch.cuda._graph_annotations import _is_tools_id_unavailable
+
+    from torchtitan.experiments.graph_trainer.common_utils import _MODULE_FQN
+    from torchtitan.experiments.graph_trainer.cudagraph import (
+        enable_cudagraph_annotations,
+    )
+
+    def _enter(annotation: dict) -> object:
+        from torch.cuda._graph_annotations import mark_kernels
+
+        ctx = mark_kernels(annotation)
+        ctx.__enter__()
+        return ctx
+
+    def _exit(ctx: object) -> None:
+        ctx.__exit__(None, None, None)  # type: ignore[union-attr]
+
+    if _is_tools_id_unavailable():
+        return gm
+
+    enable_cudagraph_annotations()
+
+    graph = gm.graph
+    current_fqn: str | None = None
+    current_ctx_node = None
+
+    for node in list(graph.nodes):
+        fqn = (node.meta.get("custom") or {}).get(_MODULE_FQN)
+
+        if fqn != current_fqn:
+            # Close previous scope
+            if current_ctx_node is not None:
+                with graph.inserting_before(node):
+                    exit_node = graph.call_function(_exit, (current_ctx_node,))
+                    exit_node.meta["custom"] = {}
+                current_ctx_node = None
+
+            # Open new scope
+            if fqn is not None:
+                with graph.inserting_before(node):
+                    enter_node = graph.call_function(
+                        _enter,
+                        ({_MODULE_FQN: fqn},),
+                    )
+                    enter_node.meta["custom"] = {}
+                current_ctx_node = enter_node
+
+            current_fqn = fqn
+
+    # Close any trailing scope (before output/return)
+    if current_ctx_node is not None:
+        output_nodes = [n for n in graph.nodes if n.op == "output"]
+        if output_nodes:
+            with graph.inserting_before(output_nodes[0]):
+                exit_node = graph.call_function(_exit, (current_ctx_node,))
+                exit_node.meta["custom"] = {}
+
+    graph.lint()
+    gm.recompile()
+    return gm
+
+
 def cudagraph_pass(
     gm: torch.fx.GraphModule,
     example_inputs: tuple,

torchtitan/experiments/graph_trainer/qwen3/parallelize.py

@@ -43,6 +43,8 @@ def annotate_qwen3(model: GraphTrainerQwen3Model) -> None:
     - AC region annotation: Tags each transformer block's forward with a unique
       ac_region_id so that apply_sac_pass can assign per-block ac_graph_id
       boundaries for the min-cut partitioner.
+    - Module FQN annotation: Tags each submodule's forward with its
+      fully-qualified name for downstream passes.
     """
     # Annotate MoE EP regions if any layer has MoE enabled
     if any(layer.moe is not None for layer in model.config.layers):