From a212d822981c184a82f6ff93a2f8306709a71979 Mon Sep 17 00:00:00 2001
From: wu-qing-157 <wuqing157@gmail.com>
Date: Mon, 17 Nov 2025 12:11:08 -0800
Subject: [PATCH 1/4] save for merge

---
 d2/planner/ilp_planner.py                     |   0
 d2/runtime/attn_kernels/ops.py                |   1 -
 d2/runtime/megatron/base_transformer_layer.py | 188 +++++++++++++++++-
 d2/runtime/megatron/forward_backward_func.py  |   2 +-
 d2/runtime/megatron/ops/fused_comm_attn.py    |   2 +-
 d2/runtime/megatron/ping_pong/tick_ops.py     | 106 +++++++++-
 .../megatron/ping_pong/transformer_block.py   |  36 +++-
 tests/test_megatron_e2e.py                    |   5 +
 tests/test_megatron_e2e_pipeline.py           |  16 +-
 tests/test_util.py                            |   4 +-
 10 files changed, 329 insertions(+), 31 deletions(-)
 create mode 100644 d2/planner/ilp_planner.py

diff --git a/d2/planner/ilp_planner.py b/d2/planner/ilp_planner.py
new file mode 100644
index 00000000..e69de29b
diff --git a/d2/runtime/attn_kernels/ops.py b/d2/runtime/attn_kernels/ops.py
index 266944bd..fbe65759 100644
--- a/d2/runtime/attn_kernels/ops.py
+++ b/d2/runtime/attn_kernels/ops.py
@@ -242,7 +242,6 @@ def fast_a2a(
         sender_send_disp, sender_transfer_sz,
         sender_recv_disp, recver_transfer_sz,
         my_rank_send_offset, my_rank_recv_offset, my_rank_send_sz,
-        True
     )
 
     return ret
diff --git a/d2/runtime/megatron/base_transformer_layer.py b/d2/runtime/megatron/base_transformer_layer.py
index 714dd949..f59f1504 100644
--- a/d2/runtime/megatron/base_transformer_layer.py
+++ b/d2/runtime/megatron/base_transformer_layer.py
@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any, Iterable, Optional
 import time
 import torch
 from torch import Tensor
@@ -38,6 +38,192 @@ def __init__(
         from megatron.core.extensions.transformer_engine import TEDotProductAttention
         self.self_attention: SelfAttention
         assert isinstance(self.self_attention.core_attention, TEDotProductAttention)
+        self.pre_attn_cuda_graph = None
+        self.post_attn_cuda_graph = None
+
+    def init_pre_attn_cuda_graph(self, prev_layer: "Optional[TransformerLayer]", seq_len: int, device: torch.device, dtype: torch.dtype):
+        if prev_layer is None:
+            static_input = torch.zeros((seq_len, 1, self.config.hidden_size), device=device, dtype=dtype, requires_grad=True)
+            self.pre_attn_cuda_graph = torch.cuda.make_graphed_callables(self.__pre_attn_cuda_graph, (static_input,))
+        else:
+            static_core_attn_out = torch.zeros((seq_len, 1, prev_layer.self_attention.query_projection_size), device=device, dtype=dtype, requires_grad=True)
+            static_residual = torch.zeros((seq_len, 1, self.config.hidden_size), device=device, dtype=dtype, requires_grad=True)
+            def post_then_pre_core_attn_cuda_graph(core_attn_out: Tensor, residual: Tensor):
+                hidden_states = prev_layer._post_attn_cuda_graph(core_attn_out, residual)
+                return self.__pre_attn_cuda_graph(hidden_states)
+            self.pre_attn_cuda_graph = torch.cuda.make_graphed_callables(post_then_pre_core_attn_cuda_graph, (static_core_attn_out, static_residual))
+    
+    def init_post_attn_cuda_graph(self, seq_len: int, device: torch.device, dtype: torch.dtype):
+        static_core_attn_out = torch.zeros((seq_len, 1, self.self_attention.query_projection_size), device=device, dtype=dtype, requires_grad=True)
+        static_residual = torch.zeros((seq_len, 1, self.config.hidden_size), device=device, dtype=dtype, requires_grad=True)
+        self.post_attn_cuda_graph = torch.cuda.make_graphed_callables(self._post_attn_cuda_graph, (static_core_attn_out, static_residual))
+
+    def _forward_pre_attn_cuda_graph(
+        self,
+        args: Iterable[torch.Tensor],
+        rotary_pos_emb: Optional[Tensor] = None,
+        rotary_pos_cos: Optional[Tensor] = None,
+        rotary_pos_sin: Optional[Tensor] = None,
+        packed_seq_params: Optional[PingPangSingleStepPackedSeqParams] = None,
+        sequence_len_offset: Optional[Tensor] = None
+    ):
+        query, key, value, residual = self.pre_attn_cuda_graph(*args)
+
+        log_memory_usage(f"(L{self.layer_number}) _forward_pre_core_attn:(init, before input layernorm)")
+
+        assert rotary_pos_cos is None and rotary_pos_sin is None
+
+        # For self attention we just duplicate the rotary_pos_emb if it isn't already
+        if rotary_pos_emb is not None and not isinstance(rotary_pos_emb, tuple):
+            rotary_pos_emb = (rotary_pos_emb,) * 2
+
+        #### Some code in core_attention. This is because we don't want the pos embedding
+        # being handled in the attention layout (the pos id will be hard to handle)
+        inference_context = None
+
+        log_memory_usage(f"(L{self.layer_number}) _forward_pre_core_attn:(after  qkv, before adjust_key_value_for_inference)")
+
+        query, key, value, rotary_pos_emb, attn_mask_type = self.self_attention._adjust_key_value_for_inference(
+            inference_context,
+            query,
+            key,
+            value,
+            rotary_pos_emb,
+            rotary_pos_cos,
+            rotary_pos_sin,
+            sequence_len_offset,
+        )
+        if packed_seq_params is not None:
+            query = query.squeeze(1)
+            key = key.squeeze(1)
+            value = value.squeeze(1)
+
+        log_memory_usage(f"(L{self.layer_number}) _forward_pre_core_attn:(after adjust_key_value_for_inference, before rope)")
+
+        # ================================================
+        # relative positional embedding (rotary embedding)
+        # ================================================
+        if rotary_pos_emb is not None and not self.config.flash_decode:
+            q_pos_emb, k_pos_emb = rotary_pos_emb
+
+            if packed_seq_params is not None:
+                if packed_seq_params.cu_seqlens_q_padded is not None:
+                    cu_seqlens_q = packed_seq_params.cu_seqlens_q_padded
+                else:
+                    cu_seqlens_q = packed_seq_params.cu_seqlens_q
+                if packed_seq_params.cu_seqlens_kv_padded is not None:
+                    cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded
+                else:
+                    cu_seqlens_kv = packed_seq_params.cu_seqlens_kv
+            else:
+                cu_seqlens_q = cu_seqlens_kv = None
+
+            if q_pos_emb is not None:
+                # TODO VIJAY: simplify
+                query = apply_rotary_pos_emb(
+                    query, q_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q
+                )
+            if k_pos_emb is not None:
+                key = apply_rotary_pos_emb(
+                    key, k_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv
+                )
+
+            # TODO, can apply positional embedding to value_layer so it has
+            # absolute positional embedding.
+            # otherwise, only relative positional embedding takes effect
+            # value_layer = apply_rotary_pos_emb(value_layer, k_pos_emb)
+        log_memory_usage(f"(L{self.layer_number}) _forward_pre_core_attn:(after rope, before return)")
+        return query, key, value, residual, attn_mask_type
+
+    def _forward_post_attn_cuda_graph(
+        self, core_attn_out: Tensor, residual: Tensor,
+        context: Optional[Tensor] = None, context_mask: Optional[Tensor] = None,
+    ):
+        assert context is None and context_mask is None, "not supported in cudagraph"
+        return self.post_attn_cuda_graph(core_attn_out, residual)
+
+    
+    def __pre_attn_cuda_graph(self, hidden_states: Tensor):
+        if self.recompute_input_layernorm:
+            self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput()
+            input_layernorm_output = self.input_layernorm_checkpoint.checkpoint(
+                self.input_layernorm, hidden_states
+            )
+        else:
+            input_layernorm_output = self.input_layernorm(hidden_states)
+
+        log_memory_usage(f"(L{self.layer_number}) _forward_pre_core_attn:(after input layernorm, before qkv)")
+
+        # q, k, v
+        query, key, value = self.self_attention.get_query_key_value_tensors(input_layernorm_output, None)
+        return query, key, value, hidden_states
+
+    def _post_attn_cuda_graph(self, core_attn_out: Tensor, residual: Tensor):
+        attention_output_with_bias = self.self_attention.linear_proj(core_attn_out)
+
+        log_memory_usage(f"(L{self.layer_number}) _forward_post_core_attn:(before layernorm)")
+        if self.recompute_input_layernorm:
+            # discard the output of the input layernorm and register the recompute
+            # as a gradient hook of attention_output_with_bias[0]
+            self.input_layernorm_checkpoint.discard_output_and_register_recompute(
+                attention_output_with_bias[0]
+            )
+
+        log_memory_usage(f"(L{self.layer_number}) _forward_post_core_attn:(after layernorm)")
+
+        # TODO: could we move `bias_dropout_add_exec_handler` itself
+        # inside the module provided in the `bias_dropout_add_spec` module?
+        with self.bias_dropout_add_exec_handler():
+            hidden_states = self.self_attn_bda(self.training, self.config.bias_dropout_fusion)(
+                attention_output_with_bias, residual, self.hidden_dropout
+            )
+
+        # Residual connection.
+        residual = hidden_states
+
+        # # Optional Layer norm after self-attention
+        # pre_cross_attn_layernorm_output = self.pre_cross_attn_layernorm(hidden_states)
+
+        # # Cross attention.
+        # log_memory_usage(f"(L{self.layer_number}) _forward_post_core_attn:(cross attention)")
+        # attention_output_with_bias = self.cross_attention(
+        #     pre_cross_attn_layernorm_output,
+        #     attention_mask=context_mask,
+        #     key_value_states=context,
+        #     inference_context=inference_context,
+        # )
+        # log_memory_usage(f"(L{self.layer_number}) _forward_post_core_attn:(after cross attention)")
+
+        # if isinstance(attention_output_with_bias, dict) and "context" in attention_output_with_bias:
+        #     context = attention_output_with_bias["context"]
+        
+        attention_output_with_bias = self.pre_cross_attn_layernorm(hidden_states)
+
+        # TODO: could we move `bias_dropout_add_exec_handler` itself
+        # inside the module provided in the `bias_dropout_add_spec` module?
+        with self.bias_dropout_add_exec_handler():
+            hidden_states = self.cross_attn_bda(self.training, self.config.bias_dropout_fusion)(
+                attention_output_with_bias, residual, self.hidden_dropout
+            )
+        log_memory_usage(f"(L{self.layer_number}) _forward_post_core_attn:(after cross attn bda)")
+
+        # Residual connection.
+        residual = hidden_states
+
+        # Optional Layer norm post the cross-attention.
+        if self.recompute_pre_mlp_layernorm:
+            self.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput()
+            pre_mlp_layernorm_output = self.pre_mlp_norm_checkpoint.checkpoint(
+                self.pre_mlp_layernorm, hidden_states
+            )
+        else:
+            pre_mlp_layernorm_output = self.pre_mlp_layernorm(hidden_states)
+
+        log_memory_usage(f"(L{self.layer_number}) _forward_post_core_attn:(after pre mlp layernorm)")
+
+        mlp_output = self._forward_mlp(pre_mlp_layernorm_output, residual)
+        log_memory_usage(f"(L{self.layer_number}) _forward_post_core_attn:(after mlp)")
+        return mlp_output
 
     def _forward_pre_core_attn(
         self,
diff --git a/d2/runtime/megatron/forward_backward_func.py b/d2/runtime/megatron/forward_backward_func.py
index e002bfc7..b2e7add6 100644
--- a/d2/runtime/megatron/forward_backward_func.py
+++ b/d2/runtime/megatron/forward_backward_func.py
@@ -90,7 +90,7 @@ def send_forward_backward_recv_forward_backward(output_tensors, input_tensor_gra
 forward_backward_pipelining_without_interleaving_first_run = True
 
 
-import wlbllm.registry
+# import wlbllm.registry
 def wlb_swap_next_forward_metadata():
     # Call this function before entering forward step.
     swap_metadata_fn = wlbllm.registry.get("swap_metadata_fn")
diff --git a/d2/runtime/megatron/ops/fused_comm_attn.py b/d2/runtime/megatron/ops/fused_comm_attn.py
index 0ff16d4c..a52a1186 100644
--- a/d2/runtime/megatron/ops/fused_comm_attn.py
+++ b/d2/runtime/megatron/ops/fused_comm_attn.py
@@ -46,7 +46,7 @@ def __post_init__(self):
         if self.deterministic is None:
             env_val = os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0")
             self.deterministic = env_val != "1"
-            print(f"[FlashAttnArgs] Set deterministic to {self.deterministic} (NVTE_ALLOW_NONDETERMINISTIC_ALGO={env_val})")
+            # print(f"[FlashAttnArgs] Set deterministic to {self.deterministic} (NVTE_ALLOW_NONDETERMINISTIC_ALGO={env_val})")
         if self.deterministic:
             print("⚠️⚠️⚠️ Using deterministic = True! This is not recommended when profiling for performance as it will degrade backward pass performance!")
 
diff --git a/d2/runtime/megatron/ping_pong/tick_ops.py b/d2/runtime/megatron/ping_pong/tick_ops.py
index 0d18ee91..2eeee17e 100644
--- a/d2/runtime/megatron/ping_pong/tick_ops.py
+++ b/d2/runtime/megatron/ping_pong/tick_ops.py
@@ -1,6 +1,10 @@
+from decimal import Context
 import os
+from turtle import forward
 from typing import Any, Dict, Optional
 
+import torch
+
 from d2.runtime.megatron.ops import FusedCommAttn, TickSync, post_a2a_attn_out_with_lse
 from d2.runtime.megatron.ops.fused_comm_attn import FlashAttnArgs
 from d2.runtime.megatron.packed_seq_params import PingPangSingleStepPackedSeqParams
@@ -27,20 +31,26 @@ def forward_pre_core_attn(layer: TransformerLayer, args: Dict[str, Any]):
     )
     log_memory_usage(f"(L{layer.layer_number}) forward_pre_core_attn:(after pre core attn)")
 
-    log_memory_usage(f"(L{layer.layer_number}) forward_pre_core_attn:(before pre mlp to attn)")
-    signal = layer._pre_mlp_to_attn(query, key, value, args["packed_seq_params"])
-    log_memory_usage(f"(L{layer.layer_number}) forward_pre_core_attn:(after pre mlp to attn)")
-
     args["query"] = query
     args["key"] = key
     args["value"] = value
     args["residual"] = residual
     args["attn_mask_type"] = attn_mask_type
-    args["signal"] = signal
+    forward_pre_core_attn_comm(layer, args)
     log_memory_usage(f"(L{layer.layer_number}) forward_pre_core_attn:(return)")
     return args
 
 
+def forward_pre_core_attn_comm(layer: TransformerLayer, args: Dict[str, Any]):
+    log_memory_usage(f"(L{layer.layer_number}) forward_pre_core_attn:(before pre mlp to attn)")
+    signal = layer._pre_mlp_to_attn(
+        args["query"], args["key"], args["value"], args["packed_seq_params"],
+    )
+    log_memory_usage(f"(L{layer.layer_number}) forward_pre_core_attn:(after pre mlp to attn)")
+    args["signal"] = signal
+    return args
+
+
 def layout_mlp_to_attn(layer: TransformerLayer, args: Dict[str, Any]):
     log_memory_usage(f"(L{layer.layer_number}) layout_mlp_to_attn:(start)")
     bwd_resend_qkv = args["packed_seq_params"].bwd_packed_seq_params is not None
@@ -108,8 +118,7 @@ def layout_attn_to_mlp(layer: TransformerLayer, args: Dict[str, Any]):
     return args
 
 
-def forward_post_core_attn(layer: TransformerLayer, args: Dict[str, Any]):
-    log_memory_usage(f"(L{layer.layer_number}) forward_post_core_attn:(start)")
+def forward_post_core_attn_comm(layer: TransformerLayer, args: Dict[str, Any]):
     signal = args.pop("signal")
     packed_seq_params: PingPangSingleStepPackedSeqParams = args["packed_seq_params"]
     bwd_resend_qkv = packed_seq_params.bwd_packed_seq_params is not None
@@ -124,10 +133,16 @@ def forward_post_core_attn(layer: TransformerLayer, args: Dict[str, Any]):
         args.pop("query"), args.pop("key"), args.pop("value")
     else:
         core_attn_out = layer._post_attn_to_mlp(signal, args["packed_seq_params"])
-    residual = args.pop("residual")
+    args["core_attn_out"] = core_attn_out
+    return args
+
+
+def forward_post_core_attn(layer: TransformerLayer, args: Dict[str, Any]):
+    log_memory_usage(f"(L{layer.layer_number}) forward_post_core_attn:(start)")
+    forward_post_core_attn_comm(layer, args)
     mlp_output, context = layer._forward_post_core_attn(
-        core_attn_out,
-        residual,
+        args.pop("core_attn_out"),
+        args.pop("residual"),
         args["context"],
         args["context_mask"],
     )
@@ -137,6 +152,67 @@ def forward_post_core_attn(layer: TransformerLayer, args: Dict[str, Any]):
     return args
 
 
+def forward_post_then_pre_core_attn_cuda_graph(layer: TransformerLayer, args: Dict[str, Any]):
+    log_memory_usage(f"(L{layer.layer_number}) forward_post_then_pre_core_attn:(start)")
+    assert args["context"] is None and args["context_mask"] is None, "not supported in cudagraph"
+    forward_post_core_attn_comm(layer, args)
+    log_memory_usage(f"(L{layer.layer_number}) forward_post_then_pre_core_attn:(before non core attn)")
+    query, key, value, residual, attn_mask_type = layer._forward_pre_attn_cuda_graph(
+        (args.pop("core_attn_out"), args.pop("residual")),
+        args["rotary_pos_emb"],
+        args["rotary_pos_cos"],
+        args["rotary_pos_sin"],
+        args["mlp_packed_seq_params"],
+        args["sequence_len_offset"],
+    )
+    log_memory_usage(f"(L{layer.layer_number}) forward_post_then_pre_core_attn:(after non core attn)")
+    args["query"] = query
+    args["key"] = key
+    args["value"] = value
+    args["residual"] = residual
+    args["attn_mask_type"] = attn_mask_type
+    forward_pre_core_attn_comm(layer, args)
+    log_memory_usage(f"(L{layer.layer_number}) forward_pre_then_pre_core_attn:(return)")
+    return args
+
+
+def forward_pre_core_attn_cuda_graph(layer: TransformerLayer, args: Dict[str, Any]):
+    log_memory_usage(f"(L{layer.layer_number}) forward_pre_core_attn:(start)")
+    log_memory_usage(f"(L{layer.layer_number}) forward_pre_core_attn:(before pre core attn)")
+    hidden_states = args.pop("hidden_states")
+    query, key, value, residual, attn_mask_type = layer._forward_pre_attn_cuda_graph(
+        (hidden_states,),
+        args["rotary_pos_emb"],
+        args["rotary_pos_cos"],
+        args["rotary_pos_sin"],
+        args["mlp_packed_seq_params"],
+        args["sequence_len_offset"],
+    )
+    log_memory_usage(f"(L{layer.layer_number}) forward_pre_core_attn:(after pre core attn)")
+
+    args["query"] = query
+    args["key"] = key
+    args["value"] = value
+    args["residual"] = residual
+    args["attn_mask_type"] = attn_mask_type
+    forward_pre_core_attn_comm(layer, args)
+    log_memory_usage(f"(L{layer.layer_number}) forward_pre_core_attn:(return)")
+    return args
+
+
+def forward_post_core_attn_cuda_graph(layer: TransformerLayer, args: Dict[str, Any]):
+    forward_post_core_attn_comm(layer, args)
+    mlp_output = layer._forward_post_attn_cuda_graph(
+        args.pop("core_attn_out"),
+        args.pop("residual"),
+        args["context"],
+        args["context_mask"],
+    )
+    args["hidden_states"] = mlp_output
+    log_memory_usage(f"(L{layer.layer_number}) forward_post_core_attn:(end)")
+    return args
+
+
 def tick_sync(compute_stream, comm_stream, arg_group_0, keys_0, arg_group_1, keys_1):
     log_memory_usage(f"(L?) tick_sync:(start)")
     if isinstance(keys_0, str):
@@ -173,3 +249,13 @@ def tick_nonca_compute(
         arg_group = forward_post_core_attn(prev_layer, arg_group)
     arg_group = forward_pre_core_attn(layer, arg_group)
     return arg_group
+
+def tick_nonca_compute_cuda_graph(
+    layer: TransformerLayer, prev_layer: Optional[TransformerLayer],
+    arg_group: Dict[str, Any], is_last_layer_post_attn: bool
+):
+    if is_last_layer_post_attn:
+        return forward_post_core_attn_cuda_graph(layer, arg_group)
+    if prev_layer is None:
+        return forward_pre_core_attn(layer, arg_group)
+    return forward_post_then_pre_core_attn_cuda_graph(layer, arg_group)
diff --git a/d2/runtime/megatron/ping_pong/transformer_block.py b/d2/runtime/megatron/ping_pong/transformer_block.py
index 1f6e6f0e..f2a0c1c8 100644
--- a/d2/runtime/megatron/ping_pong/transformer_block.py
+++ b/d2/runtime/megatron/ping_pong/transformer_block.py
@@ -18,7 +18,7 @@
 from d2.runtime.megatron.ops.fused_comm_attn import dummy_backward
 from d2.runtime.megatron.packed_seq_params import PingPangPackedSeqParams
 from d2.runtime.megatron.ping_pong.tick_ops import (
-    forward_core_attn, layout_mlp_to_attn, layout_attn_to_mlp, log_memory_usage, tick_nonca_compute,
+    forward_core_attn, layout_mlp_to_attn, layout_attn_to_mlp, log_memory_usage, tick_nonca_compute, tick_nonca_compute_cuda_graph,
     tick_sync
 )
 from d2.runtime.megatron.ping_pong.transformer_layer import TransformerLayer
@@ -47,6 +47,22 @@ def init_value(self):
         self.comm_stream: torch.cuda.Stream = None
         self._ping_pong_debug: bool = False
         self._debug_forward_impl: str = "orig"
+        self.use_cuda_graph: bool = True  # FIXME: ugly hardcode here
+        self.tick_nonca_compute = tick_nonca_compute_cuda_graph if self.use_cuda_graph else tick_nonca_compute
+    
+    def init_layer_cuda_graphs(self):
+        prev_layer = None
+        for layer in self.layers:
+            layer: TransformerLayer
+            layer.init_pre_attn_cuda_graph(prev_layer, seq_len=1024,  # FIXME: ugly hardcode here
+                                           device=layer.self_attention.linear_qkv.weight.device,
+                                           dtype=layer.self_attention.linear_qkv.weight.dtype)
+            prev_layer = layer
+        
+        # This is the last layer
+        layer.init_post_attn_cuda_graph(seq_len=1024,  # FIXME: ugly hardcode here
+                                        device=layer.self_attention.linear_qkv.weight.device,
+                                        dtype=layer.self_attention.linear_qkv.weight.dtype)
 
     def init_ping_pong_communication_ctx(self, device: torch.device):
         assert not self.ping_pong_comm_initialized
@@ -174,8 +190,8 @@ def forward_layer_ping_pong(
 
         # tick 0, compute part
         with torch.cuda.nvtx.range(linear_name + ".0"):
-            arg_group_0 = tick_nonca_compute(layer, prev_layer, arg_group_0,
-                                             is_last_layer_post_attn=False)
+            arg_group_0 = self.tick_nonca_compute(layer, prev_layer, arg_group_0,
+                                                  is_last_layer_post_attn=False)
         with torch.cuda.nvtx.range(f"forward[{l_no}].sync_tick.0"):
             tick_sync(
                 compute_stream, self.comm_stream,
@@ -190,8 +206,8 @@ def forward_layer_ping_pong(
             arg_group_0 = layout_mlp_to_attn(layer, arg_group_0)
         # compute
         with torch.cuda.nvtx.range(linear_name + ".1"):
-            arg_group_1 = tick_nonca_compute(layer, prev_layer, arg_group_1,
-                                             is_last_layer_post_attn=False)
+            arg_group_1 = self.tick_nonca_compute(layer, prev_layer, arg_group_1,
+                                                  is_last_layer_post_attn=False)
         with torch.cuda.nvtx.range(f"forward[{l_no}].sync_tick.1"):
             tick_sync(
                 compute_stream, self.comm_stream,
@@ -238,8 +254,8 @@ def forward_layer_ping_pong(
             # No next layer, do the sync here.
             # compute
             with torch.cuda.nvtx.range(f"forward[{l_no}].post_core_attn.0"):
-                arg_group_0 = tick_nonca_compute(layer, None, arg_group_0,
-                                                 is_last_layer_post_attn=True)
+                arg_group_0 = self.tick_nonca_compute(layer, None, arg_group_0,
+                                                      is_last_layer_post_attn=True)
             with torch.cuda.nvtx.range(f"forward[{l_no}].sync_tick.4"):
                 tick_sync(
                     compute_stream, self.comm_stream,
@@ -247,8 +263,8 @@ def forward_layer_ping_pong(
                     arg_group_1, "signal",          # comm out
                 )
             with torch.cuda.nvtx.range(f"forward[{l_no}].post_core_attn.1"):
-                arg_group_1 = tick_nonca_compute(layer, None, arg_group_1,
-                                                 is_last_layer_post_attn=True)
+                arg_group_1 = self.tick_nonca_compute(layer, None, arg_group_1,
+                                                      is_last_layer_post_attn=True)
             # gathering the result
             with torch.cuda.nvtx.range(f"forward[{l_no}].gather_ping_pong"):
                 hidden_states = gather_tensor([arg_group_0["hidden_states"], arg_group_1["hidden_states"]], 2)
@@ -295,6 +311,8 @@ def add_ping_pong_forward(block: MegatronTransformerBlock) -> PingPongTransforme
         PingPongTransformerBlockInterface.forward_layer_ping_pong, block)
     block.forward = types.MethodType(
         PingPongTransformerBlockInterface.forward, block)
+    block.init_layer_cuda_graphs = types.MethodType(
+        PingPongTransformerBlockInterface.init_layer_cuda_graphs, block)
 
     return block
 
diff --git a/tests/test_megatron_e2e.py b/tests/test_megatron_e2e.py
index 92e3f12d..b36d94d7 100644
--- a/tests/test_megatron_e2e.py
+++ b/tests/test_megatron_e2e.py
@@ -295,6 +295,9 @@ def test(args):
     # set again to potentially adapt to the ray launch case.
     set_random_seed(seed, set_megatron=False)
 
+    # torch.distributed.breakpoint()
+    worker.train_module[0].module.module.decoder.init_layer_cuda_graphs()  # FIXME: hardcode for now, where to put?
+
     as_rank = worker.as_rank
     as_world_size = worker.as_world_size
 
@@ -347,6 +350,7 @@ def test(args):
             normal_forward_fn=False,
             forward_only=False,
         )
+        print(f"{ref=}")
     time.sleep(1)
     torch.cuda.synchronize()
     torch.distributed.barrier()
@@ -357,6 +361,7 @@ def test(args):
             normal_forward_fn=False,
             forward_only=False,
         )
+        print(f"{ref=}")
     torch.cuda.synchronize()
     torch.distributed.barrier()
     print("=" * 20 + "forward_backward_batch attention server, done")
diff --git a/tests/test_megatron_e2e_pipeline.py b/tests/test_megatron_e2e_pipeline.py
index 9cad4831..c16fb9e8 100644
--- a/tests/test_megatron_e2e_pipeline.py
+++ b/tests/test_megatron_e2e_pipeline.py
@@ -134,6 +134,7 @@ def dummy_backward_step(model, dummy_bwd_iter, skip: bool):
                 forward_only=forward_only,
             )
         grad_sample = unwrap_model(self.train_module[0]).decoder.layers[-1].self_attention.linear_proj.weight.main_grad.clone()
+        print(f"{losses_reduced=}, {grad_sample=}")
 
         # when testing numerical correctness, instead of running optimizer step, reset grads.
         for tm in self.train_module:
@@ -293,6 +294,9 @@ def test(args):
     )
     worker.set_config(dtype=dtype)
     worker.init(model_path, seed=seed)
+
+    # torch.distributed.breakpoint()
+    worker.train_module[0].module.module.decoder.init_layer_cuda_graphs()  # FIXME: hardcode for now, where to put?
     # set again to potentially adapt to the ray launch case.
     set_random_seed(seed, set_megatron=False)
 
@@ -369,12 +373,12 @@ def test(args):
         orig_impl_microbatches.append(orig_mb)
 
     time.sleep(2)
-    loss_orig_reimpl, grad_orig_reimpl = worker.forward_backward_batch(
-        microbatches=orig_impl_microbatches,
-        forward_only=False,
-        mode="orig_reimpl",
-        with_dummy=True,
-    )
+    # loss_orig_reimpl, grad_orig_reimpl = worker.forward_backward_batch(
+    #     microbatches=orig_impl_microbatches,
+    #     forward_only=False,
+    #     mode="orig_reimpl",
+    #     with_dummy=True,
+    # )
     loss_orig, grad_orig = worker.forward_backward_batch(
         microbatches=orig_impl_microbatches,
         forward_only=False,
diff --git a/tests/test_util.py b/tests/test_util.py
index 48bf34dd..fcbb1725 100644
--- a/tests/test_util.py
+++ b/tests/test_util.py
@@ -472,7 +472,7 @@ def create_pipeline_doclens(
     if add_dummy:
         # Each rank only gets one dummy document, which is of `tp_size`
         # to avoid Sequence Parallel having an issue.
-        pp_head_new_doc_len = [[tp_size] for _ in range(dp_size)]
+        pp_head_new_doc_len = [[total_token_on_rank // dp_size] for _ in range(dp_size)]  # FIXME: using total_token_on_rank for cudagraph to work out-of-box, can be tp_size when cudagraph is disabled.
     else:
         assert total_token_on_rank % tp_size == 0, "Sequence Parallel requires total token divisible by tp_size"
         # TODO: Do not do "use_planner" to decide if we should grab batch from global batch or random.
@@ -503,7 +503,7 @@ def create_pipeline_doclens(
             other_pp_doc_len = ref_doc_lens[:-num_batches]
     else:
         dummy_fwd_num = world_size - dp_size
-        other_pp_doc_len = [[tp_size] for _ in range(dummy_fwd_num)]
+        other_pp_doc_len = [[total_token_on_rank // dp_size] for _ in range(dummy_fwd_num)]
     tick_per_rank_doc_len = pp_head_new_doc_len + other_pp_doc_len
     print(f"In util.py, finally tick_per_rank_doc_len is: {tick_per_rank_doc_len}")
 

From 1213e31a5041e065ccfde6d7458742081bebfb66 Mon Sep 17 00:00:00 2001
From: wu-qing-157 <wuqing157@gmail.com>
Date: Mon, 17 Nov 2025 13:55:12 -0800
Subject: [PATCH 2/4] clean code

---
 d2/runtime/attn_kernels/ops.py             |  1 +
 d2/runtime/megatron/ops/fused_comm_attn.py |  2 +-
 d2/runtime/megatron/ping_pong/tick_ops.py  |  5 +----
 tests/test_megatron_e2e.py                 |  2 --
 tests/test_megatron_e2e_pipeline.py        | 15 +++++++--------
 tests/test_util.py                         |  2 +-
 6 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/d2/runtime/attn_kernels/ops.py b/d2/runtime/attn_kernels/ops.py
index 27ea8f4b..54c82547 100644
--- a/d2/runtime/attn_kernels/ops.py
+++ b/d2/runtime/attn_kernels/ops.py
@@ -252,6 +252,7 @@ def fast_a2a(
         sender_send_disp, sender_transfer_sz,
         sender_recv_disp, recver_transfer_sz,
         my_rank_send_offset, my_rank_recv_offset, my_rank_send_sz,
+        True
     )
 
     return ret
diff --git a/d2/runtime/megatron/ops/fused_comm_attn.py b/d2/runtime/megatron/ops/fused_comm_attn.py
index fa6d1387..92b1d11f 100644
--- a/d2/runtime/megatron/ops/fused_comm_attn.py
+++ b/d2/runtime/megatron/ops/fused_comm_attn.py
@@ -44,7 +44,7 @@ def __post_init__(self):
         if self.deterministic is None:
             env_val = os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0")
             self.deterministic = env_val != "1"
-            # print(f"[FlashAttnArgs] Set deterministic to {self.deterministic} (NVTE_ALLOW_NONDETERMINISTIC_ALGO={env_val})")
+            print(f"[FlashAttnArgs] Set deterministic to {self.deterministic} (NVTE_ALLOW_NONDETERMINISTIC_ALGO={env_val})")
         if self.deterministic:
             print("⚠️⚠️⚠️ Using deterministic = True! This is not recommended when profiling for performance as it will degrade backward pass performance!")
 
diff --git a/d2/runtime/megatron/ping_pong/tick_ops.py b/d2/runtime/megatron/ping_pong/tick_ops.py
index cdb27094..85316868 100644
--- a/d2/runtime/megatron/ping_pong/tick_ops.py
+++ b/d2/runtime/megatron/ping_pong/tick_ops.py
@@ -2,9 +2,6 @@
 import os
 import torch
 
-import torch
-
-from d2.runtime.megatron.ops import FusedCommAttn, TickSync, post_a2a_attn_out_with_lse
 from d2.runtime.megatron.ops import FusedCommAttn, post_a2a_attn_out_with_lse
 from d2.runtime.megatron.ops.stream_sync_fn import tick_sync_with_info
 from d2.runtime.megatron.ops.fused_comm_attn import FlashAttnArgs
@@ -156,7 +153,7 @@ def forward_post_core_attn_comm(layer: TransformerLayer, args: Dict[str, Any]):
 
 def forward_post_core_attn(layer: TransformerLayer, args: Dict[str, Any]):
     log_memory_usage(f"(L{layer.layer_number}) forward_post_core_attn:(start)")
-    
+
     # Setup timing if enabled
     start_event = None
     end_event = None
diff --git a/tests/test_megatron_e2e.py b/tests/test_megatron_e2e.py
index 54269194..d69d0791 100644
--- a/tests/test_megatron_e2e.py
+++ b/tests/test_megatron_e2e.py
@@ -351,7 +351,6 @@ def test(args):
             normal_forward_fn=False,
             forward_only=False,
         )
-        print(f"{ref=}")
     time.sleep(1)
     torch.cuda.synchronize()
     torch.distributed.barrier()
@@ -362,7 +361,6 @@ def test(args):
             normal_forward_fn=False,
             forward_only=False,
         )
-        print(f"{ref=}")
     torch.cuda.synchronize()
     torch.distributed.barrier()
     print("=" * 20 + "forward_backward_batch attention server, done")
diff --git a/tests/test_megatron_e2e_pipeline.py b/tests/test_megatron_e2e_pipeline.py
index c16fb9e8..80588f59 100644
--- a/tests/test_megatron_e2e_pipeline.py
+++ b/tests/test_megatron_e2e_pipeline.py
@@ -134,7 +134,6 @@ def dummy_backward_step(model, dummy_bwd_iter, skip: bool):
                 forward_only=forward_only,
             )
         grad_sample = unwrap_model(self.train_module[0]).decoder.layers[-1].self_attention.linear_proj.weight.main_grad.clone()
-        print(f"{losses_reduced=}, {grad_sample=}")
 
         # when testing numerical correctness, instead of running optimizer step, reset grads.
         for tm in self.train_module:
@@ -373,12 +372,12 @@ def test(args):
         orig_impl_microbatches.append(orig_mb)
 
     time.sleep(2)
-    # loss_orig_reimpl, grad_orig_reimpl = worker.forward_backward_batch(
-    #     microbatches=orig_impl_microbatches,
-    #     forward_only=False,
-    #     mode="orig_reimpl",
-    #     with_dummy=True,
-    # )
+    loss_orig_reimpl, grad_orig_reimpl = worker.forward_backward_batch(
+        microbatches=orig_impl_microbatches,
+        forward_only=False,
+        mode="orig_reimpl",
+        with_dummy=True,
+    )
     loss_orig, grad_orig = worker.forward_backward_batch(
         microbatches=orig_impl_microbatches,
         forward_only=False,
@@ -414,7 +413,7 @@ def test(args):
     parser.add_argument("--num-gpus-per-node", type=int, default=4)
     parser.add_argument("--tp-size", type=int, default=1)
     parser.add_argument("--pp-size", type=int, default=4)
-    parser.add_argument("--num-microbatch", type=int, default=2)
+    parser.add_argument("--num-microbatch", type=int, default=6)
     parser.add_argument("--use-planner", action="store_true")
     args = parser.parse_args()
     test(args)
\ No newline at end of file
diff --git a/tests/test_util.py b/tests/test_util.py
index 4cb7ab3f..fe52b9f1 100644
--- a/tests/test_util.py
+++ b/tests/test_util.py
@@ -499,7 +499,7 @@ def create_pipeline_doclens(
             other_pp_doc_len = ref_doc_lens[:-num_batches]
     else:
         dummy_fwd_num = world_size - dp_size
-        other_pp_doc_len = [[total_token_on_rank // dp_size] for _ in range(dummy_fwd_num)]
+        other_pp_doc_len = [[total_token_on_rank // dp_size] for _ in range(dummy_fwd_num)]  # FIXME: using total_token_on_rank for cudagraph to work out-of-box, can be tp_size when cudagraph is disabled.
     tick_per_rank_doc_len = pp_head_new_doc_len + other_pp_doc_len
     print(f"In util.py, finally tick_per_rank_doc_len is: {tick_per_rank_doc_len}")
 

From d6c9d928baa8fa6854dc9b72e67bb620ea0e751d Mon Sep 17 00:00:00 2001
From: GindaChen <32371474+GindaChen@users.noreply.github.com>
Date: Sat, 13 Dec 2025 20:01:12 +0000
Subject: [PATCH 3/4] Fix some hardcoding

---
 .../megatron/ping_pong/transformer_block.py   | 10 ++-
 d2/utils/network_inspect.py                   | 85 ++++++++++---------
 tests/test_megatron_e2e_pipeline.py           | 23 +++--
 3 files changed, 67 insertions(+), 51 deletions(-)

diff --git a/d2/runtime/megatron/ping_pong/transformer_block.py b/d2/runtime/megatron/ping_pong/transformer_block.py
index 92679381..037db808 100644
--- a/d2/runtime/megatron/ping_pong/transformer_block.py
+++ b/d2/runtime/megatron/ping_pong/transformer_block.py
@@ -47,20 +47,24 @@ def init_value(self):
         self.comm_stream: torch.cuda.Stream = None
         self._ping_pong_debug: bool = False
         self._debug_forward_impl: str = "orig"
-        self.use_cuda_graph: bool = True  # FIXME: ugly hardcode here
+        self.use_cuda_graph: bool = False
         self.tick_nonca_compute = tick_nonca_compute_cuda_graph if self.use_cuda_graph else tick_nonca_compute
     
     def init_layer_cuda_graphs(self):
+        self.use_cuda_graph = True
         prev_layer = None
+        seq_len = int(os.environ.get("D2_SEQ_LEN", -1))
+        if seq_len == -1:
+            raise ValueError("D2_SEQ_LEN is not set")
         for layer in self.layers:
             layer: TransformerLayer
-            layer.init_pre_attn_cuda_graph(prev_layer, seq_len=1024,  # FIXME: ugly hardcode here
+            layer.init_pre_attn_cuda_graph(prev_layer, seq_len=seq_len,
                                            device=layer.self_attention.linear_qkv.weight.device,
                                            dtype=layer.self_attention.linear_qkv.weight.dtype)
             prev_layer = layer
         
         # This is the last layer
-        layer.init_post_attn_cuda_graph(seq_len=1024,  # FIXME: ugly hardcode here
+        layer.init_post_attn_cuda_graph(seq_len=seq_len,
                                         device=layer.self_attention.linear_qkv.weight.device,
                                         dtype=layer.self_attention.linear_qkv.weight.dtype)
 
diff --git a/d2/utils/network_inspect.py b/d2/utils/network_inspect.py
index 28fd1183..fa426206 100644
--- a/d2/utils/network_inspect.py
+++ b/d2/utils/network_inspect.py
@@ -100,47 +100,50 @@ def inspect_network_metadata(metadata: 'FastAllToAllMetadata_Tuple', is_ping, sa
     ).max().item()
 
     if rank == 0:
-        network_inspect_file = os.path.join(output_dir, "network_inspect.jsonl")
-        with open(network_inspect_file, "a") as f:
-            f.write(json.dumps({
-                "sample_id": sample_id,
-                "is_ping": is_ping,
-                "tolerance_factor": tolerance_factor,
-                "qkv_fwd_metadata__send_transfer_sz_mb": (qkv_fwd_metadata__send_transfer_sz // (1024 * 1024)).tolist(),
-                "qkv_fwd_metadata__recv_transfer_sz_mb": (qkv_fwd_metadata__recv_transfer_sz // (1024 * 1024)).tolist(),
-                "attn_out_fwd_metadata__send_transfer_sz_mb": (attn_out_fwd_metadata__send_transfer_sz // (1024 * 1024)).tolist(),
-                "attn_out_fwd_metadata__recv_transfer_sz_mb": (attn_out_fwd_metadata__recv_transfer_sz // (1024 * 1024)).tolist(),
-
-                "qkv_fwd_metadata__send_transfer_sz_mb_to_others": (qkv_fwd_metadata__send_transfer_sz_to_others // (1024 * 1024)).tolist(),
-                "qkv_fwd_metadata__recv_transfer_sz_mb_from_others": (qkv_fwd_metadata__recv_transfer_sz_to_others // (1024 * 1024)).tolist(),
-
-                "max_comm_budget_all_rank_mb": max_comm_budget_all_rank // (1024 * 1024),
-                "max_buffer_budget_all_rank_mb": max_buffer_budget_all_rank // (1024 * 1024),
-                "bandwidth_mb": bandwidth_mb,
-                "send_time_ms": send_time_ms.tolist(),
-                "recv_time_ms": recv_time_ms.tolist(),
-                "max_send_time_ms": max_send_time_ms,
-                "max_recv_time_ms": max_recv_time_ms,
-                "seq_len": seq_len if seq_len is not None else None,
-            }) + "\n")
-
-        network_inspect_summary_file = os.path.join(output_dir, "network_inspect.summary.jsonl")
-        with open(network_inspect_summary_file, "a") as f:
-            f.write(json.dumps({
-                "sample_id": sample_id,
-                "is_ping": is_ping,
-                "tolerance_factor": tolerance_factor,
-                "qkv_fwd_send_mb": (qkv_fwd_metadata__send_transfer_sz_to_others // (1024 * 1024)).tolist(),
-                "qkv_fwd_recv_mb": (qkv_fwd_metadata__recv_transfer_sz_to_others // (1024 * 1024)).tolist(),
-
-                "max_comm_budget_all_rank_mb": max_comm_budget_all_rank // (1024 * 1024),
-                "max_buffer_budget_all_rank_mb": max_buffer_budget_all_rank // (1024 * 1024),
-                "send_time_ms": send_time_ms.tolist(),
-                "recv_time_ms": recv_time_ms.tolist(),
-                "max_send_time_ms": max_send_time_ms,
-                "max_recv_time_ms": max_recv_time_ms,
-                "seq_len": seq_len if seq_len is not None else None,
-            }) + "\n")
+        try:
+            network_inspect_file = os.path.join(output_dir, "network_inspect.jsonl")
+            with open(network_inspect_file, "a") as f:
+                f.write(json.dumps({
+                    "sample_id": sample_id,
+                    "is_ping": is_ping,
+                    "tolerance_factor": tolerance_factor,
+                    "qkv_fwd_metadata__send_transfer_sz_mb": (qkv_fwd_metadata__send_transfer_sz // (1024 * 1024)).tolist(),
+                    "qkv_fwd_metadata__recv_transfer_sz_mb": (qkv_fwd_metadata__recv_transfer_sz // (1024 * 1024)).tolist(),
+                    "attn_out_fwd_metadata__send_transfer_sz_mb": (attn_out_fwd_metadata__send_transfer_sz // (1024 * 1024)).tolist(),
+                    "attn_out_fwd_metadata__recv_transfer_sz_mb": (attn_out_fwd_metadata__recv_transfer_sz // (1024 * 1024)).tolist(),
+
+                    "qkv_fwd_metadata__send_transfer_sz_mb_to_others": (qkv_fwd_metadata__send_transfer_sz_to_others // (1024 * 1024)).tolist(),
+                    "qkv_fwd_metadata__recv_transfer_sz_mb_from_others": (qkv_fwd_metadata__recv_transfer_sz_to_others // (1024 * 1024)).tolist(),
+
+                    "max_comm_budget_all_rank_mb": max_comm_budget_all_rank // (1024 * 1024),
+                    "max_buffer_budget_all_rank_mb": max_buffer_budget_all_rank // (1024 * 1024),
+                    "bandwidth_mb": bandwidth_mb,
+                    "send_time_ms": send_time_ms.tolist(),
+                    "recv_time_ms": recv_time_ms.tolist(),
+                    "max_send_time_ms": max_send_time_ms,
+                    "max_recv_time_ms": max_recv_time_ms,
+                    "seq_len": seq_len if seq_len is not None else None,
+                }) + "\n")
+
+            network_inspect_summary_file = os.path.join(output_dir, "network_inspect.summary.jsonl")
+            with open(network_inspect_summary_file, "a") as f:
+                f.write(json.dumps({
+                    "sample_id": sample_id,
+                    "is_ping": is_ping,
+                    "tolerance_factor": tolerance_factor,
+                    "qkv_fwd_send_mb": (qkv_fwd_metadata__send_transfer_sz_to_others // (1024 * 1024)).tolist(),
+                    "qkv_fwd_recv_mb": (qkv_fwd_metadata__recv_transfer_sz_to_others // (1024 * 1024)).tolist(),
+
+                    "max_comm_budget_all_rank_mb": max_comm_budget_all_rank // (1024 * 1024),
+                    "max_buffer_budget_all_rank_mb": max_buffer_budget_all_rank // (1024 * 1024),
+                    "send_time_ms": send_time_ms.tolist(),
+                    "recv_time_ms": recv_time_ms.tolist(),
+                    "max_send_time_ms": max_send_time_ms,
+                    "max_recv_time_ms": max_recv_time_ms,
+                    "seq_len": seq_len if seq_len is not None else None,
+                }) + "\n")
+        except Exception as e:
+            print(f"🟡 Error writing network inspect summary file: {e}")
 
     ret = dict(
         max_comm_budget_all_rank=max_comm_budget_all_rank,
diff --git a/tests/test_megatron_e2e_pipeline.py b/tests/test_megatron_e2e_pipeline.py
index 80588f59..f5f97dbf 100644
--- a/tests/test_megatron_e2e_pipeline.py
+++ b/tests/test_megatron_e2e_pipeline.py
@@ -2,6 +2,10 @@
 Debug example:
 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 torchrun --nnodes 1 --nproc_per_node 2 test_megatron_e2e_pipeline.py --num-gpus-per-node 2 --pp-size 2 --num-microbatch 2
 """
+
+from d2.utils.traceback import enable_clickable_excepthook, enable_trace_calls
+enable_clickable_excepthook()
+
 import argparse
 from functools import partial
 import os
@@ -151,12 +155,13 @@ def init_megatron_e2e_test(
     token_bytes_kv = hidden_size_kv * dtype.itemsize // tp_size
     max_tokens_query = num_tokens * (world_size // tp_size)
     max_tokens_key_value = num_tokens * (world_size // tp_size)
-    buffer_size = (
-        token_bytes_q * max_tokens_query * 3 +
-        # lse_norm. TODO: the factor of 2 might be removed
-        num_heads * torch.float32.itemsize * 2 * max_tokens_query +
-        token_bytes_kv * max_tokens_key_value * max_cp_degree * 2
-    )
+    # buffer_size = (
+    #     token_bytes_q * max_tokens_query * 3 +
+    #     # lse_norm. TODO: the factor of 2 might be removed
+    #     num_heads * torch.float32.itemsize * 2 * max_tokens_query +
+    #     token_bytes_kv * max_tokens_key_value * max_cp_degree * 2
+    # )
+    buffer_size = 512 * 1024 ** 2 # 512MB
     print(f'{buffer_size=}', flush=True)
     parallel_config = ParallelConfig(
         tensor_model_parallel_size=tp_size,
@@ -286,6 +291,10 @@ def test(args):
         hidden_size_kv = (hidden_size_kv * hf_config.num_key_value_heads //
                           hf_config.num_attention_heads)
 
+    
+    num_tokens_for_cuda_graph = num_token_per_rank
+    os.environ["D2_SEQ_LEN"] = str(num_tokens_for_cuda_graph)
+
     worker: MegatronE2eWorker = init_megatron_e2e_test(
         hidden_size_q, hidden_size_kv, hf_config.num_attention_heads, num_tokens,
         world_size, max_cp_degree, tp_size, pp_size,
@@ -400,7 +409,7 @@ def test(args):
         torch.testing.assert_close(grad_orig_reimpl, grad_sample, rtol=1.1e-3, atol=1.1e-3)
     print(f"{worker.rank} finish pingpong")
 
-    print("=" * 20 + "forward_backward_batch attention server, done")
+    print("=" * 20 + "✅ forward_backward_batch attention server, done")
 
 
 if __name__ == "__main__":

From fa9cf1ef0a1c35fb8639ece84530fcf66feb2c85 Mon Sep 17 00:00:00 2001
From: GindaChen <32371474+GindaChen@users.noreply.github.com>
Date: Sun, 14 Dec 2025 01:58:25 +0000
Subject: [PATCH 4/4] Make TP running but we have numerical error

---
 d2/runtime/megatron/base_transformer_layer.py | 15 ++++++++++++---
 tests/test_megatron_e2e_pipeline.py           |  7 ++++++-
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/d2/runtime/megatron/base_transformer_layer.py b/d2/runtime/megatron/base_transformer_layer.py
index dd6a822f..3a9a6f60 100644
--- a/d2/runtime/megatron/base_transformer_layer.py
+++ b/d2/runtime/megatron/base_transformer_layer.py
@@ -13,6 +13,7 @@
     TransformerLayer as MegatronTransformerLayer,
     TransformerLayerSubmodules,
 )
+from megatron.core import parallel_state
 
 from d2.runtime.megatron.packed_seq_params import PingPangSingleStepPackedSeqParams
 
@@ -42,11 +43,17 @@ def __init__(
         self.post_attn_cuda_graph = None
 
     def init_pre_attn_cuda_graph(self, prev_layer: "Optional[TransformerLayer]", seq_len: int, device: torch.device, dtype: torch.dtype):
+        # TODO: May need to also check `self.config.sequence_parallel_size`. If not, just assume SP == TP
+        # use_sp = self.config.sequence_parallel
+        tp = parallel_state.get_tensor_model_parallel_world_size()
         if prev_layer is None:
             static_input = torch.zeros((seq_len, 1, self.config.hidden_size), device=device, dtype=dtype, requires_grad=True)
             self.pre_attn_cuda_graph = torch.cuda.make_graphed_callables(self.__pre_attn_cuda_graph, (static_input,))
         else:
-            static_core_attn_out = torch.zeros((seq_len, 1, prev_layer.self_attention.query_projection_size), device=device, dtype=dtype, requires_grad=True)
+            hidden_size_tp = self.config.hidden_size // tp
+            # prev_layer.self_attention.query_projection_size // tp
+            static_core_attn_out = torch.zeros((seq_len * tp, 1, hidden_size_tp), device=device, dtype=dtype, requires_grad=True)
+            # static_residual = torch.zeros((seq_len // tp, 1, self.config.hidden_size), device=device, dtype=dtype, requires_grad=True)
             static_residual = torch.zeros((seq_len, 1, self.config.hidden_size), device=device, dtype=dtype, requires_grad=True)
             def post_then_pre_core_attn_cuda_graph(core_attn_out: Tensor, residual: Tensor):
                 hidden_states = prev_layer._post_attn_cuda_graph(core_attn_out, residual)
@@ -54,7 +61,10 @@ def post_then_pre_core_attn_cuda_graph(core_attn_out: Tensor, residual: Tensor):
             self.pre_attn_cuda_graph = torch.cuda.make_graphed_callables(post_then_pre_core_attn_cuda_graph, (static_core_attn_out, static_residual))
     
     def init_post_attn_cuda_graph(self, seq_len: int, device: torch.device, dtype: torch.dtype):
-        static_core_attn_out = torch.zeros((seq_len, 1, self.self_attention.query_projection_size), device=device, dtype=dtype, requires_grad=True)
+        tp = parallel_state.get_tensor_model_parallel_world_size()
+        hidden_size_tp = self.config.hidden_size // tp
+        # prev_layer.self_attention.query_projection_size // tp
+        static_core_attn_out = torch.zeros((seq_len * tp, 1, hidden_size_tp), device=device, dtype=dtype, requires_grad=True)
         static_residual = torch.zeros((seq_len, 1, self.config.hidden_size), device=device, dtype=dtype, requires_grad=True)
         self.post_attn_cuda_graph = torch.cuda.make_graphed_callables(self._post_attn_cuda_graph, (static_core_attn_out, static_residual))
 
@@ -261,7 +271,6 @@ def _forward_pre_core_attn(
         # q, k, v
         log_memory_usage(f"(L{self.layer_number}) _forward_pre_core_attn:(before qkv)")
         query, key, value = self.self_attention.get_query_key_value_tensors(input_layernorm_output, None)
-        # print(f"🟡 query: {query.shape} (query.dtype: {query.dtype}), key: {key.shape} (key.dtype: {key.dtype}), value: {value.shape} (value.dtype: {value.dtype})")
 
         #### Some code in core_attention. This is because we don't want the pos embedding
         # being handled in the attention layout (the pos id will be hard to handle)
diff --git a/tests/test_megatron_e2e_pipeline.py b/tests/test_megatron_e2e_pipeline.py
index f5f97dbf..2386d4eb 100644
--- a/tests/test_megatron_e2e_pipeline.py
+++ b/tests/test_megatron_e2e_pipeline.py
@@ -304,7 +304,12 @@ def test(args):
     worker.init(model_path, seed=seed)
 
     # torch.distributed.breakpoint()
-    worker.train_module[0].module.module.decoder.init_layer_cuda_graphs()  # FIXME: hardcode for now, where to put?
+    # FIXME: hardcode for now, where to put?
+    print("init cuda graphs")
+    worker.train_module[0].module.module.decoder.init_layer_cuda_graphs()
+    print("init cuda graphs done")
+    # return
+    
     # set again to potentially adapt to the ray launch case.
     set_random_seed(seed, set_megatron=False)