add cudagraph annotation

Author: yushangdi

.gitignore

@@ -49,3 +49,4 @@ CLAUDE.local.md
 /debug_*.py
 CLAUDE_CONTEXT/
 /.claude/settings.local.json
+agent_space/

torchtitan/experiments/graph_trainer/compile.py

@@ -240,6 +240,7 @@ def apply_compile(
         return model
 
     torch._inductor.config.reorder_for_peak_memory = False
+    torch._inductor.config.triton.cudagraph_kernel_annotations = True
     torch._dynamo.config.capture_scalar_outputs = True
 
     fsdp_reshard_after_forward = get_fsdp_reshard_after_forward_policy(

torchtitan/experiments/graph_trainer/cudagraph.py

@@ -31,6 +31,7 @@ def __init__(self) -> None:
         self._initialized = False
         self._cudagraph_wrappers: list["CUDAGraphWrapper"] = []
         self._teardown_called = False
+        self.all_annotations: dict[int, list] = {}
 
     def maybe_initialize(self) -> None:
         if self._initialized:
@@ -100,6 +101,11 @@ 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
+
+
 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
@@ -206,10 +212,16 @@ def __call__(self, *args):
                 self._cudagraph,
                 pool=_cg_manager.graph_pool,
                 stream=_cg_manager.stream,
+                enable_annotations=True,
             ):
                 # `output` is managed by pytorch's cudagraph pool
                 self._output = self._runnable(*args)
 
+            # Save kernel annotations for trace post-processing.
+            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/graph_utils.py

@@ -519,6 +519,11 @@ def get_compiler_passes_from_config(
                 )
             )
         else:
+            # Auto-insert annotation pass before cudagraph capture.
+            if pass_name == "cudagraph":
+                compiler_passes.append(
+                    AVAILABLE_COMPILER_PASSES["insert_kernel_annotations"]
+                )
             compiler_passes.append(AVAILABLE_COMPILER_PASSES[pass_name])
 
     if pass_names:

torchtitan/experiments/graph_trainer/passes.py

@@ -82,6 +82,9 @@ def construct_default_graph_passes(
         regional_inductor_pass,
     ]
 
+    # Insert kernel annotation markers before cudagraph capture.
+    passes.append(insert_kernel_annotations_pass)
+
     # cudagraph should be the last pass.
     from torchtitan.experiments.graph_trainer.cudagraph import is_cudagraph_compatible
 
@@ -260,6 +263,74 @@ def _get_fake_mode_from_gm(gm: torch.fx.GraphModule):
     return gm
 
 
+def _mark_kernels_enter(annotation: dict) -> object:
+    """Helper called from FX graph to enter a mark_kernels scope."""
+    from torch.cuda._graph_annotations import mark_kernels
+
+    ctx = mark_kernels(annotation)
+    ctx.__enter__()
+    return ctx
+
+
+def _mark_kernels_exit(ctx: object) -> None:
+    """Helper called from FX graph to exit a mark_kernels scope."""
+    ctx.__exit__(None, None, None)  # type: ignore[union-attr]
+
+
+def insert_kernel_annotations_pass(
+    gm: torch.fx.GraphModule,
+    example_inputs: tuple | None = None,
+) -> torch.fx.GraphModule:
+    """Insert mark_kernels() calls at component boundaries in the FX graph.
+
+    Reads ``node.meta["custom"]["component"]`` (set via
+    ``torch.fx.traceback.annotate``) and inserts enter/exit calls so that
+    CUDA graph capture records the annotations.
+    """
+    graph = gm.graph
+    current_component: str | None = None
+    current_ctx_node = None
+
+    for node in list(graph.nodes):
+        component = (node.meta.get("custom") or {}).get("component")
+
+        if component != current_component:
+            # Close previous scope
+            if current_ctx_node is not None:
+                with graph.inserting_before(node):
+                    exit_node = graph.call_function(
+                        _mark_kernels_exit, (current_ctx_node,)
+                    )
+                    exit_node.meta["custom"] = {}
+                current_ctx_node = None
+
+            # Open new scope
+            if component is not None:
+                with graph.inserting_before(node):
+                    enter_node = graph.call_function(
+                        _mark_kernels_enter,
+                        ({"component": component},),
+                    )
+                    enter_node.meta["custom"] = {}
+                current_ctx_node = enter_node
+
+            current_component = component
+
+    # 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(
+                    _mark_kernels_exit, (current_ctx_node,)
+                )
+                exit_node.meta["custom"] = {}
+
+    graph.lint()
+    gm.recompile()
+    return gm
+
+
 def cudagraph_pass(
     gm: torch.fx.GraphModule,
     example_inputs: tuple,
@@ -680,6 +751,7 @@ def tlparse_log_graph_pass(
     "auto_bucketing": autobucketing_reordering_pass,
     "transformer_block_bucketing": transformer_block_bucketing_reordering_pass,
     "regional_inductor": regional_inductor_pass,
+    "insert_kernel_annotations": insert_kernel_annotations_pass,
     "cudagraph": cudagraph_pass,
     "full_inductor_compilation": full_inductor_compilation_pass,
 }

torchtitan/models/common/attention.py

@@ -4,11 +4,19 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import os
+
+
+# TODO: Re-enable once we have closed
+# https://github.com/pytorch/torchtitan/issues/2722
+os.environ.setdefault("DISABLE_LLVM_OPT", "1")
+
 from collections.abc import Callable
 from dataclasses import dataclass, field
 from typing import ClassVar, NamedTuple
 
 import torch
+from torch.fx.traceback import annotate
 import torch.nn.functional as F
 from torch.distributed.tensor import DTensor, Shard
 from torch.distributed.tensor.experimental import local_map
@@ -636,7 +644,8 @@ def forward(
         positions: torch.Tensor | None = None,
     ) -> torch.Tensor:
         bs, seqlen, _ = x.shape
-        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+        with annotate({"component": "qkv_proj"}):
+            xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
 
         # Use -1 instead of `n_heads` (or `n_kv_heads`) to infer the actual
         # local heads from sizes of xq, xk, and xv as TP may have sharded them
@@ -653,25 +662,30 @@ def forward(
 
         # Apply rotary embeddings
         if self.use_rope:
-            if self.rope_backend == "cos_sin":
-                xq, xk = apply_rotary_emb_cos_sin(xq, xk, rope_cache, positions)
-            else:
-                xq, xk = apply_rotary_emb_complex(
-                    xq, xk, freqs_cis=rope_cache, positions=positions
-                )
+            with annotate({"component": "rope"}):
+                if self.rope_backend == "cos_sin":
+                    xq, xk = apply_rotary_emb_cos_sin(
+                        xq, xk, rope_cache, positions
+                    )
+                else:
+                    xq, xk = apply_rotary_emb_complex(
+                        xq, xk, freqs_cis=rope_cache, positions=positions
+                    )
 
         # Handle iRoPE dict masks (Llama4)
         if isinstance(attention_masks, dict):
             mask_key = "rope" if self.use_rope else "nope"
             attention_masks = attention_masks[mask_key]
 
-        output = self.inner_attention(
-            xq,
-            xk,
-            xv,
-            attention_masks=attention_masks,
-            scale=self.scaling,
-            enable_gqa=self.enable_gqa,
-        ).contiguous()
+        with annotate({"component": "sdpa"}):
+            output = self.inner_attention(
+                xq,
+                xk,
+                xv,
+                attention_masks=attention_masks,
+                scale=self.scaling,
+                enable_gqa=self.enable_gqa,
+            ).contiguous()
         output = output.view(bs, seqlen, -1)
-        return self.wo(output)
+        with annotate({"component": "o_proj"}):
+            return self.wo(output)

torchtitan/models/common/feed_forward.py

@@ -8,6 +8,7 @@
 
 import torch
 import torch.nn.functional as F
+from torch.fx.traceback import annotate
 
 from torchtitan.models.common.linear import Linear
 from torchtitan.protocols.module import Module
@@ -51,4 +52,7 @@ def __init__(self, config: Config):
         self.w3 = config.w3.build()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+        with annotate({"component": "silu_mul"}):
+            h = F.silu(self.w1(x)) * self.w3(x)
+        with annotate({"component": "down_proj"}):
+            return self.w2(h)

torchtitan/models/llama3/model.py

@@ -11,6 +11,7 @@
 
 import torch
 from torch import nn
+from torch.fx.traceback import annotate
 
 from torchtitan.models.common.attention import AttentionMasksType, VarlenAttention
 from torchtitan.models.common.decoder import Decoder, TransformerBlock
@@ -48,10 +49,14 @@ def forward(
         attention_masks: AttentionMasksType | None,
         positions: torch.Tensor | None = None,
     ):
-        h = x + self.attention(
-            self.attention_norm(x), freqs_cis, attention_masks, positions
-        )
-        out = h + self.feed_forward(self.ffn_norm(h))
+        with annotate({"component": "attention_norm"}):
+            h_norm = self.attention_norm(x)
+        with annotate({"component": "attention"}):
+            h = x + self.attention(h_norm, freqs_cis, attention_masks, positions)
+        with annotate({"component": "ffn_norm"}):
+            h_norm = self.ffn_norm(h)
+        with annotate({"component": "feed_forward"}):
+            out = h + self.feed_forward(h_norm)
         return out