diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
new file mode 100644
index 00000000..b15d11e0
--- /dev/null
+++ b/.github/workflows/ci.yml
@@ -0,0 +1,54 @@
+name: CI
+
+on:
+  push:
+    branches:
+      - main
+      - llm-dev
+  pull_request:
+    
+jobs:
+  tests:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          submodules: recursive
+
+      - name: Install AstraSim system dependencies
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y \
+            gcc-11 g++-11 make cmake \
+            clang-format \
+            libboost-dev libboost-program-options-dev \
+            libprotobuf-dev protobuf-compiler \
+            openmpi-bin openmpi-doc libopenmpi-dev
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v2
+
+      - name: Sync Python dependencies
+        run: uv sync
+
+      - name: Determine AstraSim commit
+        run: echo "ASTRASIM_SHA=$(git -C astra-sim rev-parse HEAD)" >> $GITHUB_ENV
+
+      - name: Restore AstraSim build cache
+        id: ascache
+        uses: actions/cache@v4
+        with:
+          path: astra-sim/build/astra_analytical/build
+          key: astra-${{ env.ASTRASIM_SHA }}
+
+      - name: Build AstraSim (cache miss)
+        if: steps.ascache.outputs.cache-hit != 'true'
+        run: |
+          chmod +x scripts/install_astrasim.sh
+          scripts/install_astrasim.sh
+
+      - name: Run pytest
+        env:
+          DEEPFLOW_VALIDATION_WORKERS: "1"
+        run: uv run pytest
diff --git a/.gitignore b/.gitignore
index 0d6ea66b..fd069542 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,8 +1,19 @@
 *.swp
 *.swo
+.cursor/
 __pycache__/
 output/
 venv/
 astra_cache/
 log/*.log
-tmp/
\ No newline at end of file
+tmp/
+output*/
+dbg_hw_conf/
+tools/fault_sweep_results*
+tools/*.png
+tools/*.tsv
+tools/*.json
+tools/gmap_artifacts/
+*.md
+*.et
+tools/comp/wrapper_outputs/*
\ No newline at end of file
diff --git a/.gitmodules b/.gitmodules
index 41a0a515..33f7ddab 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -1,3 +1,6 @@
 [submodule "astra-sim"]
 	path = astra-sim
 	url = https://github.com/gkarfakis19/astra-sim
+[submodule "extern/graph_frontend/param"]
+	path = extern/graph_frontend/param
+	url = https://github.com/facebookresearch/param.git
diff --git a/AGENTS.md b/AGENTS.md
new file mode 100644
index 00000000..ce5a1188
--- /dev/null
+++ b/AGENTS.md
@@ -0,0 +1,12 @@
+DeepFlow quick brief for future agents
+
+- Purpose: Python simulator for LSTM/GEMM/LLM performance, memory, and (rudimentary) energy. Entry point `run_perf.py` loads YAML configs (hardware + model) and builds compute/comm graphs, then simulates analytically or via AstraSim.
+- Execution backends (set in hardware config `execution_backend`): analytical (ring-only, no congestion); hybrid (AstraSim for transformer block graph, analytical pipeline); full_astrasim_hierarchical (AstraSim separately for transformer and pipeline graphs, assumes no congestion between them); full_astrasim_flattened (one big combined graph in AstraSim, slowest/most detailed).
+- Parallelism model: supports DP, TP (+optional sequence parallel), CP, GPipe-style pipeline (`lp` stages), and hybrid TP+CP. Pipeline graph assumes uniform per-layer timing and GPipe scheduling; gradient accumulation builds a no-DP variant so only the last accumulation step does DP all-reduce.
+- Communication assumptions: blocking collectives; analytical path uses ring formulas; AstraSim collectives flatten 2D topologies to 1D for collectives. Hierarchical mode requires the first active network dimension map exactly to {tp,cp}; non-ring topologies are rejected in analytical mode. Faulty links only valid in full AstraSim modes.
+- ZeRO assumptions: stage 0/1 ≈ DDP; stage 2 adds per-layer param AG + grad RS sizes; stage 3 adds per-layer param materialization (counted as comm + ephemeral memory, no overlap). MoE assumes perfect balance and TP all-to-all; multi-GPU MoE is WIP.
+- Attention: dense by default. FlashAttention supported for training forward/prefill when `attention.use_flashattention=true` and tile size is set manually; decode path ignores FlashAttention. Sliding-window/MLA are WIP.
+- Memory model: static formulas for activations/weights/optimizer; KV-cache memory estimation WIP. Peak memory compared to DRAM per device; ZeRO-3 ephemeral gather tracked separately. No fragmentation/eviction modeling.
+- Inference: `TimeCalculationLLMInference` handles prefill + sampled decode steps (linear interpolation between samples). DP acts as replica multiplier only; CP decode is noted WIP. FlashAttention disabled for decode.
+- Artifacts/logs: outputs under `output/` (mode subdir). AstraSim artifacts live in `astra_cache/` and optionally `output/LLM/astra_hier` or `astra_flat` when persistence flags are set. Current logs/marker: `log/this_folder_contains_astrasim_logs`, `log/log.log` hold prior runs. Debug hardware snapshots in `dbg_hw_conf/20251105_173209/` (e.g., `baseline_001_tp64_cp1_dp1_lp1.yaml`, `baseline_053_tp1_cp1_dp64_lp1.yaml`).
+- Key environment flags: `DEEPFLOW_VISUALIZE_GRAPHS`, `DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS`, `DEEPFLOW_PERSIST_ARTIFACT_VIZ`, `DEEPFLOW_ASTRA_CACHE_MODE` (default CACHE_READWRITE).
diff --git a/LLM_execution.py b/LLM_execution.py
index a4e59768..65933d11 100644
--- a/LLM_execution.py
+++ b/LLM_execution.py
@@ -4,11 +4,284 @@
 import os
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Dict, Tuple, Optional, List, Set
+from typing import Any, Dict, Tuple, Optional, List, Set, Sequence, Mapping
+from collections import defaultdict
+from collections import defaultdict
 
 import simulate_LLM
 from astrasim_lib import run_astra_simulation_only_onepath
+from astrasim_lib.fault_projection import FaultProjectionResult, FaultSpace
+from astrasim_lib.layout_utils import axis_layout_from_descriptor
 from simulate_LLM import Graph
+from util import log_message
+from typing import Iterable
+
+
+def _mode_label(mode: Any) -> str:
+    for attr in ("value", "name"):
+        if hasattr(mode, attr):
+            return str(getattr(mode, attr)).lower()
+    return str(mode).lower()
+
+
+def _copy_node_metadata(source: simulate_LLM.Node, target: simulate_LLM.Node) -> None:
+    for attr in (
+        "micro_batch_index",
+        "layer_index",
+        "direction",
+        "stage_id",
+        "tp_rank",
+        "cp_rank",
+    ):
+        if hasattr(source, attr):
+            setattr(target, attr, getattr(source, attr))
+
+
+def _copy_edge_metadata(source: simulate_LLM.Edge, target: simulate_LLM.Edge) -> None:
+    for attr in (
+        "local_hw_id",
+        "stage_id",
+        "micro_batch_index",
+        "layer_index",
+        "direction",
+        "tp_rank",
+        "cp_rank",
+    ):
+        if hasattr(source, attr):
+            setattr(target, attr, getattr(source, attr))
+
+
+def _detach_edge(parent: Any, child: Any) -> None:
+    children = getattr(parent, "children", None)
+    if isinstance(children, list):
+        try:
+            children.remove(child)
+        except ValueError:
+            pass
+    parents = getattr(child, "parents", None)
+    if isinstance(parents, list):
+        try:
+            parents.remove(parent)
+        except ValueError:
+            pass
+
+
+def _connect_edge(parent: Any, child: Any) -> None:
+    if child in getattr(parent, "children", []):
+        return
+    parent.add_child(child)
+
+
+def _split_tp_node(
+    node: simulate_LLM.Node,
+    tp_children: List[simulate_LLM.Edge],
+    overlap: float,
+) -> None:
+    if overlap <= 0.0 or not tp_children:
+        return
+
+    duration = float(getattr(node, "duration", 0.0) or 0.0)
+    if duration <= 0.0:
+        return
+
+    if overlap >= 1.0:
+        parents = list(getattr(node, "parents", []))
+        tp_succs: List[Any] = []
+        for tp_edge in tp_children:
+            _detach_edge(node, tp_edge)
+            for parent in parents:
+                _connect_edge(parent, tp_edge)
+            for succ in list(getattr(tp_edge, "children", []) or []):
+                tp_succs.append(succ)
+        for succ in tp_succs:
+            _connect_edge(node, succ)
+        return
+
+    head_duration = duration * (1.0 - overlap)
+    tail_duration = duration * overlap
+    if head_duration <= 0.0 or tail_duration <= 0.0:
+        return
+
+    tail = node
+    tail.duration = tail_duration
+    head = simulate_LLM.Node(
+        name=f"{tail.name}_head",
+        op_id=getattr(tail, "op_id", 0),
+        hw_id=tail.hw_id,
+        duration=head_duration,
+        fwd=tail.fwd,
+    )
+    _copy_node_metadata(tail, head)
+
+    parents = list(getattr(tail, "parents", []))
+    for parent in parents:
+        _detach_edge(parent, tail)
+        _connect_edge(parent, head)
+
+    _connect_edge(head, tail)
+
+    for tp_edge in tp_children:
+        _detach_edge(tail, tp_edge)
+        _connect_edge(head, tp_edge)
+        for succ in list(getattr(tp_edge, "children", []) or []):
+            _connect_edge(tail, succ)
+
+
+def _apply_tp_overlap_transforms(root: Any, parallelism_mode: Any, tp_overlap: float, tp_sp_overlap: float) -> Any:
+    mode = _mode_label(parallelism_mode)
+    if mode == "tensor_sequence":
+        overlap = tp_sp_overlap
+    elif mode in {"tensor", "tensor_context_hybrid"}:
+        overlap = tp_overlap
+    else:
+        return root
+    if overlap <= 0.0:
+        return root
+
+    nodes_to_process: List[simulate_LLM.Node] = []
+    visited: Set[int] = set()
+    stack: List[Any] = list(root) if isinstance(root, (list, tuple)) else [root]
+    while stack:
+        obj = stack.pop()
+        obj_id = id(obj)
+        if obj_id in visited:
+            continue
+        visited.add(obj_id)
+        if isinstance(obj, simulate_LLM.Node):
+            nodes_to_process.append(obj)
+        for child in getattr(obj, "children", []):
+            stack.append(child)
+
+    for node in nodes_to_process:
+        tp_children = [
+            child for child in getattr(node, "children", [])
+            if isinstance(child, simulate_LLM.Edge) and getattr(child, "comm_interconnect_type", None) == "tp"
+        ]
+        if not tp_children:
+            continue
+        _split_tp_node(node, tp_children, overlap)
+    return root
+
+
+def _apply_cp_overlap_transforms(root: Any, parallelism_mode: Any, cp_overlap: float) -> Any:
+    mode = _mode_label(parallelism_mode)
+    if mode not in {"context", "tensor_context_hybrid"}:
+        return root
+    overlap = cp_overlap
+    if overlap <= 0.0:
+        return root
+
+    cp_edges: List[simulate_LLM.Edge] = []
+    visited: Set[int] = set()
+    stack: List[Any] = list(root) if isinstance(root, (list, tuple)) else [root]
+    while stack:
+        obj = stack.pop()
+        obj_id = id(obj)
+        if obj_id in visited:
+            continue
+        visited.add(obj_id)
+        if isinstance(obj, simulate_LLM.Edge) and getattr(obj, "comm_interconnect_type", None) == "cp":
+            cp_edges.append(obj)
+        for child in getattr(obj, "children", []):
+            stack.append(child)
+
+    for edge in cp_edges:
+        _split_cp_edge(edge, overlap)
+    return root
+
+
+def _split_cp_edge(edge: simulate_LLM.Edge, overlap: float) -> None:
+    attention_children = [
+        child for child in getattr(edge, "children", [])
+        if isinstance(child, simulate_LLM.Node) and "attention" in str(getattr(child, "name", "")).lower()
+    ]
+    if not attention_children:
+        return
+
+    total_bytes = int(getattr(edge, "comm_size_bytes", 0) or 0)
+    preds = list(getattr(edge, "parents", []))
+    succs = list(getattr(edge, "children", []))
+
+    if overlap >= 1.0 or total_bytes <= 0:
+        for attention in attention_children:
+            _detach_edge(edge, attention)
+            for pred in preds:
+                _connect_edge(pred, attention)
+        for succ in succs:
+            if succ in attention_children:
+                continue
+            for attention in attention_children:
+                _connect_edge(attention, succ)
+            _connect_edge(edge, succ)
+        return
+
+    block_bytes = int(math.ceil(total_bytes * (1.0 - overlap)))
+    ovlp_bytes = max(0, total_bytes - block_bytes)
+    if block_bytes <= 0:
+        _split_cp_edge(edge, 1.0)
+        return
+
+    block_edge = simulate_LLM.Edge(
+        name=f"{edge.name}_block",
+        op_id=getattr(edge, "op_id", 0),
+        duration=0,
+        is_dp=edge.is_dp,
+        comm_size_bytes=block_bytes,
+        comm_type=edge.comm_type,
+        participants=edge.participants,
+        comm_interconnect_type=edge.comm_interconnect_type,
+    )
+    ovlp_edge = None
+    if ovlp_bytes > 0:
+        ovlp_edge = simulate_LLM.Edge(
+            name=f"{edge.name}_ovlp",
+            op_id=getattr(edge, "op_id", 0),
+            duration=0,
+            is_dp=edge.is_dp,
+            comm_size_bytes=ovlp_bytes,
+            comm_type=edge.comm_type,
+            participants=edge.participants,
+            comm_interconnect_type=edge.comm_interconnect_type,
+        )
+        _copy_edge_metadata(edge, ovlp_edge)
+    _copy_edge_metadata(edge, block_edge)
+
+    for pred in preds:
+        _detach_edge(pred, edge)
+        _connect_edge(pred, block_edge)
+
+    for attention in attention_children:
+        _detach_edge(edge, attention)
+        _connect_edge(block_edge, attention)
+
+    if ovlp_edge is not None:
+        _connect_edge(block_edge, ovlp_edge)
+
+    for succ in succs:
+        if succ in attention_children:
+            continue
+        _detach_edge(edge, succ)
+        for attention in attention_children:
+            _connect_edge(attention, succ)
+        if ovlp_edge is not None:
+            _connect_edge(ovlp_edge, succ)
+        else:
+            _connect_edge(block_edge, succ)
+
+
+def apply_overlap_transforms(
+    root: Any,
+    parallelism_mode: Any,
+    tp_overlap: float,
+    tp_sp_overlap: float,
+    cp_overlap: float,
+) -> Any:
+    """Apply TP/TP_SP/CP overlap rewrites to a constructed graph."""
+    if root is None:
+        return None
+    root = _apply_tp_overlap_transforms(root, parallelism_mode, tp_overlap, tp_sp_overlap)
+    root = _apply_cp_overlap_transforms(root, parallelism_mode, cp_overlap)
+    return root
 
 
 def _env_flag(name: str) -> bool:
@@ -48,6 +321,8 @@ def __init__(
         self,
         pipeline_graph: Graph,
         transformer_graph: Graph,
+        *,
+        rank_layout: Optional[Dict[str, Any]] = None,
     ) -> None:
         if transformer_graph is None:
             raise ValueError("Transformer graph is required for flattening")
@@ -63,11 +338,50 @@ def __init__(
         par_degree = transformer_graph.tp * transformer_graph.cp
         self._par_degree = max(1, int(par_degree))
         self._zero_stage = int(getattr(transformer_graph, "misc_metadata", {}).get("dp_zero_stage", 0))
+        self._layout_axis_order: Optional[List[str]] = None
+        self._layout_axis_sizes: Dict[str, int] = {}
+        self._layout_axis_strides: Dict[str, int] = {}
+        self._tp_size = int(getattr(transformer_graph, "tp", 1))
+        self._cp_size = int(getattr(transformer_graph, "cp", 1))
+        self._lp_size = int(getattr(pipeline_graph, "lp", 1))
+        if rank_layout is not None:
+            self._configure_rank_layout(rank_layout)
 
         # Track original ZeRO-3 transformer gather edges -> per-rank clones
         self._clone_cache: Dict[int, Any] = {}
         self._op_id_counter: int = 0
-        
+        self._stage_span = 1
+
+    def _configure_rank_layout(self, descriptor: Dict[str, Any]) -> None:
+        axis_order = list(descriptor.get("axis_order", []))
+        axis_sizes = dict(descriptor.get("axis_sizes", {}))
+        axis_strides = dict(descriptor.get("axis_strides", {}))
+        if not axis_order:
+            self._layout_axis_order = None
+            return
+        self._layout_axis_order = axis_order
+        self._layout_axis_sizes = axis_sizes
+        if not axis_strides:
+            span = 1
+            for axis in axis_order:
+                axis_strides[axis] = span
+                span *= axis_sizes.get(axis, 1)
+        self._layout_axis_strides = axis_strides
+        self._tp_size = max(1, axis_sizes.get("tp", self._tp_size))
+        self._cp_size = max(1, axis_sizes.get("cp", self._cp_size))
+        self._lp_size = max(1, axis_sizes.get("lp", self._lp_size))
+        if self._par_degree != self._tp_size * self._cp_size:
+            raise ValueError(
+                f"Inconsistent tensor/context parallel factors: tp={self._tp_size}, cp={self._cp_size}, "
+                f"product does not equal par_degree={self._par_degree}"
+            )
+        stage_span = 1
+        for axis in axis_order:
+            if axis == "dp":
+                continue
+            stage_span *= axis_sizes.get(axis, 1)
+        self._stage_span = stage_span
+
     def _should_shard_zero3_transformer(self, edge: Any) -> bool:
         if self._par_degree <= 1 or self._zero_stage < 3:
             return False
@@ -170,6 +484,35 @@ def _clone(self, obj: Any) -> Any:
                     obj.duration,
                     fwd=obj.fwd,
                 )
+            elif "optimizer" in obj.name:
+                # Optimizer nodes need to be expanded per TP rank
+                cloned_nodes = []
+                for tp_rank in range(self._par_degree):
+                    hw_id = self._hw_id_for_rank(obj.hw_id, tp_rank)
+                    cloned_node = simulate_LLM.Node(
+                        name=f"{obj.name}_rank{tp_rank}",
+                        op_id=self._next_op_id(),
+                        hw_id=hw_id,
+                        duration=obj.duration,
+                        fwd=obj.fwd,
+                    )
+                    cloned_nodes.append(cloned_node)
+                
+                cloned_tuple = tuple(cloned_nodes)
+                self._clone_cache[obj_id] = cloned_tuple
+                
+                # We also need to copy metadata if any
+                for cloned_node in cloned_nodes:
+                    self._copy_metadata(obj, cloned_node)
+                    
+                # Attach children (if any) to all clones? 
+                # Let's follow the pattern:
+                for child in getattr(obj, "children", []):
+                    child_clone = self._clone(child)
+                    if child_clone is not None:
+                        self._attach(cloned_tuple, child_clone)
+                        
+                return cloned_tuple
             else:
                 cloned = simulate_LLM.Node(
                     obj.name,
@@ -179,6 +522,7 @@ def _clone(self, obj: Any) -> Any:
                     fwd=obj.fwd,
                 )
 
+
             # following _expand_transformer_node logic, we need to find siblings that are zero3 tp_shard=True.
             zero3_attachments: List[simulate_LLM.Edge] = []
             for parent in getattr(obj, "parents", []):
@@ -240,6 +584,52 @@ def _clone(self, obj: Any) -> Any:
 
         raise TypeError(f"Unsupported graph element type: {type(obj)!r}")
 
+    def _propagate_local_hw_ids(self, roots: Any) -> None:
+        stack: List[Any] = []
+        if isinstance(roots, (list, tuple)):
+            stack.extend(list(roots))
+        elif roots is not None:
+            stack.append(roots)
+        visited: Set[int] = set()
+        while stack:
+            obj = stack.pop()
+            obj_id = id(obj)
+            if obj_id in visited:
+                continue
+            visited.add(obj_id)
+
+            children = getattr(obj, "children", [])
+            if isinstance(children, (list, tuple)):
+                stack.extend(children)
+            elif children is not None:
+                stack.append(children)
+
+            parents = getattr(obj, "parents", None)
+            if parents:
+                if isinstance(parents, (list, tuple)):
+                    stack.extend(parents)
+                else:
+                    stack.append(parents)
+
+            if isinstance(obj, simulate_LLM.Edge):
+                new_hw = None
+                for parent in getattr(obj, "parents", []) or []:
+                    parent_hw = getattr(parent, "hw_id", None)
+                    if parent_hw is None:
+                        parent_hw = getattr(parent, "local_hw_id", None)
+                    if parent_hw is not None and parent_hw >= 0:
+                        new_hw = parent_hw
+                        break
+                if new_hw is None:
+                    stage_id = getattr(obj, "stage_id", None)
+                    if stage_id is not None:
+                        try:
+                            new_hw = self._hw_id_for_rank(stage_id, 0)
+                        except Exception:
+                            new_hw = None
+                if new_hw is not None and new_hw >= 0:
+                    obj.local_hw_id = new_hw
+
 
     def _expand_transformer_node(self, node: simulate_LLM.Node) -> Tuple[Any, ...]:
         node_id = id(node)
@@ -421,6 +811,14 @@ def _expand_transformer_node(self, node: simulate_LLM.Node) -> Tuple[Any, ...]:
             child_clone = self._clone(child)
             if child_clone is None:
                 continue
+            
+            # Special handling for optimizer nodes: 1-to-1 connection
+            if isinstance(child, simulate_LLM.Node) and "optimizer" in child.name and isinstance(child_clone, (list, tuple)):
+                if len(child_clone) == len(rank_tails):
+                    for r in range(len(rank_tails)):
+                        rank_tails[r].add_child(child_clone[r])
+                    continue
+            
             self._attach(downstream_parents, child_clone)
 
         for zero3_edge in zero3_attachments:
@@ -501,7 +899,33 @@ def _next_op_id(self) -> int:
 
     def _hw_id_for_rank(self, stage_id: int, tp_rank: int) -> int:
         stage_int = int(stage_id) if stage_id is not None else 0
-        return stage_int * self._par_degree + tp_rank
+        tp_rank_int = int(tp_rank)
+        if tp_rank_int < 0 or tp_rank_int >= self._par_degree:
+            raise ValueError(f"tp_rank {tp_rank_int} is out of range for par_degree {self._par_degree}")
+        if not self._layout_axis_order:
+            return stage_int * self._par_degree + tp_rank_int
+
+        coords: Dict[str, int] = {}
+        if "tp" in self._layout_axis_order:
+            coords["tp"] = tp_rank_int % self._tp_size
+        if "cp" in self._layout_axis_order:
+            coords["cp"] = (tp_rank_int // self._tp_size) % self._cp_size
+        if "lp" in self._layout_axis_order:
+            if stage_int < 0 or stage_int >= self._lp_size:
+                raise ValueError(f"stage_id {stage_int} is out of range for lp={self._lp_size}")
+            coords["lp"] = stage_int % self._lp_size
+
+        linear_rank = 0
+        for axis in self._layout_axis_order:
+            coord = coords.get(axis, 0)
+            size = self._layout_axis_sizes.get(axis, 1)
+            if coord < 0 or coord >= size:
+                raise ValueError(f"Coordinate {coord} for axis '{axis}' is out of range <{size}")
+            stride = self._layout_axis_strides.get(axis)
+            if stride is None:
+                raise KeyError(f"Rank layout stride missing for axis '{axis}'")
+            linear_rank += coord * stride
+        return linear_rank
     
     
 class LLMExecutionDispatcher:
@@ -514,6 +938,7 @@ def __init__(
         transformer_graph: Optional[Graph] = None,
         transformer_forward_root: Optional[Any] = None,
         transformer_backward_root: Optional[Any] = None,
+        no_data_parallel: bool = False,
     ) -> None:
         self.time_calc = time_calc
         self.pipeline_graph = pipeline_graph
@@ -523,6 +948,359 @@ def __init__(
         self.transformer_forward_root = transformer_forward_root
         self.transformer_backward_root = transformer_backward_root
         self.flattened_root: Optional[Any] = None
+        self._transformer_rank_layout: Dict[str, Any] = {}
+        self._pipeline_rank_layout: Dict[str, Any] = {}
+        self._network_dimensions: Tuple[Any, ...] = tuple()
+        self._axis_dimension_map: Dict[str, int] = {}
+        self._transformer_stage_dp_faults: Dict[Tuple[int, int], Tuple[Tuple[int, int, float], ...]] = {}
+        self._transformer_stage_timings: Dict[Tuple[int, int], TransformerTimings] = {}
+        self._transformer_baseline_timings: Optional[TransformerTimings] = None
+        self.no_data_parallel = bool(no_data_parallel)
+        self._rank_layout = self._build_rank_layout_descriptor()
+        self._first_dim_optimize_cfg: Optional[Dict[str, Any]] = getattr(self, "_first_dim_optimize_cfg", None)
+        self._fault_space: Optional[FaultSpace] = None
+        self._fault_projections: Dict[str, FaultProjectionResult] = {}
+        self._initialize_fault_mappings()
+
+    def _build_rank_layout_descriptor(self) -> Dict[str, Any]:
+        hw_config = getattr(self.time_calc, "hw_config", None)
+        layout = getattr(hw_config, "network_layout", None)
+        dimensions = getattr(layout, "dimensions", None) if layout is not None else None
+        if not dimensions:
+            return {}
+        self._network_dimensions = tuple(dimensions)
+        optimize_cfg: Optional[Dict[str, Any]] = None
+        for idx, dim in enumerate(dimensions):
+            if getattr(dim, "optimize_2dmap", False):
+                if idx != 0:
+                    raise ValueError("optimize_2dmap is only supported on the first network dimension.")
+                if optimize_cfg is not None:
+                    raise ValueError("Multiple network dimensions requested optimize_2dmap; only one is supported.")
+                topo_type = getattr(dim, "topology_type", None)
+                if not topo_type:
+                    raise ValueError("optimize_2dmap requires a topology type on the target dimension.")
+                if str(topo_type).strip().lower() == "kingmesh2d":
+                    raise NotImplementedError("optimize_2dmap is not implemented for KingMesh2D.")
+                size_value = getattr(dim, "size", None)
+                if size_value is None:
+                    raise ValueError("optimize_2dmap requires an explicit dimension size.")
+                dims_value = getattr(dim, "size_2d", None)
+                if dims_value is not None:
+                    dims_value = (int(dims_value[0]), int(dims_value[1]))
+                optimize_cfg = {
+                    "dimension_index": idx,
+                    "topology": str(topo_type),
+                    "size": int(size_value),
+                    "parallelisms": tuple(getattr(dim, "parallelisms", ()) or ()),
+                }
+                if dims_value:
+                    optimize_cfg["dims"] = dims_value
+        self._first_dim_optimize_cfg = optimize_cfg
+
+        def _safe_int(value: Any, default: int = 1) -> int:
+            try:
+                candidate = int(value)
+            except (TypeError, ValueError):
+                candidate = default
+            return max(1, candidate)
+
+        tp_size = _safe_int(getattr(self.pipeline_graph, "tp", getattr(self.time_calc, "tp", 1)))
+        cp_size = _safe_int(getattr(self.pipeline_graph, "cp", getattr(self.time_calc, "cp", 1)))
+        lp_size = _safe_int(getattr(self.pipeline_graph, "lp", getattr(self.time_calc, "lp", 1)))
+        dp_size = _safe_int(getattr(self.time_calc, "dp", 1))
+
+        axis_sizes: Dict[str, int] = {"tp": tp_size, "cp": cp_size, "lp": lp_size, "dp": dp_size}
+        axis_order: List[str] = []
+
+        # Enforce axis ordering for hierarchical/hybrid modes: the first active
+        # dimension must contain exactly {'tp','cp'}. Subsequent active
+        # dimensions may contain 'lp' (optionally combined with 'dp'). This
+        # matches the assumptions in the hierarchical graphs where a stage is a
+        # TP/CP cluster replicated across LP (and potentially DP) axes.
+        enforce_layout = self.time_calc.execution_mode in {
+            ExecutionMode.HYBRID,
+            ExecutionMode.FULL_ASTRASIM_HIERARCHICAL,
+        }
+        first_active_checked = False
+
+        for dim in dimensions:
+            dim_axes = [str(axis).strip().lower() for axis in getattr(dim, "parallelisms", ())]
+            declared = int(getattr(dim, "size", 1))
+
+            axes_without_dp = [axis for axis in dim_axes if axis != "dp"]
+            if enforce_layout and declared > 1 and not first_active_checked and axes_without_dp:
+                canon = sorted(axes_without_dp)
+                if canon != ["cp", "tp"] and canon != ["tp", "cp"] and (axis_sizes["tp"] > 1 or axis_sizes["cp"] > 1):
+                    raise ValueError(
+                        "For hierarchical/hybrid AstraSim modes, the first active network "
+                        "dimension must contain exactly {'tp','cp'} to represent the tensor/context cluster."
+                    )
+                first_active_checked = True
+
+            for name in dim_axes:
+                if name not in axis_sizes:
+                    raise ValueError(
+                        f"Unsupported parallelism axis '{name}' in network layout. "
+                        "Supported axes for AstraSim integration are: tp, cp, lp, dp."
+                    )
+                if name not in axis_order:
+                    axis_order.append(name)
+
+            expected = 1
+            for axis_name in dim_axes:
+                expected *= axis_sizes.get(axis_name, 1)
+            if expected != declared:
+                raise ValueError(
+                    f"Network dimension '{getattr(dim, 'label', getattr(dim, 'id', '<unnamed>'))}' "
+                    f"size mismatch: declared {declared}, but parallelism factors imply {expected}."
+                )
+
+        # Ensure the layout covers active parallel axes
+        if tp_size > 1 and "tp" not in axis_order:
+            raise ValueError("Network layout must include 'tp' when tensor parallelism > 1.")
+        if cp_size > 1 and "cp" not in axis_order:
+            raise ValueError("Network layout must include 'cp' when context parallelism > 1.")
+        if lp_size > 1 and "lp" not in axis_order:
+            raise ValueError("Network layout must include 'lp' when pipeline parallelism > 1.")
+
+        axis_strides: Dict[str, int] = {}
+        span = 1
+        for axis in axis_order:
+            axis_strides[axis] = span
+            span *= axis_sizes[axis]
+
+        descriptor = {
+            "axis_order": axis_order,
+            "axis_sizes": axis_sizes,
+            "axis_strides": axis_strides,
+            "stage_span": span,
+        }
+
+        def _subset_layout(allowed: Sequence[str]) -> Optional[Dict[str, Any]]:
+            subset = [axis for axis in axis_order if axis in allowed and axis_sizes.get(axis, 1) >= 1]
+            if not subset:
+                return None
+            strides: Dict[str, int] = {}
+            span = 1
+            for axis in subset:
+                strides[axis] = span
+                span *= axis_sizes[axis]
+            return {
+                "axis_order": subset,
+                "axis_sizes": {axis: axis_sizes[axis] for axis in subset},
+                "axis_strides": strides,
+                "stage_span": span,
+            }
+
+        # Transformer graphs only encode TP/CP axes; pipeline graphs encode LP
+        # (DP replicas are handled externally). Store these subsets so callers
+        # can attach the appropriate layout before invoking AstraSim.
+        self._transformer_rank_layout = _subset_layout(["tp", "cp"])
+        self._pipeline_rank_layout = _subset_layout(["lp", "dp"])
+
+        return descriptor
+
+    def _attach_optimize_hint(self, root: Any) -> None:
+        if root is None:
+            return
+        if self._first_dim_optimize_cfg:
+            setattr(root, "_optimize_2dmap", dict(self._first_dim_optimize_cfg))
+        elif hasattr(root, "_optimize_2dmap"):
+            delattr(root, "_optimize_2dmap")
+
+    def _log_fault_summary(self, axis_order: Sequence[str], axis_sizes: Mapping[str, int]) -> None:
+        if not axis_order:
+            return
+        space = self._fault_space
+        if space is None or not space.entries:
+            return
+        network_dims = getattr(self, "_network_dimensions", tuple())
+        if not network_dims:
+            log_message("[DeepFlow][faults] Hardware dimensions unavailable; skipping fault summary.")
+            return
+
+        def coords_to_dict(coords: Tuple[Tuple[str, int], ...]) -> Dict[str, int]:
+            return {name: value for name, value in coords}
+
+        entries = space.entries
+        for dim_index, dim in enumerate(network_dims):
+            axes = [str(axis).strip().lower() for axis in getattr(dim, "parallelisms", ())]
+            if not axes:
+                continue
+            replication_axes: List[str] = []
+            for future_dim in network_dims[dim_index + 1 :]:
+                replication_axes.extend(
+                    str(axis).strip().lower() for axis in getattr(future_dim, "parallelisms", ()) if axis
+                )
+            total_clusters = 1
+            for axis in replication_axes:
+                total_clusters *= max(1, axis_sizes.get(axis, 1))
+            axes_label = ", ".join(axes) or "<none>"
+            table_rows: List[Tuple[str, str, float]] = []
+            for entry in entries:
+                if not any(axis in axes for axis in entry.affected_axes):
+                    continue
+                src_dict = coords_to_dict(entry.src_coords)
+                dst_dict = coords_to_dict(entry.dst_coords)
+                src_label = f"{entry.original[0]} (" + ", ".join(f"{axis}={src_dict.get(axis, 0)}" for axis in axes) + ")"
+                dst_label = f"{entry.original[1]} (" + ", ".join(f"{axis}={dst_dict.get(axis, 0)}" for axis in axes) + ")"
+                higher_coords = " ".join(f"{axis}={src_dict.get(axis, 0)}" for axis in replication_axes)
+                table_rows.append((higher_coords, f"{src_label} <-> {dst_label}", float(entry.original[2])))
+
+            if not table_rows:
+                log_message(f"  • dim{dim_index} ({axes_label}) → Affected HW clusters: 0 / {total_clusters}", category="faults")
+                continue
+
+            log_message(
+                f"  • dim{dim_index} ({axes_label}) → Affected HW clusters: {len(table_rows)} / {total_clusters or 1}",
+                category="faults",
+            )
+            header_axes = " ".join(replication_axes) if replication_axes else ""
+            log_message(f"  {header_axes:<12} | {'source ↔ dest':<40} | derate", category="faults")
+            log_message(f"  {'-' * max(len(header_axes),12)}-+-{'-' * 40}-+-------", category="faults")
+
+            for higher_str, pair, derate in table_rows:
+                log_message(
+                    f"  {higher_str:<12} | {pair:<40} | {derate:>6.2f}",
+                    category="faults",
+                )
+            log_message(f"  {'-' * max(len(header_axes),12)}-+-{'-' * 40}-+-------", category="faults")
+
+
+
+
+    def _initialize_fault_mappings(self) -> None:
+        hw_config = getattr(self.time_calc, "hw_config", None)
+        network_layout = getattr(hw_config, "network_layout", None)
+        faulty_links: Tuple[Tuple[int, int, float], ...] = tuple(
+            getattr(network_layout, "faulty_links", ()) or ()
+        )
+        if faulty_links and self.time_calc.execution_mode in {ExecutionMode.ANALYTICAL, ExecutionMode.HYBRID}:
+            raise ValueError("Faulty links require full AstraSim execution; analytical/hybrid modes are not supported.")
+        axis_layout = axis_layout_from_descriptor(self._rank_layout)
+        axis_dim_map = self._axis_to_dimension_map(network_layout)
+        self._axis_dimension_map = dict(axis_dim_map)
+        self._fault_space = FaultSpace(
+            axis_layout,
+            faulty_links,
+            axis_to_dimension=axis_dim_map,
+        )
+        self._fault_projections = self._build_fault_projections_from_space(self._fault_space)
+        self._validate_fault_coverage()
+        self._transformer_stage_dp_faults = self._build_transformer_stage_dp_fault_map()
+        axis_order = self._rank_layout.get("axis_order", []) if isinstance(self._rank_layout, dict) else []
+        axis_sizes = self._rank_layout.get("axis_sizes", {}) if isinstance(self._rank_layout, dict) else {}
+        self._log_fault_summary(axis_order, axis_sizes)
+
+    def _axis_to_dimension_map(self, network_layout) -> Dict[str, int]:
+        mapping: Dict[str, int] = {}
+        if network_layout is None:
+            return mapping
+        dimensions = getattr(network_layout, "dimensions", None)
+        if not dimensions:
+            return mapping
+        for idx, dim in enumerate(dimensions):
+            for axis in getattr(dim, "parallelisms", ()) or ():
+                normalized = str(axis).strip().lower()
+                if normalized:
+                    mapping[normalized] = idx
+        return mapping
+
+    def _build_fault_projections_from_space(
+        self,
+        space: FaultSpace,
+    ) -> Dict[str, FaultProjectionResult]:
+        projections: Dict[str, FaultProjectionResult] = {}
+        if not space.entries:
+            return projections
+
+        global_axes = tuple(space.layout.axis_order)
+        if global_axes:
+            projections["global"] = space.project(global_axes)
+
+        transformer_axes: Tuple[str, ...] = tuple(
+            self._transformer_rank_layout.get("axis_order", ())
+        ) if self._transformer_rank_layout else tuple()
+        if transformer_axes:
+            projections["transformer"] = space.project(transformer_axes)
+
+        pipeline_axes: Tuple[str, ...] = tuple(
+            self._pipeline_rank_layout.get("axis_order", ())
+        ) if self._pipeline_rank_layout else tuple()
+        if pipeline_axes:
+            projections["pipeline"] = space.project(pipeline_axes)
+
+        return projections
+
+    def _validate_fault_coverage(self) -> None:
+        if self._fault_space is None:
+            return
+        covered: Set[Tuple[int, int, float]] = set()
+        for label in ("transformer", "pipeline"):
+            proj = self._fault_projections.get(label)
+            if proj:
+                covered.update(proj.covered_originals)
+        uncovered = [
+            entry.original
+            for entry in self._fault_space.entries
+            if entry.original not in covered
+        ]
+        if uncovered:
+            formatted = ", ".join(str(item) for item in uncovered)
+            raise ValueError(
+                "Faulty links do not map to any transformer or pipeline axis subset: "
+                f"{formatted}"
+            )
+
+    def _axis_value_from_coords(self, coords: Tuple[Tuple[str, int], ...], axis: str) -> int:
+        for name, value in coords:
+            if name == axis:
+                return int(value)
+        layout = self._rank_layout or {}
+        if isinstance(layout, dict):
+            layout_axes = layout.get("axis_order", [])
+            if axis not in layout_axes:
+                return 0
+            axis_sizes = layout.get("axis_sizes", {})
+        else:
+            axis_sizes = {}
+        if axis_sizes.get(axis, 1) <= 1:
+            return 0
+        raise ValueError(f"Axis '{axis}' not present in coordinate tuple for faulty link.")
+
+    def _build_transformer_stage_dp_fault_map(self) -> Dict[Tuple[int, int], Tuple[Tuple[int, int, float], ...]]:
+        projection = self._fault_projection_for("transformer")
+        if not projection:
+            return {}
+        stage_dp_faults: Dict[Tuple[int, int], List[Tuple[int, int, float]]] = {}
+        for detail in projection.entries:
+            src_stage = self._axis_value_from_coords(detail.src_coords, "lp")
+            dst_stage = self._axis_value_from_coords(detail.dst_coords, "lp")
+            if src_stage != dst_stage:
+                raise ValueError(
+                    "Transformer faulty link spans multiple pipeline stages; "
+                    "hierarchical execution requires faults limited to a single stage."
+                )
+            src_dp = self._axis_value_from_coords(detail.src_coords, "dp")
+            dst_dp = self._axis_value_from_coords(detail.dst_coords, "dp")
+            affected_dp = {src_dp, dst_dp}
+            for dp_idx in affected_dp:
+                stage_dp_faults.setdefault((dp_idx, src_stage), []).append(detail.remapped)
+        return {key: tuple(links) for key, links in stage_dp_faults.items()}
+
+    def _fault_projection_for(self, label: str) -> Optional[FaultProjectionResult]:
+        return self._fault_projections.get(label)
+
+    def _fault_links_for(self, label: str) -> Tuple[Tuple[int, int, float], ...]:
+        projection = self._fault_projection_for(label)
+        if projection is None:
+            return tuple()
+        return projection.remapped_links
+
+    def _fault_override(self, label: str) -> Optional[Tuple[Tuple[int, int, float], ...]]:
+        if self._fault_space is None:
+            return None
+        return self._fault_links_for(label)
 
     def run(self, mode: ExecutionMode) -> ExecutionResult:
         if mode == ExecutionMode.ANALYTICAL:
@@ -533,19 +1311,24 @@ def run(self, mode: ExecutionMode) -> ExecutionResult:
             return self._run_full_astrasim_hierarchical()
         if mode == ExecutionMode.FULL_ASTRASIM_FLATTENED:
             return self._run_full_astrasim_flattened()
-        raise ValueError(f"Unsupported execution mode: {mode}")
 
     def _run_pipeline_with_analytical_comm(self, declared_mode: ExecutionMode) -> ExecutionResult:
         if declared_mode == ExecutionMode.HYBRID:
-            filename = "/hybrid_graph"
+            if self.no_data_parallel:
+                filename = "/hybrid_graph_no_dp"
+            else:
+                filename = "/hybrid_graph"
             timed_root = self.pipeline_root
         else: # must be "ANALYTICAL"
-            filename = "/analytical_graph"
-            timed_root = self.pipeline_graph.convert_comm_sizes_to_times(
-                self.pipeline_root,
-                self.time_calc.network_model,
-                self.interconnect_params,
-            )
+            if self.no_data_parallel:
+                filename = "/analytical_graph_no_dp"
+            else:
+                filename = "/analytical_graph"
+        timed_root = self.pipeline_graph.convert_comm_sizes_to_times(
+            self.pipeline_root,
+            self.time_calc.network_model,
+            self.interconnect_params,
+        )
             
         generate_graphs = _env_flag("DEEPFLOW_VISUALIZE_GRAPHS")
         if generate_graphs:
@@ -562,18 +1345,30 @@ def _run_pipeline_with_analytical_comm(self, declared_mode: ExecutionMode) -> Ex
 
     def _run_hybrid(self) -> ExecutionResult:
         generate_graphs = _env_flag("DEEPFLOW_VISUALIZE_GRAPHS")
+
+        transformer_time = self._run_transformer_astrasim(ExecutionMode.HYBRID)
+
+        
         if generate_graphs:
-            self.transformer_graph.save_graph(
+            transformer_timed_forward_root = self.transformer_graph.convert_comm_sizes_to_times(
                 self.transformer_forward_root,
+                self.time_calc.network_model,
+                self.interconnect_params,
+            )
+            transformer_timed_backward_root = self.transformer_graph.convert_comm_sizes_to_times(
+                self.transformer_backward_root,
+                self.time_calc.network_model,
+                self.interconnect_params,
+            )
+            self.transformer_graph.save_graph(
+                transformer_timed_forward_root,
                 self.time_calc.output_dir,
-                "/hybrid_graph_transformer",
+                "/hybrid_graph_transformer_forward",
             )
-        transformer_time = self._run_transformer_astrasim(ExecutionMode.HYBRID)
-        if generate_graphs:
             self.transformer_graph.save_graph(
-                self.transformer_forward_root,
+                transformer_timed_backward_root,
                 self.time_calc.output_dir,
-                "/hybrid_graph_transformer",
+                "/hybrid_graph_transformer_backward",
             )
 
         if transformer_time is not None:
@@ -585,6 +1380,28 @@ def _run_full_astrasim_hierarchical(self) -> ExecutionResult:
         if transformer_time is not None:
             self._apply_transformer_time(transformer_time)
 
+        if _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") and self.transformer_graph:
+            transformer_timed_forward_root = self.transformer_graph.convert_comm_sizes_to_times(
+                self.transformer_forward_root,
+                self.time_calc.network_model,
+                self.interconnect_params,
+            )
+            transformer_timed_backward_root = self.transformer_graph.convert_comm_sizes_to_times(
+                self.transformer_backward_root,
+                self.time_calc.network_model,
+                self.interconnect_params,
+            )
+            self.transformer_graph.save_graph(
+                transformer_timed_forward_root,
+                self.time_calc.output_dir,
+                "/hierarchical_graph_transformer_forward",
+            )
+            self.transformer_graph.save_graph(
+                transformer_timed_backward_root,
+                self.time_calc.output_dir,
+                "/hierarchical_graph_transformer_backward",
+            )
+
         dp_count = getattr(self.time_calc, "dp", 1) or 1
         if not self.pipeline_root:
             raise RuntimeError("Pipeline graph root is not available for AstraSim execution")
@@ -594,8 +1411,10 @@ def _run_full_astrasim_hierarchical(self) -> ExecutionResult:
         if self.time_calc.persist_astrasim_artifacts:
             artifact_dir = os.path.join(self.time_calc.output_dir, "astra_hier")
 
+        pipeline_fault_override = self._fault_override("pipeline")
         run_kwargs = {
             "persist_artifacts": self.time_calc.persist_astrasim_artifacts,
+            "faulty_links_override": pipeline_fault_override,
         }
         run_type = str(getattr(getattr(self.time_calc, "model", None), "run_type", "training")).lower()
         effective_dp = 1 if run_type == "inference" else max(1, getattr(self.time_calc, "dp", 1))
@@ -603,12 +1422,19 @@ def _run_full_astrasim_hierarchical(self) -> ExecutionResult:
             run_kwargs["dp_override"] = 1
 
         if _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") and self.pipeline_root is not None:
+            filename = "/pipeline_graph_hierarchical_no_dp" if self.no_data_parallel else "/pipeline_graph_hierarchical"
             self.pipeline_graph.save_graph(
                 self.pipeline_root,
                 self.time_calc.output_dir,
-                "/pipeline_graph_hierarchical",
+                filename,
             )
 
+        pipeline_layout = getattr(self, "_pipeline_rank_layout", None)
+        if pipeline_layout:
+            setattr(self.pipeline_root, "_astrasim_rank_layout", pipeline_layout)
+        elif hasattr(self.pipeline_root, "_astrasim_rank_layout"):
+            delattr(self.pipeline_root, "_astrasim_rank_layout")
+        self._attach_optimize_hint(self.pipeline_root)
         per_rank_sec, max_sec = run_astra_simulation_only_onepath(
             self.pipeline_root,
             self.time_calc,
@@ -630,24 +1456,37 @@ def _run_full_astrasim_flattened(self) -> ExecutionResult:
         flattener = PipelineGraphFlattener(
             pipeline_graph=self.pipeline_graph,
             transformer_graph=self.transformer_graph,
+            rank_layout=self._rank_layout,
         )
 
-        if _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") and self.pipeline_root is not None:
+        if _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") and self.pipeline_root is not None :
+            filename = "/pipeline_graph_pre_flatten_no_dp" if self.no_data_parallel else "/pipeline_graph_pre_flatten"
             self.pipeline_graph.save_graph(
                 self.pipeline_root,
                 self.time_calc.output_dir,
-                "/pipeline_graph_pre_flatten",
+                filename,
             )
         flattened_root = flattener.build(self.pipeline_root)
         if flattened_root is None:
             raise RuntimeError("Pipeline flattening produced an empty graph")
 
+        flattened_root = apply_overlap_transforms(
+            flattened_root,
+            self.time_calc.get_parallelism_mode(),
+            getattr(self.time_calc, "tp_overlap", 0.0),
+            getattr(self.time_calc, "tp_sp_overlap", 0.0),
+            getattr(self.time_calc, "cp_overlap", 0.0),
+        )
+        flattener._propagate_local_hw_ids(flattened_root)
+        setattr(flattened_root, "_astrasim_rank_layout", self._rank_layout)
+        self._attach_optimize_hint(flattened_root)
         self.time_calc.flattened_pipeline_root = flattened_root
         if _env_flag("DEEPFLOW_VISUALIZE_GRAPHS") and self.pipeline_root is not None:
+            filename = "/pipeline_graph_post_flatten_no_dp" if self.no_data_parallel else "/pipeline_graph_post_flatten"
             self.pipeline_graph.save_graph(
                 flattened_root,
                 self.time_calc.output_dir,
-                "/pipeline_graph_post_flatten",
+                filename,
             )
         self.pipeline_root = flattened_root
         # output_dir = "./astra_flattened_graph"
@@ -655,7 +1494,7 @@ def _run_full_astrasim_flattened(self) -> ExecutionResult:
         # base_path = os.path.join(output_dir, "pipeline_flattened")
         # dot = visualize_graph(flattened_root, filename=base_path)
         # try:
-        #     dot.render(base_path, format="png", cleanup=True)
+        #     dot.render(base_path, format="svg", cleanup=True)
         # except Exception as exc:  # pragma: no cover - visualization best-effort
         #     print(f"[WARN] Failed to render flattened pipeline graph: {exc}")
 
@@ -687,6 +1526,7 @@ def _run_full_astrasim_flattened(self) -> ExecutionResult:
             raise RuntimeError("AstraSim flattened execution returned no per-rank timings")
 
         expected_rank_count = effective_dp * len(unique_hw_ids)
+
         # Special case: If expected rank count is 1, then 2 is fine, but we prune the extra result
         # this is done, since astrasim backend only supports >1 ranks, so we generate extra fake result for that case.
         if expected_rank_count == 1:
@@ -696,7 +1536,6 @@ def _run_full_astrasim_flattened(self) -> ExecutionResult:
                     f"expected {expected_rank_count}, got {len(per_rank_sec)}"
                 )
             per_rank_sec = per_rank_sec[:1]
-
         if len(per_rank_sec) != expected_rank_count:
             raise RuntimeError(
                 "AstraSim rank count mismatch for flattened execution: "
@@ -752,22 +1591,78 @@ def enqueue_children(obj: Any) -> None:
     def _run_transformer_astrasim(self, mode: ExecutionMode) -> Optional[TransformerTimings]:
         del mode  # mode currently unused but kept for signature consistency
 
-        # Use hierarchical artifact directory when persisting artifacts for transformer simulation
-        artifact_dir = self.time_calc.output_dir
-        artifact_dir_fwd = artifact_dir
-        artifact_dir_bwd = artifact_dir
-        os.makedirs(artifact_dir, exist_ok=True)
-        if self.time_calc.persist_astrasim_artifacts:
-            artifact_dir = os.path.join(self.time_calc.output_dir, "astra_hier")
-            artifact_dir_fwd = os.path.join(artifact_dir, "fwd")
-            artifact_dir_bwd = os.path.join(artifact_dir, "bwd")
-            os.makedirs(artifact_dir_fwd, exist_ok=True)
-            os.makedirs(artifact_dir_bwd, exist_ok=True)
+        if not self.transformer_forward_root and not self.transformer_backward_root:
+            if getattr(self, "_transformer_stage_dp_faults", {}):
+                raise ValueError("Transformer faults require transformer graph metadata, but none is available.")
+            return None
 
+        layout = getattr(self, "_transformer_rank_layout", {})
+        if self.transformer_forward_root:
+            setattr(self.transformer_forward_root, "_astrasim_rank_layout", layout)
+        if self.transformer_backward_root:
+            setattr(self.transformer_backward_root, "_astrasim_rank_layout", layout)
+
+        persist = self.time_calc.persist_astrasim_artifacts
+        os.makedirs(self.time_calc.output_dir, exist_ok=True)
+        self._transformer_stage_timings = {}
+        self._transformer_baseline_timings = None
+
+        # Baseline run (no transformer faults)
+        baseline_fwd_dir, baseline_bwd_dir = self._transformer_artifact_dirs(label=None, persist=persist)
+        baseline_timings, baseline_fwd_per_rank, baseline_bwd_per_rank = self._execute_transformer_run(
+            baseline_fwd_dir,
+            baseline_bwd_dir,
+            faulty_links_override=(),
+        )
+        self._transformer_baseline_timings = baseline_timings
+        self.time_calc.transformer_astrasim_per_rank_forward = baseline_fwd_per_rank
+        self.time_calc.transformer_astrasim_per_rank_backward = baseline_bwd_per_rank
+        self.time_calc.transformer_astrasim_time_forward = baseline_timings.forward
+        self.time_calc.transformer_astrasim_time_backward = baseline_timings.backward
+
+        # Per-stage fault runs
+        stage_dp_faults = getattr(self, "_transformer_stage_dp_faults", {})
+        for fault_index, ((dp_idx, stage_id), fault_links) in enumerate(sorted(stage_dp_faults.items())):
+            label = f"fault{fault_index}_dp{dp_idx}_stage{stage_id}"
+            stage_fwd_dir, stage_bwd_dir = self._transformer_artifact_dirs(label=label, persist=persist)
+            stage_timings, _, _ = self._execute_transformer_run(
+                stage_fwd_dir,
+                stage_bwd_dir,
+                faulty_links_override=fault_links,
+            )
+            self._transformer_stage_timings[(dp_idx, stage_id)] = stage_timings
+
+        return baseline_timings
+
+    def _transformer_artifact_dirs(self, label: Optional[str], persist: bool) -> Tuple[str, str]:
+        if not persist:
+            base_dir = self.time_calc.output_dir
+            os.makedirs(base_dir, exist_ok=True)
+            return base_dir, base_dir
+        base_dir = os.path.join(self.time_calc.output_dir, "astra_hier")
+        os.makedirs(base_dir, exist_ok=True)
+        if label is None:
+            fwd_dir = os.path.join(base_dir, "fwd")
+            bwd_dir = os.path.join(base_dir, "bwd")
+        else:
+            fwd_dir = os.path.join(base_dir, f"{label}_fwd")
+            bwd_dir = os.path.join(base_dir, f"{label}_bwd")
+        os.makedirs(fwd_dir, exist_ok=True)
+        os.makedirs(bwd_dir, exist_ok=True)
+        return fwd_dir, bwd_dir
+
+    def _execute_transformer_run(
+        self,
+        artifact_dir_fwd: str,
+        artifact_dir_bwd: str,
+        *,
+        faulty_links_override: Optional[Tuple[Tuple[int, int, float], ...]],
+    ) -> Tuple[TransformerTimings, Optional[List[float]], Optional[List[float]]]:
         fwd_per_rank = None
         bwd_per_rank = None
         fwd_max = 0
         bwd_max = 0
+
         if self.transformer_forward_root:
             fwd_per_rank, fwd_max = run_astra_simulation_only_onepath(
                 self.transformer_forward_root,
@@ -775,7 +1670,11 @@ def _run_transformer_astrasim(self, mode: ExecutionMode) -> Optional[Transformer
                 artifact_dir_fwd,
                 dp_override=1,
                 persist_artifacts=self.time_calc.persist_astrasim_artifacts,
+                faulty_links_override=faulty_links_override,
             )
+            if fwd_max <= 0:
+                raise RuntimeError("AstraSim transformer forward execution returned non-positive duration")
+
         if self.transformer_backward_root:
             bwd_per_rank, bwd_max = run_astra_simulation_only_onepath(
                 self.transformer_backward_root,
@@ -783,28 +1682,26 @@ def _run_transformer_astrasim(self, mode: ExecutionMode) -> Optional[Transformer
                 artifact_dir_bwd,
                 dp_override=1,
                 persist_artifacts=self.time_calc.persist_astrasim_artifacts,
+                faulty_links_override=faulty_links_override,
             )
+            if bwd_max < 0:
+                raise RuntimeError("AstraSim transformer backward execution returned non-positive duration")
 
-        self.time_calc.transformer_astrasim_per_rank_forward = fwd_per_rank
-        self.time_calc.transformer_astrasim_per_rank_backward = bwd_per_rank
-        self.time_calc.transformer_astrasim_time_forward = fwd_max
-        self.time_calc.transformer_astrasim_time_backward = bwd_max
-
-        if fwd_max < 0 or bwd_max < 0:
-            raise RuntimeError("AstraSim transformer execution returned non-positive duration")
-
-        return TransformerTimings(forward=fwd_max, backward=bwd_max)
+        return TransformerTimings(forward=fwd_max, backward=bwd_max), fwd_per_rank, bwd_per_rank
 
     def _apply_transformer_time(self, timings: TransformerTimings) -> None:
         if timings.forward < 0 or timings.backward < 0:
             raise ValueError("AstraSim transformer times must be positive")
 
+        baseline_timings = self._transformer_baseline_timings or timings
+        stage_timings = getattr(self, "_transformer_stage_timings", {})
+
         comp_times = getattr(self.pipeline_graph, "comp_times", None)
         if isinstance(comp_times, dict):
             if "transformer_f" in comp_times:
-                comp_times["transformer_f"] = timings.forward
+                comp_times["transformer_f"] = baseline_timings.forward
             if "transformer_b" in comp_times:
-                comp_times["transformer_b"] = timings.backward
+                comp_times["transformer_b"] = baseline_timings.backward
 
         visited: Set[int] = set()
         roots: List[Any]
@@ -813,10 +1710,27 @@ def _apply_transformer_time(self, timings: TransformerTimings) -> None:
         else:
             roots = [self.pipeline_root]
 
+        dp_count = max(1, getattr(self.time_calc, "dp", 1))
+
         for root in roots:
-            self._assign_transformer_durations(root, visited, timings.forward, timings.backward)
+            self._assign_transformer_durations(
+                root,
+                visited,
+                stage_timings,
+                baseline_timings.forward,
+                baseline_timings.backward,
+                dp_count,
+            )
 
-    def _assign_transformer_durations(self, node: Any, visited: Set[int], forward_value: float, backward_value: float) -> None:
+    def _assign_transformer_durations(
+        self,
+        node: Any,
+        visited: Set[int],
+        stage_timings: Dict[Tuple[int, int], TransformerTimings],
+        forward_default: float,
+        backward_default: float,
+        dp_count: int,
+    ) -> None:
         if node is None:
             return
         node_id = id(node)
@@ -827,10 +1741,31 @@ def _assign_transformer_durations(self, node: Any, visited: Set[int], forward_va
         if isinstance(node, simulate_LLM.Node):
             base_name = str(getattr(node, "name", "") or "")
             if base_name.startswith("transformer_layer"):
+                hw_stage = None
+                try:
+                    hw_stage = int(getattr(node, "hw_id", None))
+                except (TypeError, ValueError):
+                    hw_stage = None
+                def _per_dp_durations(is_forward: bool) -> List[float]:
+                    values: List[float] = []
+                    for dp_idx in range(dp_count):
+                        timing_override = stage_timings.get((dp_idx, hw_stage)) if hw_stage is not None else None
+                        if timing_override:
+                            values.append(timing_override.forward if is_forward else timing_override.backward)
+                        else:
+                            default = forward_default if is_forward else backward_default
+                            values.append(default)
+                    return values
+
                 if getattr(node, "fwd", True):
-                    node.duration = forward_value
+                    per_dp_values = _per_dp_durations(is_forward=True)
+                else:
+                    per_dp_values = _per_dp_durations(is_forward=False)
+
+                if dp_count > 1:
+                    node.duration = tuple(per_dp_values)
                 else:
-                    node.duration = backward_value
+                    node.duration = per_dp_values[0]
 
         for child in getattr(node, "children", []):
-            self._assign_transformer_durations(child, visited, forward_value, backward_value)
+            self._assign_transformer_durations(child, visited, stage_timings, forward_default, backward_default, dp_count)
diff --git a/LLM_util.py b/LLM_util.py
index b8dd07ee..18bd3a00 100644
--- a/LLM_util.py
+++ b/LLM_util.py
@@ -77,12 +77,13 @@ def multihead_decoder_gemm(self, batch_size, seq_len, d_model, num_heads, kv_hea
         # assuming equal load balancing here
         num_experts = max(1, int(getattr(self, "moe_num_experts", 1)))
         top_k = max(1, int(getattr(self, "moe_top_k", 1)))
-        gemms["ffn1"] = (batch_size, seq_len * top_k / num_experts, d_model, projected_dim ) #gemm shape per expert
+        gemms["ffn1"] = (batch_size, seq_len * top_k / num_experts, d_model, projected_dim ) #gemm shape per expert assuming each dp group has all experts
         gemms["ffn2"] = (batch_size, seq_len * top_k / num_experts, intermediate_size, d_model)
     else:
         gemms["ffn1"] = (batch_size, seq_len, d_model, projected_dim)
         gemms["ffn2"] = (batch_size, seq_len, intermediate_size, d_model) 
     gemms["linear"] = (batch_size, seq_len, d_model, vocab_size)
+    gemms["router"] = (batch_size, seq_len, d_model, self.moe_num_experts) 
 
     return gemms
 
@@ -122,19 +123,25 @@ def process_gemm_shapes(self, batch_size, seq_len, d_model, num_heads, kv_heads,
 
     return processed
 
-def get_transformer_mem_layer( dp, tp, batch_size, hidden_dim, seq_len, intermediate_size, n_heads, precision, zero_stage, model_type="gpt"):#https://www.determined.ai/blog/act-mem-1.  https://arxiv.org/pdf/2205.05198. https://shjwudp.github.io/blog/2023/gpt-training-memory-estimation-nemo-training-practice/
+def get_transformer_mem_layer( dp, tp, lp, mb, batch_size, hidden_dim, seq_len, intermediate_size, n_heads, precision, zero_stage, flash_attention, full_recomputation, model_type="gpt"):#  https://arxiv.org/pdf/2205.05198. https://shjwudp.github.io/blog/2023/gpt-training-memory-estimation-nemo-training-practice/
     """ memory estimation of transformer layer for single gpu case in inference mode is supported.
     other modes or layers are work in progress."""
     #Activations refer to output activations that need to be stored
-    act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp + 5 * n_heads * seq_len/(hidden_dim * tp) ) * (precision.activations / 2) #from https://arxiv.org/pdf/2205.05198
+
+    if full_recomputation:
+        act_memory_layer = seq_len * batch_size * hidden_dim * (2 ) * (precision.activations / 2) #full recompute:
+    elif flash_attention:
+        act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp ) * (precision.activations / 2) #assuming selective recompute
+    else:
+        act_memory_layer = seq_len * batch_size * hidden_dim * (34 / tp + 5 * n_heads * seq_len/(hidden_dim * tp) ) * (precision.activations / 2) 
+    
+    act_memory_layer *= lp / mb #pipedream style pipelining activation memory
     act_memory_layer_inf = seq_len * batch_size * intermediate_size / tp * precision.activations  #inference max activation memory, no need to store for backpropagation
     ffn_proj_factor = 3 if str(model_type).lower() == "llama" else 2
     
-    
     transformer_param_layer = 4* hidden_dim * hidden_dim + intermediate_size * ffn_proj_factor * hidden_dim  # weights Wq,Wk,Wv,Wo,ffn
-    optimizer_mem = 10 * transformer_param_layer / dp # zero1 style optimizer memory
-    #TODO: which optimizer mem to use?
-    optimizer_mem = 10 * transformer_param_layer * (precision.optimizer_states / 2) / tp # don't divide by dp for DDP
+
+    optimizer_mem = (precision.optimizer_states * 2 + precision.activations) * transformer_param_layer / tp # don't divide by dp for DDP, NeMo ZeRO-1 optimizer style
     tensor_weight_memory_layer = transformer_param_layer * precision.parameters / tp #weight memory
     # master_parameters is set to 0 by default, so this works.
     master_weight_memory_layer = transformer_param_layer * precision.master_parameters / tp
@@ -200,7 +207,17 @@ def get_embedding_weight_mem(
     
 def get_tot_mem_req(exp_hw_config, exp_model_config, **kwargs):
     # Model Params
-    batch_size                   = int(kwargs.get('batch_size', exp_model_config.model_config.batch_size))
+    default_global_batch = getattr(
+        exp_model_config.model_config,
+        "global_batch_size",
+        getattr(exp_model_config.model_config, "batch_size"),
+    )
+    batch_size = int(
+        kwargs.get(
+            "global_batch_size",
+            kwargs.get("batch_size", default_global_batch),
+        )
+    )
     hidden_dim                   = int(kwargs.get('hidden_dim', exp_model_config.model_config.hidden_dim))
     vocab_size                   = int(kwargs.get('vocab_size', exp_model_config.model_config.vocab_size))
     n_layers                   = int(kwargs.get('num_layer', exp_model_config.model_config.num_layers))
diff --git a/README.md b/README.md
index 80addecd..20631f03 100644
--- a/README.md
+++ b/README.md
@@ -33,6 +33,7 @@ If uv is used, make sure to use `uv run` instead of `python` to execute python s
 - `git submodule update --remote -- astra-sim`
 - `ASTRA_SIM=$(realpath ./astra-sim)`
 - `cd ${ASTRA_SIM}`
+- `git submodule update --remote -- extern/network_backend/analytical`
 - `./build/astra_analytical/build.sh`
 - `cd ..`
 
@@ -173,7 +174,7 @@ python configs/model-config/hf_to_config.py Qwen/Qwen2.5-3B --run-type inference
 
 2. Open the provided A100 hardware file and set tensor parallel degree to 1:
 
-Edit `configs/hardware-config/a100_80GB.yaml` so that `system_hierarchy.num_devices_per_node: 1` and `scheduling_param.tp: 1`. This models one GPU without tensor parallelism.
+Edit `configs/hardware-config/a100_80GB.yaml` so that `parallelism.tp: 1`. This models one GPU without tensor parallelism.
 
 Run the inference estimation using the example configs:
 
@@ -183,7 +184,7 @@ python run_perf.py --hardware_config configs/hardware-config/a100_80GB.yaml --mo
 
 3. To see the effects of tensor-parallelism, modify the hardware file, switching to two devices per node, and tp = 2.
 
-Edit `configs/hardware-config/a100_80GB.yaml` so that `system_hierarchy.num_devices_per_node: 2` and `scheduling_param.tp: 2`. This models two GPU-inference with tensor parallelism degree of 2.
+Edit `configs/hardware-config/a100_80GB.yaml` so that `parallelism.tp: 2`. This models two GPU-inference with tensor parallelism degree of 2.
 
 Re-run the same inference command with the updated hardware config.
 
@@ -249,6 +250,10 @@ Comparing the two runs will show how increasing tensor parallelism changes the p
 - **Supported:** Very rudimentary, non validated model. Does not consider static power. Inference only for now (Training is WIP). Energy due to prefill and decode is reported separately.
 - **Work in progress:** Training support, and heavy refinement of current implementation.
 
+### Mixed Precision
+- **Supported:** Mixed precision is supported for all supported parallelism and model types.
+- **Work in progress:** KV Cache precision cannot yet be set (config value is ignored).
+
 ### Validation
 - **Status:** Validation scripts are under development. Limited inference and training validation has been performed.
 - **Supported:** For inference, we validate against Koyeb's single NVIDIA A100 data on LLama3.1-8B. The relevant script can be found in `validation_scripts/koyeb.py`. Below is the resulting plot, comparing DeepFlow's predicted decode throughput with Koyeb's reported decode throughput, available at https://www.koyeb.com/docs/hardware/gpu-benchmarks. Average absolute error is ~6%.
diff --git a/astra-sim b/astra-sim
index c13cfa86..32bec13d 160000
--- a/astra-sim
+++ b/astra-sim
@@ -1 +1 @@
-Subproject commit c13cfa86c7ba71813faf25fca3145ccc89ef307f
+Subproject commit 32bec13d1fee017c0184edba43454c1542fb1190
diff --git a/astrasim_lib/config_generation.py b/astrasim_lib/config_generation.py
index c68e15c5..9e5ac393 100644
--- a/astrasim_lib/config_generation.py
+++ b/astrasim_lib/config_generation.py
@@ -3,8 +3,10 @@
 from __future__ import annotations
 
 import os
-from typing import Any, Dict, Optional, Tuple
-
+import uuid
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+from types import SimpleNamespace
+import math
 from hw_component import Network
 
 from .bootstrap import ensure_chakra_available
@@ -16,15 +18,9 @@
 
 ASTRA_DEBUG = False
 
-_NET_YAML_CACHE: set[tuple[str, int, float, float, str]] = set()
-_JSON_WRITTEN_BY_NPUS: set[object] = set()
-
 
-def _save_json(path: str, data: Dict[str, Any], npus_key: Optional[int] = None) -> None:
-    """Write ``data`` to ``path`` once per ``npus_key`` per process."""
-    key = npus_key if npus_key is not None else path
-    if key in _JSON_WRITTEN_BY_NPUS:
-        return
+def _save_json(path: str, data: Dict[str, Any]) -> None:
+    """Write ``data`` to ``path``."""
     os.makedirs(os.path.dirname(path), exist_ok=True)
     tmp_path = path + ".tmp"
     with open(tmp_path, "w", encoding="utf-8") as handle:
@@ -32,7 +28,6 @@ def _save_json(path: str, data: Dict[str, Any], npus_key: Optional[int] = None)
 
         _json.dump(data, handle, indent=2)
     os.replace(tmp_path, path)
-    _JSON_WRITTEN_BY_NPUS.add(key)
 
 
 def _gbps_from_bps(bps: float) -> float:
@@ -55,6 +50,16 @@ def choose_collective(alg: str, topo: str, op: str) -> str:
         return "direct"
     if topo == "Switch":
         return "halvingDoubling"
+    if topo == "Torus2D":
+        return "torus2d"
+    if topo == "Mesh":
+        return "mesh"
+    if topo == "Mesh2D":
+        return "torus2d"
+    if topo == "HyperCube":
+        return "hypercube"
+    if topo == "KingMesh2D":
+        return "torus2d"
     return "ring"
 
 
@@ -68,95 +73,416 @@ def compute_intra_inter_ib_ll_from_hw(hw_obj) -> Tuple[Tuple[float, float], Tupl
 
 def derive_topology_from_hw(hw_obj) -> str:
     """Map DeepFlow network topology enums to AstraSim names."""
-    try:
-        topo = getattr(hw_obj.network_topology.intra, "topology", None)
-        topo_str = (topo or "fc").lower()
-        if topo_str in ("fc", "fullyconnected", "fully_connected", "fully-connected"):
-            return "FullyConnected"
-        if topo_str in ("ring",):
-            return "Ring"
-        if topo_str in ("switch",):
-            return "Switch"
-    except Exception:  # pragma: no cover - defensive
-        pass
+    layout = getattr(hw_obj, "network_layout", None)
+    primary = layout.primary_dimension() if layout else None
+    topo = (primary.topology_type if primary else "ring") or "ring"
+    return _normalize_topology_name(topo)
+
+
+def _normalize_topology_name(topo: str) -> str:
+    topo_str = str(topo).lower()
+    if topo_str in ("fc", "fullyconnected", "fully_connected", "fully-connected"):
+        return "FullyConnected"
+    if topo_str in ("ring",):
+        return "Ring"
+    if topo_str in ("switch",):
+        return "Switch"
+    if topo_str in ("torus2d"):
+        return "Torus2D"
+    if topo_str in ("mesh",):
+        return "Mesh"
+    if topo_str in ("hypercube",):
+        return "HyperCube"
+    if topo_str in ("mesh2d", "mesh-2d"):
+        return "Mesh2D"
+    if topo_str in ("kingmesh2d", "king-mesh2d"):
+        return "KingMesh2D"
     return "FullyConnected"
 
 
+def _expand_network_entries(entries: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Expand multi-dimensional mesh/torus entries into per-dimension entries."""
+
+    expanded: List[Dict[str, Any]] = []
+    for entry in entries:
+        topo = entry.get("topology")
+        npus_value = entry.get("npus")
+        dims_value = entry.get("dims")
+        bw_value = entry.get("bandwidth")
+        lat_value = entry.get("latency")
+
+        if isinstance(topo, str) and (topo.startswith("Torus") or topo.startswith("Mesh")):
+            dim_count: int
+            if dims_value is not None and isinstance(dims_value, (list, tuple)):
+                dim_count = len(dims_value)
+            elif isinstance(npus_value, (list, tuple)):
+                dim_count = len(npus_value)
+            elif len(topo) >= 2 and topo[-2].isdigit():
+                dim_count = int(topo[-2])
+            else:
+                dim_count = 1
+            base_name = "Ring" if topo.startswith("Torus") else "Mesh"
+            if dims_value is None and isinstance(npus_value, (list, tuple)) and len(npus_value) == dim_count:
+                dims_value = tuple(npus_value)
+            for dim_idx in range(dim_count):
+                if dims_value is not None:
+                    if not isinstance(dims_value, (list, tuple)) or len(dims_value) != dim_count:
+                        raise ValueError(f"dims for topology {topo} must match dimension count {dim_count}.")
+                    curr_npu = dims_value[dim_idx]
+                else:
+                    curr_npu = (
+                        npus_value[dim_idx]
+                        if isinstance(npus_value, (list, tuple))
+                        else npus_value
+                    )
+                    if dim_count == 2:
+                        root = int(math.isqrt(int(curr_npu)))
+                        if root * root != int(curr_npu):
+                            raise ValueError(f"npus ({curr_npu}) must be a perfect square for 2D topology {topo}.")
+                        curr_npu = root
+                topology_name = "Mesh" if dim_count == 2 and curr_npu <= 2 else base_name
+                expanded.append(
+                    {
+                        **entry,
+                        "topology": topology_name,
+                        "npus": curr_npu,
+                        "bandwidth": bw_value,
+                        "latency": lat_value,
+                    }
+                )
+        else:
+            expanded.append(dict(entry))
+    return expanded
+
+
 def generate_astrasim_configs_from_hw(
     hw_obj,
     out_dir: str = "./astra_cache",
     npus_count: Optional[int] = None,
+    *,
+    axes_filter: Optional[Sequence[str]] = None,
+    transform_2d_to_1d: bool = False,
+    faulty_links_override: Optional[Sequence[Tuple[int, int, float]]] = None,
+    ephemeral_outputs: bool = False,
+    preferred_axes_for_synthetic: Optional[Sequence[str]] = None,
 ) -> Dict[str, str]:
     """Write AstraSim network/system configs derived from ``hw_obj``."""
     if npus_count is None:
         raise ValueError("npus_count must be provided explicitly when generating AstraSim configs.")
 
-    net_yaml = os.path.join(out_dir, f"network_analytical_{int(npus_count)}.yml")
-    sys_json = os.path.join(out_dir, "system_native_collectives.json")
+    layout = getattr(hw_obj, "network_layout", None)
+    dimensions = list(getattr(layout, "dimensions", [])) if layout else []
+    layout_faults: Tuple[Tuple[int, int, float], ...]
+    if faulty_links_override is not None:
+        layout_faults = tuple(
+            (int(src), int(dst), float(weight))
+            for src, dst, weight in faulty_links_override
+        )
+    else:
+        layout_faults = tuple(getattr(layout, "faulty_links", ()) or ()) if layout else tuple()
+    if not dimensions:
+        raise ValueError("Hardware config is missing network dimensions required for AstraSim integration.")
 
-    topo = derive_topology_from_hw(hw_obj)
-    if topo == "Ring" and int(npus_count) <= 2:
-        topo = "FullyConnected"
+    exec_backend = getattr(hw_obj, "execution_backend", None)
+    astra_cfg = getattr(exec_backend, "astra", None) if exec_backend else None
+    astra_mode = str(getattr(astra_cfg, "mode", "") or "").strip().lower()
 
-    (intra_ib_bps, intra_ll_s), _ = compute_intra_inter_ib_ll_from_hw(hw_obj)
-    if not intra_ib_bps or intra_ib_bps <= 0:
-        raise ValueError("Intra-node bandwidth computed as 0. Check hardware config/network model.")
-    if not intra_ll_s or intra_ll_s <= 0:
-        raise ValueError("Intra-node latency computed as 0. Check hardware config/network model.")
-    ib_gbps = round(_gbps_from_bps(intra_ib_bps), 6)
-    ll_ns = round(_ns_from_s(intra_ll_s), 3)
+    sch_config = getattr(hw_obj, "sch_config", None)
+
+    axes_filter_original = tuple(axes_filter) if axes_filter else None
+    axes_filter_normalized: Optional[Tuple[str, ...]] = None
+    if axes_filter_original:
+        axes_filter_normalized = tuple(str(axis).strip().lower() for axis in axes_filter_original)
+
+    preferred_axes_synth: Tuple[str, ...] = tuple(
+        str(axis).strip().lower()
+        for axis in preferred_axes_for_synthetic or ()
+        if str(axis).strip()
+    )
+
+    axis_sizes_full: Dict[str, int] = {
+        "tp": int(sch_config.tp),
+        "cp": int(sch_config.cp),
+        "lp": int(sch_config.lp),
+        "dp": int(sch_config.dp),
+    }
+    synthetic_only = axes_filter_normalized is not None and set(axes_filter_normalized) == {"synthetic2"}
+    if synthetic_only:
+        preferred_set = set(preferred_axes_synth)
+        base_dim = None
+        if dimensions:
+            for dim in dimensions:
+                dim_axes = [
+                    str(axis).strip().lower()
+                    for axis in getattr(dim, "parallelisms", ()) or ()
+                    if str(axis).strip()
+                ]
+                if preferred_set and any(axis in preferred_set for axis in dim_axes):
+                    base_dim = dim
+                    break
+                if base_dim is None and dim_axes:
+                    base_dim = dim
+        if base_dim is None:
+            base_dim = dimensions[0] if dimensions else None
+        if base_dim is None:
+            raise ValueError(f"Synthetic dimension requested but no dimensions found in hardware config. Info: Called with filter {axes_filter_original}, npus_count={npus_count}, axes_sizes={axis_sizes_full}.")
+        base_bw = float(getattr(base_dim, "effective_bandwidth", getattr(base_dim, "bandwidth", None)))
+        base_latency = float(getattr(base_dim, "latency", None))
+        base_topology = getattr(base_dim, "topology_type", None)
+        base_collectives = getattr(base_dim, "collective_override", {}) or {}
+        synthetic_dim = SimpleNamespace(
+            label="synthetic2",
+            parallelisms=("synthetic2",),
+            size=2,
+            topology_type=base_topology,
+            effective_bandwidth=base_bw,
+            bandwidth=base_bw,
+            latency=base_latency,
+            collective_override=base_collectives,
+            faulty_links=(),
+        )
+        dimensions = [synthetic_dim]
+        axis_sizes_full = {"synthetic2": 2}
+        axis_sizes = {"synthetic2": 2}
+        axis_order_preference = ["synthetic2"]
+    else:
+        axis_sizes: Dict[str, int] = dict(axis_sizes_full)
+        if axes_filter_normalized:
+            axis_sizes = {axis: axis_sizes_full.get(axis, 1) for axis in axes_filter_normalized}
+            axis_sizes.setdefault("tp", 1)
+            axis_sizes.setdefault("cp", 1)
+            axis_sizes.setdefault("lp", 1)
+            axis_sizes.setdefault("dp", 1)
+        axis_order_preference = ["tp", "cp", "lp", "dp"]
+
+    allowed_axes = set(axis_sizes.keys()) if axes_filter_normalized else None
+    # print(f"Allowed axes: {allowed_axes}")
+    # print(f"Axes sizes: {axis_sizes}")
+    dim_infos: List[Tuple[Any, List[str], int, int]] = []
+    # print(f"Filter: {axes_filter}")
+    for dim_idx, dim in enumerate(dimensions):
+        axes = [str(axis).strip().lower() for axis in getattr(dim, "parallelisms", ())]
+        # print(f"All axes for dimension {dim.label}: {axes}")
+        if allowed_axes is not None and axes_filter_normalized:
+            filtered_axes = [axis for axis in axes if axis in allowed_axes]
+            axes = filtered_axes
+        # print(f"Filtered axes for dimension {dim.label}: {axes}")
+        effective = 1
+        for axis in axes:
+            if axis not in axis_sizes:
+                raise ValueError(
+                    f"Unsupported parallelism axis '{axis}' referenced by network dimension '{dim.label}'. "
+                    "Supported axes are tp, cp, lp, dp."
+                )
+            effective *= axis_sizes[axis]
+        dim_infos.append((dim, axes, effective, dim_idx))
+
+    active_dim_indices = [
+        dim_idx
+        for dim, _, _, dim_idx in dim_infos
+        if int(getattr(dim, "size", 0) or 0) > 1
+    ]
+    if active_dim_indices:
+        last_active_idx = active_dim_indices[-1]
+        dp_in_active = False
+        if axis_sizes.get("dp", 1) > 1:
+            for dim, axes, _, dim_idx in dim_infos:
+                if "dp" in axes:
+                    if dim_idx != last_active_idx:
+                        raise ValueError(
+                            f"Data-parallel axis must be assigned to the last active network dimension. "
+                            f"Dimension '{dim.label}' (index {dim_idx}) includes 'dp', but last active dimension index is {last_active_idx}."
+                        )
+                    dp_in_active = True
+        if axis_sizes.get("dp", 1) > 1 and not dp_in_active:
+            raise ValueError(
+                "Hardware config declares dp > 1, but no network dimension with size > 1 includes the 'dp' axis."
+            )
+    else:
+        if layout_faults:
+            raise ValueError(
+                "Faulty links require at least one active (size > 1) network dimension."
+            )
+    target = int(npus_count)
+    
+    axes_needed: List[str] = []
+    remaining = target
+
+    for axis in axis_order_preference:
+        size = axis_sizes.get(axis, 1)
+        
+        if size <= 1:
+            continue
+        
+        if remaining % size == 0:
+            axes_needed.append(axis)
+            remaining //= size
+    
+    if remaining != 1:
+        raise ValueError(
+            f"Unable to map requested npus_count={target} to network axes {axis_sizes}. Info: Called with filter {axes_filter_original}, npus_count={npus_count}, axes_sizes={axis_sizes}."
+        )
+
+    axes_needed_set = set(axes_needed)
+    selected_dims: List[Tuple[Any, List[str], int, int]] = []
+    accumulated = 1
+    for dim, axes, _effective, dim_idx in dim_infos:
+        selected_axes: List[str] = []
+        size_contrib = 1
+        for axis in axes:
+            if axis in axes_needed_set:
+                selected_axes.append(axis)
+                size_contrib *= axis_sizes[axis]
+        if not selected_axes:
+            continue
+        axes_needed_set.difference_update(selected_axes)
+        accumulated *= size_contrib
+        selected_dims.append((dim, selected_axes, size_contrib, dim_idx))
+    if axes_needed_set:
+        raise ValueError(
+            f"Requested npus_count={target} requires axes {axes_needed} but no network dimension covers {axes_needed_set}."
+        )
+    if accumulated != target:
+        raise ValueError(
+            f"Selected network dimensions {[d.label for d, _, _, _ in selected_dims]} "
+            f"describe {accumulated} NPUs but simulation expects {target} ranks."
+        )
+
+    topo_list: List[str] = []
+    npus_list: List[Any] = []
+    bw_list: List[float] = []
+    lat_list: List[float] = []
+    network_entries: List[Dict[str, Any]] = []
+
+    for dim, axes_selected, product_size, _ in selected_dims:
+        topo = _normalize_topology_name(dim.topology_type)
+        full_size = int(getattr(dim, "size", product_size) or product_size)
+        dims_value = getattr(dim, "size_2d", None)
+        dims_tuple: Optional[Tuple[int, int]] = None
+        use_shape = dims_value is not None and product_size == full_size
+        if use_shape:
+            dims_tuple = (int(dims_value[0]), int(dims_value[1]))  # type: ignore[index]
+            shape_product = dims_tuple[0] * dims_tuple[1]
+            if shape_product != product_size:
+                raise ValueError(
+                    f"Network dimension '{dim.label}' shape {dims_tuple} product {shape_product} "
+                    f"does not match parallelism product {product_size}."
+                )
+            size = shape_product
+        else:
+            size = product_size
+        if transform_2d_to_1d and topo in ["Torus2D", "Mesh2D", "KingMesh2D"]:
+            if topo == "Torus2D":
+                topo = "Ring"
+            if topo == "Mesh2D":
+                topo = "Mesh"
+            if topo == "KingMesh2D":
+                topo = "HyperCube"
+            # When flattening 2D, prefer the explicit shape product if available.
+            size = size
+        if topo == "Ring" and size <= 2:
+            topo = "FullyConnected"
+        effective_bw = float(getattr(dim, "effective_bandwidth", dim.bandwidth))
+        latency_s = float(dim.latency)
+        entry = {
+            "dim": dim,
+            "axes": axes_selected,
+            "npus": size,
+            "topology": topo,
+            "bandwidth": round(_gbps_from_bps(effective_bw), 6),
+            "latency": round(_ns_from_s(latency_s), 3),
+        }
+        if dims_tuple is not None and topo in ["Torus2D", "Mesh2D", "KingMesh2D"]:
+            entry["dims"] = dims_tuple
+        network_entries.append(entry)
+
+    network_entries = _expand_network_entries(network_entries)
+
+    topo_list = [entry["topology"] for entry in network_entries]
+    npus_list = [entry["npus"] for entry in network_entries]
+    bw_list = [entry["bandwidth"] for entry in network_entries]
+    lat_list = [entry["latency"] for entry in network_entries]
+
+    signature_parts = [str(size) for _dim, _axes, size, _ in selected_dims]
+    dim_signature = "_".join(signature_parts) if signature_parts else f"{target}"
+
+    unique_suffix = f"_{uuid.uuid4().hex}" if ephemeral_outputs else ""
+    net_yaml = os.path.join(out_dir, f"network_analytical_{dim_signature}{unique_suffix}.yml")
+    sys_json = os.path.join(out_dir, f"system_native_collectives_{dim_signature}{unique_suffix}.json")
+
+    topo_str = ", ".join(topo_list)
+    npus_str = ", ".join(str(v) for v in npus_list)
+    bw_str = ", ".join(str(v) for v in bw_list)
+    lat_str = ", ".join(str(v) for v in lat_list)
 
     net_content = (
-        f"topology: [ {topo} ]\n"
-        f"npus_count: [ {int(npus_count)} ]\n"
-        f"bandwidth: [ {ib_gbps} ]  # GB/s\n"
-        f"latency: [ {ll_ns} ]   # ns\n"
+        f"topology: [ {topo_str} ]\n"
+        f"npus_count: [ {npus_str} ]\n"
+        f"bandwidth: [ {bw_str} ]  # GB/s\n"
+        f"latency: [ {lat_str} ]   # ns\n"
     )
 
-    os.makedirs(os.path.dirname(net_yaml), exist_ok=True)
-    cache_key = (net_yaml, int(npus_count), ib_gbps, ll_ns, topo)
-    need_write = True
-    if cache_key in _NET_YAML_CACHE:
-        need_write = False
-    elif os.path.exists(net_yaml):
-        try:
-            with open(net_yaml, "r", encoding="utf-8") as handle:
-                existing = handle.read()
-            if existing == net_content:
-                need_write = False
-        except Exception:  # pragma: no cover - defensive
-            pass
-    if need_write:
-        tmp_path = net_yaml + ".tmp"
-        with open(tmp_path, "w", encoding="utf-8") as handle:
-            handle.write(net_content)
-        os.replace(tmp_path, net_yaml)
-    _NET_YAML_CACHE.add(cache_key)
+    faulty_links_tuple: Tuple[Tuple[int, int, float], ...] = tuple(layout_faults)
+    if faulty_links_tuple:
+        fault_entries = [
+            f"[{src}, {dst}, {float(weight):g}]"
+            for src, dst, weight in faulty_links_tuple
+        ]
+        faulty_links_str = f"[{', '.join(fault_entries)}]"
+        net_content += f"faulty_links: {faulty_links_str}\n"
 
-    exec_backend = getattr(hw_obj, "execution_backend", None)
-    if exec_backend and exec_backend.astra:
-        coll = exec_backend.astra.collectives
-        sys_opts = getattr(exec_backend.astra, "sys_options", None)
-        ag = choose_collective(coll.all_gather, topo, "all-gather")
-        ar = choose_collective(coll.all_reduce, topo, "all-reduce")
-        rs = choose_collective(coll.reduce_scatter, topo, "reduce-scatter")
-        a2a = choose_collective(coll.all_to_all, topo, "all-to-all")
+    os.makedirs(os.path.dirname(net_yaml), exist_ok=True)
+    tmp_path = net_yaml + ".tmp"
+    with open(tmp_path, "w", encoding="utf-8") as handle:
+        handle.write(net_content)
+    os.replace(tmp_path, net_yaml)
+
+    if astra_cfg:
+        coll = astra_cfg.collectives
+        sys_opts = getattr(astra_cfg, "sys_options", None)
+        default_ag = coll.all_gather
+        default_ar = coll.all_reduce
+        default_rs = coll.reduce_scatter
+        default_a2a = coll.all_to_all
     else:
-        ag = choose_collective("auto", topo, "all-gather")
-        ar = choose_collective("auto", topo, "all-reduce")
-        rs = choose_collective("auto", topo, "reduce-scatter")
-        a2a = choose_collective("auto", topo, "all-to-all")
+        default_ag = "auto"
+        default_ar = "auto"
+        default_rs = "auto"
+        default_a2a = "auto"
         sys_opts = None
 
+    def _collective_for_dimension(dim, topo_name: str, op: str, default_alg: str) -> str:
+        override = (
+            dim.collective_override.get(op)
+            or dim.collective_override.get(op.replace("-", "_"))
+            or dim.collective_override.get(op.replace("_", "-"))
+        )
+        if override:
+            return override
+        return choose_collective(default_alg, topo_name, op)
+
+    ag_impl: List[str] = []
+    ar_impl: List[str] = []
+    rs_impl: List[str] = []
+    a2a_impl: List[str] = []
+
+    for entry in network_entries:
+        dim = entry["dim"]
+        topo_name = entry["topology"]
+        ag_impl.append(_collective_for_dimension(dim, topo_name, "all-gather", default_ag))
+        ar_impl.append(_collective_for_dimension(dim, topo_name, "all-reduce", default_ar))
+        rs_impl.append(_collective_for_dimension(dim, topo_name, "reduce-scatter", default_rs))
+        a2a_impl.append(_collective_for_dimension(dim, topo_name, "all-to-all", default_a2a))
+
     system = {
         "scheduling-policy": "LIFO",
         "endpoint-delay": 10,
         "active-chunks-per-dimension": 1,
         "preferred-dataset-splits": 1,
-        "all-reduce-implementation": [ar],
-        "all-gather-implementation": [ag],
-        "reduce-scatter-implementation": [rs],
-        "all-to-all-implementation": [a2a],
+        "all-reduce-implementation": ar_impl,
+        "all-gather-implementation": ag_impl,
+        "reduce-scatter-implementation": rs_impl,
+        "all-to-all-implementation": a2a_impl,
         "collective-optimization": "localBWAware",
         "local-mem-bw": 1600,
         "boost-mode": 0,
@@ -167,13 +493,24 @@ def generate_astrasim_configs_from_hw(
         if getattr(sys_opts, "endpoint_delay", None) is not None:
             system["endpoint-delay"] = sys_opts.endpoint_delay
         if getattr(sys_opts, "active_chunks_per_dimension", None) is not None:
-            system["active-chunks-per-dimension"] = sys_opts.active_chunks_per_dimension
+            acpd = sys_opts.active_chunks_per_dimension
+            if isinstance(acpd, (list, tuple)):
+                acpd = acpd[0] if acpd else 1
+            system["active-chunks-per-dimension"] = int(acpd)
         if getattr(sys_opts, "preferred_dataset_splits", None) is not None:
-            system["preferred-dataset-splits"] = sys_opts.preferred_dataset_splits
-
-    _save_json(sys_json, system, npus_key=int(npus_count))
-
-    return {"network_yaml": net_yaml, "system_json": sys_json}
+            pds = sys_opts.preferred_dataset_splits
+            if isinstance(pds, (list, tuple)):
+                pds = pds[0] if pds else 1
+            system["preferred-dataset-splits"] = int(pds)
+
+    _save_json(sys_json, system)
+
+    return {
+        "network_yaml": net_yaml,
+        "system_json": sys_json,
+        "topology_list": topo_list,
+        "npus_per_dim": npus_list,
+    }
 
 
 __all__ = [
diff --git a/astrasim_lib/executor.py b/astrasim_lib/executor.py
index 57f71ebe..dc2d5fd2 100644
--- a/astrasim_lib/executor.py
+++ b/astrasim_lib/executor.py
@@ -20,7 +20,7 @@
 import time
 from collections import defaultdict, deque
 from itertools import count as _it_count
-from typing import Any, Dict, List, Optional, Set, Tuple
+from typing import Any, Dict, List, Mapping, Optional, Sequence, Set, Tuple
 
 sys.setrecursionlimit(10000)
 
@@ -28,6 +28,7 @@
 from graphviz import Digraph
 
 from .config_generation import ASTRA_DEBUG, generate_astrasim_configs_from_hw
+from . import gmap
 from .et_utils import (
     chakra_decode,
     chakra_encode,
@@ -40,7 +41,9 @@
     write_et_node,
 )
 from .integration import get_remote_memory_path, run_astrasim_analytical, run_cache_astrasim
+from .layout_utils import derive_axes_filter
 from simulate_LLM import visualize_graph
+from util import relpath_display
 
 
 def _env_truthy(name: str) -> bool:
@@ -64,11 +67,17 @@ def _clean_astrasim_artifacts(directory: str) -> None:
                     os.remove(path)
                 except OSError:
                     pass
-            elif entry in {
-                "manifest.json",
-                "system_native_collectives.json",
-                "comm_groups.json",
-            }:
+            elif entry == "manifest.json":
+                try:
+                    os.remove(path)
+                except OSError:
+                    pass
+            elif entry == "comm_groups.json" or entry.startswith("comm_groups_"):
+                try:
+                    os.remove(path)
+                except OSError:
+                    pass
+            elif entry.startswith("system_native_collectives"):
                 try:
                     os.remove(path)
                 except OSError:
@@ -86,7 +95,215 @@ def _clean_astrasim_artifacts(directory: str) -> None:
 
 _TP_GROUP_BASE_ID = 1000
 _LAST_TP_GROUPS: Optional[Dict[str, List[int]]] = None
+_LAST_STAGE_GROUPS: Optional[List[List[int]]] = None
 _TP_LABEL_DP_TO_ID: Dict[Tuple[str, int], str] = {}
+_TP_MEMBERS_TO_GID: Dict[Tuple[int, ...], str] = {}
+
+
+def _extract_axis_layout(rank_layout: Optional[Dict[str, Any]]) -> Tuple[List[str], Dict[str, int], Dict[str, int]]:
+    if not isinstance(rank_layout, dict):
+        return [], {}, {}
+    axis_order = list(rank_layout.get("axis_order", []))
+    raw_sizes = rank_layout.get("axis_sizes", {})
+    if isinstance(axis_order, tuple):
+        axis_order = list(axis_order)
+    axis_sizes: Dict[str, int] = {}
+    if isinstance(raw_sizes, dict):
+        for key, value in raw_sizes.items():
+            try:
+                axis_sizes[str(key)] = max(1, int(value))
+            except Exception:
+                axis_sizes[str(key)] = 1
+    raw_strides = rank_layout.get("axis_strides", {})
+    axis_strides: Dict[str, int] = {}
+    if isinstance(raw_strides, dict):
+        for key, value in raw_strides.items():
+            try:
+                axis_strides[str(key)] = int(value)
+            except Exception:
+                axis_strides[str(key)] = 0
+    else:
+        span = 1
+        for axis in axis_order:
+            axis_strides[axis] = span
+            span *= axis_sizes.get(axis, 1)
+    return axis_order, axis_sizes, axis_strides
+
+
+def _remap_stages_for_mapping(
+    *,
+    permutation: Sequence[int],
+    collector: Optional[gmap.GMapCollector],
+    stage_ids: List[int],
+    stage_to_ranks: Dict[int, List[int]],
+    stage_index: Dict[int, int],
+    rank_traces: Dict[int, "_RankTrace"],
+    rank_meta: Dict[int, Dict[str, int]],
+    dp_count: int,
+    num_stages: int,
+    output_dir: str,
+) -> Tuple[List[int], Dict[int, int]]:
+    """Reorder stages so AstraSim ranks match the SCOTCH permutation.
+
+    We apply the permutation *before* ET emission so the rest of the converter
+    remains unaware of the remap. Only stage ordering / rank bookkeeping is
+    touched here, which keeps the change localized and avoids renaming files
+    after the fact.
+    """
+
+    if collector is None or not permutation:
+        return stage_ids, stage_index
+
+    if len(permutation) != collector.vertex_count:
+        raise ValueError("Permutation length does not match first-dimension vertex count.")
+
+    def _sort_key(stage: int) -> Tuple[int, ...]:
+        coords = collector.stage_axis_coords.get(stage, {})
+        higher = tuple(int(coords.get(axis, 0)) for axis in collector.replication_axes)
+        local_idx = collector.local_index_for_stage(stage)
+        target_idx = permutation[local_idx]
+        return higher + (target_idx,)
+
+    ordered_stages = sorted(stage_ids, key=_sort_key)
+    new_stage_index = {stage: idx for idx, stage in enumerate(ordered_stages)}
+
+    new_stage_to_ranks: Dict[int, List[int]] = {stage: [0] * dp_count for stage in ordered_stages}
+    new_rank_traces: Dict[int, _RankTrace] = {}
+    new_rank_meta: Dict[int, Dict[str, int]] = {}
+
+    for stage in stage_ids:
+        for dp_idx in range(dp_count):
+            old_rank = stage_to_ranks[stage][dp_idx]
+            new_idx = new_stage_index[stage]
+            new_rank = dp_idx * num_stages + new_idx
+            trace = rank_traces[old_rank]
+            trace.rank = new_rank
+            trace.path = os.path.join(output_dir, f"llm_graph.{new_rank}.et")
+            new_rank_traces[new_rank] = trace
+            new_stage_to_ranks[stage][dp_idx] = new_rank
+            new_rank_meta[new_rank] = {"stage": stage, "dp": dp_idx}
+
+    stage_to_ranks.clear()
+    stage_to_ranks.update(new_stage_to_ranks)
+    rank_traces.clear()
+    rank_traces.update(new_rank_traces)
+    rank_meta.clear()
+    rank_meta.update(new_rank_meta)
+
+    stage_ids[:] = ordered_stages
+    stage_index = new_stage_index
+    return ordered_stages, stage_index
+
+
+def _compute_stage_axis_coords(
+    stage_ids: Sequence[int],
+    axis_order: Sequence[str],
+    axis_sizes: Mapping[str, int],
+) -> Dict[int, Dict[str, int]]:
+    coords: Dict[int, Dict[str, int]] = {}
+    if not axis_order:
+        return coords
+    for stage in stage_ids:
+        remaining = int(stage)
+        stage_coords: Dict[str, int] = {}
+        for axis in axis_order:
+            size = max(1, int(axis_sizes.get(axis, 1)))
+            if size > 1:
+                stage_coords[axis] = remaining % size
+                remaining //= size
+            else:
+                stage_coords[axis] = 0
+        coords[stage] = stage_coords
+    return coords
+
+
+def _build_axis_groups(
+    axis_order: Sequence[str],
+    axis_sizes: Mapping[str, int],
+    stage_axis_coords: Mapping[int, Dict[str, int]],
+    stage_to_ranks: Mapping[int, List[int]],
+) -> Dict[str, Dict[Tuple[int, Tuple[Tuple[str, int], ...]], List[int]]]:
+    axis_groups: Dict[
+        str, Dict[Tuple[int, Tuple[Tuple[str, int], ...]], List[int]]
+    ] = {}
+    if not axis_order:
+        return axis_groups
+    for axis in axis_order:
+        size = max(1, int(axis_sizes.get(axis, 1)))
+        if size <= 1:
+            continue
+        groups: Dict[
+            Tuple[int, Tuple[Tuple[str, int], ...]], List[int]
+        ] = defaultdict(list)
+        for stage, coords in stage_axis_coords.items():
+            key_base = tuple((ax, coords.get(ax, 0)) for ax in axis_order if ax != axis)
+            ranks = stage_to_ranks.get(stage, [])
+            for dp_idx, rank in enumerate(ranks):
+                groups[(dp_idx, key_base)].append(rank)
+        axis_groups[axis] = {
+            key: sorted(set(values)) for key, values in groups.items() if values
+        }
+    return axis_groups
+
+
+def _assign_collective_labels(
+    tp_collective_groups: Mapping[str, List[Any]],
+    collective_info: Mapping[Any, Dict[str, Any]],
+    axis_order: Sequence[str],
+    axis_sizes: Mapping[str, int],
+    stage_axis_coords: Mapping[int, Dict[str, int]],
+    axis_groups: Mapping[str, Dict[Tuple[int, Tuple[Tuple[str, int], ...]], List[int]]],
+    stage_to_ranks: Mapping[int, List[int]],
+    dp_count: int,
+) -> Tuple[Dict[Any, str], Dict[str, List[Any]], Dict[str, Set[Tuple[str, int, Tuple[int, ...]]]]]:
+    tp_collective_labels: Dict[Any, str] = {}
+    label_to_edges: Dict[str, List[Any]] = {}
+    label_group_tokens: Dict[str, Set[Tuple[str, int, Tuple[int, ...]]]] = defaultdict(set)
+    group_key_to_label: Dict[Tuple[str, Tuple[int, ...]], str] = {}
+    label_suffix_counter: Dict[str, int] = defaultdict(int)
+
+    for base_name, edges in tp_collective_groups.items():
+        if not edges:
+            continue
+        for edge in sorted(edges, key=lambda e: getattr(e, "op_id", 0)):
+            info = collective_info[edge]
+            axis = str(info.get("interconnect_type", "tp")).lower()
+            stage = info["stage"]
+            coords = stage_axis_coords.get(stage, {})
+            key_base = tuple((ax, coords.get(ax, 0)) for ax in axis_order if ax != axis)
+
+            members_per_dp: List[Tuple[int, Tuple[int, ...]]] = []
+            for dp_idx in range(dp_count):
+                members = None
+                axis_map = axis_groups.get(axis)
+                if axis_map is not None:
+                    members = axis_map.get((dp_idx, key_base))
+                if not members:
+                    ranks = stage_to_ranks.get(stage, [])
+                    if dp_idx < len(ranks):
+                        members = [ranks[dp_idx]]
+                if members:
+                    members_tuple = tuple(sorted(members))
+                    members_per_dp.append((dp_idx, members_tuple))
+
+            if not members_per_dp:
+                continue
+
+            primary_members = members_per_dp[0][1]
+            label_key = (base_name, primary_members)
+            label = group_key_to_label.get(label_key)
+            if label is None:
+                suffix = label_suffix_counter[base_name]
+                label = base_name if suffix == 0 else f"{base_name}_{suffix}"
+                label_suffix_counter[base_name] += 1
+                group_key_to_label[label_key] = label
+
+            tp_collective_labels[edge] = label
+            label_to_edges.setdefault(label, []).append(edge)
+            for dp_idx, members_tuple in members_per_dp:
+                label_group_tokens[label].add((axis, dp_idx, members_tuple))
+
+    return tp_collective_labels, label_to_edges, label_group_tokens
 
 
 class _RankTrace:
@@ -128,12 +345,19 @@ def _is_control_send(node: pb.Node) -> bool:
                 node.type == pb.COMM_SEND_NODE
                 and (node.name or "").endswith("_send_control")
             )
+        def _is_control_recv(node: pb.Node) -> bool:
+            return (
+                node.type == pb.COMM_RECV_NODE
+                and (node.name or "").endswith("_recv_control")
+            )
 
         control_nodes: List[pb.Node] = []
         regular_nodes: List[pb.Node] = []
         for node in self.nodes:
             if _is_control_send(node):
                 control_nodes.append(node)
+            elif _is_control_recv(node):
+                control_nodes.append(node)
             else:
                 regular_nodes.append(node)
 
@@ -203,6 +427,9 @@ def _visualize_et_files(et_paths: List[str]) -> None:
     if not et_paths:
         return
 
+    # prune et_paths to only include the first 10
+    et_paths = et_paths[:20]
+
     def _render_et_file(et_path: str) -> None:
         try:
             fh = chakra_open(et_path)
@@ -224,8 +451,7 @@ def _render_et_file(et_path: str) -> None:
             nodes.append(node)
         fh.close()
 
-        dot = Digraph(comment=os.path.basename(et_path))
-        dot.graph_attr.update({"rankdir": "TB", "fontsize": "10"})
+        dot = Digraph(comment=os.path.basename(et_path), format="svg",graph_attr={"rankdir": "TB", "fontsize": "10"})
 
         id_to_node = {int(node.id): node for node in nodes}
         type_color = {
@@ -271,25 +497,48 @@ def _render_et_file(et_path: str) -> None:
         dir_name, base_name = os.path.split(et_path)
         viz_base = base_name + ".viz"
         try:
-            output_path = dot.render(viz_base, directory=dir_name or None, format="png", cleanup=True)
-            final_png = et_path + ".png"
-            if output_path and output_path != final_png:
+            output_path = dot.render(viz_base, directory=dir_name or None, format="svg", cleanup=True)
+            final_svg = et_path + ".svg"
+            if output_path and output_path != final_svg:
                 try:
-                    os.replace(output_path, final_png)
+                    os.replace(output_path, final_svg)
                 except FileNotFoundError:
                     # Some graphviz versions return path without creating file when empty graph
                     pass
         except Exception as exc:
             print(f"[WARN] Failed to render Graphviz graph for {et_path}: {exc}")
 
-    for et_path in et_paths:
-        message = f" | ET Graph saved to {et_path}.png"
+    if not et_paths:
+        return
+
+    if len(et_paths) == 1:
+        et_path = et_paths[0]
+        display_path = relpath_display(f"{et_path}.svg")
+        message = f" | ET Graph saved to {display_path}"
         graphviz_async.submit(
             f"et:{os.path.basename(et_path)}",
             _render_et_file,
             et_path,
             print_message=message,
         )
+        return
+
+    first_path = et_paths[0]
+    display_first = relpath_display(f"{first_path}.svg")
+    summary_message = f" | {len(et_paths)} ET graphs saved to {display_first} (..)"
+    graphviz_async.submit(
+        f"et:{os.path.basename(first_path)}",
+        _render_et_file,
+        first_path,
+        print_message=summary_message,
+    )
+    for et_path in et_paths[1:]:
+        graphviz_async.submit(
+            f"et:{os.path.basename(et_path)}",
+            _render_et_file,
+            et_path,
+            print_message=None,
+        )
 
 
 
@@ -367,7 +616,7 @@ def fmt_attr_map(node: pb.Node) -> str:
         try:
             with open(out_path, "w") as outf:
                 outf.write("\n".join(lines) + "\n")
-            print(f"[AstraSim] Saved ET text dump to {out_path}")
+            # print(f"[AstraSim] Saved ET text dump to {out_path}")
         except OSError as exc:
             print(f"[WARN] Failed to write ET text dump for {et_path}: {exc}")
 
@@ -423,12 +672,32 @@ def convert_deepflow_graph_to_chakra_et(
     output_dir: str,
 ) -> Tuple[str, List[int], str]:
     """Convert DeepFlow graph to AstraSim ET format by scheduling per stage and DP rank.
-    This function is black magic. Don't touch it unless you know what you are doing.
-    Prefer changes *anywhere* else in the codebase if possible."""
+    This means converting the single DeepFlow DAG into a set of different AstraSim ET files.
+
+    INPUTS: ``graph_root`` is the DeepFlow DAG (after flattening). ``dp_size`` is the
+    data-parallel replication degree. ``output_dir`` is where ET traces + metadata
+    should be written.
+
+    OUTPUTS: returns ``(et_prefix, rank_ids, manifest_path)``. ``et_prefix`` is the
+    file prefix for ``llm_graph.<rank>.et``. ``rank_ids`` is the list of produced rank
+    IDs, and ``manifest_path`` points to the per-rank summary JSON.
+
+    CONCEPTS: a *stage* is a unique ``hw_id`` (flattened hardware coordinate) of the
+    compute graph. Each (stage, DP replica) pair becomes an AstraSim rank. The function
+    recovers tensor/context/pipeline axes, schedules tasks per stage and per DP rank, and
+    emits Chakra ET traces with matching communicator metadata so AstraSim can replay the
+    workload deterministically.
+    
+    This function is very complicated. Don't touch it unless you know what you are doing.
+    Prefer changes *anywhere* else in the codebase if possible. I have tried to document it
+    as much as possible, but it definitely needs further cleaning and breaking up. This is TODO."""
 
     os.makedirs(output_dir, exist_ok=True)
     debug_enabled = _env_truthy("DEEPFLOW_DUMP_CONVERTER_DEBUG")
 
+    # Step 1: Traverse the graph once to snapshot every object reachable from the
+    # provided root(s). Collecting upfront avoids mutating the graph while we
+    # analyse dependencies.
     def collect_objects(root) -> List[Any]:
         visited: Set[int] = set()
         ordered: List[Any] = []
@@ -459,6 +728,9 @@ def dfs(obj: Any) -> None:
     if not compute_nodes:
         raise ValueError("DeepFlow graph did not expose any executable compute nodes (hw_id >= 0).")
 
+    # Step 2: Each unique ``hw_id`` defines a logical stage (one flattened hardware
+    # coordinate). We build a DP-major rank ordering so AstraSim sees one rank per
+    # (stage, data-parallel replica) pair.
     stage_ids = sorted({node.hw_id for node in compute_nodes})
     dp_count = max(int(dp_size) if dp_size else 1, 1)
 
@@ -481,6 +753,33 @@ def dfs(obj: Any) -> None:
     def rank_for(stage: int, dp_idx: int) -> int:
         return stage_to_ranks[stage][dp_idx]
 
+    def _compute_duration_seconds(task: Any, dp_idx: int) -> float:
+        """Return the per-DP duration (in seconds) for ``task``."""
+        profile = getattr(task, "duration_profile", None)
+        if profile:
+            if len(profile) != dp_count:
+                raise ValueError(
+                    f"Duration profile for node '{getattr(task, 'name', '<unnamed>')}' "
+                    f"has length {len(profile)} but dp_count={dp_count}."
+                )
+            return float(profile[dp_idx])
+        duration_sec = getattr(task, "duration", 0.0) or 0.0
+        return float(duration_sec)
+
+    # Step 3: Recover the multi-dimensional axis layout (tp/cp/lp, etc.) so that
+    # communicator membership can be reconstructed deterministically later on.
+    rank_layout = getattr(graph_root, "_astrasim_rank_layout", None)
+    axis_order, axis_sizes, axis_strides = _extract_axis_layout(rank_layout)
+    stage_axis_coords = _compute_stage_axis_coords(stage_ids, axis_order, axis_sizes)
+    axis_groups = _build_axis_groups(axis_order, axis_sizes, stage_axis_coords, stage_to_ranks)
+    if hasattr(graph_root, "_optimize_2dmap"):
+        optimize_cfg = getattr(graph_root, "_optimize_2dmap")
+    else:
+        optimize_cfg = None
+    collector = gmap.begin_collection(optimize_cfg, axis_sizes, stage_axis_coords, output_dir)
+
+    # Step 4: Attribute each comm edge to the stage that owns it. This lets the
+    # dependency walkers treat collectives the same way they treat compute nodes.
     edge_stage: Dict[Any, int] = {}
     for obj in all_objects:
         comm = getattr(obj, "comm_type", None)
@@ -523,6 +822,37 @@ def rank_for(stage: int, dp_idx: int) -> int:
 
     pipeline_edge_map: Dict[Tuple[Any, Any], Any] = {}
 
+    def _record_pipeline_edges_for_first_dim(collector_obj: gmap.GMapCollector) -> None:
+        if not collector_obj.include_pipeline:
+            return
+        for obj in all_objects:
+            if getattr(obj, "comm_type", None) != "pipeline":
+                continue
+            parents = getattr(obj, "parents", [])
+            src_stage = None
+            for parent in parents:
+                if hasattr(parent, "hw_id") and parent.hw_id is not None and int(parent.hw_id) >= 0:
+                    src_stage = int(parent.hw_id)
+                    break
+            if src_stage is None:
+                raise ValueError("Unable to resolve source stage for pipeline edge during gmap collection.")
+            children = getattr(obj, "children", [])
+            dst_stage = None
+            for child in children:
+                if hasattr(child, "hw_id") and child.hw_id is not None and int(child.hw_id) >= 0:
+                    dst_stage = int(child.hw_id)
+                    break
+            if dst_stage is None:
+                raise ValueError("Unable to resolve destination stage for pipeline edge during gmap collection.")
+            if src_stage == dst_stage:
+                continue
+            if not hasattr(obj, "comm_size_bytes"):
+                raise ValueError("Pipeline edge is missing comm_size_bytes for gmap collection.")
+            size_bytes = int(getattr(obj, "comm_size_bytes"))
+            collector_obj.record_pipeline(src_stage=src_stage, dst_stage=dst_stage, size_bytes=size_bytes)
+
+    # Step 5a: Walk backwards from a compute node to classify dependencies into
+    # stage-local, cross-stage (pipeline), or collective inputs.
     def analyze_for_compute(node: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]:
         stage = node.hw_id
         stage_deps: Set[Any] = set()
@@ -685,6 +1015,8 @@ def analyze_for_compute(node: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]:
         
         return stage_deps, pipeline_deps, collective_deps
 
+    # Step 5b: Same classification for collective edges (they often use local_hw_id
+    # instead of hw_id, so we check both).
     def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]:
         stage = edge_stage[edge]
         stage_deps: Set[Any] = set()
@@ -744,6 +1076,7 @@ def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]:
 
         return stage_deps, pipeline_deps, collective_deps
 
+    # Aggregate the analysis results so later stages can fetch dependencies by node.
     compute_info: Dict[Any, Dict[str, Any]] = {}
     for node in compute_nodes:
         stage_deps, pipeline_deps, collective_deps = analyze_for_compute(node)
@@ -758,7 +1091,7 @@ def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]:
     collective_info: Dict[Any, Dict[str, Any]] = {}
     for edge, stage in edge_stage.items():
         stage_deps, pipeline_deps, collective_deps = analyze_for_collective(edge)
-        collective_info[edge] = {
+        entry = {
             "stage": stage,
             "stage_deps": stage_deps,
             "pipeline_deps": pipeline_deps,
@@ -769,6 +1102,20 @@ def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]:
             "interconnect_type": getattr(edge, "comm_interconnect_type", None),
             "participants": int(getattr(edge, "participants", 0) or 0),
         }
+        collective_info[edge] = entry
+        if collector:
+            axis_name = entry["interconnect_type"]
+            if axis_name:
+                axis_norm = str(axis_name).strip().lower()
+                if axis_norm in collector.subset_axes:
+                    should_double = entry["comm_type"] in (pb.ALL_REDUCE, pb.ALL_TO_ALL)
+                    collector.record_collective(
+                        axis=axis_norm,
+                        stage_id=stage,
+                        size_bytes=entry["size"],
+                        participant_count=entry["participants"],
+                        double_weight=should_double,
+                    )
 
     if debug_enabled:
         print("[ConverterDebug] === COMPUTE INFO (stage deps / pipeline deps / collectives) ===")
@@ -792,36 +1139,34 @@ def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]:
                 f"[ConverterDebug] collective {name}: stage={info['stage']}"
                 f" stage_deps={stage_names} pipeline_deps={pipe_names} collectives={coll_names}"
             )
+    if collector:
+        _record_pipeline_edges_for_first_dim(collector)
+    # Step 6: Group collectives by label/axis so we can reconstruct communicator
+    # memberships (tp/cp/lp) without relying on the original ordering.
     tp_collective_groups: Dict[str, List[Any]] = defaultdict(list)
     for edge, info in collective_info.items():
         interconnect_type = info.get("interconnect_type")
         if interconnect_type and interconnect_type not in {"dp", "lp", "pipeline"}:
             tp_collective_groups[info["name"]].append(edge)
 
-    tp_collective_labels: Dict[Any, str] = {}
-    # Map final labels (with suffix) to their member edges, preserving order
-    label_to_edges: Dict[str, List[Any]] = {}
-    for base_name, edges in tp_collective_groups.items():
-        if not edges:
-            continue
-        sorted_edges = sorted(edges, key=lambda edge: getattr(edge, "op_id", 0))
-        group_index = 0
-        idx = 0
-        while idx < len(sorted_edges):
-            edge = sorted_edges[idx]
-            info = collective_info[edge]
-            participants = max(1, int(info.get("participants") or 1))
-            label = base_name if group_index == 0 else f"{base_name}_{group_index}"
-            group_members = sorted_edges[idx: idx + participants]
-            for member in group_members:
-                tp_collective_labels[member] = label
-            # Record members for this final label in insertion order
-            if label not in label_to_edges:
-                label_to_edges[label] = []
-            label_to_edges[label].extend(group_members)
-            idx += len(group_members)
-            group_index += 1
+    (
+        tp_collective_labels,
+        label_to_edges,
+        label_group_tokens,
+    ) = _assign_collective_labels(
+        tp_collective_groups,
+        collective_info,
+        axis_order,
+        axis_sizes,
+        stage_axis_coords,
+        axis_groups,
+        stage_to_ranks,
+        dp_count,
+    )
 
+    # Step 7: Build a per-stage task list (compute + collectives). If a collective
+    # introduces a new stage we extend ``stage_to_ranks`` so every stage has the
+    # full (dp-major) rank list.
     stage_tasks: Dict[int, Set[Any]] = {stage: set() for stage in stage_ids}
     for node in compute_nodes:
         stage_tasks[node.hw_id].add(node)
@@ -832,6 +1177,7 @@ def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]:
             # Ensure mapping present even for stages discovered only via collectives
             stage_idx = stage_index.setdefault(stage, len(stage_index))
             stage_to_ranks[stage] = [dp_idx * num_stages + stage_idx for dp_idx in range(dp_count)]
+            stage_ids.append(stage)
         stage_tasks[stage].add(edge)
 
     stage_adj: Dict[int, Dict[Any, Set[Any]]] = {stage: defaultdict(set) for stage in stage_tasks}
@@ -957,33 +1303,8 @@ def analyze_for_collective(edge: Any) -> Tuple[Set[Any], Set[Any], Set[Any]]:
     pipeline_recv_cache: Dict[Tuple[Any, Any, int], int] = {}
     tag_counter = itertools.count(start=1)
 
-    # Build TP communicator groups (ids start at 1000) per label and dp
-    global _LAST_TP_GROUPS, _TP_LABEL_DP_TO_ID
-    _LAST_TP_GROUPS = {}
-    _TP_LABEL_DP_TO_ID = {}
-    tp_gid_counter = _it_count(start=_TP_GROUP_BASE_ID)
-
-    # Collect per-label TP members per dp index (use prebuilt mapping to avoid rescans)
-    for label in sorted(label_to_edges.keys()):
-        edges_for_label = label_to_edges.get(label, [])
-        if not edges_for_label:
-            continue
-        # Determine stage for each edge
-        stages = [collective_info[e]["stage"] for e in edges_for_label]
-        # For each dp index, map to ranks
-        for dp_idx in range(dp_count):
-            members: List[int] = []
-            for e, st in zip(edges_for_label, stages):
-                try:
-                    members.append(stage_to_ranks[st][dp_idx])
-                except Exception:
-                    continue
-            if not members:
-                continue
-            gid = str(next(tp_gid_counter))
-            _TP_LABEL_DP_TO_ID[(label, dp_idx)] = gid
-            _LAST_TP_GROUPS[gid] = sorted(members)
-
+    # Step 8 helper: materialise SEND/RECV control edges across stages whenever a
+    # pipeline dependency exists between ``parent`` and ``target``.
     def ensure_pipeline(parent: Any, target: Any, dp_idx: int) -> int:
         parent_stage = getattr(parent, "hw_id", None)
         if parent_stage == None:
@@ -1065,9 +1386,65 @@ def ensure_pipeline(parent: Any, target: Any, dp_idx: int) -> int:
         pipeline_recv_cache[key] = recv_id
         return recv_id
 
+    mapping_result = gmap.finalize_collection(collector)
+    if mapping_result:
+        stage_ids, stage_index = _remap_stages_for_mapping(
+            permutation=mapping_result.permutation,
+            collector=collector,
+            stage_ids=stage_ids,
+            stage_to_ranks=stage_to_ranks,
+            stage_index=stage_index,
+            rank_traces=rank_traces,
+            rank_meta=rank_meta,
+            dp_count=dp_count,
+            num_stages=num_stages,
+            output_dir=output_dir,
+        )
+        # Rebuild communicator layout after remapping so pg_name and comm_groups.json
+        # reflect the SCOTCH permutation.
+        axis_groups = _build_axis_groups(axis_order, axis_sizes, stage_axis_coords, stage_to_ranks)
+        (
+            tp_collective_labels,
+            label_to_edges,
+            label_group_tokens,
+        ) = _assign_collective_labels(
+            tp_collective_groups,
+            collective_info,
+            axis_order,
+            axis_sizes,
+            stage_axis_coords,
+            axis_groups,
+            stage_to_ranks,
+            dp_count,
+        )
+
+    # Build TP communicator groups (ids start at 1000) per label and dp using the
+    # (potentially remapped) label tokens.
+    global _LAST_TP_GROUPS, _TP_LABEL_DP_TO_ID, _TP_MEMBERS_TO_GID
+    _LAST_TP_GROUPS = {}
+    _TP_LABEL_DP_TO_ID = {}
+    _TP_MEMBERS_TO_GID = {}
+    tp_gid_counter = _it_count(start=_TP_GROUP_BASE_ID)
 
-    for stage in stage_order:
-        order = stage_order[stage]
+    for label in sorted(label_group_tokens.keys()):
+        tokens = sorted(label_group_tokens[label])
+        for axis, dp_idx, members_tuple in tokens:
+            members_list = list(members_tuple)
+            members_key = tuple(members_list)
+            gid = _TP_MEMBERS_TO_GID.get(members_key)
+            if gid is None:
+                gid = str(next(tp_gid_counter))
+                _TP_MEMBERS_TO_GID[members_key] = gid
+                if gid not in _LAST_TP_GROUPS:
+                    _LAST_TP_GROUPS[gid] = members_list
+            _TP_LABEL_DP_TO_ID[(label, dp_idx)] = gid
+
+    # Step 10: Emit ET nodes per stage (collectives first, then compute nodes),
+    # iterating in the topological order computed above.
+    for stage in stage_ids:
+        order = stage_order.get(stage)
+        if not order:
+            continue
         for task in order:
             if task in collective_info:
                 info = collective_info[task]
@@ -1136,7 +1513,7 @@ def ensure_pipeline(parent: Any, target: Any, dp_idx: int) -> int:
                         if dep not in unique_deps:
                             unique_deps.append(dep)
 
-                    duration_sec = getattr(task, "duration", 0.0) or 0.0
+                    duration_sec = _compute_duration_seconds(task, dp_idx)
                     duration_micros = int(round(duration_sec * 1e6)) if duration_sec else 0
                     node_id = trace.next_id
                     comp_node = new_comp_node(
@@ -1153,7 +1530,8 @@ def ensure_pipeline(parent: Any, target: Any, dp_idx: int) -> int:
                     # above to avoid duplication. Here we only reference them
                     # via dependencies when needed.
 
-    # Attach cross-stage (pipeline) dependencies after all nodes are created.
+    # Step 11: Attach cross-stage (pipeline) dependencies by wiring the RECV
+    # nodes created via ``ensure_pipeline`` into the ET nodes’ control deps.
     for stage, order in stage_order.items():
         for task in order:
             if task in collective_info:
@@ -1196,7 +1574,8 @@ def ensure_pipeline(parent: Any, target: Any, dp_idx: int) -> int:
                         node.ctrl_deps.append(rid)
 
 
-    # Build manifest ops directly from in-memory traces (graph-derived), not ET files
+    # Step 12: Build manifest ops directly from the in-memory traces. This lets the
+    # caller reuse the same metadata that will be written into the ET files.
     manifest_ranks: Dict[str, List[List]] = {}
     def _manifest_op_key(op: List) -> tuple:
         # Sort primarily by size/duration, with stable type ordering
@@ -1257,7 +1636,8 @@ def _manifest_op_key(op: List) -> tuple:
         rank_ops_sorted = sorted(rank_ops, key=_manifest_op_key)
         manifest_ranks[str(int(rank))] = rank_ops_sorted
 
-    # Now actually write ET files
+    # Step 13: Write the Chakra ET traces (with control-send renumbering) and
+    # persist the manifest/metadata alongside them.
     for trace in rank_traces.values():
         trace.close()
 
@@ -1285,6 +1665,7 @@ def run_astra_simulation_only_onepath(
     output_dir: str = "astra_comparison_output",
     dp_override: Optional[int] = None,
     persist_artifacts: Optional[bool] = None,
+    faulty_links_override: Optional[Sequence[Tuple[int, int, float]]] = None,
 ):
     """
     Run AstraSim simulation on DeepFlow graph and print results.
@@ -1293,6 +1674,7 @@ def run_astra_simulation_only_onepath(
         fwdbwd_root: Forward and backward graph root node
         time_calc_obj: TimeCalculationLLM object with hw_config and dp attributes
         output_dir: Directory for temporary files and results
+        faulty_links_override: Optional remapped faulty link list for this run
     """
     print("\n" + "="*60)
     print("ASTRASIM SIMULATION RESULTS")
@@ -1326,20 +1708,20 @@ def run_astra_simulation_only_onepath(
         rank_count = len(rank_ids)
         # Astrasim doesn't play well with only 1 rank.
         # When that happens, let's duplicate to 2 ranks. No collectives exist between the two so this should not have an effect.
+        synthetic_pair = False
         if rank_count == 1:
             # duplicate the .et file
             src = os.path.join(work_dir, "llm_graph.0.et")
             dst = os.path.join(work_dir, "llm_graph.1.et")
             shutil.copy(src, dst)
             rank_ids = [rank_ids[0], rank_ids[0]+1]
+            synthetic_pair = True
         rank_count = len(rank_ids)
 
         # Handle artifact visualization if enabled
         if persist and _env_truthy("DEEPFLOW_PERSIST_ARTIFACT_VIZ"):
             et_paths = []
             for rank in rank_ids:
-                if rank > 10:
-                    break
                 et_path = os.path.join(work_dir, f"llm_graph.{rank}.et")
                 if os.path.exists(et_path):
                     et_paths.append(et_path)
@@ -1356,22 +1738,47 @@ def run_astra_simulation_only_onepath(
         if os.environ.get("DEEPFLOW_ASTRA_SKIP_EXEC"):
             print("[AstraSim] DEEPFLOW_ASTRA_SKIP_EXEC set. Exiting after ET artifact generation.")
             exit()
-        astra_configs = generate_astrasim_configs_from_hw(time_calc_obj.hw_config, work_dir, rank_count)
 
-        # Run AstraSim simulation on forward graph (cached via manifest)
-        print(f"[AstraSim] Executing forward simulation with {rank_count} ranks...")
-        fwd_times, fwd_total = run_cache_astrasim(
+        rank_layout = getattr(fwdbwd_root, "_astrasim_rank_layout", None)
+        axis_order, axis_sizes, _ = _extract_axis_layout(rank_layout)
+        preferred_axes_for_synthetic = tuple(axis_order) if axis_order else tuple()
+        axes_filter = derive_axes_filter(axis_order, axis_sizes, dp_count)
+        if not axes_filter:
+            axes_filter = None
+        if synthetic_pair:
+            axes_filter = ["synthetic2"]
+        astra_configs = generate_astrasim_configs_from_hw(
             time_calc_obj.hw_config,
-            comm="graph",
-            npus_count=rank_count,
-            size_bytes=0,
-            astra_config_dir="./astra_cache",
-            cache_path="./astra_cache/cache.json",
-            manifest_json_path=fwd_manifest,
-            workload_prefix=fwd_et_prefix,
-            comm_group_json=comm_groups_path,
+            work_dir,
+            rank_count,
+            axes_filter=axes_filter,
+            faulty_links_override=faulty_links_override,
+            preferred_axes_for_synthetic=preferred_axes_for_synthetic,
         )
 
+        # Run AstraSim simulation on forward graph (cached via manifest)
+        print(f"[AstraSim] Executing forward simulation with {rank_count} ranks...")
+        cache_override_env = os.environ.pop("ASTRA_CACHE_DIR", None)
+        local_cache_dir = work_dir
+        local_cache_path = os.path.join(work_dir, "cache.json")
+        try:
+            fwd_times, fwd_total = run_cache_astrasim(
+                time_calc_obj.hw_config,
+                comm="graph",
+                npus_count=rank_count,
+                size_bytes=0,
+                astra_config_dir=local_cache_dir,
+                cache_path=local_cache_path,
+                manifest_json_path=fwd_manifest,
+                workload_prefix=fwd_et_prefix,
+                comm_group_json=comm_groups_path,
+                axes_filter=axes_filter,
+                files=astra_configs,
+            )
+        finally:
+            if cache_override_env is not None:
+                os.environ["ASTRA_CACHE_DIR"] = cache_override_env
+
 
         conversion_and_sim_time = time.time() - astrasim_start
 
@@ -1408,8 +1815,8 @@ def run_astra_simulation_only_onepath(
 
     def my_save_graph(roots, output_folder = "output_graph/", filename="graph"):
         dot_fw = visualize_graph(roots, filename=output_folder + filename)
-        dot_fw.render(output_folder + filename , format="png", cleanup=True)
-        print("graph saved to %s%s.png" % (output_folder , filename ))
+        dot_fw.render(output_folder + filename , format="svg", cleanup=True)
+        print("graph saved to %s%s.svg" % (output_folder , filename ))
 
     import config
     import pickle
diff --git a/astrasim_lib/fault_projection.py b/astrasim_lib/fault_projection.py
new file mode 100644
index 00000000..54a2e12b
--- /dev/null
+++ b/astrasim_lib/fault_projection.py
@@ -0,0 +1,159 @@
+"""Utilities for projecting faulty link endpoints onto axis subsets."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, Iterable, Mapping, Optional, Sequence, Tuple, List
+
+from .layout_utils import AxisLayout, decode_axis_coordinates, build_subset_strides
+
+
+def _coords_dict(coords: Tuple[Tuple[str, int], ...]) -> Dict[str, int]:
+    return {axis: value for axis, value in coords}
+
+
+def _normalize_axes(axes: Sequence[str]) -> Tuple[str, ...]:
+    normalized: List[str] = []
+    for entry in axes:
+        axis = str(entry).strip().lower()
+        if not axis:
+            continue
+        if axis not in normalized:
+            normalized.append(axis)
+    return tuple(normalized)
+
+
+@dataclass(frozen=True)
+class FaultSpaceEntry:
+    original: Tuple[int, int, float]
+    src_coords: Tuple[Tuple[str, int], ...]
+    dst_coords: Tuple[Tuple[str, int], ...]
+    affected_axes: Tuple[str, ...]
+    affected_dimensions: Tuple[int, ...]
+
+
+@dataclass(frozen=True)
+class FaultProjectionDetail:
+    original: Tuple[int, int, float]
+    remapped: Tuple[int, int, float]
+    subset_axes: Tuple[str, ...]
+    affected_axes: Tuple[str, ...]
+    affected_dimensions: Tuple[int, ...]
+    src_coords: Tuple[Tuple[str, int], ...]
+    dst_coords: Tuple[Tuple[str, int], ...]
+
+
+@dataclass(frozen=True)
+class FaultProjectionResult:
+    subset_axes: Tuple[str, ...]
+    remapped_links: Tuple[Tuple[int, int, float], ...]
+    entries: Tuple[FaultProjectionDetail, ...]
+    covered_originals: Tuple[Tuple[int, int, float], ...]
+    skipped_originals: Tuple[Tuple[int, int, float], ...]
+
+
+class FaultSpace:
+    """Decoded view of global faulty links for downstream projections."""
+
+    def __init__(
+        self,
+        layout: AxisLayout,
+        faulty_links: Iterable[Tuple[int, int, float]],
+        axis_to_dimension: Optional[Mapping[str, int]] = None,
+    ) -> None:
+        entries: List[FaultSpaceEntry] = []
+        layout_order = layout.axis_order
+        for idx, link in enumerate(faulty_links):
+            if len(link) != 3:
+                raise ValueError(f"faulty_links[{idx}] must be a (src, dst, weight) tuple.")
+            src, dst, weight = link
+            src_coords = decode_axis_coordinates(int(src), layout)
+            dst_coords = decode_axis_coordinates(int(dst), layout)
+            src_dict = _coords_dict(src_coords)
+            dst_dict = _coords_dict(dst_coords)
+            affected_axes = tuple(
+                axis for axis in layout_order if src_dict.get(axis) != dst_dict.get(axis)
+            )
+            if axis_to_dimension is not None:
+                dims = {
+                    axis_to_dimension[axis]
+                    for axis in affected_axes
+                    if axis in axis_to_dimension
+                }
+                affected_dimensions = tuple(sorted(dims))
+            else:
+                affected_dimensions = tuple()
+            entries.append(
+                FaultSpaceEntry(
+                    original=(int(src), int(dst), float(weight)),
+                    src_coords=src_coords,
+                    dst_coords=dst_coords,
+                    affected_axes=affected_axes,
+                    affected_dimensions=affected_dimensions,
+                )
+            )
+        self._entries: Tuple[FaultSpaceEntry, ...] = tuple(entries)
+        self._layout = layout
+
+    @property
+    def layout(self) -> AxisLayout:
+        return self._layout
+
+    @property
+    def entries(self) -> Tuple[FaultSpaceEntry, ...]:
+        return self._entries
+
+    def project(self, subset_axes: Sequence[str]) -> FaultProjectionResult:
+        normalized = _normalize_axes(subset_axes)
+        if not normalized:
+            raise ValueError("subset_axes must not be empty when projecting faults.")
+        subset_axes_set = set(normalized)
+        subset_strides = build_subset_strides(normalized, self._layout.axis_sizes)
+        details: List[FaultProjectionDetail] = []
+        remapped_links: List[Tuple[int, int, float]] = []
+        covered: List[Tuple[int, int, float]] = []
+        unmapped: List[Tuple[int, int, float]] = []
+
+        for entry in self._entries:
+            if not any(axis in subset_axes_set for axis in entry.affected_axes):
+                unmapped.append(entry.original)
+                continue
+            src_dict = _coords_dict(entry.src_coords)
+            dst_dict = _coords_dict(entry.dst_coords)
+            src_linear = _linearize_subset(src_dict, normalized, subset_strides)
+            dst_linear = _linearize_subset(dst_dict, normalized, subset_strides)
+            remapped = (src_linear, dst_linear, entry.original[2])
+            details.append(
+                FaultProjectionDetail(
+                    original=entry.original,
+                    remapped=remapped,
+                    subset_axes=normalized,
+                    affected_axes=entry.affected_axes,
+                    affected_dimensions=entry.affected_dimensions,
+                    src_coords=entry.src_coords,
+                    dst_coords=entry.dst_coords,
+                )
+            )
+            remapped_links.append(remapped)
+            covered.append(entry.original)
+
+        return FaultProjectionResult(
+            subset_axes=normalized,
+            remapped_links=tuple(remapped_links),
+            entries=tuple(details),
+            covered_originals=tuple(covered),
+            skipped_originals=tuple(unmapped),
+        )
+
+
+def _linearize_subset(
+    coords: Mapping[str, int],
+    subset_axes: Sequence[str],
+    subset_strides: Mapping[str, int],
+) -> int:
+    linear = 0
+    for axis in subset_axes:
+        coord = coords.get(axis, 0)
+        stride = subset_strides[axis]
+        linear += coord * stride
+    return linear
diff --git a/astrasim_lib/gmap.py b/astrasim_lib/gmap.py
new file mode 100644
index 00000000..b7d0470b
--- /dev/null
+++ b/astrasim_lib/gmap.py
@@ -0,0 +1,532 @@
+"""First-dimension traffic aggregation and SCOTCH mapping helpers."""
+
+from __future__ import annotations
+
+import math
+import os
+import tempfile
+from dataclasses import dataclass
+from typing import Dict, List, Mapping, Optional, Sequence, Tuple
+
+import numpy as np
+
+from .layout_utils import build_subset_strides
+
+import scotchpy
+from scotchpy.common import proper_int as sct_i
+from scotchpy.maporder import Mapping as ScotchMapping
+from scotchpy import strat as scotch_strat
+
+_EDGE_WEIGHT_SCALE_BYTES = 1024 * 1024 * 1024  # convert bytes -> gigabytes for SCOTCH stability
+
+
+@dataclass(frozen=True)
+class MappingResult:
+    permutation: Tuple[int, ...]
+    graph_path: str
+    map_path: str
+    metrics_path: str
+
+
+class GMapCollector:
+    """Accumulates per-pair byte weights for the first network dimension."""
+
+    def __init__(
+        self,
+        *,
+        topology: str,
+        subset_axes: Tuple[str, ...],
+        axis_sizes: Mapping[str, int],
+        subset_strides: Mapping[str, int],
+        stage_axis_coords: Mapping[int, Mapping[str, int]],
+        output_dir: str,
+        vertex_count: int,
+        include_pipeline: bool,
+        dims: Optional[Tuple[int, int]] = None,
+    ) -> None:
+        if not subset_axes:
+            raise ValueError("optimize_2dmap requires at least one axis in the first network dimension")
+
+        self.topology = str(topology).strip().lower()
+        self.subset_axes = tuple(subset_axes)
+        self.axis_sizes = dict(axis_sizes)
+        self.subset_strides = dict(subset_strides)
+        self.stage_axis_coords = stage_axis_coords
+        self.output_dir = output_dir
+        self.vertex_count = int(vertex_count)
+        self.include_pipeline = include_pipeline
+        self.dims: Optional[Tuple[int, int]] = tuple(int(v) for v in dims) if dims else None
+        self._edge_weights: Dict[Tuple[int, int], float] = {}
+
+        missing = [axis for axis in self.subset_axes if axis not in self.axis_sizes]
+        if missing:
+            raise ValueError(f"Subset axes missing from axis_sizes: {missing}")
+
+        product = 1
+        for axis in self.subset_axes:
+            product *= int(self.axis_sizes[axis])
+        if product != self.vertex_count:
+            raise ValueError(
+                f"Vertex count {self.vertex_count} does not match product of subset axis sizes {product}."
+            )
+
+        self.replication_axes: Tuple[str, ...] = tuple(
+            axis
+            for axis in self.axis_sizes
+            if axis not in self.subset_axes and int(self.axis_sizes[axis]) > 1
+        )
+        self._axis_groups = self._build_axis_groups()
+        self._canonical_signature = self._select_canonical_signature()
+
+    # ------------------------------------------------------------------
+    # Canonical cluster helpers (only one LP/DP slice contributes weights)
+    # ------------------------------------------------------------------
+
+    def _select_canonical_signature(self) -> Optional[Tuple[Tuple[str, int], ...]]:
+        if not self.replication_axes:
+            return None
+        zero_signature = tuple((axis, 0) for axis in self.replication_axes)
+        for coords in self.stage_axis_coords.values():
+            if all(int(coords.get(axis, 0)) == 0 for axis in self.replication_axes):
+                return zero_signature
+        signatures = sorted(
+            {
+                tuple((axis, int(coords.get(axis, 0))) for axis in self.replication_axes)
+                for coords in self.stage_axis_coords.values()
+            }
+        )
+        if not signatures:
+            raise ValueError("Unable to determine canonical cluster for optimize_2dmap.")
+        return signatures[0]
+
+    def _is_canonical_stage(self, stage_id: int) -> bool:
+        if not self.replication_axes:
+            return True
+        coords = self.stage_axis_coords.get(stage_id)
+        if coords is None:
+            raise ValueError(f"Unknown stage id {stage_id} when checking canonical membership.")
+        signature = tuple((axis, int(coords.get(axis, 0))) for axis in self.replication_axes)
+        return signature == self._canonical_signature
+
+    # ------------------------------------------------------------------
+    # Axis indexing helpers
+    # ------------------------------------------------------------------
+
+    def _coord(self, coords: Mapping[str, int], axis: str) -> int:
+        if axis not in coords:
+            raise ValueError(f"Axis '{axis}' missing from coordinate tuple {coords}.")
+        return int(coords[axis])
+
+    def _build_axis_groups(self) -> Dict[str, Dict[Tuple[Tuple[str, int], ...], Tuple[int, ...]]]:
+        axis_groups: Dict[str, Dict[Tuple[Tuple[str, int], ...], List[int]]] = {
+            axis: {} for axis in self.subset_axes
+        }
+        for stage_id, coords in self.stage_axis_coords.items():
+            for axis in self.subset_axes:
+                key = self._axis_key_entries(axis, coords)
+                axis_groups[axis].setdefault(key, []).append(stage_id)
+        normalized: Dict[str, Dict[Tuple[Tuple[str, int], ...], Tuple[int, ...]]] = {}
+        for axis, entries in axis_groups.items():
+            normalized[axis] = {
+                key: tuple(sorted(stages))
+                for key, stages in entries.items()
+            }
+        return normalized
+
+    def local_index_for_stage(self, stage_id: int) -> int:
+        coords = self.stage_axis_coords.get(stage_id)
+        if coords is None:
+            raise ValueError(f"Unknown stage id {stage_id} when computing local index.")
+        local = 0
+        for axis in self.subset_axes:
+            coord = self._coord(coords, axis)
+            stride = self.subset_strides[axis]
+            local += coord * stride
+        return local
+
+    def coords_from_local_index(self, local_index: int) -> Dict[str, int]:
+        coords: Dict[str, int] = {}
+        for axis in reversed(self.subset_axes):
+            stride = self.subset_strides[axis]
+            size = int(self.axis_sizes[axis])
+            if stride == 0:
+                coords[axis] = 0
+                continue
+            coords[axis] = (local_index // stride) % size
+        return coords
+
+    # ------------------------------------------------------------------
+    # Byte aggregation
+    # ------------------------------------------------------------------
+
+    def _per_pair_bytes(self, total_bytes: int, participants: int, double_weight: bool) -> float:
+        if participants <= 1:
+            raise ValueError("participants must be > 1 to compute per-pair bytes")
+        if total_bytes < 0:
+            raise ValueError("total_bytes must be non-negative")
+        per_pair = float(total_bytes) / (participants * (participants - 1))
+        if double_weight:
+            per_pair *= 2.0
+        return per_pair
+
+    def _accumulate_pairs(self, local_indices: Sequence[int], per_pair_bytes: float) -> None:
+        unique = sorted(set(int(idx) for idx in local_indices))
+        if len(unique) <= 1 or per_pair_bytes <= 0:
+            return
+        for i, src in enumerate(unique):
+            for dst in unique[i + 1 :]:
+                key = (src, dst)
+                self._edge_weights[key] = self._edge_weights.get(key, 0.0) + per_pair_bytes
+
+    def _axis_key_entries(self, axis: str, coords: Mapping[str, int]) -> Tuple[Tuple[str, int], ...]:
+        key_entries: List[Tuple[str, int]] = []
+        for other_axis in self.subset_axes:
+            if other_axis == axis:
+                continue
+            key_entries.append((other_axis, self._coord(coords, other_axis)))
+        for repl_axis in self.replication_axes:
+            key_entries.append((repl_axis, self._coord(coords, repl_axis)))
+        return tuple(key_entries)
+
+    def _participants_for_axis(self, axis: str, coords: Mapping[str, int]) -> Optional[Tuple[int, ...]]:
+        axis_map = self._axis_groups.get(axis)
+        if not axis_map:
+            return None
+        key = self._axis_key_entries(axis, coords)
+        return axis_map.get(key)
+
+    def record_collective(
+        self,
+        *,
+        axis: str,
+        stage_id: int,
+        size_bytes: int,
+        participant_count: int,
+        double_weight: bool,
+    ) -> None:
+        if not self._is_canonical_stage(stage_id):
+            return
+        axis_norm = str(axis).strip().lower()
+        coords = self.stage_axis_coords.get(stage_id)
+        if coords is None:
+            raise ValueError(f"Unknown stage id {stage_id} when resolving participants.")
+        axis_candidates = [axis_norm] + [ax for ax in self.subset_axes if ax != axis_norm]
+        participants: Optional[Tuple[int, ...]] = None
+        chosen_axis = axis_norm
+        for candidate in axis_candidates:
+            group = self._participants_for_axis(candidate, coords)
+            if not group:
+                continue
+            if participant_count > 1 and len(group) == participant_count:
+                participants = group
+                chosen_axis = candidate
+                break
+            if participants is None:
+                participants = group
+                chosen_axis = candidate
+        if not participants:
+            raise ValueError(f"Failed to resolve participants for axis '{axis_norm}' at stage {stage_id}.")
+        axis_norm = chosen_axis
+        effective_participants = participant_count if participant_count > 1 else len(participants)
+        effective_participants = max(effective_participants, len(participants))
+        if effective_participants <= 1:
+            return
+        local_indices = [self.local_index_for_stage(stage) for stage in participants]
+        per_pair = self._per_pair_bytes(size_bytes, effective_participants, double_weight)
+        self._accumulate_pairs(local_indices, per_pair)
+
+    def record_pipeline(self, src_stage: int, dst_stage: int, size_bytes: int) -> None:
+        if not self.include_pipeline:
+            return
+        if not (self._is_canonical_stage(src_stage) and self._is_canonical_stage(dst_stage)):
+            return
+        if size_bytes < 0:
+            raise ValueError("Pipeline edge size must be non-negative")
+        local_indices = [self.local_index_for_stage(src_stage), self.local_index_for_stage(dst_stage)]
+        per_pair = self._per_pair_bytes(size_bytes, 2, False)
+        self._accumulate_pairs(local_indices, per_pair)
+
+    # ------------------------------------------------------------------
+    # Graph emission helpers
+    # ------------------------------------------------------------------
+
+    def emit_graph(self) -> str:
+        os.makedirs(self.output_dir, exist_ok=True)
+        path = os.path.join(self.output_dir, "first_dim_comm.grf")
+        adjacency: List[List[Tuple[int, int]]] = [[] for _ in range(self.vertex_count)]
+        for (src, dst), weight in self._edge_weights.items():
+            weight /= _EDGE_WEIGHT_SCALE_BYTES
+            w_int = max(1, int(round(weight)))
+            adjacency[src].append((dst, w_int))
+            adjacency[dst].append((src, w_int))
+        num_arcs = sum(len(neigh) for neigh in adjacency)
+        with open(path, "w", encoding="ascii") as handle:
+            handle.write("0\n")
+            handle.write(f"{self.vertex_count} {num_arcs}\n")
+            handle.write("0 010\n")
+            for neighbors in adjacency:
+                neighbors.sort(key=lambda item: item[0])
+                line_parts: List[str] = [str(len(neighbors))]
+                for neighbor, weight in neighbors:
+                    line_parts.extend([str(weight), str(neighbor)])
+                handle.write(" ".join(line_parts) + "\n")
+        return path
+
+    def build_csr(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        nbrs = [[] for _ in range(self.vertex_count)]
+        weights = [[] for _ in range(self.vertex_count)]
+        for (src, dst), weight in self._edge_weights.items():
+            if src == dst or weight <= 0:
+                continue
+            # print(f"src: {src}, dst: {dst}, weight: {weight}")
+            w_int = _scotch_weight(weight)
+            # new weight
+            # print(f"w_int: {w_int}")
+            nbrs[src].append(dst)
+            weights[src].append(w_int)
+            nbrs[dst].append(src)
+            weights[dst].append(w_int)
+        dtype = sct_i if sct_i is not None else np.int64
+        xadj = np.zeros(self.vertex_count + 1, dtype=dtype)
+        adjncy: List[int] = []
+        adjwgt: List[int] = []
+        for idx in range(self.vertex_count):
+            xadj[idx + 1] = xadj[idx] + len(nbrs[idx])
+            adjncy.extend(nbrs[idx])
+            adjwgt.extend(weights[idx])
+        return xadj, np.array(adjncy, dtype=dtype), np.array(adjwgt, dtype=dtype)
+
+
+def begin_collection(
+    optimize_cfg: Optional[Mapping[str, object]],
+    axis_sizes: Mapping[str, int],
+    stage_axis_coords: Mapping[int, Mapping[str, int]],
+    output_dir: str,
+) -> Optional[GMapCollector]:
+    if not optimize_cfg:
+        return None
+    subset_axes_raw = optimize_cfg.get("parallelisms")
+    if not subset_axes_raw:
+        raise ValueError("optimize_2dmap requires at least one parallelism axis in the first dimension.")
+    subset_axes: Tuple[str, ...] = tuple(
+        str(axis).strip().lower()
+        for axis in subset_axes_raw
+        if str(axis).strip()
+    )
+    subset_strides = build_subset_strides(subset_axes, axis_sizes)
+    dims_raw = optimize_cfg.get("dims")
+    if dims_raw is None and isinstance(optimize_cfg.get("size"), (tuple, list)):
+        dims_raw = optimize_cfg.get("size")
+    dims: Optional[Tuple[int, int]] = None
+    if dims_raw is not None:
+        if not isinstance(dims_raw, (tuple, list)) or len(dims_raw) != 2:
+            raise ValueError("optimize_2dmap dims must be a two-item tuple/list when provided.")
+        try:
+            dims = (int(dims_raw[0]), int(dims_raw[1]))
+        except (TypeError, ValueError) as exc:
+            raise ValueError("optimize_2dmap dims entries must be integers.") from exc
+        if dims[0] <= 0 or dims[1] <= 0:
+            raise ValueError("optimize_2dmap dims entries must be > 0.")
+
+    subset_product = 1
+    for axis in subset_axes:
+        if axis not in axis_sizes:
+            raise ValueError(f"Axis '{axis}' from optimize_2dmap is missing in axis_sizes.")
+        subset_product *= int(axis_sizes[axis])
+
+    if dims:
+        vertex_count = int(dims[0]) * int(dims[1])
+        if vertex_count != subset_product:
+            raise ValueError(
+                f"optimize_2dmap dims product {vertex_count} does not match product of subset axes {subset_product}."
+            )
+    else:
+        vertex_count = subset_product
+    declared_size_raw = optimize_cfg.get("size", vertex_count)
+    if isinstance(declared_size_raw, (tuple, list)):
+        try:
+            declared_size = int(declared_size_raw[0]) * int(declared_size_raw[1])
+        except Exception:
+            declared_size = vertex_count
+    else:
+        declared_size = int(declared_size_raw)
+    if declared_size != vertex_count:
+        raise ValueError(
+            f"optimize_2dmap size {declared_size} does not match product of subset axes {vertex_count}."
+        )
+    topology = str(optimize_cfg.get("topology", "")).strip()
+    if not topology:
+        raise ValueError("optimize_2dmap requires a topology string.")
+    include_pipeline = "lp" in subset_axes
+    return GMapCollector(
+        topology=topology,
+        subset_axes=subset_axes,
+        axis_sizes=axis_sizes,
+        subset_strides=subset_strides,
+        stage_axis_coords=stage_axis_coords,
+        output_dir=output_dir,
+        vertex_count=vertex_count,
+        include_pipeline=include_pipeline,
+        dims=dims,
+    )
+
+
+def finalize_collection(collector: Optional[GMapCollector]) -> Optional[MappingResult]:
+    if collector is None:
+        return None
+    graph_path = collector.emit_graph()
+    xadj, adjncy, adjwgt = collector.build_csr()
+    g = scotchpy.Graph()
+    g.build(
+        vertnbr=int(collector.vertex_count),
+        verttab=xadj,
+        vendtab=None,
+        velotab=None,
+        vlbltab=None,
+        edgenbr=int(len(adjncy)),
+        edgetab=adjncy,
+        edlotab=adjwgt,
+    )
+    dims = _infer_target_dims(collector)
+    arch = _build_architecture(collector.topology, dims)
+    baseline_map = np.arange(collector.vertex_count, dtype=np.int64)
+    before_metrics = _collect_map_metrics(g, arch, baseline_map.copy(), run_map=False)
+    maptab = np.empty(collector.vertex_count, dtype=np.int64)
+    after_metrics = _collect_map_metrics(g, arch, maptab, run_map=True)
+    before_commexp, after_commexp = _commexpan_values(before_metrics, after_metrics)
+    should_apply_mapping = True
+    if before_commexp is not None and after_commexp is not None:
+        should_apply_mapping = after_commexp < before_commexp
+    map_path = os.path.join(collector.output_dir, "first_dim_comm.out")
+    final_map = maptab if should_apply_mapping else baseline_map.copy()
+    _write_mapping_file(map_path, final_map)
+    metrics_path = os.path.join(collector.output_dir, "first_dim_comm.metrics")
+    _write_metrics_report(metrics_path, before_metrics, after_metrics)
+    _maybe_print_commexpan_delta(before_commexp, after_commexp, should_apply_mapping)
+    if _env_truthy("DEEPFLOW_GMAP_ONLY"):
+        status = "applied" if should_apply_mapping else "skipped"
+        print(
+            "[GMapDebug] DEEPFLOW_GMAP_ONLY=1; "
+            f"{status} mapping. Artifacts: map={map_path}, metrics={metrics_path}"
+        )
+        raise SystemExit(0)
+    permutation = tuple(int(val) for val in final_map)
+    return MappingResult(
+        permutation=permutation,
+        graph_path=graph_path,
+        map_path=map_path,
+        metrics_path=metrics_path,
+    )
+
+
+def _infer_target_dims(collector: GMapCollector) -> Tuple[int, int]:
+    if collector.dims:
+        dim_x, dim_y = collector.dims
+        if dim_x <= 0 or dim_y <= 0:
+            raise ValueError("optimize_2dmap requires positive dims entries.")
+        return dim_x, dim_y
+    total = int(collector.vertex_count)
+    if total <= 0:
+        raise ValueError("optimize_2dmap requires a positive vertex count.")
+    side = int(math.isqrt(total))
+    if side * side != total:
+        raise ValueError(
+            f"optimize_2dmap first-dimension size {total} is not a perfect square; Mesh2D/Torus2D mapping now requires square layouts."
+        )
+    return side, side
+
+
+def _build_architecture(topology: str, dims: Tuple[int, int]):
+    dim_x, dim_y = dims
+    topo = topology.lower()
+    arch = scotchpy.Arch()
+    if topo == "mesh2d":
+        arch.mesh2(int(dim_x), int(dim_y))
+        return arch
+    if topo == "torus2d":
+        arch.torus2(int(dim_x), int(dim_y))
+        return arch
+    if topo == "kingmesh2d":
+        raise NotImplementedError("optimize_2dmap does not yet support KingMesh2D mappings.")
+    raise ValueError(f"Unsupported topology '{topology}' for optimize_2dmap.")
+
+
+def _scotch_weight(weight_bytes: float) -> int:
+    if weight_bytes <= 0:
+        return 0
+    scaled = weight_bytes / float(_EDGE_WEIGHT_SCALE_BYTES)
+    if scaled < 1.0:
+        scaled = 1.0
+    return max(1, int(round(scaled)))
+
+
+def _write_mapping_file(path: str, maptab: np.ndarray) -> None:
+    with open(path, "w", encoding="ascii") as handle:
+        handle.write(f"{len(maptab)}\n")
+        for idx, target in enumerate(maptab):
+            handle.write(f"{idx} {int(target)}\n")
+
+
+def _write_metrics_report(path: str, before: str, after: str) -> None:
+    with open(path, "w", encoding="ascii") as handle:
+        handle.write("SCOTCH Mapping Metrics (graphMapView output)\n\n")
+        handle.write("Before:\n")
+        handle.write(before.rstrip() + "\n\n")
+        handle.write("After:\n")
+        handle.write(after.rstrip() + "\n")
+
+
+def _collect_map_metrics(
+    graph: scotchpy.Graph, arch: scotchpy.Arch, parttab: np.ndarray, *, run_map: bool
+) -> str:
+    mapping = ScotchMapping()
+    graph.map_init(mapping, arch, parttab=parttab)
+    try:
+        if run_map:
+            strategy = scotch_strat.Strat()
+            graph.map_compute(mapping, strategy)
+        with tempfile.TemporaryFile(mode="w+") as tmp:
+            graph.map_view(mapping, tmp)
+            tmp.seek(0)
+            return tmp.read()
+    finally:
+        graph.map_exit(mapping)
+
+
+def _maybe_print_commexpan_delta(
+    before_val: Optional[float], after_val: Optional[float], applied: bool
+) -> None:
+    if before_val is None or after_val is None or before_val == 0:
+        return
+    if not (_env_truthy("DEEPFLOW_VISUALIZE_GRAPHS") or _env_truthy("DEEPFLOW_GMAP_ONLY")):
+        return
+    delta_pct = (after_val - before_val) / before_val * 100.0
+    status = "applied" if applied else "skipped"
+    print(
+        "[GMapDebug] CommExpan before "
+        f"{before_val:.6f} after {after_val:.6f} change {delta_pct:+.2f}% ({status})"
+    )
+
+
+def _extract_commexpan(metrics: str) -> Optional[float]:
+    for line in metrics.splitlines():
+        if "CommExpan=" in line:
+            try:
+                prefix = line.split("CommExpan=")[1]
+                value_str = prefix.split("\t", 1)[0]
+                return float(value_str)
+            except (ValueError, IndexError):
+                return None
+    return None
+
+
+def _commexpan_values(before: str, after: str) -> Tuple[Optional[float], Optional[float]]:
+    return _extract_commexpan(before), _extract_commexpan(after)
+
+
+def _env_truthy(name: str) -> bool:
+    value = os.environ.get(name)
+    if value is None:
+        return False
+    normalized = value.strip().lower()
+    return normalized not in {"", "0", "false", "no"}
diff --git a/astrasim_lib/integration.py b/astrasim_lib/integration.py
index d70a8317..22bd1997 100644
--- a/astrasim_lib/integration.py
+++ b/astrasim_lib/integration.py
@@ -6,9 +6,11 @@
 import json
 import os
 import re
+import shutil
 import subprocess
+import tempfile
 import time
-from typing import Any, Dict, Iterable, List, Optional, Tuple
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
 
 from .bootstrap import ensure_chakra_available
 from .config_generation import (
@@ -34,6 +36,8 @@
 
 ensure_chakra_available()
 
+_DEFAULT_ASTRA_CACHE_DIR = os.path.abspath("./astra_cache")
+
 _CACHE_MODES = {
     "NO_CACHE",
     "CACHE_READONLY",
@@ -133,6 +137,19 @@ def _hash_file_bundle(paths: Iterable[str]) -> str:
     return h.hexdigest()
 
 
+def _path_within_dir(path: Optional[str], directory: str) -> bool:
+    """Return True when ``path`` is within ``directory``."""
+
+    if not path:
+        return False
+    try:
+        path_abs = os.path.abspath(path)
+        dir_abs = os.path.abspath(directory)
+        return os.path.commonpath([path_abs, dir_abs]) == dir_abs
+    except (OSError, ValueError):
+        return False
+
+
 def _load_cache(cache_path: str) -> Dict[str, Any]:
     """Load a JSON cache manifest if available, otherwise return an empty dict."""
 
@@ -315,158 +332,216 @@ def run_cache_astrasim(
     manifest_json_path: Optional[str] = None,
     workload_prefix: Optional[str] = None,
     comm_group_json: Optional[str] = None,
+    axes_filter: Optional[Sequence[str]] = None,
+    transform_2d_to_1d: bool = False,
+    files = None,
 ) -> Tuple[List[float], float]:
     """Run AstraSim with caching to avoid recomputation when inputs match."""
 
     if isinstance(hw_obj, str):
         raise TypeError("run_cache_astrasim expects a parsed HWConfig object, not a path")
 
-    override_dir = os.environ.get("ASTRA_CACHE_DIR")
-    if override_dir:
-        astra_config_dir = override_dir
-        cache_path = os.path.join(override_dir, "cache.json")
+    comm_lower = str(comm).lower()
 
     cache_mode = _cache_mode()
     allow_read = cache_mode in {"CACHE_READONLY", "CACHE_READWRITE"}
     allow_write = cache_mode == "CACHE_READWRITE"
 
-    if allow_read or allow_write:
-        ensure_cache_file_exists(cache_path)
-
-    (intra_ib_bps, intra_ll_s), _ = compute_intra_inter_ib_ll_from_hw(hw_obj)
-    if not intra_ib_bps or intra_ib_bps <= 0 or not intra_ll_s or intra_ll_s <= 0:
-        raise ValueError("Invalid intra-node IB/LL computed from HW config")
-    ib_gbps = round((intra_ib_bps) / float(1 << 30), 6)
-    ll_ns = round(intra_ll_s * 1e9, 3)
-    topo = derive_topology_from_hw(hw_obj)
-    colls = _collectives_from_hw(hw_obj, topo)
-    sys_opts_sig = _sys_options_from_hw(hw_obj)
-
-    sig: Dict[str, Any] = {
-        "comm": comm.lower(),
-        "npus": int(npus_count),
-        "size_bytes": int(size_bytes),
-        "topology": topo,
-        "physical_topology": topo,
-        "ib_gbps": ib_gbps,
-        "ll_ns": ll_ns,
-        "collectives": colls,
-        "backend": "analytical",
-    }
-    if sys_opts_sig is not None:
-        sig["sys_options"] = sys_opts_sig
-    if manifest_json_path:
-        sig["multinode"] = True
-    canonical = _canonical_sig(sig)
-
-    files = generate_astrasim_configs_from_hw(hw_obj, out_dir=astra_config_dir, npus_count=npus_count)
-
-    if workload_prefix:
-        pass
-    elif not manifest_json_path:
-        label = size_label(size_bytes)
-        base_dir = os.path.join(astra_config_dir, "workload", comm.lower(), f"{npus_count}npus_{label}")
-        workload_prefix = os.path.join(base_dir, f"{comm.lower()}_{label}")
-        expected = [f"{workload_prefix}.{r}.et" for r in range(npus_count)]
-        if not all(os.path.exists(p) for p in expected):
-            generate_workload_et(comm, npus_count, size_bytes, astra_config_dir=astra_config_dir)
+    temp_cache_dir: Optional[str] = None
+    if cache_mode == "NO_CACHE":
+        temp_cache_dir = tempfile.mkdtemp(prefix="astrasim_cache_")
+        astra_config_dir = temp_cache_dir
+        cache_path = os.path.join(temp_cache_dir, "cache.json")
     else:
-        workload_prefix = os.path.splitext(manifest_json_path)[0]
+        override_dir = os.environ.get("ASTRA_CACHE_DIR")
+        if override_dir:
+            astra_config_dir = override_dir
+            cache_path = os.path.join(override_dir, "cache.json")
+        if allow_read or allow_write:
+            ensure_cache_file_exists(cache_path)
 
-    remote_mem_path = get_remote_memory_path()
+    ephemeral_paths: List[str] = []
 
-    if manifest_json_path:
-        bundle_list: List[str] = [manifest_json_path]
-        for path in (files["system_json"], files["network_yaml"]):
-            if path and os.path.exists(path):
-                bundle_list.append(path)
-        if remote_mem_path and os.path.exists(remote_mem_path):
-            bundle_list.append(remote_mem_path)
-        if comm_group_json and os.path.exists(comm_group_json):
-            bundle_list.append(comm_group_json)
-        cache_key = _hash_file_bundle(bundle_list)
-    else:
-        cache_key = _hash_sig(canonical)
-
-    if allow_read:
-        cache = _load_cache(cache_path)
-        entry = cache.get(cache_key) if cache else None
-        if entry:
-            cached_per = entry.get("per_node_sec", [])
-            cached_max = float(entry.get("max_sec", 0.0))
-            if cached_per and cached_max > 0 and len(cached_per) == int(npus_count) and all((t > 0 for t in cached_per)):
-                if str(comm).lower() == "graph":
+    try:
+        (intra_ib_bps, intra_ll_s), _ = compute_intra_inter_ib_ll_from_hw(hw_obj)
+        if not intra_ib_bps or intra_ib_bps <= 0 or not intra_ll_s or intra_ll_s <= 0:
+            raise ValueError("Invalid intra-node IB/LL computed from HW config")
+        ib_gbps = round((intra_ib_bps) / float(1 << 30), 6)
+        ll_ns = round(intra_ll_s * 1e9, 3)
+        topo = derive_topology_from_hw(hw_obj)
+        colls = _collectives_from_hw(hw_obj, topo)
+        sys_opts_sig = _sys_options_from_hw(hw_obj)
+        axes_filter_tuple: Optional[Tuple[str, ...]] = tuple(axes_filter) if axes_filter else None
+
+        cache_dir_abs = os.path.abspath(astra_config_dir)
+        ephemeral_configs = (
+            files is None
+            and comm_lower != "graph"
+            and cache_dir_abs == _DEFAULT_ASTRA_CACHE_DIR
+        )
+
+        sig: Dict[str, Any] = {
+            "comm": comm_lower,
+            "npus": int(npus_count),
+            "size_bytes": int(size_bytes),
+            "topology": topo,
+            "physical_topology": topo,
+            "ib_gbps": ib_gbps,
+            "ll_ns": ll_ns,
+            "collectives": colls,
+            "backend": "analytical",
+        }
+        if sys_opts_sig is not None:
+            sig["sys_options"] = sys_opts_sig
+        if manifest_json_path:
+            sig["multinode"] = True
+        if axes_filter_tuple is not None:
+            sig["axes_filter"] = axes_filter_tuple
+        if not files:
+            files = generate_astrasim_configs_from_hw(
+                hw_obj,
+                out_dir=astra_config_dir,
+                npus_count=npus_count,
+                axes_filter=axes_filter_tuple,
+                transform_2d_to_1d=transform_2d_to_1d,
+                ephemeral_outputs=ephemeral_configs,
+            )
+            if ephemeral_configs:
+                ephemeral_paths.extend(
+                    [
+                        files.get("system_json"),
+                        files.get("network_yaml"),
+                    ]
+                )
+        topo_list = files.get("topology_list")
+        npus_per_dim = files.get("npus_per_dim")
+        if topo_list:
+            sig["topology"] = list(topo_list)
+            sig["physical_topology"] = list(topo_list)
+        if npus_per_dim:
+            sig["npus_per_dim"] = list(npus_per_dim)
+
+        canonical = _canonical_sig(sig)
+
+        if workload_prefix:
+            pass
+        elif not manifest_json_path:
+            label = size_label(size_bytes)
+            base_dir = os.path.join(astra_config_dir, "workload", comm_lower, f"{npus_count}npus_{label}")
+            workload_prefix = os.path.join(base_dir, f"{comm_lower}_{label}")
+            expected = [f"{workload_prefix}.{r}.et" for r in range(npus_count)]
+            if not all(os.path.exists(p) for p in expected):
+                generate_workload_et(comm, npus_count, size_bytes, astra_config_dir=astra_config_dir)
+        else:
+            workload_prefix = os.path.splitext(manifest_json_path)[0]
+
+        remote_mem_path = get_remote_memory_path()
+
+        if manifest_json_path:
+            bundle_list: List[str] = [manifest_json_path]
+            for path in (files["system_json"], files["network_yaml"]):
+                if path and os.path.exists(path):
+                    bundle_list.append(path)
+            if remote_mem_path and os.path.exists(remote_mem_path):
+                bundle_list.append(remote_mem_path)
+            if comm_group_json and os.path.exists(comm_group_json):
+                bundle_list.append(comm_group_json)
+            cache_key = _hash_file_bundle(bundle_list)
+        else:
+            cache_key = _hash_sig(canonical)
+
+        if allow_read:
+            cache = _load_cache(cache_path)
+            entry = cache.get(cache_key) if cache else None
+            if entry:
+                cached_per = entry.get("per_node_sec", [])
+                cached_max = float(entry.get("max_sec", 0.0))
+                if cached_per and cached_max > 0 and len(cached_per) == int(npus_count) and all((t > 0 for t in cached_per)):
+                    if comm_lower == "graph":
+                        try:
+                            print(
+                                f"[AstraSim] Cache HIT: comm={comm}, npus={npus_count}, size={size_bytes}, total={cached_max:.6f}s"
+                            )
+                        except Exception:
+                            pass
+                    return cached_per, cached_max
+                if allow_write and cache_key in cache:
                     try:
-                        print(
-                            f"[AstraSim] Cache HIT: comm={comm}, npus={npus_count}, size={size_bytes}, total={cached_max:.6f}s"
-                        )
+                        cache.pop(cache_key)
+                        _save_cache(cache_path, cache)
                     except Exception:
                         pass
-                return cached_per, cached_max
-            if allow_write and cache_key in cache:
-                try:
-                    cache.pop(cache_key)
-                    _save_cache(cache_path, cache)
-                except Exception:
-                    pass
-
-    attempts = 0
-    per_node_sec: List[float] = []
-    max_sec: float = 0.0
-    last_outcome_zero = False
-
-    # Capture command details for debugging
-    bin_path = _astrasim_binary_path()
-    debug_cmd = [
-        bin_path,
-        f"--workload-configuration={workload_prefix}",
-        f"--system-configuration={files['system_json']}",
-        f"--network-configuration={files['network_yaml']}",
-    ]
-    if remote_mem_path:
-        debug_cmd.append(f"--remote-memory-configuration={remote_mem_path}")
-    if comm_group_json and os.path.exists(comm_group_json):
-        debug_cmd.append(f"--comm-group-configuration={comm_group_json}")
-
-    while attempts < 5:
-        attempts += 1
-        per_node_sec, max_sec = run_astrasim_analytical(
-            workload_prefix=workload_prefix,
-            system_json=files["system_json"],
-            network_yaml=files["network_yaml"],
-            remote_memory_json=remote_mem_path,
-            comm_group_json=comm_group_json,
-        )
-        if per_node_sec and max_sec > 0 and len(per_node_sec) == int(npus_count) and all((t > 0 for t in per_node_sec)):
-            last_outcome_zero = False
-            break
-        last_outcome_zero = True
-        time.sleep(0.5)
-    if last_outcome_zero:
-        cmd_str = " ".join(debug_cmd)
-        raise RuntimeError(
-            f"AstraSim returned zero time for {comm} size={size_bytes} npus={npus_count} after 5 retries.\n"
-            f"Full AstraSim command:\n{cmd_str}"
-        )
-
-    cache_entry = {
-        "signature": sig,
-        "canonical": canonical,
-        "per_node_sec": per_node_sec,
-        "max_sec": max_sec,
-        "workload_prefix": workload_prefix,
-        "system_json": files["system_json"],
-        "network_yaml": files["network_yaml"],
-    }
-    if manifest_json_path:
-        cache_entry["manifest_path"] = manifest_json_path
 
-    if allow_write:
-        cache = _load_cache(cache_path)
-        cache[cache_key] = cache_entry
-        _save_cache(cache_path, cache)
-
-    return per_node_sec, max_sec
+        attempts = 0
+        per_node_sec: List[float] = []
+        max_sec: float = 0.0
+        last_outcome_zero = False
+
+        # Capture command details for debugging
+        bin_path = _astrasim_binary_path()
+        debug_cmd = [
+            bin_path,
+            f"--workload-configuration={workload_prefix}",
+            f"--system-configuration={files['system_json']}",
+            f"--network-configuration={files['network_yaml']}",
+        ]
+        if remote_mem_path:
+            debug_cmd.append(f"--remote-memory-configuration={remote_mem_path}")
+        if comm_group_json and os.path.exists(comm_group_json):
+            debug_cmd.append(f"--comm-group-configuration={comm_group_json}")
+        cmd_str = " ".join(debug_cmd)
+        # print(f"Full AstraSim command:\n{cmd_str}")
+
+        while attempts < 5:
+            attempts += 1
+            per_node_sec, max_sec = run_astrasim_analytical(
+                workload_prefix=workload_prefix,
+                system_json=files["system_json"],
+                network_yaml=files["network_yaml"],
+                remote_memory_json=remote_mem_path,
+                comm_group_json=comm_group_json,
+            )
+            if per_node_sec and max_sec > 0 and len(per_node_sec) == int(npus_count) and all((t > 0 for t in per_node_sec)):
+                last_outcome_zero = False
+                break
+            last_outcome_zero = True
+            time.sleep(0.1)
+        if last_outcome_zero:
+            raise RuntimeError(
+                f"AstraSim returned zero time for {comm} size={size_bytes} npus={npus_count} after 5 retries.\n"
+                f"Full AstraSim command:\n{cmd_str}"
+            )
+
+        cache_entry = {
+            "signature": sig,
+            "canonical": canonical,
+            "per_node_sec": per_node_sec,
+            "max_sec": max_sec,
+            "workload_prefix": workload_prefix,
+            "system_json": files["system_json"],
+            "network_yaml": files["network_yaml"],
+        }
+        if manifest_json_path:
+            cache_entry["manifest_path"] = manifest_json_path
+
+        if allow_write:
+            cache = _load_cache(cache_path)
+            cache[cache_key] = cache_entry
+            _save_cache(cache_path, cache)
+
+        return per_node_sec, max_sec
+    finally:
+        if ephemeral_paths:
+            for cleanup_path in {p for p in ephemeral_paths if p}:
+                if _path_within_dir(cleanup_path, astra_config_dir):
+                    try:
+                        os.remove(cleanup_path)
+                    except FileNotFoundError:
+                        pass
+                    except OSError:
+                        pass
+        if temp_cache_dir is not None:
+            shutil.rmtree(temp_cache_dir, ignore_errors=True)
 
 
 __all__ = [
diff --git a/astrasim_lib/layout_utils.py b/astrasim_lib/layout_utils.py
new file mode 100644
index 00000000..7cea167a
--- /dev/null
+++ b/astrasim_lib/layout_utils.py
@@ -0,0 +1,154 @@
+"""Shared helpers for reasoning about DeepFlow axis layouts."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Mapping, Sequence, Tuple, List
+
+
+def _normalize_axis_name(name: str) -> str:
+    if not isinstance(name, str):
+        raise TypeError(f"Axis name must be a string, got {type(name)!r}")
+    normalized = name.strip().lower()
+    if not normalized:
+        raise ValueError("Axis name cannot be empty or whitespace.")
+    return normalized
+
+
+@dataclass(frozen=True)
+class AxisLayout:
+    """Immutable view of the flattened axis layout that defines hw_id numbering."""
+
+    axis_order: Tuple[str, ...]
+    axis_sizes: Dict[str, int]
+    axis_strides: Dict[str, int]
+    total_ranks: int
+
+    def __post_init__(self) -> None:
+        if any(size < 1 for size in self.axis_sizes.values()):
+            raise ValueError("Axis sizes must all be >= 1.")
+        for axis in self.axis_order:
+            if axis not in self.axis_sizes:
+                raise ValueError(f"Axis '{axis}' is missing a declared size.")
+            if axis not in self.axis_strides:
+                raise ValueError(f"Axis '{axis}' is missing a declared stride.")
+
+
+def axis_layout_from_descriptor(descriptor: Mapping[str, Any]) -> AxisLayout:
+    """Create an :class:`AxisLayout` from `_build_rank_layout_descriptor` output."""
+    if descriptor is None:
+        raise ValueError("Axis layout descriptor must not be None.")
+
+    axis_order_raw = descriptor.get("axis_order", [])
+    if not isinstance(axis_order_raw, Sequence):
+        raise TypeError("Axis layout descriptor axis_order must be a sequence.")
+    axis_order: Tuple[str, ...] = tuple(_normalize_axis_name(entry) for entry in axis_order_raw)
+    if not axis_order:
+        raise ValueError("Axis layout descriptor axis_order is empty.")
+
+    sizes_raw = descriptor.get("axis_sizes", {})
+    if not isinstance(sizes_raw, Mapping):
+        raise TypeError("Axis layout descriptor axis_sizes must be a mapping.")
+    axis_sizes: Dict[str, int] = {}
+    for key, value in sizes_raw.items():
+        axis = _normalize_axis_name(str(key))
+        try:
+            size = int(value)
+        except (TypeError, ValueError) as exc:
+            raise ValueError(f"Axis '{axis}' size must be an integer.") from exc
+        if size < 1:
+            raise ValueError(f"Axis '{axis}' size must be >= 1.")
+        axis_sizes[axis] = size
+
+    strides_raw = descriptor.get("axis_strides", {})
+    if strides_raw and not isinstance(strides_raw, Mapping):
+        raise TypeError("Axis layout descriptor axis_strides must be a mapping when provided.")
+
+    axis_strides: Dict[str, int] = {}
+    if strides_raw:
+        for key, value in strides_raw.items():
+            axis = _normalize_axis_name(str(key))
+            try:
+                stride = int(value)
+            except (TypeError, ValueError) as exc:
+                raise ValueError(f"Axis '{axis}' stride must be an integer.") from exc
+            if stride < 0:
+                raise ValueError(f"Axis '{axis}' stride must be >= 0.")
+            axis_strides[axis] = stride
+
+    if not axis_strides:
+        span = 1
+        for axis in axis_order:
+            axis_strides[axis] = span
+            span *= axis_sizes.get(axis, 1)
+
+    total = 1
+    for axis in axis_order:
+        total *= axis_sizes[axis]
+
+    return AxisLayout(
+        axis_order=axis_order,
+        axis_sizes=axis_sizes,
+        axis_strides=axis_strides,
+        total_ranks=total,
+    )
+
+
+def decode_axis_coordinates(hw_id: int, layout: AxisLayout) -> Tuple[Tuple[str, int], ...]:
+    """Return ordered axis coordinates for ``hw_id``."""
+    if hw_id < 0:
+        raise ValueError(f"Hardware id must be >= 0, got {hw_id}.")
+    if hw_id >= layout.total_ranks:
+        raise ValueError(
+            f"Hardware id {hw_id} exceeds layout span ({layout.total_ranks} ranks)."
+        )
+    coords: List[Tuple[str, int]] = []
+    for axis in layout.axis_order:
+        stride = layout.axis_strides[axis]
+        size = layout.axis_sizes[axis]
+        if size == 1:
+            coords.append((axis, 0))
+            continue
+        coord = (hw_id // stride) % size
+        coords.append((axis, coord))
+    return tuple(coords)
+
+
+def build_subset_strides(
+    subset_axes: Sequence[str],
+    axis_sizes: Mapping[str, int],
+) -> Dict[str, int]:
+    """Compute strides for a subset axis ordering."""
+    strides: Dict[str, int] = {}
+    span = 1
+    for raw_axis in subset_axes:
+        axis = _normalize_axis_name(raw_axis)
+        if axis not in axis_sizes:
+            raise ValueError(f"Axis '{axis}' is not defined in the layout axis_sizes.")
+        size = axis_sizes[axis]
+        if size < 1:
+            raise ValueError(f"Axis '{axis}' size must be >= 1.")
+        strides[axis] = span
+        span *= size
+    return strides
+
+
+def derive_axes_filter(
+    axis_order: Sequence[str],
+    axis_sizes: Mapping[str, int],
+    dp_count: int,
+) -> List[str]:
+    """Return the axes filter used to select network dimensions for AstraSim."""
+    axes: List[str] = []
+    for entry in axis_order:
+        axis = _normalize_axis_name(entry)
+        try:
+            size = max(1, int(axis_sizes.get(axis, 1)))
+        except Exception:
+            size = 1
+        if size > 1 and axis not in axes:
+            axes.append(axis)
+    if dp_count > 1:
+        axes = [axis for axis in axes if axis != "dp"]
+        axes.append("dp")
+    return axes
diff --git a/config.py b/config.py
index 716d8064..27c64aa2 100644
--- a/config.py
+++ b/config.py
@@ -1,10 +1,10 @@
 from dataclasses import dataclass, field
-import ruamel as _ruamel
-import ruamel.yaml as _yaml
-from ruamel.yaml import YAMLError as _YAMLError
 import math
 from collections import namedtuple as _namedtuple
-from typing import Optional
+from typing import Dict, List, Optional, Sequence, Tuple
+
+import yaml as _yaml
+from yaml import YAMLError as _YAMLError
 
 
 _PRECISION_DTYPE_BYTES = {
@@ -234,50 +234,6 @@ def from_dict(cls, sram_config_dict):
         )
 
 
-@dataclass
-class SubNetworkConfig:
-    latency: float
-    nominal_freq: float
-    nominal_voltage: float
-    nominal_energy_per_link: float
-    nominal_area_per_link: float
-    threshold_voltage: float
-    margin_voltage: float
-    num_links_per_mm: int
-    util: float
-    operating_frequency: float = None  # Optional: overrides voltage-scaled frequency
-    bandwidth: float = None  # Optional: master parameter that directly sets throughput
-
-    @classmethod
-    def from_dict(cls, config_dict):
-        return cls(
-            latency=config_dict["latency"],
-            nominal_freq=config_dict.get("nominal_frequency", 0.1),
-            nominal_voltage=config_dict.get("nominal_voltage", 0.1),
-            nominal_energy_per_link=config_dict.get("nominal_energy_per_link", 0.1),
-            nominal_area_per_link=config_dict.get("nominal_area_per_link", 0.1),
-            threshold_voltage=config_dict.get("threshold_voltage", 0.1),
-            margin_voltage=config_dict.get("margin_voltage", 0.1),
-            num_links_per_mm=config_dict.get("num_links_per_mm", 1),
-            util=config_dict["util"],
-            operating_frequency=config_dict.get("operating_frequency", None),
-            bandwidth=parse_bandwidth_string(config_dict.get("bandwidth", None)),
-        )
-
-
-@dataclass
-class NetworkConfig:
-    intra_node: SubNetworkConfig
-    inter_node: SubNetworkConfig
-
-    @classmethod
-    def from_dict(cls, d):
-        return cls(
-            intra_node=SubNetworkConfig.from_dict(d["intra_node"]),
-            inter_node=SubNetworkConfig.from_dict(d["inter_node"]),
-        )
-
-
 @dataclass
 class TechConfig:
     core: CoreConfig
@@ -285,7 +241,6 @@ class TechConfig:
     SRAML2: SRAMConfig
     SRAML1: SRAMConfig
     SRAMR: SRAMConfig
-    network: NetworkConfig
 
     @classmethod
     def from_dict(cls, tech_config_dict):
@@ -295,7 +250,6 @@ def from_dict(cls, tech_config_dict):
             SRAML2=SRAMConfig.from_dict(tech_config_dict["SRAM-L2"]),
             SRAML1=SRAMConfig.from_dict(tech_config_dict["SRAM-L1"]),
             SRAMR=SRAMConfig.from_dict(tech_config_dict["SRAM-R"]),
-            network=NetworkConfig.from_dict(tech_config_dict["network"]),
         )
 
 
@@ -388,69 +342,535 @@ def from_dict(cls, network_power_config_dict):
         )
 
 
-@dataclass
-class SystemHierarchyConfig:
-    num_nodes_per_wafer: int
-    num_wafers: int
-    num_workers: int
-    inter_derate: float
-    intra_derate: float
-    kp1_inter: float
-    kp2_inter: float
-    dp_inter: float
-    lp_inter: float
-    tp_inter: float
-    par2cross: dict
+@dataclass(frozen=True)
+class NetworkDimensionLayout:
+    id: str
+    label: str
+    size: int
+    topology_type: str
+    bandwidth: float
+    util: float
+    latency: float
+    size_2d: Optional[Tuple[int, int]] = None
+    collective_override: Dict[str, str] = field(default_factory=dict)
+    parallelisms: Tuple[str, ...] = field(default_factory=tuple)
+    energy_per_bit: float = 0.0
+    optimize_2dmap: bool = False
 
     @classmethod
-    def from_dict(cls, system_config_dict):
+    def from_raw(
+        cls,
+        raw: dict,
+        *,
+        parallelism_params: Dict[str, object],
+        index: int,
+    ) -> "NetworkDimensionLayout":
+        if not isinstance(raw, dict):
+            raise TypeError("each network dimension must be a mapping")
+
+        raw_id = raw.get("id")
+        dim_id = str(raw_id) if raw_id is not None else f"dim{index}"
+        label = str(raw.get("label", dim_id))
+
+        if "size" not in raw:
+            raise ValueError(f"network dimension '{label}' is missing required field 'size'")
+        size_raw = raw["size"]
+        size_mode: str
+        size_value: Optional[int] = None
+        size_2d: Optional[Tuple[int, int]] = None
+        tuple_entries: Optional[Tuple[object, object]] = None
+        if isinstance(size_raw, str):
+            normalized_size = size_raw.strip()
+            if normalized_size.startswith("(") and normalized_size.endswith(")"):
+                tuple_entries = tuple(
+                    part.strip() for part in normalized_size[1:-1].split(",")
+                )  # type: ignore[assignment]
+                size_mode = "tuple"
+            elif normalized_size.lower() == "auto":
+                size_mode = "auto_scalar"
+            else:
+                try:
+                    size_value = int(size_raw)
+                except (TypeError, ValueError) as exc:
+                    raise ValueError(f"network dimension '{label}' size must be an integer or tuple") from exc
+                if size_value < 1:
+                    raise ValueError(f"network dimension '{label}' size must be >= 1")
+                size_mode = "scalar"
+        elif isinstance(size_raw, (list, tuple)):
+            if len(size_raw) != 2:
+                raise ValueError(f"network dimension '{label}' 2D size must have exactly two entries")
+            tuple_entries = tuple(size_raw)  # type: ignore[assignment]
+            size_mode = "tuple"
+        else:
+            size_mode = "auto_scalar"
+
+        topo_dict = raw.get("topology")
+        if not isinstance(topo_dict, dict):
+            raise ValueError(f"network dimension '{label}' requires a 'topology' mapping")
+        if "type" not in topo_dict:
+            raise ValueError(f"network dimension '{label}' topology missing required 'type'")
+        topo_type = str(topo_dict["type"])
+
+        if "bandwidth" not in topo_dict:
+            raise ValueError(f"network dimension '{label}' topology missing required 'bandwidth'")
+        bandwidth = parse_bandwidth_string(topo_dict["bandwidth"])
+
+        try:
+            util = float(topo_dict.get("util", 1.0))
+        except (TypeError, ValueError) as exc:
+            raise ValueError(
+                f"network dimension '{label}' topology util must be numeric"
+            ) from exc
+        if util <= 0:
+            raise ValueError(f"network dimension '{label}' topology util must be > 0")
+
+        if "energy_per_bit" not in topo_dict:
+            raise ValueError(
+                f"network dimension '{label}' topology missing required 'energy_per_bit'"
+            )
+        try:
+            energy_per_bit = float(topo_dict["energy_per_bit"])
+        except (TypeError, ValueError) as exc:
+            raise ValueError(
+                f"network dimension '{label}' energy_per_bit must be numeric"
+            ) from exc
+        if energy_per_bit < 0:
+            raise ValueError(
+                f"network dimension '{label}' energy_per_bit must be >= 0"
+            )
+
+        if "latency" not in topo_dict:
+            raise ValueError(f"network dimension '{label}' topology missing required 'latency'")
+        try:
+            latency = float(topo_dict["latency"])
+        except (TypeError, ValueError) as exc:
+            raise ValueError(f"network dimension '{label}' latency must be numeric") from exc
+
+        raw_optimize = topo_dict.get("optimize_2dmap", False)
+        if raw_optimize not in (True, False):
+            raise ValueError(
+                f"network dimension '{label}' topology optimize_2dmap must be a boolean when provided"
+            )
+        optimize_2dmap = bool(raw_optimize)
+        if optimize_2dmap:
+            normalized_topo = topo_type.lower()
+            if normalized_topo not in {"mesh2d", "torus2d", "kingmesh2d"}:
+                raise ValueError(
+                    f"network dimension '{label}' sets optimize_2dmap but topology type '{topo_type}'"
+                    " is not Mesh2D/Torus2D/KingMesh2D"
+                )
+
+        collectives_raw = raw.get("collective_override")
+        if collectives_raw is None and "collectives" in raw:
+            collectives_raw = raw.get("collectives")
+        if not collectives_raw:
+            collectives_raw = {}
+        if not isinstance(collectives_raw, dict):
+            raise ValueError(f"network dimension '{label}' collective_override must be a mapping if provided")
+        collective_override = {str(k): str(v) for k, v in collectives_raw.items()}
+
+        parallelisms_raw = raw.get("parallelisms", [])
+        if parallelisms_raw is None:
+            parallelisms_raw = []
+        if not isinstance(parallelisms_raw, Sequence) or isinstance(parallelisms_raw, (str, bytes)):
+            raise ValueError(
+                f"network dimension '{label}' parallelisms must be a sequence of names"
+            )
+
+        normalized_parallelisms: List[str] = []
+        alias_map: Dict[str, str] = {}
+        for entry in parallelisms_raw:
+            name = str(entry).strip()
+            if not name:
+                raise ValueError(f"network dimension '{label}' has an empty parallelism name")
+            normalized = name.lower()
+            normalized_parallelisms.append(normalized)
+            alias_map[normalized] = name
+
+        computed_product: Optional[int] = None
+        auto_product: Optional[int] = None
+
+        if size_mode in {"auto_scalar", "tuple"}:
+            auto_product = _compute_dimension_parallelism_product(
+                dimension_label=label,
+                normalized_names=tuple(normalized_parallelisms),
+                alias_map=alias_map,
+                parallelism_params=parallelism_params,
+            )
+            if auto_product < 1:
+                raise ValueError(f"network dimension '{label}' inferred size must be >= 1")
+
+        if size_mode == "auto_scalar":
+            size_value = auto_product
+            computed_product = auto_product
+        elif size_mode == "tuple":
+            assert tuple_entries is not None
+            resolved: List[Optional[int]] = [None, None]
+            auto_entries = 0
+            for idx, entry in enumerate(tuple_entries):
+                if isinstance(entry, str):
+                    normalized = entry.strip().lower()
+                    if normalized == "auto":
+                        auto_entries += 1
+                        if auto_entries > 1:
+                            raise ValueError(
+                                f"network dimension '{label}' 2D size may include at most one 'auto' entry"
+                            )
+                        resolved[idx] = None
+                        continue
+                    if normalized not in parallelism_params:
+                        alias = alias_map.get(normalized, normalized)
+                        raise ValueError(
+                            f"network dimension '{label}' 2D size references parallelism '{alias}' "
+                            "which is not defined in parallelism"
+                        )
+                    try:
+                        factor = int(parallelism_params[normalized])
+                    except (TypeError, ValueError) as exc:
+                        alias = alias_map.get(normalized, normalized)
+                        raise ValueError(
+                            f"network dimension '{label}' 2D size parallelism '{alias}' must be an integer"
+                        ) from exc
+                    if factor < 1:
+                        alias = alias_map.get(normalized, normalized)
+                        raise ValueError(
+                            f"network dimension '{label}' 2D size parallelism '{alias}' must be >= 1"
+                        )
+                    resolved[idx] = factor
+                else:
+                    try:
+                        factor = int(entry)
+                    except (TypeError, ValueError) as exc:
+                        raise ValueError(
+                            f"network dimension '{label}' 2D size entries must be integers, parallelism names, or 'auto'"
+                        ) from exc
+                    if factor < 1:
+                        raise ValueError(f"network dimension '{label}' 2D size entries must be >= 1")
+                    resolved[idx] = factor
+
+            known_product = 1
+            for val in resolved:
+                if val:
+                    known_product *= val
+
+            if auto_entries:
+                if auto_product is None:
+                    auto_product = _compute_dimension_parallelism_product(
+                        dimension_label=label,
+                        normalized_names=tuple(normalized_parallelisms),
+                        alias_map=alias_map,
+                        parallelism_params=parallelism_params,
+                    )
+                if known_product == 0:
+                    raise ValueError(f"network dimension '{label}' 2D size known entries must be > 0")
+                if auto_product % known_product != 0:
+                    raise ValueError(
+                        f"network dimension '{label}' 2D size mismatch: auto product {auto_product} "
+                        f"is not divisible by provided factors {tuple_entries}"
+                    )
+                auto_value = auto_product // known_product
+                if auto_value < 1:
+                    raise ValueError(
+                        f"network dimension '{label}' 2D size auto-resolved entry must be >= 1 (got {auto_value})"
+                    )
+                for idx in range(2):
+                    if resolved[idx] is None:
+                        resolved[idx] = auto_value
+                        break
+            else:
+                if auto_product is None:
+                    auto_product = _compute_dimension_parallelism_product(
+                        dimension_label=label,
+                        normalized_names=tuple(normalized_parallelisms),
+                        alias_map=alias_map,
+                        parallelism_params=parallelism_params,
+                    )
+                if known_product != auto_product:
+                    raise ValueError(
+                        f"network dimension '{label}' 2D size mismatch: provided shape {tuple_entries} "
+                        f"product {known_product} does not match parallelism product {auto_product}"
+                    )
+
+            if resolved[0] is None or resolved[1] is None:
+                raise ValueError(f"network dimension '{label}' 2D size could not be fully resolved")
+
+            size_2d = (int(resolved[0]), int(resolved[1]))
+            size_value = int(size_2d[0]) * int(size_2d[1])
+            computed_product = auto_product
+        else:
+            computed_product = None
+
+        if size_value is None:
+            raise ValueError(f"network dimension '{label}' size could not be resolved")
+
+        _validate_dimension_parallelisms(
+            dimension_label=label,
+            dimension_size=int(size_value) if size_value is not None else 0,
+            normalized_names=tuple(normalized_parallelisms),
+            alias_map=alias_map,
+            parallelism_params=parallelism_params,
+            expected_product=computed_product,
+        )
+
         return cls(
-            num_nodes_per_wafer=system_config_dict["num_devices_per_node"],
-            num_wafers=system_config_dict["num_nodes"],
-            num_workers=int(
-                system_config_dict["num_nodes"]
-                * system_config_dict["num_devices_per_node"]
-            ),
-            inter_derate=system_config_dict.get("inter_derate", 1.0),
-            intra_derate=system_config_dict.get("intra_derate", 1.0),
-            kp1_inter=system_config_dict["kp1_inter"],
-            kp2_inter=system_config_dict["kp2_inter"],
-            dp_inter=system_config_dict["dp_inter"],
-            lp_inter=system_config_dict["lp_inter"],
-            tp_inter=system_config_dict["tp_inter"],
-            par2cross={
-                "kp1": system_config_dict["kp1_inter"],
-                "kp2": system_config_dict["kp2_inter"],
-                "dp": system_config_dict["dp_inter"],
-                "lp": system_config_dict["lp_inter"],
-                "tp": system_config_dict["tp_inter"],
-            },
+            id=dim_id,
+            label=label,
+            size=int(size_value) if size_value is not None else 0,
+            size_2d=size_2d,
+            topology_type=topo_type,
+            bandwidth=bandwidth,
+            util=util,
+            latency=latency,
+            collective_override=collective_override,
+            parallelisms=tuple(normalized_parallelisms),
+            energy_per_bit=energy_per_bit,
+            optimize_2dmap=optimize_2dmap,
         )
 
+    @property
+    def effective_bandwidth(self) -> float:
+        return float(self.bandwidth) * float(self.util)
+
+
+def _validate_dimension_parallelisms(
+    *,
+    dimension_label: str,
+    dimension_size: int,
+    normalized_names: Tuple[str, ...],
+    alias_map: Dict[str, str],
+    parallelism_params: Dict[str, object],
+    expected_product: Optional[int] = None,
+) -> None:
+    if not normalized_names:
+        return
+
+    product = expected_product
+    if product is None:
+        product = _compute_dimension_parallelism_product(
+            dimension_label=dimension_label,
+            normalized_names=normalized_names,
+            alias_map=alias_map,
+            parallelism_params=parallelism_params,
+        )
 
-@dataclass
-class TopologyConfig:
-    topology: str = None
+    if product != dimension_size:
+        readable = [alias_map.get(name, name) for name in normalized_names]
+        raise ValueError(
+            f"Network dimension '{dimension_label}' size mismatch: declared size {dimension_size} "
+            f"but parallelism factors ({readable}) imply {product}"
+        )
 
-    @classmethod
-    def from_dict(cls, d):
-        if d == "hybrid":
-            NotImplemented()
-        else:
-            return cls(topology=d)
 
+def _compute_dimension_parallelism_product(
+    *,
+    dimension_label: str,
+    normalized_names: Tuple[str, ...],
+    alias_map: Dict[str, str],
+    parallelism_params: Dict[str, object],
+) -> int:
+    product = 1
+    for name in normalized_names:
+        if name not in parallelism_params:
+            alias = alias_map.get(name, name)
+            raise ValueError(
+                f"network dimension '{dimension_label}' references parallelism '{alias}' "
+                "which is not defined in parallelism"
+            )
+        value = parallelism_params[name]
+        if value in (None, False):
+            alias = alias_map.get(name, name)
+            raise ValueError(
+                f"network dimension '{dimension_label}' parallelism '{alias}' must have a "
+                "positive parallelism factor"
+            )
+        try:
+            factor = int(value)
+        except (TypeError, ValueError) as exc:
+            alias = alias_map.get(name, name)
+            raise ValueError(
+                f"parallelism.{alias} must be an integer to compute network dimension sizes"
+            ) from exc
+        if factor < 1:
+            alias = alias_map.get(name, name)
+            raise ValueError(
+                f"parallelism.{alias} must be >= 1 to compute network dimension sizes"
+            )
+        product *= factor
 
-@dataclass
-class NetworkTopologyConfig:
-    inter: TopologyConfig
-    intra: TopologyConfig
+    return product
 
-    @classmethod
-    def from_dict(cls, d):
-        return cls(
-            inter=TopologyConfig.from_dict(d["inter_node"]),
-            intra=TopologyConfig.from_dict(d["intra_node"]),
+
+def _parse_network_layout(
+    network_spec,
+    parallelism_params: Dict[str, object],
+) -> Tuple[Tuple[NetworkDimensionLayout, ...], Tuple[Tuple[int, int, float], ...], "NetworkOverlapConfig"]:
+    if network_spec is None:
+        raise ValueError("network section must be specified and include overlap settings")
+
+    if not isinstance(network_spec, dict):
+        raise ValueError("network must be provided as a mapping to supply overlap settings")
+
+    faulty_links: Tuple[Tuple[int, int, float], ...] = _parse_faulty_links("network", network_spec.get("faulty_links", []))
+    overlap_config = _parse_network_overlap(network_spec.get("overlap"))
+    dimensions_spec = network_spec.get("dimensions")
+    if dimensions_spec is None:
+        raise ValueError("network.dimensions must be specified when network is a mapping")
+
+    if not isinstance(dimensions_spec, Sequence) or isinstance(dimensions_spec, (str, bytes)):
+        raise ValueError("network.dimensions must be a sequence of dimension mappings")
+
+    dimensions: List[NetworkDimensionLayout] = []
+    for index, entry in enumerate(dimensions_spec):
+        dimensions.append(
+            NetworkDimensionLayout.from_raw(
+                entry,
+                parallelism_params=parallelism_params,
+                index=index,
+            )
         )
+    return tuple(dimensions), faulty_links, overlap_config
+
+
+@dataclass(frozen=True)
+class NetworkOverlapConfig:
+    tp_overlap: float
+    tp_sp_overlap: float
+    cp_overlap: float
+
+
+@dataclass(frozen=True)
+class NetworkLayoutConfig:
+    dimensions: Tuple[NetworkDimensionLayout, ...]
+    faulty_links: Tuple[Tuple[int, int, float], ...] = field(default_factory=tuple)
+    parallelism_map: Dict[str, NetworkDimensionLayout] = field(default_factory=dict)
+    overlap_config: "NetworkOverlapConfig" = None
+
+    def primary_dimension(self) -> Optional[NetworkDimensionLayout]:
+        return self.dimensions[0] if self.dimensions else None
+
+    def dimension_for_parallelism(self, name: str) -> Optional[NetworkDimensionLayout]:
+        normalized = str(name).strip().lower()
+        if normalized in self.parallelism_map:
+            return self.parallelism_map[normalized]
+        return self.primary_dimension()
+
+    def link_for_parallelism(self, name: str) -> Tuple[float, float]:
+        dim = self.dimension_for_parallelism(name)
+        if dim is None:
+            return 0.0, 0.0
+        return dim.effective_bandwidth, dim.latency
+
+
+def _parse_faulty_links(owner_label: str, faulty_links_raw) -> Tuple[Tuple[int, int, float], ...]:
+    entries: List[Tuple[int, int, float]] = []
+    if not faulty_links_raw:
+        return tuple()
+    if not isinstance(faulty_links_raw, Sequence) or isinstance(faulty_links_raw, (str, bytes)):
+        raise ValueError(
+            f"{owner_label} faulty_links must be a sequence of [src, dst, weight] entries"
+        )
+    for idx, entry in enumerate(faulty_links_raw):
+        if not isinstance(entry, Sequence) or isinstance(entry, (str, bytes)) or len(entry) != 3:
+            raise ValueError(
+                f"{owner_label} faulty_links[{idx}] must be a three-item sequence [src, dst, weight]"
+            )
+        src_raw, dst_raw, weight_raw = entry
+        try:
+            src = int(src_raw)
+        except (TypeError, ValueError) as exc:
+            raise ValueError(
+                f"{owner_label} faulty_links[{idx}][0] must be an integer endpoint"
+            ) from exc
+        try:
+            dst = int(dst_raw)
+        except (TypeError, ValueError) as exc:
+            raise ValueError(
+                f"{owner_label} faulty_links[{idx}][1] must be an integer endpoint"
+            ) from exc
+        if src < 0 or dst < 0:
+            raise ValueError(
+                f"{owner_label} faulty_links[{idx}] endpoints must be >= 0"
+            )
+        try:
+            weight = float(weight_raw)
+        except (TypeError, ValueError) as exc:
+            raise ValueError(
+                f"{owner_label} faulty_links[{idx}][2] must be a numeric reliability weight"
+            ) from exc
+        if weight < 0.0 or weight > 1.0:
+            raise ValueError(
+                f"{owner_label} faulty_links[{idx}] weight must be between 0.0 and 1.0"
+            )
+        entries.append((src, dst, weight))
+    return tuple(entries)
+
+def _parse_network_overlap(overlap_raw) -> "NetworkOverlapConfig":
+    if not isinstance(overlap_raw, dict):
+        raise ValueError("network.overlap must be a mapping with tp_overlap, tp_sp_overlap, and cp_overlap")
+    required_fields = ("tp_overlap", "tp_sp_overlap", "cp_overlap")
+    values = {}
+    for field in required_fields:
+        if field not in overlap_raw:
+            raise ValueError(f"network.overlap missing required field '{field}'")
+        try:
+            val = float(overlap_raw[field])
+        except (TypeError, ValueError) as exc:
+            raise ValueError(f"network.overlap.{field} must be numeric") from exc
+        if val < 0.0 or val > 1.0:
+            raise ValueError(f"network.overlap.{field} must be between 0.0 and 1.0")
+        values[field] = val
+    return NetworkOverlapConfig(
+        tp_overlap=values["tp_overlap"],
+        tp_sp_overlap=values["tp_sp_overlap"],
+        cp_overlap=values["cp_overlap"],
+    )
+
+
+def _build_network_layout_config(
+    dimensions: Sequence[NetworkDimensionLayout],
+    faulty_links: Sequence[Tuple[int, int, float]] = (),
+    overlap_config: Optional["NetworkOverlapConfig"] = None,
+) -> NetworkLayoutConfig:
+    parallelism_map: Dict[str, NetworkDimensionLayout] = {}
+    for dim in dimensions:
+        for pname in dim.parallelisms:
+            if pname in parallelism_map:
+                raise ValueError(
+                    f"parallelism '{pname}' assigned to multiple network dimensions"
+                )
+            parallelism_map[pname] = dim
+    return NetworkLayoutConfig(
+        dimensions=tuple(dimensions),
+        faulty_links=tuple(faulty_links),
+        parallelism_map=parallelism_map,
+        overlap_config=overlap_config,
+    )
+
+
+_PARALLELISM_DEFAULTS: Dict[str, object] = {
+    "auto": False,
+    "dp": 1,
+    "lp": 1,
+    "mb": 1,
+    "kp_hidden_dim1": 1,
+    "kp_softmax_dim1": 1,
+    "kp_embedding_dim1": 1,
+    "kp_projection_dim1": 1,
+    "kp_hidden_dim2": 1,
+    "kp_softmax_dim2": 1,
+    "kp_embedding_dim2": 1,
+    "kp_projection_dim2": 1,
+    "kp_hidden_type": -1,
+    "kp_softmax_type": -1,
+    "kp_embedding_type": -1,
+    "kp_projection_type": -1,
+    "t": "CR",
+    "tp": 1,
+    "cp": 1,
+    "tp_sp": False,
+    "kp1": 1,
+    "kp2": 1,
+}
 
 
 @dataclass
@@ -528,7 +948,8 @@ class LLMConfig:
     tied_embeddings: bool
     num_layers: int
     hidden_dim: int
-    batch_size: int
+    global_batch_size: int
+    gradient_accumulation_steps: int
     seq_len: int
     decode_len: Optional[int]
     intermediate_size: Optional[int]
@@ -550,6 +971,12 @@ def use_flashattention(self) -> bool:
     @property
     def use_moe(self) -> bool:
         return self.num_experts > 1
+
+
+    @property
+    def grad_accumulation_steps(self) -> int:
+        """Backward-compatible alias for gradient accumulation steps."""
+        return self.gradient_accumulation_steps
 LLMInferenceConfig = _namedtuple(
     "inference_param",
     [
@@ -557,11 +984,19 @@ def use_moe(self) -> bool:
     ],
 )
 SWConfig = _namedtuple(
-    "sw_param", ["kernel_launch_overhead", "precision", "h2d_bandwidth", "dp_zero_stage"]
+    "sw_param",
+    [
+        "kernel_launch_overhead",
+        "precision",
+        "h2d_bandwidth",
+        "dp_zero_stage",
+        "full_recomputation",
+        "dp_microbatch",
+    ],
 )
 
 SchedulingConfig = _namedtuple(
-    "scheduling_param",
+    "parallelism",
     [
         "auto",
         "dp",
@@ -598,9 +1033,8 @@ def use_moe(self) -> bool:
         "sch_config",
         "area_breakdown",
         "perimeter_breakdown",
-        "system_config",
         "memory_hierarchy",
-        "network_topology",
+        "network_layout",
     ],
 )
 
@@ -645,7 +1079,6 @@ def use_moe(self) -> bool:
     "InferenceHWConfig",
     [
         "kvcache_type",
-        "kvcache_fetch_overlap",
     ],
 )
 
@@ -658,16 +1091,15 @@ def use_moe(self) -> bool:
         "sch_config",
         "area_breakdown",
         "perimeter_breakdown",
-        "system_config",
         "memory_hierarchy",
-        "network_topology",
+        "network_layout",
         "execution_backend",
         "inference_config",
     ],
 )
 
-MODELConfig = _namedtuple(
-    "MODELConfig",
+ModelConfig = _namedtuple(
+    "ModelConfig",
     [
         "model_config",
         "inference_config",
@@ -675,50 +1107,60 @@ def use_moe(self) -> bool:
 )
 
 
+def _convert_scalar_string(value: str):
+    try:
+        return float(value)
+    except ValueError:
+        pass
+
+    digit = [int(s) for s in value.split() if s.isdigit()]
+    order = [str(s) for s in value.split() if not s.isdigit()]
+    if not (order and digit):
+        return value
+
+    prefix = order[0][0]
+    bit = order[0][1] if len(order[0]) > 1 else "B"
+    mult = 1
+
+    if prefix == "K":
+        mult = 1024
+    elif prefix == "M":
+        mult = 1024 * 1024
+    elif prefix == "G":
+        mult = 1024 * 1024 * 1024
+    elif prefix == "T":
+        mult = 1024 * 1024 * 1024 * 1024
+    else:
+        raise ValueError(f"Unknown prefix '{prefix}' while parsing value '{value}'")
+
+    if bit == "b":
+        mult = mult / 8  # Capacity is expected in Bytes
+    elif bit != "B":
+        raise ValueError(f"Unknown type '{bit}' while parsing value '{value}'")
+
+    return digit[0] * mult
+
+
+def _convert_value(value):
+    if isinstance(value, dict):
+        convert(value)
+        return value
+    if isinstance(value, list):
+        return [_convert_value(item) for item in value]
+    if isinstance(value, str):
+        try:
+            return _convert_scalar_string(value)
+        except ValueError:
+            return value
+    return value
+
+
 def convert(d):
-    for key1, val1 in d.items():
-        for key2, val2 in val1.items():
-            if isinstance(val2, dict):
-                for key3, val3 in val2.items():
-                    if isinstance(val3, str):
-                        digit = [int(s) for s in val3.split() if s.isdigit()]
-                        order = [str(s) for s in val3.split() if not s.isdigit()]
-                        if order and digit:
-                            assert len(order) >= 1
-                            assert len(digit) >= 1
-
-                            prefix = order[0][0]
-                            bit = order[0][1]
-                            mult = 1
-
-                            if prefix == "K":
-                                mult = 1024
-                            elif prefix == "M":
-                                mult = 1024 * 1024
-                            elif prefix == "G":
-                                mult = 1024 * 1024 * 1024
-                            elif prefix == "T":
-                                mult = 1024 * 1024 * 1024 * 1024
-                            else:
-                                print(
-                                    "Unknown prefix: {} at {}: {}".format(
-                                        prefix, key3, val3
-                                    )
-                                )
-                                exit(0)
-
-                            if bit == "b":
-                                mult = mult / 8  # Capacity is expected in Bytes
-                            elif bit == "B":
-                                mult = mult
-                            else:
-                                print(
-                                    "Unknown type: {} at {}: {}".format(bit, key3, val3)
-                                )
-                                exit(0)
-
-                            new_val = digit[0] * mult
-                            d[key1][key2][key3] = new_val
+    if not isinstance(d, dict):
+        return d
+    for key, val in list(d.items()):
+        d[key] = _convert_value(val)
+    return d
 
 
 def parse_bandwidth_string(value):
@@ -781,14 +1223,14 @@ def parse_config(filename, config_type):
     """
     with open(filename, "r") as f:
         try:
-            config_dict = _yaml.load(f, Loader=_ruamel.yaml.Loader)
+            config_dict = _yaml.safe_load(f)
         except _YAMLError as exc:
             hint = (
                 f"Failed to parse YAML config '{filename}'. "
                 "Please check indentation and required sections like 'attention' and parameters such as 'num_experts'."
             )
             raise ValueError(hint) from exc
-        # print(config_dict) 
+        # print(config_dict)
         convert(config_dict)
     if config_type == "hardware":
         sw_block = dict(config_dict["sw_param"])
@@ -797,18 +1239,42 @@ def parse_config(filename, config_type):
         kernel_launch_overhead = sw_block["kernel_launch_overhead"]
         h2d_bandwidth = sw_block["h2d_bandwidth"]
         dp_zero_stage = sw_block["dp_zero_stage"]
+        full_recomp_raw = sw_block.get("full_recomputation", False)
+        if isinstance(full_recomp_raw, str):
+            full_recomputation = full_recomp_raw.strip().lower() in {"1", "true", "yes", "y", "on", "full"}
+        else:
+            full_recomputation = bool(full_recomp_raw)
+        dp_microbatch_raw = sw_block.get("dp_microbatch", "every_mb")
+        if isinstance(dp_microbatch_raw, str):
+            dp_microbatch = dp_microbatch_raw.strip().lower()
+        else:
+            dp_microbatch = str(dp_microbatch_raw).strip().lower()
+        if dp_microbatch not in {"every_mb", "last_mb"}:
+            raise ValueError("sw_param.dp_microbatch must be 'every_mb' or 'last_mb'")
         sw_config = SWConfig(
             kernel_launch_overhead=kernel_launch_overhead,
             precision=precision_config,
             h2d_bandwidth=h2d_bandwidth,
             dp_zero_stage=dp_zero_stage,
+            full_recomputation=full_recomputation,
+            dp_microbatch=dp_microbatch,
         )
-        sch_params = dict(config_dict["scheduling_param"])
-        if "tp" not in sch_params:
-            sch_params["tp"] = None
-        if "cp" not in sch_params:
-            sch_params["cp"] = None
-        sch_config = SchedulingConfig(**sch_params)
+        raw_parallelism_block = config_dict.get("parallelism", {})
+        if not isinstance(raw_parallelism_block, dict):
+            raise ValueError("parallelism must be a mapping")
+        parallelism_params = dict(_PARALLELISM_DEFAULTS)
+        parallelism_params.update(raw_parallelism_block)
+        parallelism_params["auto"] = bool(parallelism_params["auto"])
+        parallelism_params["tp_sp"] = bool(parallelism_params["tp_sp"])
+        sch_config = SchedulingConfig(**parallelism_params)
+        scheduling_for_network = {str(k).lower(): v for k, v in parallelism_params.items()}
+
+        network_spec = config_dict.get("network")
+        network_dimensions, network_faults, network_overlap = _parse_network_layout(network_spec, scheduling_for_network)
+        if not network_dimensions:
+            raise ValueError("network section must define at least one dimension")
+        network_layout_config = _build_network_layout_config(network_dimensions, network_faults, network_overlap)
+
         tech_config = TechConfig.from_dict(config_dict["tech_param"])
 
         # Optional breakdown configs (not needed for simplified configs)
@@ -836,13 +1302,9 @@ def parse_config(filename, config_type):
             # Create dummy perimeter config with zeros
             perimeter_config = PerimeterBreakdownConfig(DRAM=0.1, inter_node=0.1, intra_node=0.1)
 
-        system_config = SystemHierarchyConfig.from_dict(config_dict["system_hierarchy"])
         memory_hierarchy_config = MemoryHierarchyConfig.from_dict(
             config_dict["memory_hierarchy"]
         )
-        network_topology_config = NetworkTopologyConfig.from_dict(
-            config_dict["network_topology"]
-        )
         # execution backend (optional)
         eb_dict = config_dict.get("execution_backend", {})
         eb_model = eb_dict.get("model", "analytical")
@@ -881,7 +1343,6 @@ def parse_config(filename, config_type):
         inference_dict = config_dict.get("inference", {}) or {}
         inference_cfg = InferenceHWConfig(
             kvcache_type=inference_dict.get("kvcache_type", "hbm_only"),
-            kvcache_fetch_overlap=bool(inference_dict.get("kvcache_fetch_overlap", False)),
         )
 
         config = HWConfig(
@@ -891,21 +1352,20 @@ def parse_config(filename, config_type):
             sch_config=sch_config,
             area_breakdown=area_config,
             perimeter_breakdown=perimeter_config,
-            system_config=system_config,
             memory_hierarchy=memory_hierarchy_config,
-            network_topology=network_topology_config,
+            network_layout=network_layout_config,
             execution_backend=exec_backend,
             inference_config=inference_cfg,
         )
     elif config_type == "LSTM":
         model_config = ModelLSTMConfig(**config_dict["model_param"])
-        config = MODELConfig(model_config=model_config, inference_config=None)
+        config = ModelConfig(model_config=model_config, inference_config=None)
     elif config_type == "GEMM":
         mp = dict(config_dict["model_param"])  # copy
         if "backward" not in mp:
             mp["backward"] = False
         model_config = GEMMConfig(**mp)
-        config = MODELConfig(model_config=model_config, inference_config=None)
+        config = ModelConfig(model_config=model_config, inference_config=None)
     elif config_type == "LLM":
         mp = dict(config_dict["model_param"])
         if "attention" not in mp:
@@ -1016,8 +1476,7 @@ def parse_config(filename, config_type):
                 ) from exc
             if top_k <= 0:
                 raise ValueError("model_param.top_k must be a positive integer when provided")
-        if top_k > num_experts:
-            raise ValueError("model_param.top_k cannot exceed model_param.num_experts")
+
 
         if "run_type" not in mp:
             raise ValueError("model_param.run_type must be specified")
@@ -1082,7 +1541,25 @@ def _pop_required_int(field: str) -> int:
 
         num_layers = _pop_required_int("num_layers")
         hidden_dim = _pop_required_int("hidden_dim")
-        batch_size = _pop_required_int("batch_size")
+
+        global_batch_size = _pop_required_int("global_batch_size")
+
+
+        grad_accum_field = mp.pop("gradient_accumulation_steps", None)
+        if grad_accum_field is None:
+            grad_accum_field = mp.pop("gradient_accumulation_step", None)
+        if grad_accum_field is None:
+            gradient_accumulation_steps = 1
+        else:
+            try:
+                gradient_accumulation_steps = int(grad_accum_field)
+            except (TypeError, ValueError) as exc:
+                raise ValueError(
+                    f"model_param.gradient_accumulation_steps must be an integer (got {grad_accum_field!r})"
+                ) from exc
+            if gradient_accumulation_steps <= 0:
+                raise ValueError("model_param.gradient_accumulation_steps must be a positive integer")
+
         seq_len = _pop_required_int("seq_len")
         vocab_size = _pop_required_int("vocab_size")
 
@@ -1103,7 +1580,8 @@ def _pop_required_int(field: str) -> int:
             tied_embeddings=tied_embeddings,
             num_layers=num_layers,
             hidden_dim=hidden_dim,
-            batch_size=batch_size,
+            global_batch_size=global_batch_size,
+            gradient_accumulation_steps=gradient_accumulation_steps,
             seq_len=seq_len,
             decode_len=decode_len,
             intermediate_size=intermediate_size,
@@ -1114,13 +1592,13 @@ def _pop_required_int(field: str) -> int:
             num_experts=num_experts,
             top_k=top_k,
         )
-        config = MODELConfig(model_config=model_config, inference_config=inference_config)
+        config = ModelConfig(model_config=model_config, inference_config=inference_config)
     else:
         raise ValueError("Invalid config type: {}".format(config_type))
     
     # model_config = ModelConfig(**config_dict["model_param"])
     # sw_config = SWConfig(**config_dict["sw_param"])
-    # sch_config = SchedulingConfig(**config_dict["scheduling_param"])
+    # sch_config = SchedulingConfig(**config_dict["parallelism"])
     # tech_config = TechConfig.from_dict(config_dict["tech_param"])
     # power_config = PowerBreakdownConfig.from_dict(config_dict["power_breakdown"])
     # area_config = AreaBreakdownConfig.from_dict(config_dict["area_breakdown"])
diff --git a/configs/hardware-config/H100_SXM5_80GB.yaml b/configs/hardware-config/H100_SXM5_80GB.yaml
index bfadaf2c..f5e72fbc 100644
--- a/configs/hardware-config/H100_SXM5_80GB.yaml
+++ b/configs/hardware-config/H100_SXM5_80GB.yaml
@@ -5,10 +5,11 @@
 # [Semi] https://semianalysis.com/2025/06/23/nvidia-tensor-core-evolution-from-volta-to-blackwell/#4th-generation-tensor-core-warpgroup-level-asynchronous-mma
 
 sw_param:
+  full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation
   kernel_launch_overhead: 9e-6   # (s)
   # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3
   # only used if data parallelism (dp) > 1
-  dp_zero_stage: 0
+  dp_zero_stage: 1
   precision:
     # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32).
     # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats.
@@ -19,7 +20,7 @@ sw_param:
     # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes.
     # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode.
     # if true, optimizer states and stats are always FP32.
-    mixed_precision: false
+    mixed_precision: true
     # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params.
     # -> as_tensor_format makes weights/gradients use tensor_format bytes.
     # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and
@@ -28,8 +29,9 @@ sw_param:
     param_storage_mode: as_tensor_format
     # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format".
     kv_cache: as_tensor_format
-  h2d_bandwidth: 13297674240    # Bytes/s; set -1 to disable H2D timing
-
+  h2d_bandwidth: -1  # Bytes/s; set -1 to disable H2D timing
+  # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb)
+  dp_microbatch: last_mb
 tech_param:
   core:
     operating_frequency: 1.98e9     # (Hz) GPU boost clock frequency [TePU]
@@ -37,13 +39,13 @@ tech_param:
     nominal_flop_rate_per_mcu: 1024     # F16 FLOP / cycle / SM is 4096 [Semi]
     nominal_energy_per_flop: 1.8e-13
     num_mcu_per_bundle: 4     # number of tensorcores per SM [TePU]
-    FMA_d1: 8                # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock
-    FMA_d2: 4
-    dataflow: best              #{wst. ast, ost, best, none}: wst: weight stationary, ast: activation stationary, ost: output stationary
-    util: 0.75     #util should be 0.75 (12/16) for tensorcore
+    FMA_d1: 16                # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock
+    FMA_d2: 8
+    dataflow: best
+    util: 1.0
   DRAM:
     size: 80 GB     # [TePU]
-    bandwidth: 3440 GB     # (/s) [TePU]
+    bandwidth: 3430 GB     # (/s) [TePU]
     dynamic_energy_per_bit: 4.5e-12     # (J) [https://passlab.github.io/mchpc/mchpc2019/presentations/MCHPC_Pawlowski_keynote.pdf, slide 36]
     latency: 100e-9     # (s)
     util: 1
@@ -61,33 +63,10 @@ tech_param:
     util: 1
   SRAM-R:
     size: 33 MB     # register file memory per SM * number of SMs [NVIDIA]
-    bandwidth: 244 TB     # traffic for 2 reads and 1 write for GEMMs computed per clock
+    bandwidth: 122 TB     # traffic for 2 reads and 1 write for GEMMs computed per clock
     dynamic_energy_per_bit: 110e-15
     latency: 0
     util: 1
-  network:
-    intra_node:
-      latency: 5e-6
-      util: 1
-      bandwidth: 400 Gb
-    inter_node:
-      latency: 5e-6
-      util: 0.96
-      bandwidth: 400 Gb
-
-system_hierarchy:
-  num_devices_per_node: 1
-  num_nodes: 1
-  kp1_inter: false
-  kp2_inter: false
-  dp_inter: false
-  lp_inter: false
-  tp_inter: false
-
-network_topology:
-  # Allowed values: ring | fc | switch
-  inter_node: ring
-  intra_node: ring
 
 memory_hierarchy:
   l0:   #Register Memory
@@ -103,48 +82,78 @@ memory_hierarchy:
     type: DRAM
     scope: global
 
-scheduling_param:
+parallelism:
   auto: false
   ### LLM PARAMS ###
   dp: 1 # data parallelism. Acts as replica count during inference
+  tp: 8 # tensor parallelism
+  tp_sp: false # tensor-sequence parallelism
+  cp: 1 # context parallelism. Not supported for inference.
   lp: 1 # pipeline parallelism. Number of pipeline stages
-  mb: 1 # micro-batches (for best results set to >= lp)
-  tp: 1 # tensor parallelism
-  tp_sp: true # tensor-sequence parallelism
-  cp : 1 # context parallelism
-  ### LSTM/GEMM LEGACY PARAMS ###
-  t: CR
-  kp1: 1
-  kp2: 1
-  kp_hidden_dim1: 1
-  kp_hidden_dim2: 1
-  kp_softmax_dim1: 1
-  kp_softmax_dim2: 1
-  kp_embedding_dim1: 1
-  kp_embedding_dim2: 1
-  kp_projection_dim1: 1
-  kp_projection_dim2: 1
-  kp_hidden_type: -1
-  kp_softmax_type: -1
-  kp_embedding_type: -1
-  kp_projection_type: -1
+  # mb is not a true parallelism parameter, does not affect num of GPUs.
+  mb: 1 # micro-batches (for best results set to >= lp). 
+
+network:
+  dimensions:
+    - id: dim0
+      label: torus2d_tp
+      size: auto
+      topology:
+        type: Ring
+        bandwidth: 100 GB # We assume 8 25GB/s NVLINKs in the torus. Torus2D is hardcoded bidirectional, so it's actually 4 links, 25GB/s each (because they get replicated). So 100GB/s total. 
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+        optimize_2dmap: false
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [tp, cp]
+    - id: dim1
+      label: inter_torus_infiniband
+      size: auto
+      topology:
+        type: Ring 
+        bandwidth: 25 GB
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [lp,dp]
+    - id: dim2
+      label: inter_torus_infiniband_dp
+      size: auto
+      topology:
+        type: Ring
+        bandwidth: 25 GB
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+      parallelisms: []
+      collective_override: {} # overrides execution_backend.collectives when provided
 
+  overlap:
+    # fraction of TP+SP compute that can overlap TP/SP comm
+    tp_sp_overlap: 0.63
+    # fraction of TP compute that can overlap TP comm
+    tp_overlap: 0.6
+    # fraction of CP comm that can overlap CP compute
+    cp_overlap: 0.0
 execution_backend:
-  model: astra # analytical | astra
+  model: analytical # analytical | astra
   astra:
     backend: astra
-    mode: hybrid  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
+    mode: full_astrasim_flattened  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
     collectives:
-      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree
+      # Default collective algorithms. Per-dimension collective_override entries can override these.
+      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d
       # Note: leave to auto unless you have a specific reason to change it.
       all_gather: auto
       all_reduce: auto
       reduce_scatter: auto
       all_to_all: auto
     sys_options:
-      endpoint_delay: 10
+      endpoint_delay: 5
       active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism
-      preferred_dataset_splits: 32 # Should be same as active_chunks_per_dimension
+      preferred_dataset_splits: 4 # Should be same as active_chunks_per_dimension
 
 inference:
   kvcache_type: hbm_only
diff --git a/configs/hardware-config/a100_80GB.yaml b/configs/hardware-config/a100_80GB.yaml
index 36373ca5..2f4503d7 100644
--- a/configs/hardware-config/a100_80GB.yaml
+++ b/configs/hardware-config/a100_80GB.yaml
@@ -1,8 +1,9 @@
 sw_param:
+  full_recomputation: true # Enable full activation recompute during backward instead of selective recomputation
   kernel_launch_overhead: 9e-6   # (s)
   # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3
   # only used if data parallelism (dp) > 1
-  dp_zero_stage: 2
+  dp_zero_stage: 1
   precision:
     # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32).
     # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats.
@@ -13,7 +14,7 @@ sw_param:
     # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes.
     # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode.
     # if true, optimizer states and stats are always FP32.
-    mixed_precision: false
+    mixed_precision: true
     # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params.
     # -> as_tensor_format makes weights/gradients use tensor_format bytes.
     # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and
@@ -23,6 +24,8 @@ sw_param:
     # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format".
     kv_cache: as_tensor_format
   h2d_bandwidth: -1  # Bytes/s; set -1 to disable H2D timing
+  # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb)
+  dp_microbatch: last_mb
 tech_param:
   core:
     operating_frequency: 1.41e9     # (Hz)
@@ -58,27 +61,7 @@ tech_param:
     dynamic_energy_per_bit: 110e-15 # (J)
     latency: 0
     util: 1
-  network:
-    intra_node:
-      latency: 5e-6
-      util: 1
-      bandwidth: 400 Gb
-    inter_node:
-      latency: 5e-6
-      util: 0.96
-      bandwidth: 400 Gb
-system_hierarchy:
-  num_devices_per_node: 8 # currently only support astrasim integration with all devices in one node.
-  num_nodes: 1
-  kp1_inter: false # these currently are not support for LLM mode.
-  kp2_inter: false
-  dp_inter: false
-  lp_inter: false
-  tp_inter: false
-network_topology:
-  # Allowed values: ring | fc | switch
-  inter_node: ring
-  intra_node: ring
+
 memory_hierarchy:
   l0:   #Register Memory
     type: SRAM-R
@@ -92,49 +75,77 @@ memory_hierarchy:
   l3:   #Global Memory
     type: DRAM
     scope: global
-scheduling_param:
-  auto: false # Not supported for LLM mode.
-  ### LLM PARAMS ###
+
+parallelism:
   dp: 2 # data parallelism. Acts as replica count during inference
-  lp: 2 # pipeline parallelism. Number of pipeline stages
-  mb: 1 # micro-batches (for best results set to >= lp)
   tp: 2 # tensor parallelism
   tp_sp: true # tensor-sequence parallelism
-  cp : 1 # context parallelism
-  ### LSTM/GEMM LEGACY PARAMS ###
-  t: "CR"
-  kp1: 1
-  kp2: 1
-  kp_hidden_dim1: 1
-  kp_hidden_dim2: 1
-  kp_softmax_dim1: 1
-  kp_softmax_dim2: 1
-  kp_embedding_dim1: 1
-  kp_embedding_dim2: 1
-  kp_projection_dim1: 1
-  kp_projection_dim2: 1
-  kp_hidden_type: -1
-  kp_softmax_type: -1
-  kp_embedding_type: -1
-  kp_projection_type: -1
+  cp: 2 # context parallelism. Not supported for inference.
+  lp: 4 # pipeline parallelism. Number of pipeline stages
+  # mb is not a true parallelism parameter, does not affect num of GPUs.
+  mb: 4 # micro-batches (for best results set to >= lp). 
 
+network:
+  overlap:
+    tp_overlap: 0.6
+    tp_sp_overlap: 0.63
+    cp_overlap: 1.0
+  # Global faulty links apply to the active AstraSim dimension: [[src, dst, weight]]
+  # faulty_links: [[0,1,0.6],[16,17,0.1]]
+  dimensions:
+    # currently only support astrasim integration with all devices in one node.
+    - id: dim0
+      label: torus2d_tp
+      size: auto
+      topology:
+        type: Ring
+        bandwidth: 100 GB # We assume 8 25GB/s NVLINKs in the torus. Torus2D is hardcoded bidirectional, so it's actually 4 links, 25GB/s each (because they get replicated). So 100GB/s total. 
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+        optimize_2dmap: false
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [tp, cp]
+    - id: dim1
+      label: inter_torus_infiniband
+      size: auto
+      topology:
+        type: Ring 
+        bandwidth: 25 GB
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [lp,dp]
+    - id: dim2
+      label: inter_torus_infiniband_dp
+      size: auto
+      topology:
+        type: Ring
+        bandwidth: 25 GB
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+      parallelisms: []
+      collective_override: {} # overrides execution_backend.collectives when provided
 
 execution_backend:
   model: astra # analytical | astra
   astra:
     backend: astra
-    mode: full_astrasim_flattened  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
+    mode: full_astrasim_hierarchica  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
     collectives:
-      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree
+      # Default collective algorithms. Per-dimension collective_override entries can override these.
+      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d
       # Note: leave to auto unless you have a specific reason to change it.
       all_gather: auto
       all_reduce: auto
       reduce_scatter: auto
       all_to_all: auto
     sys_options:
-      endpoint_delay: 10
+      endpoint_delay: 5
       active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism
-      preferred_dataset_splits: 32 # Should be same as active_chunks_per_dimension
+      preferred_dataset_splits: 4 # Should be same as active_chunks_per_dimension
 
 inference:
   kvcache_type: hbm_only # Only supported mode for now.
diff --git a/configs/hardware-config/a100_80GB_no_parallelism.yaml b/configs/hardware-config/a100_80GB_no_parallelism.yaml
index aecc7e1f..f0476141 100644
--- a/configs/hardware-config/a100_80GB_no_parallelism.yaml
+++ b/configs/hardware-config/a100_80GB_no_parallelism.yaml
@@ -2,6 +2,7 @@
 # DO NOT MODIFY.
 
 sw_param:
+  full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation
   kernel_launch_overhead: 9e-6
   # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3
   # only used if data parallelism (dp) > 1
@@ -26,6 +27,8 @@ sw_param:
     # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format".
     kv_cache: as_tensor_format
   h2d_bandwidth: -1  # Bytes/s; set -1 to disable H2D timing
+  # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb)
+  dp_microbatch: last_mb
 tech_param:
   core:
     operating_frequency: 1.41e9
@@ -70,18 +73,6 @@ tech_param:
       latency: 5e-6
       util: 0.96
       bandwidth: 400 Gb
-system_hierarchy:
-  num_devices_per_node: 1
-  num_nodes: 1
-  kp1_inter: false
-  kp2_inter: false
-  dp_inter: false
-  lp_inter: false
-  tp_inter: false
-network_topology:
-  # Allowed values: ring | fc | switch
-  inter_node: ring
-  intra_node: ring
 memory_hierarchy:
   l0:
     type: SRAM-R
@@ -95,8 +86,9 @@ memory_hierarchy:
   l3:
     type: DRAM
     scope: global
-scheduling_param:
-  auto: false # Not supported for LLM mode.
+
+
+parallelism:
   ### LLM PARAMS ###
   dp: 1 # acts as replica count during inference
   lp: 1 # pipeline stages
@@ -104,31 +96,38 @@ scheduling_param:
   tp: 1 # tensor parallelism
   tp_sp: false # tensor-sequence parallelism
   cp : 1 # context parallelism
-  ### LSTM/GEMM LEGACY PARAMS ###
-  t: "CR"
-  kp1: 1
-  kp2: 1
-  kp_hidden_dim1: 1
-  kp_hidden_dim2: 1
-  kp_softmax_dim1: 1
-  kp_softmax_dim2: 1
-  kp_embedding_dim1: 1
-  kp_embedding_dim2: 1
-  kp_projection_dim1: 1
-  kp_projection_dim2: 1
-  kp_hidden_type: -1
-  kp_softmax_type: -1
-  kp_embedding_type: -1
-  kp_projection_type: -1
+
+network:
+  dimensions:
+    - id: dim0
+      label: device_fabric
+      size: 1
+      topology:
+        type: Ring
+        bandwidth: 400 Gb
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [tp, cp, lp, dp]
+
 
 
+  overlap:
+    # fraction of TP+SP compute that can overlap TP/SP comm
+    tp_sp_overlap: 0.63
+    # fraction of TP compute that can overlap TP comm
+    tp_overlap: 0.6
+    # fraction of CP comm that can overlap CP compute
+    cp_overlap: 0.0
 execution_backend:
   model: analytical # analytical | astra
   astra:
     backend: astra
     mode: full_astrasim_flattened  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
     collectives:
-      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree
+      # Default collective algorithms. Per-dimension collective_override entries can override these.
+      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d
       # Note: leave to auto unless you have a specific reason to change it.
       all_gather: auto
       all_reduce: auto
diff --git a/configs/hardware-config/a100_80GB_problem.yaml b/configs/hardware-config/a100_80GB_problem.yaml
new file mode 100644
index 00000000..3b1bc65c
--- /dev/null
+++ b/configs/hardware-config/a100_80GB_problem.yaml
@@ -0,0 +1,144 @@
+sw_param:
+  full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation
+  kernel_launch_overhead: 9e-6   # (s)
+  # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3
+  # only used if data parallelism (dp) > 1
+  dp_zero_stage: 2
+  precision:
+    # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32).
+    # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats.
+    # For FP32 or other formats they need to be manually adjusted.
+    # The following are equivalent: (tensor_format: 2.0, tensor_format: fp16, tensor_format: bf16).
+    tensor_format: bf16  # baseline storage width for activations/gradients.
+    # mixed_precision: true | false
+    # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes.
+    # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode.
+    # if true, optimizer states and stats are always FP32.
+    mixed_precision: false
+    # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params.
+    # -> as_tensor_format makes weights/gradients use tensor_format bytes.
+    # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and
+    #    keeps a separate FP32 master copy for weights.
+    # -> fp32_params makes weights *and* gradients use FP32 bytes without a master copy.
+    param_storage_mode: as_tensor_format
+    # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format".
+    kv_cache: as_tensor_format
+  h2d_bandwidth: -1  # Bytes/s; set -1 to disable H2D timing
+  # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb)
+  dp_microbatch: last_mb
+tech_param:
+  core:
+    operating_frequency: 1.41e9     # (Hz)
+    num_bundles: 108
+    nominal_flop_rate_per_mcu: 512
+    nominal_energy_per_flop: 1.8e-13 # (J)
+    num_mcu_per_bundle: 4     # number of tensorcores per SM
+    FMA_d1: 8                 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock
+    FMA_d2: 4
+    dataflow: best            # {wst | ast | ost | best}: wst: weight stationary, ast: activation stationary, ost: output stationary, best: best of all 3
+    util: 100.0
+  DRAM:
+    size: 80 GB
+    bandwidth: 198600 GB     # (/s)
+    dynamic_energy_per_bit: 4.5e-12     # (J)
+    latency: 100e-9     # (s)
+    util: 1
+  SRAM-L2:
+    size: 40 MB
+    bandwidth: 7050 GB
+    dynamic_energy_per_bit: 130e-15 # (J)
+    latency: 0
+    util: 1
+  SRAM-L1:
+    size: 20736 KB
+    bandwidth: 18 TB
+    dynamic_energy_per_bit: 130e-15 # (J)
+    latency: 0
+    util: 1
+  SRAM-R:
+    size: 27 MB
+    bandwidth: 122 TB     
+    dynamic_energy_per_bit: 110e-15 # (J)
+    latency: 0
+    util: 1
+
+memory_hierarchy:
+  l0:   #Register Memory
+    type: SRAM-R
+    scope: mcu
+  l1:   #Shared Memory
+    type: SRAM-L1
+    scope: mcu-bundle
+  l2:   #L2
+    type: SRAM-L2
+    scope: global
+  l3:   #Global Memory
+    type: DRAM
+    scope: global
+
+parallelism:
+  dp: 2 # data parallelism. Acts as replica count during inference
+  tp: 3 # tensor parallelism
+  tp_sp: true # tensor-sequence parallelism
+  cp: 2 # context parallelism. Not supported for inference.
+  lp: 3 # pipeline parallelism. Number of pipeline stages
+  # mb is not a true parallelism parameter, does not affect num of GPUs.
+  mb: 2 # micro-batches (for best results set to >= lp). 
+
+network:
+  # currently only support astrasim integration with all devices in one node.
+  - id: dim0
+    label: intra_wafer_mesh_silicon
+    size: 6
+    topology:
+      type: Ring
+      bandwidth: 1000 Gb
+      latency: 5e-6
+      energy_per_bit: 8e-12 # (J/bit)
+      util: 1.0
+    collective_override: {} # overrides execution_backend.collectives when provided
+    parallelisms: [tp, cp]
+  - id: dim1
+    label: inter_wafer_ring_infiniband
+    size: 3
+    topology:
+      type: Ring
+      bandwidth: 200 Gb
+      latency: 5e-6
+      energy_per_bit: 8e-12 # (J/bit)
+      util: 1.0
+    collective_override: {} # overrides execution_backend.collectives when provided
+    parallelisms: [lp]
+  - id: dim2
+    label: wafer_dp_chains_ring_infiniband
+    size: 2
+    topology:
+      type: ring
+      bandwidth: 200 Gb
+      latency: 5e-6
+      energy_per_bit: 8e-12 # (J/bit)
+      util: 0.96
+    collective_override: {} # overrides execution_backend.collectives when provided
+    parallelisms: [dp]
+
+
+execution_backend:
+  model: astra # analytical | astra
+  astra:
+    backend: astra
+    mode: full_astrasim_flattened  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
+    collectives:
+      # Default collective algorithms. Per-dimension collective_override entries can override these.
+      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d
+      # Note: leave to auto unless you have a specific reason to change it.
+      all_gather: auto
+      all_reduce: auto
+      reduce_scatter: auto
+      all_to_all: auto
+    sys_options:
+      endpoint_delay: 5
+      active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism
+      preferred_dataset_splits: 4 # Should be same as active_chunks_per_dimension
+
+inference:
+  kvcache_type: hbm_only # Only supported mode for now.
diff --git a/configs/hardware-config/gmap_mesh2d_demo.yaml b/configs/hardware-config/gmap_mesh2d_demo.yaml
new file mode 100644
index 00000000..cb62cc48
--- /dev/null
+++ b/configs/hardware-config/gmap_mesh2d_demo.yaml
@@ -0,0 +1,177 @@
+sw_param:
+  full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation
+  kernel_launch_overhead: 5e-6
+  dp_zero_stage: 0
+  precision:
+    tensor_format: bf16
+    mixed_precision: true
+    param_storage_mode: as_tensor_format
+    kv_cache: as_tensor_format
+  h2d_bandwidth: -1
+
+  # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb)
+  dp_microbatch: last_mb
+tech_param:
+  core:
+    operating_frequency: 1.41e9
+    num_bundles: 108
+    nominal_flop_rate_per_mcu: 512
+    nominal_energy_per_flop: 1.8e-13
+    nominal_voltage: 0.8
+    threshold_voltage: 0.6
+    margin_voltage: 0.1
+    operating_area_per_mcu: 1.0
+    num_mcu_per_bundle: 4
+    FMA_d1: 8
+    FMA_d2: 4
+    dataflow: best
+    util: 1.0
+  DRAM:
+    size: 80 GB
+    bandwidth: 1986 GB
+    dynamic_energy_per_bit: 4.5e-12
+    static_power_per_bit: 0.1
+    area_per_bit: 0.1
+    stack_capacity: 80
+    area_per_stack: 0.1
+    latency: 100e-9
+    mem_ctrl_area: 0.1
+    nominal_voltage: 0.8
+    threshold_voltage: 0.6
+    margin_voltage: 0.1
+    num_links_per_mm: 1
+    num_links_per_stack: 1
+    max_voltage: 0.8
+    util: 1
+  SRAM-L2:
+    size: 40 MB
+    bandwidth: 7050 GB
+    dynamic_energy_per_bit: 130e-15
+    static_power_per_bit: 0.1
+    area_per_bit: 0.1
+    bank_capacity: 1.0
+    controller_area_per_link: 0.1
+    latency: 0
+    overhead: 0.1
+    util: 1
+  SRAM-L1:
+    size: 20736 KB
+    bandwidth: 18 TB
+    dynamic_energy_per_bit: 130e-15
+    static_power_per_bit: 0.1
+    area_per_bit: 0.1
+    bank_capacity: 1.0
+    controller_area_per_link: 0.1
+    latency: 0
+    overhead: 0.1
+    util: 1
+  SRAM-R:
+    size: 27 MB
+    bandwidth: 122 TB
+    dynamic_energy_per_bit: 110e-15
+    static_power_per_bit: 0.1
+    area_per_bit: 0.1
+    bank_capacity: 1.0
+    controller_area_per_link: 0.1
+    latency: 0
+    overhead: 0.1
+    util: 1
+
+memory_hierarchy:
+  l0:
+    type: SRAM-R
+    scope: mcu
+  l1:
+    type: SRAM-L1
+    scope: mcu-bundle
+  l2:
+    type: SRAM-L2
+    scope: global
+  l3:
+    type: DRAM
+    scope: global
+
+parallelism:
+  dp: 1
+  tp: 4
+  tp_sp: false
+  cp: 4
+  lp: 1
+  mb: 1
+
+network:
+  dimensions:
+    - id: dim0
+      label: mesh2d_cluster
+      size: auto
+      topology:
+        type: Torus2D
+        bandwidth: 100 GB
+        latency: 5e-6
+        energy_per_bit: 8e-12
+        util: 1.0
+        optimize_2dmap: true
+      parallelisms: [tp, cp, lp]
+    - id: dim1
+      label: pipeline_ring
+      size: auto
+      topology:
+        type: Ring
+        bandwidth: 25 GB
+        latency: 5e-6
+        energy_per_bit: 8e-12
+        util: 1.0
+      parallelisms: [dp]
+
+  overlap:
+    # fraction of TP+SP compute that can overlap TP/SP comm
+    tp_sp_overlap: 0.63
+    # fraction of TP compute that can overlap TP comm
+    tp_overlap: 0.6
+    # fraction of CP comm that can overlap CP compute
+    cp_overlap: 0.0
+power_breakdown:
+  TDP: 400
+  core: 200
+  DRAM: 100
+  L2: 30
+  L1: 20
+  reg_mem: 10
+  network:
+    inter_node: 5
+    intra_node: 5
+
+area_breakdown:
+  proc_chip_area_budget: 1.0
+  core: 0.5
+  DRAM: 0.1
+  L2: 0.1
+  L1: 0.1
+  reg_mem: 0.1
+  device_area_budget: 1.0
+  network:
+    inter_node: 0.1
+    intra_node: 0.1
+
+perimeter_breakdown:
+  DRAM: 0.1
+  inter_node: 0.1
+  intra_node: 0.1
+
+execution_backend:
+  model: astra
+  astra:
+    backend: astra
+    mode: full_astrasim_flattened
+    collectives:
+      all_gather: auto
+      all_reduce: auto
+      reduce_scatter: auto
+      all_to_all: auto
+    sys_options:
+      endpoint_delay: 5
+      active_chunks_per_dimension: 8
+      preferred_dataset_splits: 2
+
+inference:
+  kvcache_type: hbm_only
diff --git a/configs/hardware-config/v100.yaml b/configs/hardware-config/v100.yaml
index a68ff884..782ddac6 100644
--- a/configs/hardware-config/v100.yaml
+++ b/configs/hardware-config/v100.yaml
@@ -1,5 +1,6 @@
 
 sw_param:
+  full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation
   kernel_launch_overhead: 9e-6   #9us
   # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3
   # only used if data parallelism (dp) > 1
@@ -25,6 +26,8 @@ sw_param:
     kv_cache: as_tensor_format
   h2d_bandwidth: 13297674240    # Bytes/s; set -1 to disable H2D timing
 
+  # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb)
+  dp_microbatch: last_mb
 tech_param:
   core:
     operating_frequency: 1.33e9     #Hz
@@ -70,20 +73,6 @@ tech_param:
       util: 0.96
       bandwidth: 300 Gb
 
-system_hierarchy:
-  num_devices_per_node: 1
-  num_nodes: 1
-  kp1_inter: false
-  kp2_inter: false
-  dp_inter: false
-  lp_inter: false
-  tp_inter: false
-
-network_topology:
-  # Allowed values: ring | fc | switch
-  inter_node: ring
-  intra_node: ring
-
 memory_hierarchy:
   l0:   #Register Memory
     type: SRAM-R
@@ -98,8 +87,8 @@ memory_hierarchy:
     type: DRAM
     scope: global
 
-scheduling_param:
-  auto: false
+
+parallelism:
   ### LLM PARAMS ###
   dp: 1 # acts as replica count during inference
   lp: 1 # pipeline stages
@@ -107,30 +96,37 @@ scheduling_param:
   tp: 1 # tensor parallelism
   tp_sp: true # tensor-sequence parallelism
   cp : 1 # context parallelism
-  ### LSTM/GEMM LEGACY PARAMS ###
-  t: CR
-  kp1: 1
-  kp2: 1
-  kp_hidden_dim1: 1
-  kp_hidden_dim2: 1
-  kp_softmax_dim1: 1
-  kp_softmax_dim2: 1
-  kp_embedding_dim1: 1
-  kp_embedding_dim2: 1
-  kp_projection_dim1: 1
-  kp_projection_dim2: 1
-  kp_hidden_type: -1
-  kp_softmax_type: -1
-  kp_embedding_type: -1
-  kp_projection_type: -1
 
+network:
+  dimensions:
+    - id: dim0
+      label: device_fabric
+      size: 1
+      topology:
+        type: Mesh
+        bandwidth: 300 Gb
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [tp, cp, lp, dp]
+
+
+  overlap:
+    # fraction of TP+SP compute that can overlap TP/SP comm
+    tp_sp_overlap: 0.63
+    # fraction of TP compute that can overlap TP comm
+    tp_overlap: 0.6
+    # fraction of CP comm that can overlap CP compute
+    cp_overlap: 0.0
 execution_backend:
   model: astra # analytical | astra
   astra:
     backend: astra
     mode: hybrid  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
     collectives:
-      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree
+      # Default collective algorithms. Per-dimension collective_override entries can override these.
+      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d
       # Note: leave to auto unless you have a specific reason to change it.
       all_gather: auto
       all_reduce: auto
diff --git a/configs/hardware-config/waferscale_20v100_80hbm.yaml b/configs/hardware-config/waferscale_20v100_80hbm.yaml
index 0ffdf7b3..c4d4bbfe 100644
--- a/configs/hardware-config/waferscale_20v100_80hbm.yaml
+++ b/configs/hardware-config/waferscale_20v100_80hbm.yaml
@@ -3,6 +3,7 @@
 
 
 sw_param:
+  full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation
   kernel_launch_overhead: 9e-6   #9us
   # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3
   # only used if data parallelism (dp) > 1
@@ -28,6 +29,8 @@ sw_param:
     kv_cache: as_tensor_format
   h2d_bandwidth: 13297674240    # Bytes/s; set -1 to disable H2D timing
 
+  # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb)
+  dp_microbatch: last_mb
 tech_param:
   core:
     operating_frequency: 1.33e9     #Hz
@@ -88,19 +91,6 @@ tech_param:
       util: 0.96
       bandwidth: 300 Gb
 
-system_hierarchy:
-  num_devices_per_node: 1   # devices per package essentially
-  num_nodes: 20   # one package
-  kp1_inter: true
-  kp2_inter: true
-  dp_inter: false
-  lp_inter: false
-  tp_inter: true
-
-network_topology:
-  inter_node: none     #mesh, torus, crossbar, custom, hierarchical
-  intra_node: mesh
-
 memory_hierarchy:
   l0:   #Register Memory
     type: SRAM-R
@@ -115,7 +105,7 @@ memory_hierarchy:
     type: DRAM
     scope: global
 
-scheduling_param:
+parallelism:
   auto: false
   dp: 1
   lp: 1
@@ -123,28 +113,36 @@ scheduling_param:
   tp: 20
   tp_sp: false
   cp: 1
-  t: CR
-  kp1: 20
-  kp2: 1
-  kp_hidden_dim1: 1
-  kp_hidden_dim2: 1
-  kp_softmax_dim1: 1
-  kp_softmax_dim2: 1
-  kp_embedding_dim1: 1
-  kp_embedding_dim2: 1
-  kp_projection_dim1: 1
-  kp_projection_dim2: 1
-  kp_hidden_type: -1
-  kp_softmax_type: -1
-  kp_embedding_type: -1
-  kp_projection_type: -1
+
+network:
+  dimensions:
+    - id: dim0
+      label: wafer_mesh
+      size: 20
+      topology:
+        type: Mesh
+        bandwidth: 5000 Gb
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 0.95
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [tp, cp, lp, dp]
+
+  overlap:
+    # fraction of TP+SP compute that can overlap TP/SP comm
+    tp_sp_overlap: 0.63
+    # fraction of TP compute that can overlap TP comm
+    tp_overlap: 0.6
+    # fraction of CP comm that can overlap CP compute
+    cp_overlap: 0.0
 execution_backend:
   model: astra # analytical | astra
   astra:
     backend: astra
     mode: hybrid  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
     collectives:
-      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree
+      # Default collective algorithms. Per-dimension collective_override entries can override these.
+      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d
       # Note: leave to auto unless you have a specific reason to change it.
       all_gather: auto
       all_reduce: auto
diff --git a/configs/model-config/LLM.yaml b/configs/model-config/LLM.yaml
index 7ced3ae6..73ea9b04 100644
--- a/configs/model-config/LLM.yaml
+++ b/configs/model-config/LLM.yaml
@@ -3,7 +3,8 @@ model_param:
     run_type: "training" # Type of run: "training" or "inference"
     tied_embeddings: false  # Share input and output embedding weights (WIP: affects memory estimation only)
     model_type: "gpt"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
-    batch_size: 1024  # Sequences per global batch
+    global_batch_size: 1024  # Sequences per global batch
+    gradient_accumulation_steps: 1  # Number of gradient accumulation micro-steps per optimizer update
     seq_len: 16384  # Sequence length (tokens per sequence)
     hidden_dim: 8192 # Transformer hidden dimension
     attention: # Attention-related parameters
@@ -16,4 +17,4 @@ model_param:
     top_k: 1   # Number of active experts per token in MoE
     intermediate_size: 32768 # Output dimension of the first FFN layer (typically 4×hidden_dim)
     vocab_size: 800000  # Vocabulary size
-    num_layers: 2 # Number of transformer layers
+    num_layers: 1 # Number of transformer layers
diff --git a/configs/model-config/LLM_inf.yaml b/configs/model-config/LLM_inf.yaml
index f66a6db9..34e92bcc 100644
--- a/configs/model-config/LLM_inf.yaml
+++ b/configs/model-config/LLM_inf.yaml
@@ -3,18 +3,19 @@ model_param:
     run_type: "inference" # Type of run: "training" or "inference"
     tied_embeddings: true  # Share input and output embedding weights (WIP: affects memory estimation only)
     model_type: "gpt"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
-    batch_size: 64  # Sequences per global batch
-    seq_len: 4096  # Total context length (prefill + decode_len for inference)
+    global_batch_size: 2048  # Sequences per global batch
+    gradient_accumulation_steps: 1  # Number of gradient accumulation micro-steps per optimizer update
+    seq_len: 2048  # Total context length (prefill + decode_len for inference)
     decode_len: 1024  # Number of autoregressive decode steps.
     hidden_dim: 8192 # Transformer hidden dimension
     attention: # Attention-related parameters
         attention_type: "mha"  # Supported attention types: "mha" or "gqa" ("mla" not implemented)
         num_heads: 64 # Total number of attention heads
-        kv_heads: 8  # Number of key/value heads (used only when attention_type = "gqa")
+        kv_heads: 64  # Number of key/value heads (used only when attention_type = "gqa")
         use_flashattention: false  # Enable FlashAttention optimization
         attention_tile_size: 256  # Tile size used in attention computation
     num_experts: 1   # Number of experts in MoE layers. Set to 1 to disable MoE.
     top_k: 1   # Number of active experts per token in MoE
     intermediate_size: 32768 # Output dimension of the first FFN layer (typically 4×hidden_dim)
     vocab_size: 129280  # Vocabulary size
-    num_layers: 96 # Number of transformer layers
+    num_layers: 32 # Number of transformer layers
diff --git a/configs/model-config/LLM_problem.yaml b/configs/model-config/LLM_problem.yaml
new file mode 100644
index 00000000..146d22e8
--- /dev/null
+++ b/configs/model-config/LLM_problem.yaml
@@ -0,0 +1,20 @@
+model_param:
+    mode: "LLM"  # Mode of operation
+    run_type: "training" # Type of run: "training" or "inference"
+    tied_embeddings: false  # Share input and output embedding weights (WIP: affects memory estimation only)
+    model_type: "gpt"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
+    global_batch_size: 1024  # Sequences per global batch
+    gradient_accumulation_steps: 1  # Number of gradient accumulation micro-steps per optimizer update
+    seq_len: 16384  # Sequence length (tokens per sequence)
+    hidden_dim: 8192 # Transformer hidden dimension
+    attention: # Attention-related parameters
+        attention_type: "mha"  # Supported attention types: "mha" or "gqa" ("mla" not implemented)
+        num_heads: 64 # Total number of attention heads
+        kv_heads: 8  # Number of key/value heads (used only when attention_type = "gqa")
+        use_flashattention: false  # Enable FlashAttention optimization
+        attention_tile_size: 256  # Tile size used in attention computation
+    num_experts: 1   # Number of experts in MoE layers. Set to 1 to disable MoE.
+    top_k: 1   # Number of active experts per token in MoE
+    intermediate_size: 32768 # Output dimension of the first FFN layer (typically 4×hidden_dim)
+    vocab_size: 800000  # Vocabulary size
+    num_layers: 3 # Number of transformer layers
diff --git a/configs/model-config/Llama2-13B_inf.yaml b/configs/model-config/Llama2-13B_inf.yaml
new file mode 100644
index 00000000..a5105fa0
--- /dev/null
+++ b/configs/model-config/Llama2-13B_inf.yaml
@@ -0,0 +1,21 @@
+model_param:
+  mode: "LLM"  # Mode of operation
+  run_type: "inference" # Type of run: "training" or "inference"
+  tied_embeddings: true  # Share input and output embedding weights (WIP: affects memory estimation only)
+  model_type: "llama"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
+  global_batch_size: 16  # Sequences per global batch
+  gradient_accumulation_steps: 1  # Number of gradient accumulation micro-steps per optimizer update
+  seq_len: 2048  # Total context length (prefill + decode_len for inference)
+  decode_len: 1024 # Number of tokens to decode during inference
+  hidden_dim: 5120 # Transformer hidden dimension
+  attention: # Attention-related parameters
+    attention_type: "mha"  # Supported attention types: "mha" or "gqa" ("mla" not implemented)
+    num_heads: 40 # Total number of attention heads
+    kv_heads: 8  # Number of key/value heads (used only when attention_type = "gqa")
+    use_flashattention: true  # Enable FlashAttention optimization
+    attention_tile_size: 256  # Tile size used in attention computation
+  num_experts: 1   # Number of experts in MoE layers. Set to 1 to disable MoE.
+  top_k: 1   # Number of active experts per token in MoE
+  intermediate_size: 13824 # Output dimension of the first FFN layer (typically 4×hidden_dim)
+  vocab_size: 32000  # Vocabulary size
+  num_layers: 40 # Number of transformer layers
diff --git a/configs/model-config/Llama2-70B_inf.yaml b/configs/model-config/Llama2-70B_inf.yaml
new file mode 100644
index 00000000..24facac3
--- /dev/null
+++ b/configs/model-config/Llama2-70B_inf.yaml
@@ -0,0 +1,21 @@
+model_param:
+  mode: "LLM"  # Mode of operation
+  run_type: "inference" # Type of run: "training" or "inference"
+  tied_embeddings: true  # Share input and output embedding weights (WIP: affects memory estimation only)
+  model_type: "llama"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
+  global_batch_size: 16  # Sequences per global batch
+  gradient_accumulation_steps: 1  # Number of gradient accumulation micro-steps per optimizer update
+  seq_len: 2048  # Total context length (prefill + decode_len for inference)
+  decode_len: 1024 # Number of tokens to decode during inference
+  hidden_dim: 8192 # Transformer hidden dimension
+  attention: # Attention-related parameters
+    attention_type: "gqa"  # Supported attention types: "mha" or "gqa" ("mla" not implemented)
+    num_heads: 64 # Total number of attention heads
+    kv_heads: 8  # Number of key/value heads (used only when attention_type = "gqa")
+    use_flashattention: true  # Enable FlashAttention optimization
+    attention_tile_size: 256  # Tile size used in attention computation
+  num_experts: 1   # Number of experts in MoE layers. Set to 1 to disable MoE.
+  top_k: 1   # Number of active experts per token in MoE
+  intermediate_size: 28672 # Output dimension of the first FFN layer (typically 4×hidden_dim)
+  vocab_size: 32000  # Vocabulary size
+  num_layers: 80 # Number of transformer layers
diff --git a/configs/model-config/Llama2-7B.yaml b/configs/model-config/Llama2-7B.yaml
index bf26bb5b..d862e9e4 100644
--- a/configs/model-config/Llama2-7B.yaml
+++ b/configs/model-config/Llama2-7B.yaml
@@ -3,7 +3,8 @@ model_param:
   run_type: "training" # Type of run: "training" or "inference"
   tied_embeddings: true  # Share input and output embedding weights (WIP: affects memory estimation only)
   model_type: "llama"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
-  batch_size: 1024  # Sequences per global batch
+  global_batch_size: 1024  # Sequences per global batch
+  gradient_accumulation_steps: 1  # Number of gradient accumulation micro-steps per optimizer update
   seq_len: 2048  # Sequence length (tokens per sequence)
   hidden_dim: 4096 # Transformer hidden dimension
   attention: # Attention-related parameters
diff --git a/configs/model-config/Llama2-7B_inf.yaml b/configs/model-config/Llama2-7B_inf.yaml
index 092477c1..0bed4a63 100644
--- a/configs/model-config/Llama2-7B_inf.yaml
+++ b/configs/model-config/Llama2-7B_inf.yaml
@@ -3,7 +3,8 @@ model_param:
   run_type: "inference" # Type of run: "training" or "inference"
   tied_embeddings: true  # Share input and output embedding weights (WIP: affects memory estimation only)
   model_type: "llama"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
-  batch_size: 16  # Sequences per global batch
+  global_batch_size: 16  # Sequences per global batch
+  gradient_accumulation_steps: 1  # Number of gradient accumulation micro-steps per optimizer update
   seq_len: 2048  # Total context length (prefill + decode_len for inference)
   decode_len: 1024 # Number of tokens to decode during inference
   hidden_dim: 4096 # Transformer hidden dimension
diff --git a/configs/model-config/Llama3.1-405B.yaml b/configs/model-config/Llama3.1-405B.yaml
index 3b277e7c..6894dea3 100644
--- a/configs/model-config/Llama3.1-405B.yaml
+++ b/configs/model-config/Llama3.1-405B.yaml
@@ -3,7 +3,8 @@ model_param:
   run_type: "training" # Type of run: "training" or "inference"
   tied_embeddings: true  # Share input and output embedding weights (WIP: affects memory estimation only)
   model_type: "llama"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
-  batch_size: 1024  # Sequences per global batch
+  global_batch_size: 1024  # Global batch size
+  gradient_accumulation_steps: 1  # Gradient accumulation steps (only for training)
   seq_len: 2048  # Sequence length (tokens per sequence)
   hidden_dim: 16384 # Transformer hidden dimension
   attention: # Attention-related parameters
diff --git a/configs/model-config/Llama3.1-405B_inf.yaml b/configs/model-config/Llama3.1-405B_inf.yaml
index 57094ac2..cbad8a6e 100644
--- a/configs/model-config/Llama3.1-405B_inf.yaml
+++ b/configs/model-config/Llama3.1-405B_inf.yaml
@@ -3,7 +3,8 @@ model_param:
   run_type: "inference" # Type of run: "training" or "inference"
   tied_embeddings: true  # Share input and output embedding weights (WIP: affects memory estimation only)
   model_type: "llama"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
-  batch_size: 1  # Sequences per global batch
+  global_batch_size: 1  # Sequences per global batch
+  gradient_accumulation_steps: 1  # Number of gradient accumulation micro-steps per optimizer update
   seq_len: 2048  # Total context length (prefill + decode_len for inference)
   decode_len: 1024 # Number of tokens to decode during inference
   hidden_dim: 16384 # Transformer hidden dimension
diff --git a/configs/model-config/Llama3.1-70B_inf.yaml b/configs/model-config/Llama3.1-70B_inf.yaml
index 4f133ac5..4bee4285 100644
--- a/configs/model-config/Llama3.1-70B_inf.yaml
+++ b/configs/model-config/Llama3.1-70B_inf.yaml
@@ -3,7 +3,8 @@ model_param:
   run_type: "inference" # Type of run: "training" or "inference"
   tied_embeddings: true  # Share input and output embedding weights (WIP: affects memory estimation only)
   model_type: "llama"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
-  batch_size: 1  # Sequences per global batch
+  global_batch_size: 1  # Sequences per global batch
+  gradient_accumulation_steps: 1  # Number of gradient accumulation micro-steps per optimizer update
   seq_len: 2048  # Total context length (prefill + decode_len for inference)
   decode_len: 1024 # Number of tokens to decode during inference
   hidden_dim: 8192 # Transformer hidden dimension
diff --git a/configs/model-config/Llama3.1-8B_inf.yaml b/configs/model-config/Llama3.1-8B_inf.yaml
index d69e5f7a..439048b9 100644
--- a/configs/model-config/Llama3.1-8B_inf.yaml
+++ b/configs/model-config/Llama3.1-8B_inf.yaml
@@ -3,7 +3,8 @@ model_param:
   run_type: "inference" # Type of run: "training" or "inference"
   tied_embeddings: true  # Share input and output embedding weights (WIP: affects memory estimation only)
   model_type: "llama"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
-  batch_size: 1  # Sequences per global batch
+  global_batch_size: 1  # Sequences per global batch
+  gradient_accumulation_steps: 1  # Number of gradient accumulation micro-steps per optimizer update
   seq_len: 2048  # Total context length (prefill + decode_len for inference)
   decode_len: 1024 # Number of tokens to decode during inference
   hidden_dim: 4096 # Transformer hidden dimension
diff --git a/configs/model-config/adding_new_models.md b/configs/model-config/adding_new_models.md
index 07373397..3305cb67 100644
--- a/configs/model-config/adding_new_models.md
+++ b/configs/model-config/adding_new_models.md
@@ -2,10 +2,10 @@ The hf_to_config.py utility can be used to convert a Hugging Face model config.j
 
 Usage:
 ```
-python hf_to_config.py <model_id> --revision <revision> --run-type <run_type>  --batch-size <batch_size> --seq-len <seq_len> --decode-len <decode_len>  --use-flashattention <bool> --flash-tile-size <int> -o <output_file>
+python hf_to_config.py <model_id> --revision <revision> --run-type <run_type>  --global-batch-size <global-batch-size> --gradient-accumulation-steps <gradient-accumulation-steps> --seq-len <seq_len> --decode-len <decode_len>  --use-flashattention <bool> --flash-tile-size <int> -o <output_file>
 ```
 
 Example:
 ```
-python hf_to_config.py NousResearch/Hermes-3-Llama-3.1-405B --run-type inference --batch-size 32 --seq-len 65536 --decode-len 1024 --use-flashattention true --flash-tile-size 256 -o Hermes_Llama3.1-405B.yaml
+python hf_to_config.py NousResearch/Hermes-3-Llama-3.1-405B --run-type inference --global-batch-size 32 --gradient-accumulation-steps 2 --seq-len 65536 --decode-len 1024 --use-flashattention true --flash-tile-size 256 -o Hermes_Llama3.1-405B.yaml
 ```
\ No newline at end of file
diff --git a/configs/model-config/gmap_mesh2d_demo.yaml b/configs/model-config/gmap_mesh2d_demo.yaml
new file mode 100644
index 00000000..e1b4e406
--- /dev/null
+++ b/configs/model-config/gmap_mesh2d_demo.yaml
@@ -0,0 +1,20 @@
+model_param:
+  mode: "LLM"
+  run_type: "training"
+  tied_embeddings: true
+  model_type: "llama"
+  global_batch_size: 8
+  gradient_accumulation_steps: 1
+  seq_len: 32
+  hidden_dim: 512
+  attention:
+    attention_type: "mha"
+    num_heads: 8
+    kv_heads: 8
+    use_flashattention: false
+    attention_tile_size: 16
+  num_experts: 1
+  top_k: 1
+  intermediate_size: 2048
+  vocab_size: 8192
+  num_layers: 2
diff --git a/configs/model-config/hf_to_config.py b/configs/model-config/hf_to_config.py
index 6a2ca66c..8efe902e 100644
--- a/configs/model-config/hf_to_config.py
+++ b/configs/model-config/hf_to_config.py
@@ -206,7 +206,8 @@ def _build_yaml_config(cfg: dict, args: argparse.Namespace, model_type: str) ->
             "run_type": args.run_type,
             "tied_embeddings": tied_embeddings,
             "model_type": model_type,
-            "batch_size": args.batch_size,
+            "global_batch_size": args.global_batch_size,
+            "gradient_accumulation_steps": args.gradient_accumulation_steps,
             "seq_len": int(seq_len),
             "decode_len": args.decode_len,
             "hidden_dim": int(hidden_dim),
@@ -226,7 +227,13 @@ def parse_args(argv: Optional[List[str]]) -> argparse.Namespace:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("model_id", help="Hugging Face model identifier, e.g. meta-llama/Llama-2-7b-hf")
     parser.add_argument("--revision", default="main", help="Model revision/branch to use (default: main)")
-    parser.add_argument("--batch-size", type=int, default=1, help="Training batch size to encode in the config (default: 1)")
+    parser.add_argument("--global-batch-size", type=int, default=1, help="Training global batch size to encode in the config (default: 1)")
+    parser.add_argument(
+        "--gradient-accumulation-steps",
+        type=int,
+        default=1,
+        help="Number of gradient accumulation micro-steps per optimizer update (default: 1)",
+    )
     parser.add_argument("--seq-len", type=int, default=None, help="Override sequence length. Defaults to HF config max position embeddings if not provided")
     parser.add_argument("--decode-len", type=int, default=0, help="Decode sequence length (default: 0)")
     parser.add_argument("--run-type", default="training", help="Config run_type field ('training' or 'inference')")
diff --git a/configs/old-detailed-hw-configs/H100_SXM5_80GB.yaml b/configs/old-detailed-hw-configs/H100_SXM5_80GB.yaml
index 44319ea2..3c4f0450 100644
--- a/configs/old-detailed-hw-configs/H100_SXM5_80GB.yaml
+++ b/configs/old-detailed-hw-configs/H100_SXM5_80GB.yaml
@@ -99,7 +99,7 @@ tech_param:
       latency: 5e-6
       nominal_frequency: 3e9
       nominal_voltage: 1
-      nominal_energy_per_link: 1e-12
+      energy_per_bit: 1e-12
       nominal_area_per_link: 100e-6     # (mm^2)
       num_links_per_mm: 400
       threshold_voltage: 0.25
@@ -109,7 +109,7 @@ tech_param:
       latency: 5e-6
       nominal_frequency: 20e9
       nominal_voltage: 1.2
-      nominal_energy_per_link: 13e-12
+      energy_per_bit: 13e-12
       nominal_area_per_link: 0.045
       num_links_per_mm: 20     #Assuming 4 layers on PCB and 200 um pitch for high speed signalling
       threshold_voltage: 0.35
@@ -204,7 +204,7 @@ execution_backend:
   astra:
     backend: astra
     mode: hybrid  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
-    collectives:
+    collective_override:
       # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree
       # Note: leave to auto unless you have a specific reason to change it.
       all_gather: auto
diff --git a/configs/old-detailed-hw-configs/a100_80GB.yaml b/configs/old-detailed-hw-configs/a100_80GB.yaml
index 75a5a3c6..debc8ce6 100644
--- a/configs/old-detailed-hw-configs/a100_80GB.yaml
+++ b/configs/old-detailed-hw-configs/a100_80GB.yaml
@@ -89,7 +89,7 @@ tech_param:
       latency: 5e-6
       nominal_frequency: 3e9
       nominal_voltage: 1
-      nominal_energy_per_link: 8e-12
+      energy_per_bit: 8e-12
       nominal_area_per_link: 100e-6
       num_links_per_mm: 400
       threshold_voltage: 0.25
@@ -99,7 +99,7 @@ tech_param:
       latency: 5e-6
       nominal_frequency: 20e9
       nominal_voltage: 1.2
-      nominal_energy_per_link: 13e-12
+      energy_per_bit: 13e-12
       nominal_area_per_link: 0.045
       num_links_per_mm: 20
       threshold_voltage: 0.35
@@ -189,7 +189,7 @@ execution_backend:
   astra:
     backend: astra
     mode: full_astrasim_flattened  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
-    collectives:
+    collective_override:
       # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree
       # Note: leave to auto unless you have a specific reason to change it.
       all_gather: auto
diff --git a/configs/old-detailed-hw-configs/v100.yaml b/configs/old-detailed-hw-configs/v100.yaml
index 71721296..f2a59f9c 100644
--- a/configs/old-detailed-hw-configs/v100.yaml
+++ b/configs/old-detailed-hw-configs/v100.yaml
@@ -91,7 +91,7 @@ tech_param:
       latency: 5e-6
       nominal_frequency: 3e9
       nominal_voltage: 1
-      nominal_energy_per_link: 1e-12
+      energy_per_bit: 1e-12
       nominal_area_per_link: 100e-6     #mm^2
       num_links_per_mm: 400
       threshold_voltage: 0.25
@@ -101,7 +101,7 @@ tech_param:
       latency: 5e-6
       nominal_frequency: 20e9
       nominal_voltage: 1.2
-      nominal_energy_per_link: 13e-12
+      energy_per_bit: 13e-12
       nominal_area_per_link: 0.045
       num_links_per_mm: 20     #Assuming 4 layers on PCB and 200 um pitch for high speed signalling
       threshold_voltage: 0.35
@@ -197,7 +197,7 @@ execution_backend:
   astra:
     backend: astra
     mode: hybrid  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
-    collectives:
+    collective_override:
       # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree
       # Note: leave to auto unless you have a specific reason to change it.
       all_gather: auto
diff --git a/configs/old-detailed-hw-configs/waferscale_20v100_80hbm.yaml b/configs/old-detailed-hw-configs/waferscale_20v100_80hbm.yaml
index 7e99ef83..2cc8fa6f 100644
--- a/configs/old-detailed-hw-configs/waferscale_20v100_80hbm.yaml
+++ b/configs/old-detailed-hw-configs/waferscale_20v100_80hbm.yaml
@@ -111,7 +111,7 @@ tech_param:
       nominal_frequency: 12.48e9     # THIS IS BITS/SECOND PER LINK!!!! Look at lines 502 in hw_component.py, 678 in perf.py
       nominal_voltage: 1
           # energy per link 
-      nominal_energy_per_link: 8e-12     # This is probably pessimistic, as it is likely stated for 300mm. Our avg distance is 150mm. 8pj/Bit based on https://www.techpowerup.com/forums/threads/nvidia-is-preparing-co-packaged-photonics-for-nvlink.276139/#:~:text=While%20the%20current%20NVLink%202.0,(4%20pJ%2Fb), https://www.techpowerup.com/img/fpXXsSk543ZUFWcx_thm.jpg
+      energy_per_bit: 8e-12     # This is probably pessimistic, as it is likely stated for 300mm. Our avg distance is 150mm. 8pj/Bit based on https://www.techpowerup.com/forums/threads/nvidia-is-preparing-co-packaged-photonics-for-nvlink.276139/#:~:text=While%20the%20current%20NVLink%202.0,(4%20pJ%2Fb), https://www.techpowerup.com/img/fpXXsSk543ZUFWcx_thm.jpg
       nominal_area_per_link: 1e-12     #mm^2. We are setting this to an extremely low value on purpose to ignore mem ctrl area. This is basically 0 but we cant set to 0 because of div errors.
       num_links_per_mm: 3.363     #increased a tiny bit to account for integer division and the fact that I dropped digits
       threshold_voltage: 0.25
@@ -121,7 +121,7 @@ tech_param:
       latency: 5e-6
       nominal_frequency: 20e9
       nominal_voltage: 1.2
-      nominal_energy_per_link: 13e-12
+      energy_per_bit: 13e-12
       nominal_area_per_link: 0.045
       num_links_per_mm: 20     #Assuming 4 layers on PCB and 200 um pitch for high speed signalling
       threshold_voltage: 0.35
@@ -212,7 +212,7 @@ execution_backend:
   astra:
     backend: astra
     mode: hybrid  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
-    collectives:
+    collective_override:
       # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree
       # Note: leave to auto unless you have a specific reason to change it.
       all_gather: auto
diff --git a/GD_search.py b/deprecated/GD_search.py
similarity index 99%
rename from GD_search.py
rename to deprecated/GD_search.py
index d496160b..d8ddde80 100644
--- a/GD_search.py
+++ b/deprecated/GD_search.py
@@ -12,8 +12,7 @@
 import util
 import perf
 
-import ruamel as _ruamel
-import ruamel.yaml as _yaml
+import yaml as _yaml
 
 import math
 
@@ -184,7 +183,7 @@ def initialize(self):
         random.seed(random_seed)
 
         with open(self.exp_config, "r") as f:
-            config_dict = _yaml.load(f, Loader=_ruamel.yaml.Loader)
+            config_dict = _yaml.safe_load(f)
 
         for param_class in self.parameters:
           if param_class == 'perimeter_breakdown': 
@@ -263,7 +262,7 @@ def populate_config(self, params, exp_dir):
         
         config_dict = {}
         with open(exp_config, "r") as f:
-            config_dict = _yaml.load(f, Loader=_ruamel.yaml.Loader)
+            config_dict = _yaml.safe_load(f)
         
         config_dict['area_breakdown']['proc_chip_area_budget'] = self.chip_area_budget if self.chip_area_budget != -1 else config_dict['area_breakdown']['proc_chip_area_budget']#in mm^2
         for param_class in params:
@@ -305,7 +304,7 @@ def populate_config(self, params, exp_dir):
             config_dict['system_hierarchy'][param] = self.system_hierarchy_params[param]
         
         with open(config_file, 'w') as yaml_file:
-            _yaml.dump(config_dict, yaml_file, default_flow_style=False)
+            _yaml.safe_dump(config_dict, yaml_file, default_flow_style=False, sort_keys=False)
       
         #print(config_file)
         #if self.debug:
diff --git a/deviceMapping.py b/deprecated/deviceMapping.py
similarity index 100%
rename from deviceMapping.py
rename to deprecated/deviceMapping.py
diff --git a/energy.py b/deprecated/energy.py
similarity index 100%
rename from energy.py
rename to deprecated/energy.py
diff --git a/main.py b/deprecated/main.py
similarity index 100%
rename from main.py
rename to deprecated/main.py
diff --git a/topology.py b/deprecated/topology.py
similarity index 100%
rename from topology.py
rename to deprecated/topology.py
diff --git a/topology_hack.py b/deprecated/topology_hack.py
similarity index 100%
rename from topology_hack.py
rename to deprecated/topology_hack.py
diff --git a/examples/korthi.sh b/examples/korthi.sh
new file mode 100755
index 00000000..2cb3c7f7
--- /dev/null
+++ b/examples/korthi.sh
@@ -0,0 +1,11 @@
+#!/bin/bash
+
+# Runs the first Korthikanti training case (GPT-22B, bs=4, mb=1, dp=1, tp=8, pp=1, cp=1, tp_sp=true)
+# with all debug/visualization flags enabled.
+
+DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS=1 \
+DEEPFLOW_VISUALIZE_GRAPHS=1 \
+DEEPFLOW_PERSIST_ARTIFACT_VIZ=1 \
+uv run run_perf.py \
+  --hardware_config validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml \
+  --model_config validation_scripts/validation_configs/model-config/GPT_22_B.yaml
diff --git a/examples/llm_astra_inference.sh b/examples/llm_astra_inference.sh
index c8cb3551..553f31f1 100755
--- a/examples/llm_astra_inference.sh
+++ b/examples/llm_astra_inference.sh
@@ -4,4 +4,4 @@
 # Setup model parameter in configs/model-config/Llama2-7B_inf.yaml
 # Setup hardware parameters (incl. parallelism/network) in configs/hardware-config/a100_80GB.yaml
 
-python run_perf.py --hardware_config configs/hardware-config/a100_80GB.yaml --model_config configs/model-config/Llama2-7B_inf.yaml
\ No newline at end of file
+python run_perf.py --hardware_config configs/hardware-config/a100_80GB.yaml --model_config configs/model-config/LLM_inf.yaml
\ No newline at end of file
diff --git a/examples/llm_astra_inference_debug_graphviz.sh b/examples/llm_astra_inference_debug_graphviz.sh
index eda86b7a..70426991 100755
--- a/examples/llm_astra_inference_debug_graphviz.sh
+++ b/examples/llm_astra_inference_debug_graphviz.sh
@@ -4,4 +4,4 @@
 # Setup model parameter in configs/model-config/Llama2-7B.yaml
 # Setup hardware parameters (incl. parallelism/network) in configs/hardware-config/a100_80GB.yaml
 
-DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS=1 DEEPFLOW_VISUALIZE_GRAPHS=1 DEEPFLOW_PERSIST_ARTIFACT_VIZ=1 python run_perf.py --hardware_config configs/hardware-config/a100_80GB.yaml --model_config configs/model-config/Llama2-7B_inf.yaml
\ No newline at end of file
+DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS=1 DEEPFLOW_VISUALIZE_GRAPHS=1 DEEPFLOW_PERSIST_ARTIFACT_VIZ=1 python run_perf.py --hardware_config configs/hardware-config/a100_80GB.yaml --model_config configs/model-config/LLM_inf.yaml
\ No newline at end of file
diff --git a/examples/problem.sh b/examples/problem.sh
new file mode 100755
index 00000000..01bb28fa
--- /dev/null
+++ b/examples/problem.sh
@@ -0,0 +1,7 @@
+#!/bin/bash
+
+# Runs performance model in LLM mode
+# Setup model parameter in configs/model-config/Llama2-7B.yaml
+# Setup hardware parameters (incl. parallelism/network) in configs/hardware-config/a100_80GB.yaml
+
+DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS=1 DEEPFLOW_VISUALIZE_GRAPHS=1 DEEPFLOW_PERSIST_ARTIFACT_VIZ=1 python run_perf.py --hardware_config configs/hardware-config/a100_80GB_problem.yaml --model_config configs/model-config/LLM_problem.yaml
\ No newline at end of file
diff --git a/gemm_calc.py b/gemm_calc.py
new file mode 100644
index 00000000..8b7478bb
--- /dev/null
+++ b/gemm_calc.py
@@ -0,0 +1,79 @@
+import sys
+from tile import TiledGEMM
+
+
+class GEMMCalc:
+    def __init__(self, core, mem_hierarchy, dtype_bytes, flashattn_enable):
+        self.core = core
+        self.mem_hierarchy = mem_hierarchy
+        self.dtype_bytes = dtype_bytes
+        self.flashattn_enable = flashattn_enable
+
+    def run(self, M : int, K : int, N : int, name = "", throughput: float | None = None):
+        best_result = self._search_tile_space(M, K, N, throughput=throughput)
+        return best_result
+    
+    def _search_tile_space(self, M : int, K : int, N : int, throughput: float | None = None):
+        best_gemm = None
+        best_time = float("inf")
+
+        for gemm in TiledGEMM.enumerate_candidates(
+            self.core,
+            self.mem_hierarchy,
+            M, K, N,
+            self.dtype_bytes
+        ):
+            time = self._compute_time(gemm, throughput=throughput)
+            if time < best_time:
+                best_time = time
+                best_gemm = gemm
+
+        return best_gemm, best_time
+
+    def _compute_time(self, gemm : TiledGEMM, throughput: float | None = None):
+        times,_,_ = self._roofline(
+            gemm.GEMM_flop,
+            gemm.mem_accesses.totals(),
+            gemm.bundle_util,
+            flashattn_enable=self.flashattn_enable,
+            throughput=throughput
+        )
+        return max(times)
+    
+    def _roofline(self, total_flops, mem_access : tuple, util : float,
+                   flashattn_enable : bool, throughput: float | None = None):
+
+        if not throughput:
+            throughput = self.core.getThroughput()
+        throughput *= util
+        num_level = len(mem_access)
+        try:
+            if not flashattn_enable:
+                assert mem_access[num_level - 1] > 0, f"mem_access: {mem_access}"
+        except Exception as e:
+            print(
+                "Number of accesses to the last level of memory hierarchy cannot be zero:\n {}".format(e),
+                flush=True,
+            )
+            sys.exit(0)
+        levels_to_compute = [(i, mem_access[i]) for i in range(num_level)]
+
+        # Compute roofline time for each level
+        times = []
+        inf_pts = []
+        comp_ints = []
+        for level_idx, num_mem in levels_to_compute:
+            mem_bw = self.mem_hierarchy[level_idx].getThroughput()
+            mem_latency = self.mem_hierarchy[level_idx].getLatency()
+            inflection_point = float("inf") if mem_bw == 0 else throughput / mem_bw
+            # print(f"Level {level_idx}: mem_bw={mem_bw}, mem_latency={mem_latency}, inflection_point={inflection_point}", flush=True)
+            comp_int = float("inf") if num_mem == 0 else total_flops / num_mem
+            if comp_int < inflection_point:  # mem-bound
+                level_time = (float("inf") if (mem_bw == 0 or num_mem == 0) else (num_mem / mem_bw)) + mem_latency
+            else:  # compute-bound
+                level_time = float("inf") if (throughput == 0) else (total_flops / throughput)
+            times.append(level_time)
+            inf_pts.append(inflection_point)
+            comp_ints.append(comp_int)
+        # print("Roofline: exited {}".format(name))
+        return times, inf_pts, comp_ints
\ No newline at end of file
diff --git a/hw_component.py b/hw_component.py
index 9e2f9133..602a5685 100644
--- a/hw_component.py
+++ b/hw_component.py
@@ -1,17 +1,13 @@
 import math
+from typing import Tuple
+
 import numpy as np
 import util
-from topology_hack import Topology
 
 kilo = 1024.0
 giga = 1024.0 * 1024.0 * 1024.0
 
 
-class System:
-    def __init__(self, exp_config):
-        self.num_nodes_per_wafer = exp_config.system_hierarchy.num_nodes_per_wafer
-
-
 class Base:
     def __init__(self, exp_config):
         self.exp_config = exp_config
@@ -64,6 +60,29 @@ def getTileDims(self):
         return self.getPower2TileDims()
         # return self.getArbitraryTileDims(self):
 
+    def calc_waves_per_sm(self, M, K, N, m, k, n):
+        bytes_required = self.precision * (m * k + k * n)
+        grid_size = math.ceil(M / m) * math.ceil(N / n)
+
+        compute_time = 2 * (m * n * k) / 312e12
+        load_time = (self.precision * (m * k + k * n)) / (7050 * 1024 ** 3)
+
+        stages = 3 if load_time < compute_time else 5
+        smem_required = bytes_required * stages
+        if smem_required > self.size_per_bundle:
+            return -1
+        
+        reg_accum = m * n
+        reg_input = 16 * (m + n)
+        if reg_accum + reg_input > 65536:
+            return -1
+
+        smem_blocks = self.size_per_bundle // smem_required
+        reg_blocks = 65536 // (reg_accum + reg_input)
+
+        return grid_size / min(smem_blocks, reg_blocks) / 108
+
+
     def getGEMMBasedTileDims(self, M, K, N):
         m_dims = [ M >> i for i in range(M.bit_length()) if (M >> i) >= 1 ]
         k_dims = [ K >> i for i in range(K.bit_length()) if (K >> i) >= 1 ]
@@ -708,188 +727,48 @@ def printStats(self, f):
 class Network(Base):
     def __init__(self, exp_config):
         super().__init__(exp_config)
-        self.intra_network = SubNetwork(
-            exp_config,
-            exp_config.tech_config.network.intra_node,
-            exp_config.power_breakdown.network.intra_node,
-            exp_config.area_breakdown.network.intra_node,
-            "intra",
-        )
-        self.inter_network = SubNetwork(
-            exp_config,
-            exp_config.tech_config.network.inter_node,
-            exp_config.power_breakdown.network.inter_node,
-            exp_config.area_breakdown.network.inter_node,
-            "inter",
-        )
+        layout = getattr(exp_config, "network_layout", None)
+        if layout is None:
+            raise ValueError("hardware config is missing network layout information")
+        self.layout = layout
+        self.energies_per_bit = [
+            float(getattr(dim, "energy_per_bit", 0.0))
+            for dim in layout.dimensions
+        ]
 
     def calcThroughput(self):
-        inter_throughput = self.inter_network.calcThroughput()
-        intra_throughput = self.intra_network.calcThroughput()
-
-        return intra_throughput, inter_throughput
+        primary = self.layout.primary_dimension()
+        if primary is None:
+            return 0.0, 0.0
+        bw = float(primary.effective_bandwidth)
+        return bw, bw
 
     def calcLatency(self):
-        inter_latency = self.inter_network.latency
-        intra_latency = self.intra_network.latency
-
-        return intra_latency, inter_latency
-
-    def printStats(self, f):
-        self.inter_network.printStats(f, "inter")
-        self.intra_network.printStats(f, "intra")
+        primary = self.layout.primary_dimension()
+        if primary is None:
+            return 0.0, 0.0
+        latency = float(primary.latency)
+        return latency, latency
 
+    def get_link(self, parallelism: str) -> Tuple[float, float]:
+        bandwidth, latency = self.layout.link_for_parallelism(parallelism)
+        return float(bandwidth), float(latency)
 
-class SubNetwork(Base):
-    def __init__(
-        self, exp_config, net_config, power_breakdown, area_breakdown, netLevel
-    ):
-        super().__init__(exp_config)
-        self.tot_power = power_breakdown  # * self.TDP
-        self.tot_area = area_breakdown  # * self.proc_chip_area_budget
-        # TODO: Rename core_perimeter to proc_chip_perimeter
-        self.latency = net_config.latency
-        self.nominal_freq = net_config.nominal_freq
-        self.nominal_voltage = net_config.nominal_voltage
-        self.nominal_energy_per_link = net_config.nominal_energy_per_link
-        self.nominal_area_per_link = net_config.nominal_area_per_link
-        self.threshold_voltage = net_config.threshold_voltage
-        self.margin_voltage = net_config.margin_voltage
-        self.num_links_per_mm = net_config.num_links_per_mm
-        self.util = net_config.util
-        self.config_operating_frequency = net_config.operating_frequency  # Master parameter
-        self.config_bandwidth = net_config.bandwidth  # Master parameter - overrides all
-        self.inter = True if netLevel == "inter" else False
-        self.intra = True if netLevel == "intra" else False
-
-        inter_frac = exp_config.perimeter_breakdown.inter_node
-        intra_frac = exp_config.perimeter_breakdown.intra_node
-        perimeter_fraction = inter_frac if self.inter else intra_frac
-        self.tot_perimeter = perimeter_fraction * self.core_perimeter
-        self.num_links = int(
-            min(
-                self.tot_area / self.nominal_area_per_link,
-                perimeter_fraction * self.core_perimeter * self.num_links_per_mm,
-            )
-        )
+    def dimension_for_parallelism(self, parallelism: str):
+        return self.layout.dimension_for_parallelism(parallelism)
 
-        self.throughput = 0
-        self.operating_freq = 0
-        self.operating_voltage = 0
-
-        # If bandwidth is explicitly set, use it directly (master parameter - overrides everything)
-        if self.config_bandwidth:
-            self.throughput = self.config_bandwidth * self.util
-            # Set nominal values when bandwidth is overridden (actual freq/voltage don't matter)
-            self.operating_freq = self.nominal_freq
-            self.operating_voltage = self.nominal_voltage
-        elif self.num_links > 0:
-            self.calcOperatingVoltageFrequency()
-            # Calculate throughput based on topology
-            # 4 edges for intra (mesh), 1 edge for inter (mesh extension)
-            degree = 4 if self.intra else 1
-            self.throughput = (self.num_links * self.operating_freq * self.util) / (8 * degree)
-
-        # self.energy_per_bit             = self.calcEnergyPerBit()
-
-    def calcOperatingVoltageFrequency(self):
-        # If operating_frequency is explicitly set, use it directly (master parameter)
-        if self.config_operating_frequency:
-            self.operating_freq = self.config_operating_frequency
-            self.operating_voltage = self.nominal_voltage  # Assume nominal voltage when freq is overridden
-            return
-
-        # Otherwise, calculate from voltage scaling based on power constraints
-        self.nominal_power = (
-            self.nominal_energy_per_link * self.num_links * self.nominal_freq
-        )
-        # minimum voltage to meet the power constraint
-        self.operating_voltage = self.solve_poly(
-            p0=1,
-            p1=-2 * self.threshold_voltage,
-            p2=self.threshold_voltage**2,
-            p3=-1
-            * self.tot_power
-            / self.nominal_power
-            * self.nominal_voltage
-            * (self.nominal_voltage - self.threshold_voltage) ** 2,
-        )
-        # operating frequency at minimum voltage
-        self.operating_freq = self.nominal_freq * (
-            (
-                (self.operating_voltage - self.threshold_voltage) ** 2
-                / (self.operating_voltage)
-            )
-            / (
-                (self.nominal_voltage - self.threshold_voltage) ** 2
-                / self.nominal_voltage
-            )
-        )
-        self.frequency_scaling_factor = 1
-        if self.operating_voltage < (self.threshold_voltage + self.margin_voltage):
-            self.scaled_voltage = self.threshold_voltage + self.margin_voltage
-            self.frequency_scaling_factor = (
-                self.operating_voltage / self.scaled_voltage
-            ) ** 2
-
-        self.operating_freq = self.frequency_scaling_factor * self.operating_freq
-
-    def calcEnergyPerBit(self):
-        self.operating_energy_per_link = (
-            self.nominal_energy_per_link
-            * (self.operating_voltage / self.nominal_voltage) ** 2
-        )
-        energy_per_bit = self.operating_energy_per_link
-        return energy_per_bit
-
-    # Return P2P bw
-    def calcThroughput(self):
-        return self.throughput
-
-    def printStats(self, f, name):
-        self.calcEnergyPerBit()
-        self.eff_power = (
-            self.num_links * self.operating_freq * self.operating_energy_per_link
-        )
-        self.eff_area = self.num_links * self.nominal_area_per_link
-        self.eff_perimeter = self.num_links / self.num_links_per_mm
-
-        if self.eff_power > 0 and self.eff_area > 0:
+    def printStats(self, f):
+        for dim in self.layout.dimensions:
             f.write("\n\n=============\n")
-            f.write("Network: {}\n".format(name))
+            f.write(f"Network Dimension: {dim.label} ({dim.id})\n")
             f.write("=============\n")
+            f.write(f"Topology: {dim.topology_type}\n")
+            f.write(f"Size: {dim.size}\n")
             f.write(
-                "operating_voltage: {0:.2f}, operating_freq: {1:.2f} (Ghz)\n".format(
-                    self.operating_voltage, self.operating_freq / 1e9
-                )
-            )
-            f.write(
-                "voltage_lowerbound: {0:.2f}\n".format(
-                    self.threshold_voltage + self.margin_voltage
+                "Effective Bandwidth: {0:.4f} (GB/s)\n".format(
+                    dim.effective_bandwidth / 1024 / 1024 / 1024
                 )
             )
-            f.write("#links: {0:5d}\n".format(self.num_links))
-            if self.tot_area != 0:
-                f.write(
-                    "eff_area: {0:.2f} (mm2), tot_area: {1:.2f} (mm2), util: {2:.2f}%\n".format(
-                        self.eff_area,
-                        self.tot_area,
-                        self.eff_area / self.tot_area * 100,
-                    )
-                )
-            if self.tot_power != 0:
-                f.write(
-                    "eff_power: {0:.2f} (watt), tot_power: {1:.2f} (watt), util: {2:.2f}%\n".format(
-                        self.eff_power,
-                        self.tot_power,
-                        self.eff_power / self.tot_power * 100,
-                    )
-                )
-            if self.tot_perimeter != 0:
-                f.write(
-                    "eff_perimeter: {0:.2f} (mm), tot_perimeter: {1:.2f} (mm), util: {2:.2f}%\n".format(
-                        self.eff_perimeter,
-                        self.tot_perimeter,
-                        self.eff_perimeter / self.tot_perimeter * 100,
-                    )
-                )
+            f.write(f"Latency: {dim.latency * 1e9:.4f} (ns)\n")
+            f.write(f"Util: {dim.util:.4f}\n")
+            f.write(f"Parallelisms: {', '.join(dim.parallelisms) or 'None'}\n")
diff --git a/ispass_plot.py b/ispass_plot.py
new file mode 100644
index 00000000..3e39bbc1
--- /dev/null
+++ b/ispass_plot.py
@@ -0,0 +1,263 @@
+from gemm_calc import GEMMCalc
+from hw_component import Core, MemoryHierarchy
+from validate_gemm import extract_gemm_shapes, parse_llm_shapes_file
+import config
+from typing import List, Dict
+
+import os
+import csv
+from pathlib import Path
+import pprint
+import matplotlib.pyplot as plt
+import numpy as np
+
+STACK_ORDER = [
+    'qkv_proj',
+    'attention_score',
+    'attention_output',
+    'output_proj',
+    'ffn1',
+    'ffn2',
+    'linear'
+]
+
+def configs():
+    """Load hardware and model configurations from YAML files."""
+    hw_config_path = "configs/hardware-config/a100_80GB_no_parallelism.yaml"
+    model_config_path = "configs/model-config/GEMM.yaml"
+    
+    # Make paths absolute relative to the repository root
+    # base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    # hw_config_path = os.path.join(base_dir, hw_config_path)
+    # model_config_path = os.path.join(base_dir, model_config_path)
+    
+    # Load configurations using config module
+    hw_config = config.parse_config(hw_config_path, config_type="hardware")
+    model_config = config.parse_config(model_config_path, config_type="GEMM")
+    
+    return hw_config, model_config
+
+def csv_validation_data():
+    """Load validation data from ncu_metrics_results.csv."""
+    csv_path = Path("validation_scripts") / "ncu_metrics_results.csv"
+    data = []
+    with open(csv_path, 'r') as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            parsed_row = {
+                'm': int(row['m']),
+                'k': int(row['k']),
+                'n': int(row['n']),
+                'latency_ms': float(row['latency_ms']),
+                'clock': float(row['clock_Ghz']) * 1e9,  # Convert GHz to Hz
+            }
+            data.append(parsed_row)
+    return data
+
+def _create_gemm_calc(configs):
+    """Helper function to create GEMMCalc instance from configs.
+    
+    Args:
+        configs: Tuple of (hw_config, model_config)
+        
+    Returns:
+        Initialized GEMMCalc instance
+    """
+    hw_config, model_config = configs
+    core = Core(hw_config)
+    mem_hierarchy = MemoryHierarchy(hw_config)
+    dtype_bytes = hw_config.sw_config.precision.activations
+    flashattn_enable = False
+    return GEMMCalc(core, mem_hierarchy.memLayer, dtype_bytes, flashattn_enable)
+
+def find_row(
+    csv_data: List[Dict],
+    m: int,
+    k: int,
+    n: int
+) -> Dict | None:
+    """Find a row in CSV data matching given m, k, n dimensions."""
+    for row in csv_data:
+        if row['m'] == m and row['k'] == k and row['n'] == n:
+            return row
+    return None
+
+def compare(
+    gemm_calc: GEMMCalc,
+    csv_validation_data: List[Dict],
+    shapes_file: str,
+):
+    models_data = parse_llm_shapes_file(shapes_file)
+    gemm_configs = extract_gemm_shapes(models_data)
+
+    data = {}
+
+    for cfg in gemm_configs:
+        if cfg['op_name'] == "router":
+            continue  # Skip router operations
+        ncu_data = find_row(
+            csv_validation_data,
+            cfg['m'],
+            cfg['k'],
+            cfg['n']
+        )
+
+        num_heads = cfg.get('batch', 1)
+
+        ncu_time = ncu_data['latency_ms']
+        clock = ncu_data['clock']
+        derate = clock / gemm_calc.core.operating_freq
+
+        # Run GEMMCalc
+        best_gemm, best_time = gemm_calc.run(cfg['m'], cfg['k'], cfg['n'], name=f"{cfg['model']}_{cfg['op_name']}", throughput=gemm_calc.core.getThroughput() * derate)
+
+        best_time = (best_time + 9e-6) * 1e3
+
+        model_dict = data.get(f"{cfg['model']}", {})
+        model_dict[f"{cfg['op_name']}"] = {
+            'm': cfg['m'],
+            'k': cfg['k'],
+            'n': cfg['n'],
+            'ncu_time_ms': ncu_time * num_heads,
+            'gemmcalc_time_ms': best_time * num_heads,
+            'err_pct': abs(ncu_time - best_time) / ncu_time * 100,
+        }
+        data[f"{cfg['model']}"] = model_dict
+    return data
+
+def plot_stacked_bar_chart(results: Dict[str, Dict[str, Dict]]):
+    """Create a double stacked bar chart comparing NCU and GEMMCalc times.
+    
+    Args:
+        results: Dictionary with structure {'model': {'op_name': {'ncu_time_ms', 'gemmcalc_time_ms', ...}}}
+    """
+    # Split models into llama2 and llama3.1 groups
+    llama2_models = {k: v for k, v in results.items() if 'llama2' in k.lower()}
+    llama31_models = {k: v for k, v in results.items() if 'llama3.1' in k.lower() or 'llama3_1' in k.lower()}
+    
+    # Get all unique operation names across all models
+    all_ops = sorted(set(op for ops in results.values() for op in ops.keys()))
+    colors = plt.cm.tab20(np.linspace(0, 1, len(all_ops)))
+    
+    # Create figure with two subplots
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 8))
+    
+    # Plot llama2 models
+    _plot_model_group(ax1, llama2_models, all_ops, colors, 'Llama2 Models (Batch Size 16)')
+    
+    # Plot llama3.1 models
+    _plot_model_group(ax2, llama31_models, all_ops, colors, 'Llama3.1 Models (Batch Size 1)')
+    
+    # Create shared legend
+    legend_elements = [plt.Rectangle((0,0),1,1, fc=colors[i], label=op) 
+                      for i, op in enumerate(STACK_ORDER)]
+    legend_elements.insert(0, plt.Rectangle((0,0),1,1, fc='white', ec='black', 
+                                            label='Actual'))
+    legend_elements.insert(1, plt.Rectangle((0,0),1,1, fc='white', ec='black', hatch='///',
+                                            label='DeepFlowSim'))
+    fig.legend(handles=legend_elements, loc='upper center', bbox_to_anchor=(0.5, 0.96), 
+               ncol=len(legend_elements)//2 + len(legend_elements)%2, frameon=True)
+    
+    fig.suptitle('aggregated error weighed by time per GEMM', fontsize=16, fontweight='bold', y=0.99)
+    
+    plt.subplots_adjust(top=0.85)
+    
+    # Save the plot
+    output_path = Path("validation_scripts") / "comparison_plot.png"
+    plt.savefig(output_path, dpi=300, bbox_inches='tight')
+    print(f"Plot saved to {output_path}")
+
+def _plot_model_group(ax, model_dict: Dict[str, Dict[str, Dict]], all_ops: List[str], colors, title: str):
+    """Helper function to plot a group of models on a given axis.
+    
+    Args:
+        ax: Matplotlib axis to plot on
+        model_dict: Dictionary of models and their operations
+        all_ops: List of all operation names
+        colors: Color map for operations
+        title: Title for the subplot
+    """
+    if not model_dict:
+        ax.set_visible(False)
+        return
+    
+    models = list(model_dict.keys())
+    num_models = len(models)
+    
+    # Prepare data structures
+    ncu_times = {}
+    gemmcalc_times = {}
+    
+    for model, ops in model_dict.items():
+        ncu_times[model] = {op: data['ncu_time_ms'] for op, data in ops.items()}
+        gemmcalc_times[model] = {op: data['gemmcalc_time_ms'] for op, data in ops.items()}
+    
+    x = np.arange(num_models)
+    width = 0.25
+    gap = 0.05
+    
+    # Create stacked bars for each operation
+    ncu_bottoms = np.zeros(num_models)
+    gemmcalc_bottoms = np.zeros(num_models)
+    
+    for idx, op in enumerate(reversed(STACK_ORDER)):
+        ncu_values = [ncu_times[model].get(op, 0) for model in models]
+        gemmcalc_values = [gemmcalc_times[model].get(op, 0) for model in models]
+        
+        ax.bar(x - (width+gap)/2, ncu_values, width, bottom=ncu_bottoms, color=colors[len(all_ops) - idx - 1])
+        ax.bar(x + (width+gap)/2, gemmcalc_values, width, bottom=gemmcalc_bottoms, 
+               color=colors[len(all_ops) - idx - 1], hatch='///', edgecolor='white', linewidth=0.5)
+        
+        ncu_bottoms += ncu_values
+        gemmcalc_bottoms += gemmcalc_values
+    
+    # Add overall error annotations with vertical brackets
+    for i, model in enumerate(models):
+        total_ncu = ncu_bottoms[i]
+        total_gemmcalc = gemmcalc_bottoms[i]
+        overall_error = abs(total_ncu - total_gemmcalc) / total_ncu * 100 if total_ncu > 0 else 0
+        
+        # Draw vertical bracket showing difference
+        min_height = min(total_ncu, total_gemmcalc)
+        max_height = max(total_ncu, total_gemmcalc)
+        
+        # Bracket position (to the right of the bars)
+        bracket_x = x[i]
+        bracket_width = 0.02
+        
+        # Draw vertical line
+        ax.plot([bracket_x, bracket_x], [min_height, max_height], 
+                'k-', linewidth=1.5, solid_capstyle='butt')
+        # Draw top horizontal line
+        ax.plot([bracket_x - bracket_width/2, bracket_x + bracket_width/2], [max_height, max_height], 
+                'k-', linewidth=1.5)
+        # Draw bottom horizontal line
+        ax.plot([bracket_x - bracket_width/2, bracket_x + bracket_width/2], [min_height, min_height], 
+                'k-', linewidth=1.5)
+        
+        # Position error text beside the bracket
+        mid_height = (min_height + max_height) / 2
+        ax.text(bracket_x, max_height, f'err: {overall_error:.1f}%', 
+                ha='center', va='bottom', fontsize=10, fontweight='bold')
+    
+    # Customize the subplot
+    # ax.set_xlabel('Model', fontsize=12, fontweight='bold')
+    ax.set_ylabel('Time (ms)', fontsize=12, fontweight='bold')
+    ax.set_title(title, fontsize=14, fontweight='bold')
+    ax.set_xticks(x)
+    ax.set_xticklabels(models)
+    ax.grid(axis='y', alpha=0.3)
+
+if __name__ == "__main__":
+    hw_config, model_config = configs()
+    gemm_calc = _create_gemm_calc((hw_config, model_config))
+    csv_data = csv_validation_data()
+    shapes_file = Path("validation_scripts") / "llm_shapes.txt"
+    
+    results = compare(gemm_calc, csv_data, shapes_file)
+    
+    # pp = pprint.PrettyPrinter(indent=4)
+    # pp.pprint(results)
+    
+    # Create visualization
+    plot_stacked_bar_chart(results)
\ No newline at end of file
diff --git a/model.py b/model.py
index f1678270..e913e67b 100644
--- a/model.py
+++ b/model.py
@@ -20,7 +20,8 @@ def __init__(self, exp_config):
       
 class Model_LLM:
   def __init__(self, exp_config):
-      self.batch_size       = exp_config.model_config.batch_size
+      self.global_batch_size = exp_config.model_config.global_batch_size
+      self.gradient_accumulation_steps = exp_config.model_config.gradient_accumulation_steps
       self.vocab_size       = exp_config.model_config.vocab_size
       self.num_layers      = exp_config.model_config.num_layers
       self.hidden_dim       = exp_config.model_config.hidden_dim
@@ -41,7 +42,7 @@ def __init__(self, exp_config):
 
       self.moe_num_experts = int(getattr(exp_config.model_config, "num_experts", 1))
       self.moe_top_k = int(getattr(exp_config.model_config, "top_k", 1))
-      if self.moe_top_k > self.moe_num_experts:
+      if self.moe_top_k > self.moe_num_experts and self.moe_num_experts > 1:
           raise ValueError("model_param.top_k cannot exceed model_param.num_experts")
       self.use_moe = self.moe_num_experts > 1
       
diff --git a/pyproject.toml b/pyproject.toml
index 0bee730a..ecfe9c9c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,12 +12,18 @@ dependencies = [
     "pandas==2.3.1",
     "protobuf==5.28.2",
     "pydot>=3.0.0",
+    "pytest>=9.0.0",
     "python-dateutil==2.9.0.post0",
     "pytz==2025.2",
     "pyyaml==6.0.2",
-    "ruamel-yaml==0.17.40",
-    "ruamel-yaml-clib==0.2.8",
+    "scotchpy64>=1.0.1",
     "seaborn>=0.13.2",
     "six==1.17.0",
+    "tqdm>=4.67.1",
     "tzdata==2025.2",
 ]
+
+[tool.pytest.ini_options]
+testpaths = [
+    "tests",
+]
diff --git a/requirements.txt b/requirements.txt
index 524f5a5e..03d8cc60 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,12 +4,11 @@ pandas==2.3.1
 python-dateutil==2.9.0.post0
 pytz==2025.2
 pyyaml==6.0.2
-ruamel-yaml==0.17.40
-ruamel-yaml-clib==0.2.8
 six==1.17.0
 tzdata==2025.2
 protobuf==5.28.2
 pydot==3.0.0
 graphviz==0.21
 joblib>=1.5.2
-seaborn
\ No newline at end of file
+seaborn
+scotchpy64>=1.0.1
diff --git a/run_perf.py b/run_perf.py
index 28020b25..65ee8a9a 100755
--- a/run_perf.py
+++ b/run_perf.py
@@ -10,6 +10,8 @@
 import pandas as pd
 import yaml
 import shutil
+import util
+from util import log_message, flush_log_queue, extend_log
 
 import graphviz_async
 from tile import TiledGEMM, formatBytes
@@ -17,11 +19,6 @@
 from time_calculation_LLM import TimeCalculationLLM
 from time_calculation_inf import TimeCalculationLLMInference
 
-algByte = False  # algorithmic ops false
-proj = False  # consider projection layer, turn off for end-2-end validation, as baeline model does not have projection layer
-validating_v100 = True
-
-
 # Cache handling policy for AstraSim integration.
 # Options: "NO CACHE", "CACHE READONLY", "CACHE READWRITE"
 cache_handling = "CACHE READWRITE"
@@ -48,6 +45,7 @@ def _report_total_wall_time() -> None:
         # Best-effort only
         pass
 
+atexit.register(flush_log_queue)
 atexit.register(_report_total_wall_time)
 
 def parse_arguments():
@@ -85,17 +83,17 @@ def _validate_astrasim_dependencies(hw_config) -> None:
 def _validate_network_topology(hw_config) -> None:
     backend = getattr(hw_config, "execution_backend", None)
     model = getattr(backend, "model", "analytical") if backend else "analytical"
-    network_topology = getattr(hw_config, "network_topology", None)
+    network_layout = getattr(hw_config, "network_layout", None)
 
-    if str(model).lower() == "analytical" and network_topology:
-        inter_topology = getattr(network_topology.inter, "topology", "ring")
-        intra_topology = getattr(network_topology.intra, "topology", "ring")
+    if str(model).lower() == "analytical" and network_layout:
+        for dim in getattr(network_layout, "dimensions", ()):
+            topo = str(getattr(dim, "topology_type", "ring")).lower()
+            if topo != "ring":
+                raise RuntimeError(
+                    "Non-ring network topologies are not supported in analytical mode. "
+                    "Only execution_backend.model='astra' (requires a valid AstraSim install) supports non-ring networks."
+                )
 
-        if str(inter_topology).lower() != "ring" or str(intra_topology).lower() != "ring":
-            raise RuntimeError(
-                "Non-ring network topologies are not supported in analytical mode. "
-                "Only execution_backend.model='astra' (requires a valid AstraSim install) supports non-ring networks."
-            )
 
 def run_LSTM(
     exp_hw_config_path,
@@ -238,6 +236,7 @@ def _run_llm_training(exp_hw_config, exp_model_config, exp_dir, mode):
     output_file = os.path.join(exp_dir, "LLM_training_results.txt")
     tc_llm = TimeCalculationLLM(exp_hw_config, exp_model_config, mode, output_dir=exp_dir)
     total_time = tc_llm.calc_time_llm()
+    topology_lines = util.network_topology_summary_training(exp_hw_config)
 
     with open(output_file, "a+") as handle:
         handle.write("\n\n==============================================\n")
@@ -246,13 +245,16 @@ def _run_llm_training(exp_hw_config, exp_model_config, exp_dir, mode):
         handle.write("Execution Mode: {}\n".format(tc_llm.execution_mode.value))
         handle.write("Total Time: {0:.8f}\n".format(total_time))
         handle.write("\n")
-        handle.write("For more info, turn on debug flags. See examples/llm_astra_inference_debug_graphviz.sh")
+        handle.write("For more info, turn on debug flags. See examples/llm_astra_inference_debug_graphviz.sh\n")
+        handle.write("\n".join(topology_lines))
+        handle.write("\n")
 
-    print("Training time for batch: {:.2f}s".format(tc_llm.get_time()))
-    print("LLM training results written to {}".format(output_file))
+    log_message("Training time for batch: {:.2f}s".format(tc_llm.get_time()), category="results")
+    extend_log(topology_lines, category="network")
+    log_message("LLM training results written to {}".format(output_file), category="results")
     warning_message = tc_llm.memory_capacity_warning()
     if warning_message:
-        print(warning_message)
+        log_message(warning_message)
 
 
 def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode):
@@ -264,20 +266,22 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode):
     total_time = inference_timing["total_inference_time"]
     decode_rates = inference_timing.get("decode_tokens_per_s") or {}
 
-    print(
+    log_message(
         "LLM inference time: {:.2f}s (mode={})".format(
             total_time, tc_inf.execution_mode.value
-        )
+        ),
+        category="results",
     )
-    print(
+    log_message(
         "LLM time to first token: {:.2f}s".format(
             inference_timing["time_to_first_token"],
-        )
+        ),
+        category="results",
     )
     dp_replicas = max(1, getattr(tc_inf, "dp", 1))
     batch_size = getattr(tc_inf, "batch_size", 1)
     if dp_replicas > 1:
-        print(f"Data parallel replicas: {dp_replicas}")
+        log_message(f"Data parallel replicas: {dp_replicas}", category="results")
     if decode_rates:
         # decode_rates are per-generation rates (tokens per second per generation)
         start_gen_rate = decode_rates.get("start", 0.0)
@@ -286,17 +290,18 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode):
         mid_step = int(decode_rates.get("midpoint_step", 0.0))
         
         # Print per-generation rates
-        print(
+        log_message(
             "Decode sequences/s: start={:.2f}, mid(token {})={:.2f}, end={:.2f}".format(
                 start_gen_rate,
                 mid_step,
                 mid_gen_rate,
                 end_gen_rate,
-            )
+            ),
+            category="results",
         )
 
         # Print aggregate decode throughput (with batch_size and dp multipliers)
-        print(
+        log_message(
             "Aggregate decode throughput tok/s (batch={}, dp={}): start={:.2f}, mid(token {})={:.2f}, end={:.2f}".format(
                 batch_size,
                 dp_replicas,
@@ -304,9 +309,12 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode):
                 mid_step,
                 mid_gen_rate * batch_size * dp_replicas,
                 end_gen_rate * batch_size * dp_replicas,
-            )
+            ),
+            category="results",
         )
 
+    topology_lines = util.network_topology_summary_inference(exp_hw_config)
+
     output_path = os.path.join(exp_dir, "LLM_inference_results.txt")
     os.makedirs(exp_dir, exist_ok=True)
     with open(output_path, "w") as handle:
@@ -329,12 +337,15 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode):
             handle.write(f"Decode Generations per Second: start={start_gen_rate:.2f}, mid(token {mid_step})={mid_gen_rate:.2f}, end={end_gen_rate:.2f}\n")
             handle.write(f"Aggregate Decode Throughput Tok/s (batch={batch_size}, dp={dp_replicas}): start={start_gen_rate * batch_size * dp_replicas:.2f}, mid(token {mid_step})={mid_gen_rate * batch_size * dp_replicas:.2f}, end={end_gen_rate * batch_size * dp_replicas:.2f}\n")
         handle.write("\n")
-        handle.write("For more info, turn on debug flags. See examples/llm_astra_inference_debug_graphviz.sh")
+        handle.write("For more info, turn on debug flags. See examples/llm_astra_inference_debug_graphviz.sh\n")
+        handle.write("\n".join(topology_lines))
+        handle.write("\n")
 
-    print("LLM inference results written to {}".format(output_path))
+    extend_log(topology_lines, category="network")
+    log_message("LLM inference results written to {}".format(output_path), category="results")
     warning_message = tc_inf.memory_capacity_warning()
     if warning_message:
-        print(warning_message)
+        log_message(warning_message)
 
 if __name__ == "__main__":
     args = parse_arguments()
@@ -352,7 +363,6 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode):
     os.makedirs(exp_dir, exist_ok=True)
     
     if mode == "LLM":
-        print("Using LLM parameters for computation...")
         run_LLM(
             exp_hw_config_path=config_hardware_path,
             exp_model_config_path=config_model_path,
@@ -361,7 +371,7 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode):
         )
     
     elif mode == "LSTM":
-        print("Using LSTM parameters for computation...")
+        log_message("Using LSTM parameters for computation...")
         run_LSTM(
             exp_hw_config_path=config_hardware_path,
             exp_model_config_path=config_model_path,
@@ -370,7 +380,7 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode):
         
         )
     elif mode == "GEMM":
-        print("Using GEMM parameters for computation...")
+        log_message("Using GEMM parameters for computation...")
         run_GEMM(
             exp_hw_config_path=config_hardware_path,
             exp_model_config_path=config_model_path,
@@ -379,3 +389,7 @@ def _run_llm_inference(exp_hw_config, exp_model_config, exp_dir, mode):
         )
     else:
         print("Invalid mode selected. Please choose 'LLM', 'LSTM', or 'GEMM'.")
+        flush_log_queue()
+        sys.exit(1)
+
+    flush_log_queue()
diff --git a/scripts/install_astrasim.sh b/scripts/install_astrasim.sh
new file mode 100755
index 00000000..cc3074ae
--- /dev/null
+++ b/scripts/install_astrasim.sh
@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
+ASTRA_SIM_DIR="${REPO_ROOT}/astra-sim"
+
+cd "${REPO_ROOT}"
+
+echo "[astrasim] Updating DeepFlow submodules..."
+git submodule update --init --recursive
+git submodule update --remote -- astra-sim
+
+if [[ ! -d "${ASTRA_SIM_DIR}" ]]; then
+  echo "[astrasim] ERROR: expected submodule directory at ${ASTRA_SIM_DIR}" >&2
+  exit 1
+fi
+
+echo "[astrasim] Updating AstraSim nested submodules..."
+git -C "${ASTRA_SIM_DIR}" submodule update --remote -- extern/network_backend/analytical
+
+BUILD_SCRIPT="${ASTRA_SIM_DIR}/build/astra_analytical/build.sh"
+if [[ ! -x "${BUILD_SCRIPT}" ]]; then
+  echo "[astrasim] ERROR: build script not found or not executable at ${BUILD_SCRIPT}" >&2
+  exit 1
+fi
+
+echo "[astrasim] Building AstraSim (this may take a while)..."
+pushd "${ASTRA_SIM_DIR}" >/dev/null
+"${BUILD_SCRIPT}"
+popd >/dev/null
+
+echo "[astrasim] AstraSim build completed successfully."
diff --git a/simulate_LLM.py b/simulate_LLM.py
index 8ee46c0a..85c6310e 100644
--- a/simulate_LLM.py
+++ b/simulate_LLM.py
@@ -7,6 +7,7 @@
 import os
 
 import graphviz_async
+import util
 debug = False
 BYTES_PER_GIB = 1024 ** 3
 
@@ -15,7 +16,7 @@ def __init__(self, name, op_id, hw_id, duration, fwd=True):
         self.name = name
         self.op_id = op_id
         self.hw_id = hw_id
-        self.duration = duration
+        self._duration_data = self._normalize_duration(duration)
         self.done = False
         self.finish_time = -1
         self.parents = []
@@ -30,6 +31,39 @@ def add_child(self, obj):
 
     def __repr__(self):
         return f"Node({self.name},op={self.op_id},hw={self.hw_id},{round(self.duration, 2)})"
+
+    @staticmethod
+    def _normalize_duration(value):
+        if isinstance(value, (list, tuple)):
+            entries = tuple(float(v) for v in value)
+            if not entries:
+                raise ValueError("Duration tuple must contain at least one entry.")
+            return entries
+        return float(value)
+
+    @property
+    def duration(self):
+        data = self._duration_data
+        if isinstance(data, tuple):
+            return data[0]
+        return data
+
+    @duration.setter
+    def duration(self, value):
+        self._duration_data = self._normalize_duration(value)
+
+    @property
+    def duration_profile(self) -> Optional[Tuple[float, ...]]:
+        data = self._duration_data
+        if isinstance(data, tuple):
+            return data
+        return None
+
+    def duration_storage(self):
+        data = self._duration_data
+        if isinstance(data, tuple):
+            return tuple(data)
+        return data
     
 class Data_batch:
     def __init__(self, name, batch_id, duration):
@@ -95,19 +129,53 @@ def __init__(
         self.comp_times = comp_times or {}
         self.comm_metadata = comm_metadata or {}
         self.misc_metadata = misc_metadata or {}
+        self.full_recomputation = bool(self.misc_metadata.get("full_recomputation", False))
+        self.flattened_mode = bool(self.misc_metadata.get("flattened_mode", False))
+        self.pipeline_style_recompute = bool(self.misc_metadata.get("pipeline_style_recompute", False))
 
         self.num_batch = self.misc_metadata.get("num_batch", 0)
         self.num_layer = self.misc_metadata.get("num_layer", 0)
-        self.layer_per_device = max(1, math.ceil(self.num_layer / self.lp)) if self.lp else self.num_layer
+        self.layer_per_device = self.num_layer / self.lp if self.lp else self.num_layer
         self.all_reduce = self.misc_metadata.get("all_reduce", "every layer")
+        self.dp_microbatch_mode = str(self.misc_metadata.get("dp_microbatch_mode", "every_mb")).lower()
+        self.dp_zero_stage = int(self.misc_metadata.get("dp_zero_stage", 0) or 0)
 
         self.transformer_cfg = self.comp_times.get("transformer", {})
         self.T_grad_transformer = self.comp_times.get("grad_transformer", 0.0)
+        self.layers_per_stage = self._compute_layers_per_stage()
+        self.layer_to_stage = self._build_layer_to_stage()
 
     def _time(self, key: str, default: float = 0.0) -> float:
         value = self.comp_times.get(key)
         return float(value) if value is not None else default
 
+    def _compute_layers_per_stage(self) -> List[int]:
+        stage_count = max(1, int(self.lp) if self.lp else 1)
+        if self.num_layer <= 0:
+            return [0 for _ in range(stage_count)]
+        base = self.num_layer // stage_count
+        remainder = self.num_layer % stage_count
+        layers_per_stage: List[int] = []
+        for stage_idx in range(stage_count):
+            count = base + (1 if stage_idx < remainder else 0)
+            layers_per_stage.append(count)
+        return layers_per_stage
+
+    def _build_layer_to_stage(self) -> List[int]:
+        mapping: List[int] = []
+        for stage_idx, count in enumerate(self.layers_per_stage):
+            if count <= 0:
+                continue
+            mapping.extend([stage_idx] * count)
+        return mapping
+
+    def _stage_for_layer(self, layer_idx: int) -> int:
+        if layer_idx < 0 or layer_idx >= self.num_layer:
+            raise IndexError(f"Layer index {layer_idx} is out of bounds for {self.num_layer} layers.")
+        if self.layer_to_stage:
+            return self.layer_to_stage[layer_idx]
+        return 0
+
     def create_comm_edge(self, name, op_id, comm_key, is_dp=False, local_hw_id=None):
         """Create a communication edge with optional local computation node.
 
@@ -200,7 +268,7 @@ def ensure_clone(obj: Any, cls: str) -> Any:
                 return cached
 
             if isinstance(obj, Node):
-                clone = Node(obj.name, obj.op_id, obj.hw_id, obj.duration, fwd=obj.fwd)
+                clone = Node(obj.name, obj.op_id, obj.hw_id, obj.duration_storage(), fwd=obj.fwd)
                 clone.memory = getattr(obj, "memory", 0)
             elif isinstance(obj, Edge):
                 clone = Edge(
@@ -376,7 +444,7 @@ def ensure_clone(obj: Any) -> Any:
                 return cached
 
             if isinstance(obj, Node):
-                clone = Node(obj.name, obj.op_id, obj.hw_id, obj.duration, fwd=obj.fwd)
+                clone = Node(obj.name, obj.op_id, obj.hw_id, obj.duration_storage(), fwd=obj.fwd)
                 clone.memory = getattr(obj, "memory", 0)
             elif isinstance(obj, Edge):
                 clone = Edge(
@@ -515,7 +583,7 @@ def traverse_and_convert(node, visited=None):
         traverse_and_convert(roots)
         return roots
 
-    def construct_fwd_bwd_graph(self, include_backward: bool = True):
+    def construct_fwd_bwd_graph(self, include_backward: bool = True, include_optimizer: bool = False):
         embedding_node = []
         data_batch_node = []
         softmax_node = []
@@ -523,9 +591,6 @@ def construct_fwd_bwd_graph(self, include_backward: bool = True):
         layer_entry_nodes = [[] for _ in range(self.num_batch)]  # lists of entry nodes per layer
         layer_exit_nodes = [[] for _ in range(self.num_batch)]   # transformer nodes per layer
 
-        embedding_b_time = self._time("embedding_b")
-        linear_softmax_b_time = self._time("linear_softmax_b")
-        transformer_b_time = self._time("transformer_b")
         linear_softmax_f_time = self._time("linear_softmax_f")
         linear_softmax_b_time = self._time("linear_softmax_b")
         transformer_f_time = self._time("transformer_f")
@@ -533,6 +598,11 @@ def construct_fwd_bwd_graph(self, include_backward: bool = True):
         embedding_f_time = self._time("embedding_f")
         embedding_b_time = self._time("embedding_b")
         cross_layer_time = self._time("cross_layer_f")
+        recompute_enabled = include_backward and self.full_recomputation and (
+            self.flattened_mode or self.pipeline_style_recompute
+        )
+        transformer_b_base_time = transformer_b_time
+
 
         def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_comm_children=None):
             parents = list(getattr(target, "parents", []))
@@ -559,6 +629,9 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_
             transformer_nodes_b = [[] for _ in range(self.num_batch)]  #
             for b in range(self.num_batch):
                 transformer_nodes_b[b] = [[] for _ in range(self.num_layer)]
+            recompute_nodes_b: Optional[List[List[Optional[Node]]]] = None
+            if recompute_enabled:
+                recompute_nodes_b = [[None for _ in range(self.num_layer)] for _ in range(self.num_batch)]
 
         op_id = 0  # operation ID, used to distinguish nodes and edges
         batch_id = 0  # batch ID, used to distinguish data batches
@@ -585,7 +658,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_
             layer_exit_nodes[b] = []
 
             for l in range(self.num_layer):
-                hw_id = min(l // self.layer_per_device , self.lp - 1)#assign hw_id for transformer
+                hw_id = self._stage_for_layer(l)
                 transformer_node = Node(f"transformer_layer{l}", op_id, hw_id, transformer_f_time)
                 transformer_node.micro_batch_index = b
                 transformer_node.layer_index = l
@@ -663,6 +736,16 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_
         if not include_backward:
             return embedding_node[0]
 
+        def _bwd_entry_node(mb_idx: int, layer_idx: int) -> Any:
+            if recompute_enabled and recompute_nodes_b is not None:
+                candidate = recompute_nodes_b[mb_idx][layer_idx]
+                if candidate is not None:
+                    return candidate
+            return transformer_nodes_b[mb_idx][layer_idx]
+
+        def _bwd_exit_node(mb_idx: int, layer_idx: int) -> Any:
+            return transformer_nodes_b[mb_idx][layer_idx]
+
         for b in reversed(range(self.num_batch)): #connect each data batch node with corresponding nodes
             emb_b = Node("embedding_b", op_id, 0, embedding_b_time, fwd=False)      # hw_id = 0
             op_id += 1
@@ -674,18 +757,42 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_
 
 
             for l in reversed(range(self.num_layer)):
-                hw_id = min(l // self.layer_per_device , self.lp - 1)
-                transformer_node_b = Node(f"transformer_layer{l}_b", op_id, hw_id, transformer_b_time, fwd=False)
+                hw_id = self._stage_for_layer(l)
+                recompute_node = None
+                if recompute_enabled and recompute_nodes_b is not None:
+                    recompute_node = Node(
+                        f"transformer_layer{l}_recompute",
+                        op_id,
+                        hw_id,
+                        transformer_f_time,
+                        fwd=True,
+                    )
+                    recompute_node.micro_batch_index = b
+                    recompute_node.layer_index = l
+                    recompute_node.direction = "forward"
+                    recompute_node.stage_id = hw_id
+                    op_id += 1
+                    recompute_nodes_b[b][l] = recompute_node
+
+                transformer_node_b = Node(
+                    f"transformer_layer{l}_b",
+                    op_id,
+                    hw_id,
+                    transformer_b_base_time,
+                    fwd=False,
+                )
                 transformer_node_b.micro_batch_index = b
                 transformer_node_b.layer_index = l
                 transformer_node_b.direction = "backward"
                 transformer_node_b.stage_id = hw_id
                 op_id += 1
                 transformer_nodes_b[b][l] = transformer_node_b
+                if recompute_node is not None:
+                    recompute_node.add_child(transformer_node_b)
 
             for l in reversed(range(1, self.num_layer)):
-                curr_node = transformer_nodes_b[b][l]         # current layer's qkv_proj.
-                next_ffn2  = transformer_nodes_b[b][l-1]            # next layer's layernorm.
+                curr_node = _bwd_exit_node(b, l)         # current layer's qkv_proj.
+                next_ffn2  = _bwd_entry_node(b, l-1)            # next layer's layernorm.
                 if curr_node.hw_id == next_ffn2.hw_id:
                     edge = Edge("cross_layer", op_id, 0, comm_type="pipeline")  
                 else:
@@ -694,7 +801,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_
 
                 curr_node.add_child(edge); edge.add_child(next_ffn2)
 
-            qkv_0_b = transformer_nodes_b[b][0]     # first layer's qkv_proj
+            qkv_0_b = _bwd_exit_node(b, 0)     # first layer's qkv_proj
             if qkv_0_b.hw_id == emb_b.hw_id:
                 edge = Edge("Emb_node0", op_id, 0, comm_type="pipeline")
             else:
@@ -704,7 +811,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_
             edge.add_child(emb_b)
 
 
-            prev_layer_norm2 = transformer_nodes_b[b][self.num_layer-1] # last layer's layernorm2
+            prev_layer_norm2 = _bwd_entry_node(b, self.num_layer-1) # last layer's layernorm2
             if prev_layer_norm2.hw_id == softmax_node_b[b].hw_id:
                 layernorm_Softmax = Edge("layernorm2_Softmax", op_id, 0, comm_type="pipeline")  # same GPU
             else:
@@ -731,6 +838,9 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_
         for b in range(self.num_batch):
             # Data parallel collectives (all reduce for DDP/ZeRO-1, reduce-scatter + all-gather for ZeRO-2/3)
             if self.dp > 1:
+                apply_dp_all_mbs = self.dp_microbatch_mode != "last_mb" or self.dp_zero_stage >= 3
+                if not apply_dp_all_mbs and b != 0: # mb is backwards in backpass
+                    continue
                 zero2_embedding_key = "zero2_embedding_gather"
                 zero2_transformer_key = "zero2_transformer_gather"
                 zero2_softmax_key = "zero2_softmax_gather"
@@ -768,7 +878,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_
                         op_id,
                         "transformer",
                         is_dp=True,
-                        local_hw_id=transformer_nodes_b[b][layer_idx].hw_id,
+                        local_hw_id=_bwd_exit_node(b, layer_idx).hw_id,
                     )
                     R_edge[b].append(reducer)
                     op_id += 1
@@ -779,7 +889,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_
                             op_id,
                             zero2_transformer_key,
                             is_dp=True,
-                            local_hw_id=transformer_nodes_b[b][layer_idx].hw_id,
+                            local_hw_id=_bwd_exit_node(b, layer_idx).hw_id,
                         )
                         op_id += 1
                         reducer.add_child(gather_edge)
@@ -863,7 +973,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_
                 softmax_node_b[b].add_child(R_edge[b][-1])
                 embedding_node_b[b].add_child(R_edge[b][0])
                 for layer_idx in range(self.num_layer):
-                    transformer_nodes_b[b][layer_idx].add_child(R_edge[b][layer_idx + 1])
+                    _bwd_exit_node(b, layer_idx).add_child(R_edge[b][layer_idx + 1])
 
 
         last_transformer_layer = [-1] * self.lp  # Initialize with -1 for all GPUs
@@ -872,8 +982,8 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_
         # first_transformer_layer.append(0)
         gpu_index = self.lp - 1
         for l in range(self.num_layer - 1, 0, -1):
-            if transformer_nodes_b[0][l].hw_id != transformer_nodes_b[0][l-1].hw_id:  # Check if on different GPU
-                # print("Layer ", l, " is on GPU ", transformer_nodes_b[0][l].hw_id)
+            if _bwd_exit_node(0, l).hw_id != _bwd_exit_node(0, l-1).hw_id:  # Check if on different GPU
+                # print("Layer ", l, " is on GPU ", _bwd_exit_node(0, l).hw_id)
                 first_transformer_layer[gpu_index-1] = l-1  # Record first layer on each GPU
                 last_transformer_layer[gpu_index] = l  # Record last layer on each GPU
                 gpu_index -= 1
@@ -885,18 +995,18 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_
         for b in range(self.num_batch-1, 0, -1):
             gpu_index = self.lp - 1
             for l in range(self.num_layer - 1, 0, -1):
-                if transformer_nodes_b[b][l].hw_id != transformer_nodes_b[b][l-1].hw_id:  # Check if on different GPUs
+                if _bwd_exit_node(b, l).hw_id != _bwd_exit_node(b, l-1).hw_id:  # Check if on different GPUs
                     # last_transformer_layer.append(l)  # Record last layer on each GPU
                     # first_transformer_layer.append(l-1)  # Record first layer on each GPU
                     
-                    if transformer_nodes_b[b][l].hw_id == self.lp - 1:
-                        transformer_nodes_b[b][l].add_child(softmax_node_b[b-1])  # Add dependency edge
+                    if _bwd_exit_node(b, l).hw_id == self.lp - 1:
+                        _bwd_exit_node(b, l).add_child(softmax_node_b[b-1])  # Add dependency edge
                     else:
-                        transformer_nodes_b[b][l].add_child(transformer_nodes_b[b-1][first_transformer_layer[gpu_index]])  # Add dependency edge
+                        _bwd_exit_node(b, l).add_child(_bwd_entry_node(b-1, first_transformer_layer[gpu_index]))  # Add dependency edge
                     gpu_index -= 1
             # Ensure embedding_node_b[b] is connected to the correct transformer node
             # if first_transformer_layer:
-            embedding_node_b[b].add_child(transformer_nodes_b[b-1][first_transformer_layer[0]])  # Add dependency edge
+            embedding_node_b[b].add_child(_bwd_entry_node(b-1, first_transformer_layer[0]))  # Add dependency edge
 
         zero3_embedding_key = "zero3_embedding_gather"
         zero3_transformer_key = "zero3_transformer_gather"
@@ -930,8 +1040,8 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_
 
             if zero3_transformer_key in self.comm_metadata and self.num_layer > 0:
                 for layer_idx in reversed(range(self.num_layer)):
-                    host = softmax_node_b[b] if layer_idx == self.num_layer - 1 else transformer_nodes_b[b][layer_idx + 1]
-                    target = transformer_nodes_b[b][layer_idx]
+                    host = softmax_node_b[b] if layer_idx == self.num_layer - 1 else _bwd_exit_node(b, layer_idx + 1)
+                    target = _bwd_entry_node(b, layer_idx)
                     gather_edge = self.create_comm_edge(
                         f"{zero3_transformer_key}_b{b}_layer{layer_idx}_bwd",
                         op_id,
@@ -951,7 +1061,7 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_
 
 
             if zero3_embedding_key in self.comm_metadata and self.num_layer > 0:
-                host = transformer_nodes_b[b][0]
+                host = _bwd_exit_node(b, 0)
                 gather_edge = self.create_comm_edge(
                     f"{zero3_embedding_key}_b{b}_bwd",
                     op_id,
@@ -962,6 +1072,37 @@ def attach_parallel_edge(target, gather_edge, skip_non_comm_children=None, skip_
                 op_id += 1
                 attach_parallel_edge(host, gather_edge)
 
+        if include_backward and include_optimizer:
+            optimizer_time = self._time("optimizer")
+            if optimizer_time > 0:
+                # Find the last node for each stage to attach the optimizer
+                # For stage 0, it's the embedding backward of the last microbatch
+                # For other stages, it's the first transformer layer (lowest index) of that stage for the last microbatch
+                
+                # Helper to find min layer for each stage
+                stage_min_layer = {}
+                for l in range(self.num_layer):
+                    stage = self._stage_for_layer(l)
+                    if stage not in stage_min_layer:
+                        stage_min_layer[stage] = l
+                    else:
+                        stage_min_layer[stage] = min(stage_min_layer[stage], l)
+                                
+                # bwd flips mbs, so we attach to first one, not last.
+                for stage in range(self.lp):
+                    last_node = None
+                    if stage == 0:
+                        last_node = embedding_node_b[0]
+                    else:
+                        min_l = stage_min_layer.get(stage)
+                        if min_l is not None:
+                            last_node = _bwd_exit_node(0, min_l)
+                    
+                    if last_node:
+                        opt_node = Node(f"optimizer_stage{stage}", op_id, stage, optimizer_time, fwd=False)
+                        op_id += 1
+                        last_node.add_child(opt_node)
+
         return root_forward_entry
 
     def construct_transformer_graph(self, direction: str = "both"):
@@ -970,85 +1111,96 @@ def construct_transformer_graph(self, direction: str = "both"):
         if not gemm_entries:
             raise ValueError("Transformer GEMM times not provided")
 
-        tp_degree = self.tp
+        tp_degree = max(1, int(self.tp))
+        cp_degree = max(1, int(self.cp))
 
         root = Data_batch("transformer_root", 0, 0)
         op_id = 0
 
-        for rank in range(tp_degree):
-            previous = root
-
-            if direction in {"forward", "both"}:
-                for idx, entry in enumerate(gemm_entries):
-                    entry_name = entry.get("name", f"g{idx}")
-                    forward_cfg = entry.get("forward", {})
-                    fwd_duration = forward_cfg.get("duration")
-                    if fwd_duration is None:
-                        raise ValueError("Transformer GEMM entry missing forward duration")
-
-                    node = Node(
-                        name=f"{entry_name}_fwd_rank{rank}",
-                        op_id=op_id,
-                        hw_id=rank,
-                        duration=fwd_duration,
-                        fwd=True,
-                    )
-                    op_id += 1
-                    previous.add_child(node)
-                    previous = node
-                    comm_key = forward_cfg.get("comm_keys",None)
-                    if comm_key:
-                        if len(comm_key) > 1:
-                            print(f"*CONSTRUCT INFO: Multiple comm keys for {entry_name}")
-                            print(f"*CONSTRUCT INFO: Comm keys: {comm_key}")
-                            raise ValueError(f"Multiple comm keys for {entry_name}")
-
-                        comm_key = comm_key[0]
-                        if comm_key not in self.comm_metadata:
-                            raise KeyError(f"Missing transformer comm metadata for key '{comm_key}'")
-                        comm_edge = self.create_comm_edge(
-                            name=comm_key,
+        for cp_idx in range(cp_degree):
+            for tp_idx in range(tp_degree):
+                rank = tp_idx + cp_idx * tp_degree
+                previous = root
+
+                if direction in {"forward", "both"}:
+                    for idx, entry in enumerate(gemm_entries):
+                        entry_name = entry.get("name", f"g{idx}")
+                        forward_cfg = entry.get("forward", {})
+                        fwd_duration = forward_cfg.get("duration")
+                        if fwd_duration is None:
+                            raise ValueError("Transformer GEMM entry missing forward duration")
+
+                        node = Node(
+                            name=f"{entry_name}_fwd_rank{rank}",
                             op_id=op_id,
-                            comm_key=comm_key,
-                            is_dp=False,
-                            local_hw_id=rank,
+                            hw_id=rank,
+                            duration=fwd_duration,
+                            fwd=True,
                         )
+                        node.tp_rank = tp_idx
+                        node.cp_rank = cp_idx
                         op_id += 1
-                        previous.add_child(comm_edge)
-                        previous = comm_edge
-
-            if direction in {"backward", "both"}:
-                for idx, entry in enumerate(reversed(gemm_entries)):
-                    entry_name = entry.get("name", f"g{idx}")
-                    backward_cfg = entry.get("backward", {})
-                    bwd_duration = backward_cfg.get("duration")
-                    if bwd_duration is None:
-                        raise ValueError("Transformer GEMM entry missing backward duration")
-
-                    node = Node(
-                        name=f"{entry_name}_bwd_rank{rank}",
-                        op_id=op_id,
-                        hw_id=rank,
-                        duration=bwd_duration,
-                        fwd=False,
-                    )
-                    op_id += 1
-                    previous.add_child(node)
-                    previous = node
-
-                    for comm_idx, comm_key in enumerate(backward_cfg.get("comm_keys", [])):
-                        if comm_key not in self.comm_metadata:
-                            raise KeyError(f"Missing transformer comm metadata for key '{comm_key}'")
-                        comm_edge = self.create_comm_edge(
-                            name=comm_key,
+                        previous.add_child(node)
+                        previous = node
+                        comm_key = forward_cfg.get("comm_keys", None)
+                        if comm_key:
+                            if len(comm_key) > 1:
+                                print(f"*CONSTRUCT INFO: Multiple comm keys for {entry_name}")
+                                print(f"*CONSTRUCT INFO: Comm keys: {comm_key}")
+                                raise ValueError(f"Multiple comm keys for {entry_name}")
+
+                            comm_key = comm_key[0]
+                            if comm_key not in self.comm_metadata:
+                                raise KeyError(f"Missing transformer comm metadata for key '{comm_key}'")
+                            comm_edge = self.create_comm_edge(
+                                name=comm_key,
+                                op_id=op_id,
+                                comm_key=comm_key,
+                                is_dp=False,
+                                local_hw_id=rank,
+                            )
+                            comm_edge.tp_rank = tp_idx
+                            comm_edge.cp_rank = cp_idx
+                            op_id += 1
+                            previous.add_child(comm_edge)
+                            previous = comm_edge
+
+                if direction in {"backward", "both"}:
+                    for idx, entry in enumerate(reversed(gemm_entries)):
+                        entry_name = entry.get("name", f"g{idx}")
+                        backward_cfg = entry.get("backward", {})
+                        bwd_duration = backward_cfg.get("duration")
+                        if bwd_duration is None:
+                            raise ValueError("Transformer GEMM entry missing backward duration")
+
+                        node = Node(
+                            name=f"{entry_name}_bwd_rank{rank}",
                             op_id=op_id,
-                            comm_key=comm_key,
-                            is_dp=False,
-                            local_hw_id=rank,
+                            hw_id=rank,
+                            duration=bwd_duration,
+                            fwd=False,
                         )
+                        node.tp_rank = tp_idx
+                        node.cp_rank = cp_idx
                         op_id += 1
-                        previous.add_child(comm_edge)
-                        previous = comm_edge
+                        previous.add_child(node)
+                        previous = node
+
+                        for comm_idx, comm_key in enumerate(backward_cfg.get("comm_keys", [])):
+                            if comm_key not in self.comm_metadata:
+                                raise KeyError(f"Missing transformer comm metadata for key '{comm_key}'")
+                            comm_edge = self.create_comm_edge(
+                                name=comm_key,
+                                op_id=op_id,
+                                comm_key=comm_key,
+                                is_dp=False,
+                                local_hw_id=rank,
+                            )
+                            comm_edge.tp_rank = tp_idx
+                            comm_edge.cp_rank = cp_idx
+                            op_id += 1
+                            previous.add_child(comm_edge)
+                            previous = comm_edge
 
         return root
         
@@ -1121,7 +1273,7 @@ def simulate(self, root):
                         print("child {}  ready at time {} ".format(child.name, time))
 
             if isinstance(event, Node):
-                GPU_list[event.hw_id] = True
+                GPU_list[int(event.hw_id)] = True
 
                 
 
@@ -1143,7 +1295,7 @@ def simulate(self, root):
                     ready_list.remove(event)
 
                 elif isinstance(event, Node): 
-                    if GPU_list[event.hw_id] == True:
+                    if GPU_list[int(event.hw_id)] == True:
                         new_time = time + event.duration
                         heappush(event_queue, (new_time, counter, event))
                         event.scheduled = True
@@ -1151,7 +1303,7 @@ def simulate(self, root):
                         if debug:
                             print("{}.{} enqueued at time {} at device {}".format(event.name, event.op_id, time, event.hw_id))
                         counter = counter + 1
-                        GPU_list[event.hw_id] = False
+                        GPU_list[int(event.hw_id)] = False
                         ready_list.remove(event)
                 elif isinstance(event, Edge): 
                     new_time = time + event.duration
@@ -1422,20 +1574,20 @@ def _count_transformer_layers_per_device(root_obj: Any, num_devices: int) -> Lis
                         print("child {}  ready at time {} ".format(child.name, time))
 
             if isinstance(event, Node):
-                GPU_list[event.hw_id] = True
+                GPU_list[int(event.hw_id)] = True
                 is_transformer_block = _is_transformer_block(event) #TODO: embedding and softmax layers
                 if mode == "training":
                     if not event.fwd:
                         if is_transformer_block:
-                            memory_snapshot.release_activation(event.hw_id, event, time)
+                            memory_snapshot.release_activation(int(event.hw_id), event, time)
                     elif event.fwd:
                         if is_transformer_block:
-                            memory_snapshot.allocate_activation(event.hw_id, event, time)
+                            memory_snapshot.allocate_activation(int(event.hw_id), event, time)
                 elif mode == "inference":
                     if event.fwd:
                         if is_transformer_block:
-                            memory_snapshot.allocate_activation(event.hw_id, event, time)
-                            memory_snapshot.release_activation(event.hw_id, event, time)
+                            memory_snapshot.allocate_activation(int(event.hw_id), event, time)
+                            memory_snapshot.release_activation(int(event.hw_id), event, time)
 
                 
             for event in ready_list[:]:
@@ -1454,7 +1606,7 @@ def _count_transformer_layers_per_device(root_obj: Any, num_devices: int) -> Lis
                     ready_list.remove(event)
 
                 elif isinstance(event, Node): 
-                    if GPU_list[event.hw_id] == True:
+                    if GPU_list[int(event.hw_id)] == True:
                         new_time = time + event.duration
                         heappush(event_queue, (new_time, counter, event))
                         event.scheduled = True
@@ -1462,7 +1614,7 @@ def _count_transformer_layers_per_device(root_obj: Any, num_devices: int) -> Lis
                         if debug:
                             print("{}.{} enqueued at time {} at device {}".format(event.name, event.op_id, time, event.hw_id))
                         counter = counter + 1
-                        GPU_list[event.hw_id] = False
+                        GPU_list[int(event.hw_id)] = False
                         ready_list.remove(event)
                 elif isinstance(event, Edge): 
                     new_time = time + event.duration
@@ -1484,25 +1636,52 @@ def _count_transformer_layers_per_device(root_obj: Any, num_devices: int) -> Lis
     def save_graph(self, roots, output_folder = "output/LLM/", filename="graph"):
         os.makedirs(output_folder, exist_ok=True)
 
-        printstr = " | Graph saved to    %s%s.png" % (output_folder, filename)
+        base_path = os.path.normpath(f"{output_folder}{filename}")
+        svg_path = f"{base_path}.svg"
+        display_path = util.relpath_display(svg_path)
+        printstr = f" | Graph saved to    {display_path}"
         def _render_graph() -> None:
-            dot_fw = visualize_graph(roots, filename=output_folder + filename)
-            dot_fw.render(output_folder + filename, format="png", cleanup=True)
+            dot_fw = visualize_graph(roots, filename=base_path)
+            dot_fw.render(base_path, format="svg", cleanup=True)
 
-        graphviz_async.submit(f"{filename}.png", _render_graph, print_message=printstr)
+        graphviz_async.submit(f"{filename}.svg", _render_graph, print_message=printstr)
 
 
 def visualize_graph(roots, filename="graph"):
     _ = filename  # unused, kept for backwards compatibility with callers
 
-    dot = Digraph(comment="Computation Graph")
+    dot = Digraph(comment="Computation Graph", format="svg")
     visited = set()
 
-    def _format_duration(value: float) -> str:
-        ms = value * 1e3
-        if abs(ms) > 1000:
-            return f"{value:.2f}s"
-        return f"{ms:.2f}ms"
+    def _format_duration(value: float, profile: Optional[Tuple[float, ...]] = None) -> str:
+        def _format_single(entry: float) -> str:
+            ms = entry * 1e3
+            if abs(ms) > 1000:
+                return f"{entry:.2f}s"
+            return f"{ms:.2f}ms"
+
+        if not profile:
+            return _format_single(value)
+
+        groups: List[Dict[str, Any]] = []
+        for idx, entry in enumerate(profile):
+            matched = False
+            for group in groups:
+                if math.isclose(entry, group["value"], rel_tol=1e-9, abs_tol=1e-12):
+                    group["indices"].append(idx)
+                    matched = True
+                    break
+            if not matched:
+                groups.append({"value": entry, "indices": [idx]})
+
+        if len(groups) == 1:
+            return _format_single(groups[0]["value"])
+
+        parts = []
+        for group in groups:
+            indices = ",".join(str(i) for i in group["indices"])
+            parts.append(f"{{{indices}}}: {_format_single(group['value'])}")
+        return ", ".join(parts)
 
     def _node_color(node) -> str:
         if isinstance(node, Node):
@@ -1523,17 +1702,19 @@ def _node_label(node) -> str:
         if isinstance(node, Data_batch):
             return f"{node.name}\n(batch_id={node.batch_id}, dur={_format_duration(node.duration)})"
         if isinstance(node, Node):
+            duration_display = _format_duration(node.duration, node.duration_profile)
             return (
                 f"{node.name}\n(op_id={node.op_id}, hw_id={node.hw_id}, "
-                f"dur={_format_duration(node.duration)})"
+                f"dur={duration_display})"
             )
         if isinstance(node, Edge):
+            duration_display = _format_duration(node.duration)
             if getattr(node, "local_hw_id", None) is not None:
                 return (
                     f"{node.name}\n(op_id={node.op_id}, local_hw_id={node.local_hw_id}, "
-                    f"dur={_format_duration(node.duration)})"
+                    f"dur={duration_display})"
                 )
-            return f"{node.name}\n(op_id={node.op_id}, dur={_format_duration(node.duration)})"
+            return f"{node.name}\n(op_id={node.op_id}, dur={duration_display})"
         return str(node)
 
     def _visit(node):
diff --git a/simulate_inf.py b/simulate_inf.py
index e97afad3..744bcb3b 100644
--- a/simulate_inf.py
+++ b/simulate_inf.py
@@ -196,6 +196,8 @@ def _execute_decode_step(
         (
             pipeline_graph,
             pipeline_root,
+            _,
+            _,
             transformer_graph,
             transformer_forward_root,
             transformer_backward_root,
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 00000000..0607f997
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,10 @@
+import sys
+from pathlib import Path
+
+# This is magic file that pytest loads. We need it unfortunately to hook the validation_helpers module into pytest.
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+  sys.path.insert(0, str(PROJECT_ROOT))
+
+pytest_plugins = ["validation_scripts.validation_helpers"]
diff --git a/tests/test_IMEC_inf_no_network.py b/tests/test_IMEC_inf_no_network.py
new file mode 100644
index 00000000..7462dfe2
--- /dev/null
+++ b/tests/test_IMEC_inf_no_network.py
@@ -0,0 +1,115 @@
+import math
+from pathlib import Path
+
+import pandas as pd
+import pytest
+
+from validation_scripts import nvidia_inf
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+DATA_DIR = PROJECT_ROOT / "validation_scripts/imec_data"
+
+
+DEFAULT_THRESHOLD = 10.0
+# Override per (device, model) average pct error.
+# Error is actual + 0.5% margin of error, rounded up to nearest 0.5%.
+MODEL_THRESHOLDS = {
+  ("A100", "Llama 2-7B"): 5.5,   # 4.62 + 0.5 -> ceil to 5.5
+  ("A100", "Llama 2-13B"): 3.0,  # 2.44 + 0.5 -> ceil to 3.0
+  ("A100", "Llama 2-70B"): 13.0, # 12.09 + 0.5 -> ceil to 13.0
+}
+# Override per (device, TP) average pct error across all models at that TP.
+TP_THRESHOLDS = {
+  ("A100", 1): 4.5,  # 3.98 + 0.5 -> ceil to 4.5
+  ("A100", 2): 8.5,  # 7.60 + 0.5 -> ceil to 8.5
+  ("A100", 4): 5.5,  # 4.91 + 0.5 -> ceil to 5.5
+  ("A100", 8): 7.0,  # 6.34 + 0.5 -> ceil to 7.0
+}
+
+
+def _load_params(by: str):
+  params = []
+  for device in ("A100", "H100"):
+    csv_path = DATA_DIR / f"{device}_inf.csv"
+    if not csv_path.exists():
+      continue
+    df = pd.read_csv(csv_path, comment="/")
+    df["TP"] = df["TP"].astype(int)
+    df["actual"] = pd.to_numeric(df["actual"], errors="coerce")
+    df = df.dropna(subset=["device", "model", "TP", "actual"])
+    if by == "model":
+      grouped = df.groupby("model")
+    else:
+      grouped = df.groupby("TP")
+    for key, group in grouped:
+      marks = ()
+      if device == "H100":
+        marks = (pytest.mark.xfail(reason="Pending model refinements", strict=False),)
+      params.append(
+        pytest.param(
+          device,
+          key,
+          marks=marks,
+          id=f"{device}-{by}-{key}",
+        )
+      )
+  return params
+
+
+MODEL_PARAMS = _load_params("model")
+TP_PARAMS = _load_params("TP")
+
+
+@pytest.fixture(scope="module")
+def cached_rows():
+  cache = {}
+  for device in ("A100", "H100"):
+    try:
+      result = nvidia_inf.run(
+        enable_plot=False,
+        network_ignored=True,
+        device=device,
+        emit_logs=False,
+      )
+    except FileNotFoundError:
+      continue
+    cache[device] = result["rows"]
+  return cache
+
+
+@pytest.mark.parametrize("device,model", MODEL_PARAMS)
+def test_avg_abs_error_by_model(device, model, cached_rows, record_validation):
+  label = f"{device} | {model}"
+  threshold = MODEL_THRESHOLDS.get((device, model), DEFAULT_THRESHOLD)
+  rows = [r for r in cached_rows.get(device, []) if r["model"] == model]
+  assert rows, f"{label} produced no validation rows"
+
+  pct_errors = []
+  for row in rows:
+    pct_error = row["pct_error"]
+    assert not math.isnan(pct_error), f"{label}-TP{row['tp']} produced no valid measurements"
+    pct_errors.append(pct_error)
+
+  avg_error = sum(pct_errors) / len(pct_errors)
+  val_label = f"{label} avg_abs_error"
+  record_validation(val_label, avg_error, expected_pct=threshold)
+  assert avg_error <= threshold, f"{label} avg error {avg_error:.2f}% exceeds {threshold}%"
+
+
+@pytest.mark.parametrize("device,tp", TP_PARAMS)
+def test_avg_abs_error_by_tp(device, tp, cached_rows, record_validation):
+  label = f"{device} | TP{tp}"
+  threshold = TP_THRESHOLDS.get((device, int(tp)), DEFAULT_THRESHOLD)
+  rows = [r for r in cached_rows.get(device, []) if int(r["tp"]) == int(tp)]
+  assert rows, f"{label} produced no validation rows"
+
+  pct_errors = []
+  for row in rows:
+    pct_error = row["pct_error"]
+    assert not math.isnan(pct_error), f"{label}-{row['model']} produced no valid measurements"
+    pct_errors.append(pct_error)
+
+  avg_error = sum(pct_errors) / len(pct_errors)
+  val_label = f"{label} avg_abs_error"
+  record_validation(val_label, avg_error, expected_pct=threshold)
+  assert avg_error <= threshold, f"{label} avg error {avg_error:.2f}% exceeds {threshold}%"
diff --git a/tests/test_IMEC_train.py b/tests/test_IMEC_train.py
new file mode 100644
index 00000000..24466d57
--- /dev/null
+++ b/tests/test_IMEC_train.py
@@ -0,0 +1,43 @@
+import math
+
+import pytest
+
+from validation_scripts import nvidia_train
+
+DEFAULT_THRESHOLD = 20.0
+# Optional overrides per (device, tp_sp_enabled)
+GROUP_THRESHOLDS = {
+  # ("A100_korthi", True): 10.0,
+}
+
+
+def _run_and_check(device: str, tp_sp: bool, record_validation):
+  result = nvidia_train.run(
+    device=device,
+    tp_sp=tp_sp,
+    enable_plot=False,
+    emit_logs=False,
+  )
+  rows = result["rows"]
+  assert rows, f"{device} tp_sp={tp_sp} produced no validation rows"
+  pct_errors = []
+  for row in rows:
+    pct = row["pct_error"]
+    assert not math.isnan(pct), f"{device} tp_sp={tp_sp} produced NaN pct_error for {row['model']}"
+    pct_errors.append(pct)
+  avg_error = sum(pct_errors) / len(pct_errors)
+  threshold = GROUP_THRESHOLDS.get((device, tp_sp), DEFAULT_THRESHOLD)
+  record_validation(f"{device} tp_sp={tp_sp}", avg_error, expected_pct=threshold)
+  assert avg_error <= threshold, f"{device} tp_sp={tp_sp} avg error {avg_error:.2f}% exceeds {threshold}%"
+
+
+# def test_train_a100_korthi_tp_sp_disabled(record_validation):
+#   _run_and_check("A100_korthi", False, record_validation)
+
+
+# def test_train_a100_korthi_tp_sp_enabled(record_validation):
+#   _run_and_check("A100_korthi", True, record_validation)
+
+
+# def test_train_a100_selene_tp_sp_disabled(record_validation):
+#   _run_and_check("A100_selene", False, record_validation)
diff --git a/tests/test_koyeb_inf_single_gpu.py b/tests/test_koyeb_inf_single_gpu.py
new file mode 100644
index 00000000..02d1a8fd
--- /dev/null
+++ b/tests/test_koyeb_inf_single_gpu.py
@@ -0,0 +1,19 @@
+import math
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+  sys.path.insert(0, str(PROJECT_ROOT))
+
+from validation_scripts import koyeb
+
+
+def test_koyeb_average_absolute_error_below_threshold(record_validation):
+  results = koyeb.run_all_and_plot(enable_plot=False, verbose=False)
+  avg_abs_error = results["avg_abs_error"]
+  record_validation("avg_abs_error", avg_abs_error, expected_pct=6.5)
+  assert not math.isnan(avg_abs_error), "Koyeb validation produced no valid measurements"
+  # Historically the average absolute error bottoms out at roughly 5.97%, so we allow
+  # a little slack (6.5%) to absorb minor DeepFlow changes while still flagging regressions.
+  assert avg_abs_error <= 6.5, f"Average absolute error {avg_abs_error:.2f}% exceeds 6.5%"
diff --git a/tile.py b/tile.py
index ff60af5c..0208dbf7 100644
--- a/tile.py
+++ b/tile.py
@@ -31,7 +31,20 @@
 from math import ceil
 from functools import lru_cache
 from enum import IntEnum
-from typing import Tuple, Union
+import math
+from typing import Tuple, Union, NamedTuple
+
+class AccessBytes(NamedTuple):
+    reads: tuple[int, ...]
+    writes: tuple[int, ...]
+    def totals(self) -> tuple[int, ...]:
+        return tuple(r + w for r, w in zip(self.reads, self.writes))
+
+def describe_accesses(profile: AccessBytes) -> str:
+    return ", ".join(
+        f"L{lvl}: R={formatBytes(r)} W={formatBytes(w)}"
+        for lvl, (r, w) in enumerate(zip(profile.reads, profile.writes))
+    )
 
 
 class Dataflow(IntEnum):
@@ -152,7 +165,7 @@ def generate_tile_space(memLayer, num_levels: int, dim1: int, dim2: int, dim3: i
 
     Largely unchanged from original function as it is not perf-critical.
     """
-    original = True
+    original = False
     tile_space = []
     tiles = [None] * num_levels
     for level in range(0, num_levels - 1):
@@ -214,18 +227,69 @@ def _sysarray_accesses_sig(M: int, N: int, K: int,
     return load_bytes + store_bytes
 
 
+def _l2_accesses_sig(M: int, N: int, K: int,
+                     l1_M: int, l1_K: int, l1_N: int,
+                     dtype_size: int) -> tuple[int, int]:
+    """Compute L2 memory accesses (reads, writes) in bytes."""
+    grid_x = -(-N // l1_N)
+    grid_y = -(-M // l1_M)
+
+    read_bytes = grid_x * (M * K * dtype_size) + grid_y * (K * N * dtype_size)
+    write_bytes = M * N * dtype_size
+    return (read_bytes, write_bytes)
+
+# Simplified approximate HBM model
+@lru_cache(maxsize=65536)
+def _hbm_accesses_sig_simplified(M: int, N: int, K: int,
+                                 l1_M: int, l1_K: int, l1_N: int,
+                                 dtype_size: int, total_bytes: int, l2_read: int,
+                                 capacity: int) -> tuple[int, int]:
+    """Compute HBM/DRAM memory accesses (reads, writes) in bytes.
+
+    Simplified reuse-distance model:
+    - Keep cold misses (first touches of A/B tiles).
+    - Approximate capacity pressure via a coarse working-set estimate per k-block.
+    - Avoid enumerating detailed reuse-distance patterns; replace with a ratio-based
+      penalty that scales cold misses when working set exceeds capacity.
+    """
+    tile_A = int(l1_M * l1_K)
+    tile_B = int(l1_K * l1_N)
+    tile_C = int(l1_M * l1_N)
+
+    capacity -= (108 * tile_C)
+
+    rows = -(-M // l1_M)
+    cols = -(-N // l1_N)
+    k_blocks = -(-K // l1_K)
+
+    total_accesses = rows * cols * k_blocks
+
+    cold_misses = rows * k_blocks * tile_A + cols * k_blocks * tile_B
+
+    working_set = rows * tile_A + cols * tile_B
+
+    if working_set > capacity:
+        excess_ratio = min(1.0, (working_set - capacity) / max(1, working_set))
+        capacity_misses = int(cold_misses * excess_ratio)
+    else:
+        capacity_misses = 0
+
+    total_misses = cold_misses + capacity_misses
+
+    write_bytes = M * N * dtype_size if total_bytes > capacity * dtype_size else 0
+
+    per_transfer_bytes = max(1, (tile_A + tile_B))
+    read_bytes = l2_read * total_misses / (total_accesses * per_transfer_bytes)
+
+    return (int(read_bytes), int(write_bytes))
+
 @lru_cache(maxsize=131072)
 def _simulate_accesses_sig(inner_code: int, M: int, K: int, N: int,
                            l2_M: int, l2_K: int, l2_N: int,
                            l1_M: int, l1_K: int, l1_N: int,
                            dtype_size: int, FMA_x: int, FMA_y: int, dataflow_code: Union[int, Dataflow],
                            capacity: int, total_bytes: int) -> Tuple[float, int, int, int]:
-    """Compute memory traffic at each level (sys, L1/shared, L2, DRAM).
-
-    Parameters are integers. The function is pure and amenable to
-    caching. It uses closed-form counts derived from tile reuse rather than
-    iterating through actual tiles, which keeps evaluation fast.
-    """
+    """Compute memory traffic at each level (sys, L1/shared, L2, DRAM)."""
     sys_bytes = _sysarray_accesses_sig(M, N, K, FMA_x, FMA_y, dataflow_code, dtype_size)
 
     num_tiles_M = -(-M // l2_M) # == math.ceil(M/l2_M)
@@ -240,58 +304,16 @@ def _simulate_accesses_sig(inner_code: int, M: int, K: int, N: int,
     write_bytes = (l1_M * l1_N * dtype_size) * (reuse_M * reuse_N)
     l1_shared_total = (read_bytes + write_bytes) * max_reload
 
-    # L2 accesses
-    # LOGIC BELOW THIS LINE CAN BE REWRITTEN IN A MORE EFFICIENT WAY
-    # IT HAS BEEN LEFT VERBOSE FOR READABILITY
-    factor_M = reuse_M
-    factor_K = reuse_K
-    factor_N = reuse_N
-    if inner_code == InnerLoop.M:
-        mk_load = max_reload
-        kn_load = num_tiles_K * num_tiles_N
-        mn_load = max_reload
-        factor_M = 1
-        dfk = Dataflow.WST
-    elif inner_code == InnerLoop.K:
-        mk_load = max_reload
-        kn_load = max_reload
-        mn_load = num_tiles_M * num_tiles_N
-        factor_K = 1
-        dfk = Dataflow.OST
-    elif inner_code == InnerLoop.N:
-        mk_load = num_tiles_M * num_tiles_K
-        kn_load = max_reload
-        mn_load = max_reload
-        factor_N = 1
-        dfk = Dataflow.AST
-
-    l2_bytes = (
-        mk_load * (l2_M * l2_K * dtype_size) * factor_N
-        + kn_load * (l2_K * l2_N * dtype_size) * factor_M
-        + 2 * mn_load * (l2_M * l2_N * dtype_size) * factor_K
-    )
-
-    l2f0 = M // l2_M # are we sure these are floor and should not be ceil? (ceil is -(-M // l2_M))
-    l2f1 = K // l2_K
-    l2f2 = N // l2_N
-    reload_mk = 1
-    reload_kn = 1
-    reload_mn = 1
-    if dfk == Dataflow.WST:
-        reload_mk = l2f2
-        reload_mn = l2f1
-    elif dfk == Dataflow.OST:
-        reload_mk = l2f2
-        reload_kn = l2f0
-    elif dfk == Dataflow.AST:
-        reload_kn = l2f0
-        reload_mn = l2f1
+    # L2 accesses (reads+writes lumped to match return contract)
+    l2_read, l2_write = _l2_accesses_sig(M, N, K, l1_M, l1_K, l1_N, dtype_size)
 
-    dram_counts = reload_mk * (M * K * dtype_size) + reload_kn * (K * N * dtype_size)
-    if total_bytes > capacity:
-        dram_counts += reload_mn * (M * N * dtype_size)
+    # HBM/DRAM accesses via selected model (reads+writes lumped)
+    hbm_read, hbm_write = _hbm_accesses_sig_simplified(M, N, K, l1_M, l1_K, l1_N, dtype_size, total_bytes, l2_read, capacity)
 
-    return (sys_bytes, l1_shared_total, l2_bytes, dram_counts)
+    return AccessBytes(
+        reads=(sys_bytes, read_bytes * max_reload, l2_read, hbm_read),
+        writes=(0, write_bytes * max_reload, l2_write, hbm_write)
+    )
 
 
 
@@ -308,7 +330,7 @@ def __init__(self, order_dims, tile_dims, core, memLayer, dtype_size=2):
         self.memLayer = memLayer
         self.dtype_size = dtype_size
         self.M, self.K, self.N = tile_dims[3]
-        self.capacity = memLayer[3].size_per_bundle
+        self.capacity = memLayer[2].size_per_bundle # L2 capacity in bytes
         self.order_dims = order_dims if isinstance(order_dims, str) else None
         # Precompute inner code once (accepts str or int code)
         self._inner_code = (
@@ -362,6 +384,27 @@ def GEMM_flop(self):
         """Theoretical FLOP count for the full GEMM of size MxKxN."""
         return self.M * self.N * (2 * self.K - 1)
 
+    @property
+    def bundle_util(self):
+        """Wave utilization of bundles"""
+        reuse_M = -(-self.M // self.l2_M)
+        reuse_K = -(-self.K // self.l2_K)
+        reuse_N = -(-self.N // self.l2_N)
+
+        gemm_waves = self.memLayer[1].calc_waves_per_sm(
+            self.M, self.K, self.N,
+            self.l2_M, self.l2_K, self.l2_N
+        )
+        
+        if gemm_waves == -1:
+            return 1.0
+        
+        full_waves = math.floor(gemm_waves)
+        partial_wave = gemm_waves - full_waves # fractional wave
+        
+        return (full_waves + partial_wave * partial_wave) / gemm_waves
+
+
     def sysarray_accesses(self):
         """Return systolic-array traffic (bytes) using cached pure function.
 
diff --git a/time_calculation.py b/time_calculation.py
index 6e05bdef..03c84d55 100644
--- a/time_calculation.py
+++ b/time_calculation.py
@@ -8,13 +8,12 @@
 # import numpy as np
 
 from parallelism import Parallelism
-from topology import Topology
 from simulate import Graph
 import util
 from hw_component import Core, MemoryHierarchy, Network, DRAM
 from astrasim_lib import run_cache_astrasim
 from model import Model_LSTM, Model_GEMM, Model_LLM
-from tile import TiledGEMM, formatBytes
+from tile import AccessBytes, TiledGEMM, formatBytes
 
 algByte = False  # algorithmic ops false
 proj = False  # consider projection layer, turn off for end-2-end validation, as baeline model does not have projection layer
@@ -28,18 +27,23 @@ def __init__(self, hw_config, precision, kernel_overhead, roofline_cb, astra_pol
         self._roofline = roofline_cb
         self._astra_policy = astra_policy or "analytical"
 
-    def _astra_collective(self, kind: str, participants: int, size_bytes: int) -> float:
+    def _astra_collective(self, kind: str, participants: int, size_bytes: int, axis: Optional[str] = None) -> float:
         part = int(participants)
         if part <= 1 or size_bytes <= 0:
             return 0.0
         byte_count = int(math.ceil(size_bytes))
+        axes_filter = [str(axis).lower()] if axis else None
+        if os.environ.get("DEEPFLOW_ASTRA_SKIP_EXEC"):
+            return 0.0
+        # for collectives ONLY, we cannot use 2D topologies (Mesh2D, Torus2D, KingMesh2D)
+        # transform them to their 1D equivalents (Ring, Mesh, HyperCube (?))
         _, max_sec = run_cache_astrasim(
             self.hw_config,
             comm=kind,
             npus_count=part,
             size_bytes=byte_count,
-            astra_config_dir="./astra_cache",
-            cache_path="./astra_cache/cache.json",
+            axes_filter=axes_filter,
+            transform_2d_to_1d=True,
         )
         return float(max_sec)
 
@@ -55,6 +59,7 @@ def collective(
         local_bytes: float = 0.0,
         local_ops: float = 0.0,
         debug_label: str = "",
+        axis: Optional[str] = None,
     ) -> float:
         if size_bytes <= 0:
             return 0.0
@@ -77,7 +82,8 @@ def collective(
             if kind in collective_ops or kind == "pipeline":
                 # Pipeline uses 2 NPUs for point-to-point, others use part
                 npus = 2 if kind == "pipeline" else part
-                network_time = self._astra_collective(kind, npus, network_bytes)
+                axis_filter = axis if kind != "pipeline" else None
+                network_time = self._astra_collective(kind, npus, network_bytes, axis_filter)
             else:
                 raise ValueError(f"Unsupported collective operation: {kind}")
         else:
@@ -136,8 +142,10 @@ def _analytical_all_reduce(self, size_bytes, participants, ib, ll, label):
         data_transfer = ((per_rank / ib) + mem_access + ll) * 2 * (participants - 1)
         prep_comp = per_rank
         prep_mem = int(math.ceil(3 * size_bytes / participants))
+
+        # NOTE: removed overhead because it made it impossible to eliminate network effects even with inf bandwidth and zero latency
         data_prep = (
-            self._roofline(prep_comp, prep_mem, name=f"{label}-prep") + self.O
+            self._roofline(prep_comp, prep_mem, name=f"{label}-prep") # + self.O
         ) * (participants - 1)
         return data_transfer + data_prep
 
@@ -195,6 +203,8 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti
         # If we want proper backported LSTM support we need to fix this at *some* point.
         self.h2d_bandwidth = getattr(hw_config.sw_config, "h2d_bandwidth", -1)
         self.zero_stage = getattr(hw_config.sw_config, "dp_zero_stage", 0)
+        self.full_recomputation = getattr(hw_config.sw_config, "full_recomputation", False)
+        self.dp_microbatch = getattr(hw_config.sw_config, "dp_microbatch", "every_mb")
         self.attached = True
 
         # Hardware Parameters
@@ -208,83 +218,20 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti
         self.num_levels = self.memoryHierarchy.num_levels
         self.memLayer = self.memoryHierarchy.memLayer
         self.tileSpace = None
-
-        # TODO: move this to config file
         self.H2Dbw = self.h2d_bandwidth
-
-
-        # System Parameters
-        self.num_wafer = hw_config.system_config.num_wafers
-        self.num_workers = hw_config.system_config.num_workers
-
-
+        self.num_workers = self._derive_num_workers(hw_config)
 
         level = 0
         mem_config = hw_config.memory_hierarchy.mem_hr[level]
         self.DRAM = DRAM(hw_config, mem_config, level)
-        self.memory_capacity = self.DRAM.size * self.num_workers# in bytes
-
-        # self.memory_capacity = hw_config.perimeter_breakdown.DRAM
-
         self.network = Network(hw_config)
 
-        intra_throughput, inter_throughput = self.network.calcThroughput()
-        intra_latency, inter_latency = self.network.calcLatency()
-
-        inter_derate = hw_config.system_config.inter_derate
-        intra_derate = hw_config.system_config.intra_derate
-        par2cross = hw_config.system_config.par2cross
-
-        derated_inter_throughput = -1
-        derated_intra_throughput = -1
-
-        # inter-wafer communications will pass through intra links too
-        if self.num_wafer > 1 and self.num_workers > 1:
-            if intra_derate != 0:
-                derated_inter_throughput = min(
-                    intra_throughput / intra_derate, inter_throughput / inter_derate
-                )
-                # print(f'intra_throughput / intra_derate: {intra_throughput / intra_derate}, inter_throughput / inter_derate: {inter_throughput / inter_derate}')
-            else:
-                derated_inter_throughput = inter_throughput / inter_derate
-                # print(f'inter_throughput / inter_derate: {inter_throughput / inter_derate}')
-        else:
-            derated_inter_throughput = 0
-
-        if self.num_workers > 1 and intra_derate != 0:
-            derated_intra_throughput = intra_throughput / intra_derate
-        else:
-            derated_intra_throughput = 0
-
-        self.IBK1, self.LLK1 = (
-            (derated_inter_throughput, inter_latency)
-            if par2cross["kp1"]
-            else (derated_intra_throughput, intra_latency)
-        )
-        self.IBK2, self.LLK2 = (
-            (derated_inter_throughput, inter_latency)
-            if par2cross["kp2"]
-            else (derated_intra_throughput, intra_latency)
-        )
-        self.IBD, self.LLD = ( #interconnect bandwidth and latency for data parallelism
-            (derated_inter_throughput, inter_latency)
-            if par2cross["dp"]
-            else (derated_intra_throughput, intra_latency)
-        )
-        self.IBL, self.LLL = (
-            (derated_inter_throughput, inter_latency)
-            if par2cross["lp"]
-            else (derated_intra_throughput, intra_latency)
-        )
-
-        self.IBTP, self.LLTP = (
-            (derated_inter_throughput, inter_latency)
-            if par2cross["tp"]
-            else (derated_intra_throughput, intra_latency)
-        )
+        self.IBK1, self.LLK1 = self.network.get_link("kp1")
+        self.IBK2, self.LLK2 = self.network.get_link("kp2")
+        self.IBD, self.LLD = self.network.get_link("dp")
+        self.IBL, self.LLL = self.network.get_link("lp")
+        self.IBTP, self.LLTP = self.network.get_link("tp")
         
-
-
         # Scheduling Parameters
         par = Parallelism(hw_config)
         par.findParallelStrategy()
@@ -312,17 +259,17 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti
         try:
             self.tp = int(tp_value)
         except (TypeError, ValueError) as exc:
-            raise ValueError("scheduling_param.tp must be an integer") from exc
+            raise ValueError("parallelism.tp must be an integer") from exc
         if self.tp < 1:
-            raise ValueError("scheduling_param.tp must be >= 1")
+            raise ValueError("parallelism.tp must be >= 1")
 
         cp_value = par.cp if par.cp not in (None, 0) else 1
         try:
             self.cp = int(cp_value)
         except (TypeError, ValueError) as exc:
-            raise ValueError("scheduling_param.cp must be an integer") from exc
+            raise ValueError("parallelism.cp must be an integer") from exc
         if self.cp < 1:
-            raise ValueError("scheduling_param.cp must be >= 1")
+            raise ValueError("parallelism.cp must be >= 1")
 
         self.tp_sp = par.tp_sp
 
@@ -333,30 +280,18 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti
         else:
             self.kp1 = None
             self.kp2 = None
-        run_type = model_config.model_config.run_type
-        if run_type == "inference":
-            if self.cp > 1:
-                raise ValueError(
-                    "Context parallelism (cp) is not supported for LLM inference. "
-                    "Please set scheduling_param.cp to 1 for inference runs."
-                )
-            if self.mb > 1:
-                print(f"[WARNING]: LLM inference configured with mb={self.mb} (>1). \n Pipeline micro-batching is ill-defined for autoregressive decode and should be avoided.")
 
-        if self.mode == "LLM":
-            expected_workers = self.tp * self.cp * self.dp * self.lp
-            label = f"tp({self.tp}) * cp({self.cp}) * dp({self.dp}) * lp({self.lp})"
-        else:
-            kp1 = int(self.kp1) if self.kp1 else 1
-            kp2 = int(self.kp2) if self.kp2 else 1
-            expected_workers = self.dp * self.lp * kp1 * kp2
-            label = f"dp({self.dp}) * lp({self.lp}) * kp1({kp1}) * kp2({kp2})"
-
-        if expected_workers != self.num_workers:
-            raise ValueError(
-                f"Parallelism mismatch: {label} = {expected_workers}, "
-                f"but system_hierarchy reports num_workers={self.num_workers}."
-            )
+        if self.mode != "GEMM":
+            run_type = model_config.model_config.run_type
+            if run_type == "inference":
+                if self.cp > 1:
+                    raise ValueError(
+                        "Context parallelism (cp) is not supported for LLM inference. "
+                        "Please set parallelism.cp to 1 for inference runs."
+                    )
+                if self.mb > 1:
+                    print(f"[WARNING]: LLM inference configured with mb={self.mb} (>1). \n Pipeline micro-batching is ill-defined for autoregressive decode and should be avoided.")
+
         
         # Statistics Param
         self.tot_flop = 0
@@ -377,12 +312,11 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti
             self.roofline,
             astra_policy=self._astra_policy,
         )
-        
+
         # Dynamically select and instantiate the model class
         model_class = self.get_model_class(mode)
         self.model = model_class(model_config)  # Instantiate the model class
 
-
         # Model Parameters
         # self.model = self.get_model_class(mode)
         if mode == "LSTM":
@@ -406,7 +340,14 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti
             self.N = self.model.N
             
         if mode == "LLM":
-            self.batch_size = self.model.batch_size
+            self.global_batch_size = self.model.global_batch_size
+            self.gradient_accumulation_steps = self.model.gradient_accumulation_steps
+
+            if self.global_batch_size % self.gradient_accumulation_steps != 0:
+                raise ValueError(
+                    "Global batch size must be divisible by gradient accumulation steps"
+                )
+            self.batch_size = self.global_batch_size // self.gradient_accumulation_steps
             self.vocab_size = self.model.vocab_size
             self.num_layers = self.model.num_layers
             self.hidden_dim = self.model.hidden_dim
@@ -429,9 +370,38 @@ def __init__(self, hw_config, model_config, mode, *, astra_policy_override: Opti
             raw_top_k = getattr(self.model, "moe_top_k", 1)
             self.moe_num_experts = max(1, int(raw_num_experts))
             self.moe_top_k = max(1, int(raw_top_k))
-            if self.moe_top_k > self.moe_num_experts:
-                raise ValueError("model_param.top_k cannot exceed model_param.num_experts")
             self.use_moe = self.moe_num_experts > 1
+
+
+    def _derive_num_workers(self, hw_config) -> int:
+        layout = getattr(hw_config, "network_layout", None)
+        if layout and getattr(layout, "dimensions", None):
+            total = 1
+            for dim in layout.dimensions:
+                try:
+                    size = int(dim.size)
+                except (TypeError, ValueError):
+                    size = 1
+                if size < 1:
+                    size = 1
+                total *= size
+            if total >= 1:
+                return int(total)
+
+        factors = []
+        for name in ("dp", "lp", "tp", "cp", "kp1", "kp2"):
+            value = getattr(hw_config.sch_config, name, 1)
+            try:
+                factor = int(value) if value else 1
+            except (TypeError, ValueError):
+                factor = 1
+            if factor < 1:
+                factor = 1
+            factors.append(factor)
+        total = 1
+        for factor in factors:
+            total *= factor
+        return max(1, total)
    
 
     def get_model_class(self, model_type):
@@ -770,7 +740,7 @@ def printSysConfig(self, exp_hw_config, exp_model_config, output_file):
                 self.memLayer[i].printStats(f)
 
             self.network.printStats(f)
-    def roofline(self, flop, mem_access_, name="", info=False, mem_level=None, flashattn_enable=False):
+    def roofline(self, flop, mem_access_, name="", util=1, info=False, mem_level=None, flashattn_enable=False):
         # print("Roofline: entered {}".format(name))
 
         # Parse mem_access_ into consistent format
@@ -785,6 +755,8 @@ def roofline(self, flop, mem_access_, name="", info=False, mem_level=None, flash
             print("mem_access_ should be integer or list, wrong input", flush=True)
             sys.exit(0)
 
+        throughput = self.th * util
+
         # Determine which levels to compute
         if mem_level is not None:
             # Single level mode
@@ -794,7 +766,7 @@ def roofline(self, flop, mem_access_, name="", info=False, mem_level=None, flash
             num_level = len(mem_access)
             try:
                 if not flashattn_enable:
-                    assert mem_access[num_level - 1] > 0, "last_level_mem = 0"
+                    assert mem_access[num_level - 1] > 0, f"mem_access: {mem_access_}"
             except Exception as e:
                 print(
                     "{}: Number of accesses to the last level of memory hierarchy cannot be zero:\n {}".format(
@@ -810,13 +782,13 @@ def roofline(self, flop, mem_access_, name="", info=False, mem_level=None, flash
         for level_idx, num_mem in levels_to_compute:
             mem_bw = self.memLayer[level_idx].getThroughput()
             mem_latency = self.memLayer[level_idx].getLatency()
-            inflection_point = float("inf") if mem_bw == 0 else self.th / mem_bw
+            inflection_point = float("inf") if mem_bw == 0 else throughput / mem_bw
             # print(f"Level {level_idx}: mem_bw={mem_bw}, mem_latency={mem_latency}, inflection_point={inflection_point}", flush=True)
             comp_int = float("inf") if num_mem == 0 else flop / num_mem
             if comp_int < inflection_point:  # mem-bound
                 level_time = (float("inf") if (mem_bw == 0 or num_mem == 0) else (num_mem / mem_bw)) + mem_latency
             else:  # compute-bound
-                level_time = float("inf") if (self.th == 0) else (flop / self.th)
+                level_time = float("inf") if (throughput == 0) else (flop / throughput)
 
             times.append(level_time)
 
@@ -824,12 +796,14 @@ def roofline(self, flop, mem_access_, name="", info=False, mem_level=None, flash
 
         # print("Roofline: exited {}".format(name))
         return max_time
-    def getGEMMTime(self, dim1, dim2, dim3, name, 
+    
+    def getGEMMTime(self, dim1, dim2, dim3, name="", 
                     flashattn_enable=False, disable_overhead=False, read_bytes_l2=0, write_bytes_l2=0, original=False):
         # Streaming best selection to avoid building large dicts
         best_time = float("inf")
         best_choice = None  # type: Optional[tuple]
         best_mem_access = None  # type: Optional[tuple]
+        best_rw_access = None  # type: Optional[AccessBytes]
         best_gemm = None  # type: Optional[TiledGEMM]
         best_metric = float("inf")
 
@@ -841,8 +815,9 @@ def getGEMMTime(self, dim1, dim2, dim3, name,
                 print("===============================================================")
 
             GEMM_flop = gemm.GEMM_flop
-            mem_access = gemm.mem_accesses
+            rw_accesses = gemm.mem_accesses
 
+            mem_access = rw_accesses.totals() # (L0, L1, L2, DRAM)
             if flashattn_enable:
                 mem_access = list(mem_access)
                 mem_access[3] = 0 # no HBM accesses
@@ -859,7 +834,17 @@ def getGEMMTime(self, dim1, dim2, dim3, name,
             if eff_sm > 0:
                 mem_access_per_sm[1] = mem_access_per_sm[1] / eff_sm
 
-            GEMM_time = self.roofline(GEMM_flop, mem_access_per_sm, name, flashattn_enable=flashattn_enable) 
+            if flashattn_enable:
+                sm_util = 0.25 # assume 25% SM utilization for FlashAttention (Dao, 2023)
+            else:
+                sm_util = 1
+                gemm_waves = self.memLayer[1].calc_waves_per_sm(dim1, dim2, dim3, gemm.l2_M, gemm.l2_K, gemm.l2_N)
+                if gemm_waves != -1:
+                    full_waves = math.floor(gemm_waves)
+                    partial_wave = gemm_waves - full_waves
+                    sm_util = (full_waves + partial_wave * partial_wave) / gemm_waves
+
+            GEMM_time = self.roofline(GEMM_flop, mem_access_per_sm, name, util=sm_util, flashattn_enable=flashattn_enable) 
             if flashattn_enable or disable_overhead:
                 pass
             else:
@@ -873,12 +858,14 @@ def getGEMMTime(self, dim1, dim2, dim3, name,
             key = (gemm._inner_code, tile_dims)
 
             # Tie-breaker metric identical to previous selection: hypot(dram, l2)
+            
             metric = math.hypot(mem_access[3], mem_access[2])
 
             if (GEMM_time < best_time) or (GEMM_time == best_time and metric < best_metric):
                 best_time = GEMM_time
                 best_choice = key
                 best_mem_access = mem_access
+                best_rw_access = rw_accesses
                 best_gemm = gemm
                 best_metric = metric
 
@@ -897,7 +884,7 @@ def getGEMMTime(self, dim1, dim2, dim3, name,
         # 2 -> inner 'n' (activation stationary)
         best_inner_code = best_choice[0]  # type: ignore[index]
         best_tile_dims = best_choice[1]  # type: ignore[index]
-        return best_time, best_inner_code, best_tile_dims, mem_access
+        return best_time, best_inner_code, best_tile_dims, mem_access #, best_rw_access
     
     def generateTileSpace(self, dim1=None, dim2=None, dim3=None, original=False):
         tile_space = []
@@ -1488,86 +1475,38 @@ def getCb(self):
 
     # Reduction and all-gather time estimation
 
-    def gradClipping(self, Dim0=None, Dim1=None, name=None):
-        if Dim0 == None:
-            Dim0 = 2 * self.D
-        if Dim1 == None:
-            Dim1 = self.G * self.D
-        if name == None:
-            name = "Hidden"
-        # t_list[i] * clip_norm / max(global_norm, clip_norm)
-        # where:
-        # global_norm = sqrt(sum([l2norm(t)**2 for t in t_list]))
-
-        norm_comp = Dim0 * Dim1 * 2
-        # 1: power 2
-        # 1: summ
-        norm_mem = (Dim0 * Dim1 * 1) * self.precision_bytes
-        # 1: one read per element and power it by 2 in local registers  anfd
-        # summing to local acc
-
-        clip_comp = Dim0 * Dim1 * 2
-        # 1: pointwise mul
-        # 1: pointwise div
-
-        clip_mem = (Dim0 * Dim1 * 2) * self.precision_bytes
-        # 1: one read for pointwise mul
-        # 1: one write for pointwise div
-
-        gradclip_mem = norm_mem + clip_mem
-        gradclip_comp = norm_comp + clip_comp
-
-        gradclip_time = self.roofline(
-            gradclip_comp, gradclip_mem, name="pointwise-grad-clipping"
-        )
-
-        if self.debug:
-            print(
-                "({}) gradclip_flop: {:,}, gradclip_mem: {:,}".format(
-                    name, gradclip_comp, gradclip_mem
-                )
-            )
-            print("({}) gradclip_time: {:,}\n".format(name, gradclip_time))
+    def grad_clipping(self, num_params: int) -> float:
+        """Gradient clipping (L2 norm + scale) cost."""
+        norm_comp = num_params * 2  # square + sum
+        clip_comp = num_params * 2  # scale + divide
+        clip_mem = num_params * 2 * self.precision.gradients  # read + write
 
-        return gradclip_time
+        gradclip_mem = clip_mem
+        gradclip_comp = norm_comp + clip_comp
 
-    def apply_grad(self, Dim0=None, Dim1=None, name=None):
-        if Dim0 == None:
-            Dim0 = 2 * self.D
-        if Dim1 == None:
-            Dim1 = self.G * self.D
-        if name == None:
-            name = "Hidden"
+        return self.roofline(gradclip_comp, gradclip_mem, name="pointwise-grad-clipping")
 
-        applyGrad_comp = Dim0 * Dim1 * 3
-        # 3: one pointwise division by  scalar after reducing all the gradients,
-        #   one final addition of gradients to the weights
-        #   one multiply by learning rate
-        applyGrad_mem = (
-            (1 * Dim0 * Dim1 * self.precision_bytes)
-            + (2 * Dim0 * Dim1 * self.precision_bytes)
-            + (1 * Dim0 * Dim1 * self.precision_bytes)
+    def apply_grad(self, num_params: int) -> float:
+        """Approximate optimizer update cost (Adam/AdamW-style) per parameter tensor."""
+        OPT_FLOPS_PER_PARAM = 14.0 # typical for Adam
+        bytes_per_param_read = (
+            self.precision.parameters
+            + self.precision.gradients
+            + 2 * self.precision.optimizer_states
+            + self.precision.master_parameters
         )
-        # 1: read for pointiwse div
-        # 2: 1 reads and one write for pointwise add
-        # 1: one write for multiplication by lr
-        applyGrad_time = self.roofline(
-            applyGrad_comp, applyGrad_mem, name="pointwise-applyGrad"
+        bytes_per_param_write = (
+            self.precision.parameters
+            + 2 * self.precision.optimizer_states
+            + self.precision.master_parameters
         )
 
-        clip_time = self.gradClipping(Dim0, Dim1, name)
-
-        grad_time = applyGrad_time + clip_time
-
-        if self.debug:
-            print(
-                "({}) applyGrad_flop: {:,}, applyGrad_mem: {:,}".format(
-                    name, applyGrad_comp, applyGrad_mem
-                )
-            )
-            print("({}) applyGrad_time: {:,}\n".format(name, applyGrad_time))
+        apply_grad_comp = num_params * OPT_FLOPS_PER_PARAM
+        apply_grad_mem = num_params * (bytes_per_param_read + bytes_per_param_write)
 
-        return grad_time
+        apply_grad_time = self.roofline(apply_grad_comp, apply_grad_mem, name="apply_grad", mem_level=self.num_levels-1)
+        clip_time = self.grad_clipping(num_params)
+        return apply_grad_time + clip_time
 
 
 
@@ -1749,7 +1688,7 @@ def getDataParallelReduction(self, k, n, dim1, dim2, name):
                 local_bytes=0.0,
                 debug_label=name or "comm",
             )
-            apply_grad_time = self.apply_grad(Dim0=k / dim1, Dim1=n, name=name)
+            apply_grad_time = self.apply_grad(int(math.ceil((k / dim1) * n)))
 
         elif self.kp_hidden_type == 2:  # RC
             total_bytes = math.ceil(self.precision_bytes * (k / dim1) * (n / dim2))
@@ -1775,7 +1714,7 @@ def getDataParallelReduction(self, k, n, dim1, dim2, name):
                 local_bytes=3 * total_bytes,
                 debug_label=name or "comm",
             )
-            apply_grad_time = self.apply_grad(Dim0=k, Dim1=n / dim2, name=name)
+            apply_grad_time = self.apply_grad(int(math.ceil(k * (n / dim2))))
         else:
             reduction_time_wt_kp = 0
             total_bytes = math.ceil(self.precision_bytes * k * n)
@@ -1788,7 +1727,7 @@ def getDataParallelReduction(self, k, n, dim1, dim2, name):
                 local_bytes=0.0,
                 debug_label=name or "comm",
             )
-            apply_grad_time = self.apply_grad(Dim0=k, Dim1=n, name=name)
+            apply_grad_time = self.apply_grad(int(math.ceil(k * n)))
         if self.debug:
             print(f"reduction_time_wt_kp: {reduction_time_wt_kp}")
             print(f"reduction_time_wt_dp: {reduction_time_wt_dp}")
diff --git a/time_calculation_LLM.py b/time_calculation_LLM.py
index 7ba1c8dd..4a8cf15b 100644
--- a/time_calculation_LLM.py
+++ b/time_calculation_LLM.py
@@ -1,15 +1,18 @@
 import math
 import os
 import json
+import warnings
 from enum import Enum
-from typing import Any, Dict, Tuple, Optional, List
+from typing import Any, Dict, Tuple, Optional, List, Mapping, Sequence, Set
+from collections.abc import Mapping as MappingABC, Sequence as SequenceABC
 import simulate_LLM
-from LLM_execution import ExecutionMode, LLMExecutionDispatcher
+from LLM_execution import ExecutionMode, LLMExecutionDispatcher, apply_overlap_transforms
 from simulate_LLM import Graph
 import LLM_util
 from time_calculation import TimeCalculation
 from itertools import zip_longest  # for element-wise aggregation of memory access lists
-
+from timing_model import CommSpec, DirectionTiming, OperationTiming, OperationGroup
+import yaml
 def _env_flag(name: str) -> bool:
     value = os.environ.get(name)
     if value is None:
@@ -41,73 +44,147 @@ class ParallelismMode(Enum):
 
 # Map each parallelism mode to operation-level collective specs used across the
 # metadata pipeline. Each spec records the collective kind, the participant
-# scope (tp/cp/seq/etc.), the interconnect label, and an identifying suffix.
+# scope (tp/cp/seq/etc.), and the interconnect label.
 COMM_RULE_DEFAULT_KEY = "__default__"
 COMMUNICATION_RULES: Dict[
     ParallelismMode, Dict[str, Dict[str, Optional[Dict[str, str]]]]
 ] = {
     ParallelismMode.TENSOR: {
         COMM_RULE_DEFAULT_KEY: {
-            'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_reduce'},
-            'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_reduce'},
+            'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'},
+        },
+        'layernorm1': {
+            'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'},
+        },
+        'layernorm2': {
+            'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'},
+        },
+        'qkv_proj': {
+            'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'},
+        },
+        'attention': {
+            'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'},
+        },
+        'output_proj': {
+            'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'},
+        },
+        'MLP': {
+            'forward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'all_reduce', 'participants': 'tp', 'interconnect': 'tp'},
         },
     },
     ParallelismMode.TENSOR_SEQUENCE: {
-        COMM_RULE_DEFAULT_KEY: {'forward': None, 'backward': None},
+        COMM_RULE_DEFAULT_KEY: {'forward': None, 'backward': None}, # <- dangerous
         'layernorm1': {
-            'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'},
-            'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'},
+            'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'},
         },
         'layernorm2': {
-            'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'},
-            'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'},
+            'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'},
+        },
+        'qkv_proj': {
+            'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'},
         },
-        'MHA': {
-            'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'reduce_scatter'},
-            'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'reduce_scatter'},
+        'attention': {
+            'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'},
+        },
+        'output_proj': {
+            'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'},
         },
         'MLP': {
-            'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'reduce_scatter'},
-            'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'reduce_scatter'},
+            'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'},
         },
     },
     ParallelismMode.CONTEXT: {
         COMM_RULE_DEFAULT_KEY: {
-            'forward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp', 'suffix': 'all_gather'},
-            'backward': {'kind': 'reduce_scatter', 'participants': 'cp', 'interconnect': 'cp', 'suffix': 'reduce_scatter'},
+            'forward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp'},
+            'backward': {'kind': 'reduce_scatter', 'participants': 'cp', 'interconnect': 'cp'},
         },
-        'attention': {'forward': None, 'backward': {'kind': 'reduce_scatter', 'participants': 'cp', 'interconnect': 'cp', 'suffix': 'reduce_scatter'}},
-        'output_proj': {'forward': None, 'backward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp', 'suffix': 'all_gather'}},
-        'qkv_proj': {'forward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp', 'suffix': 'all_gather'}, 'backward': None},
+        'attention': {'forward': None, 'backward': {'kind': 'reduce_scatter', 'participants': 'cp', 'interconnect': 'cp'}},
+        'output_proj': {'forward': None, 'backward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp'}},
+        'qkv_proj': {'forward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp'}, 'backward': None},
     },
     ParallelismMode.TENSOR_CONTEXT_HYBRID: {
         COMM_RULE_DEFAULT_KEY: {'forward': None, 'backward': None},
         'layernorm1': {
-            'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'},
-            'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'},
+            'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'},
         },
         'layernorm2': {
-            'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'},
-            'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'all_gather'},
+            'forward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'all_gather', 'participants': 'tp', 'interconnect': 'tp'},
         },
         'MLP': {
-            'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'reduce_scatter'},
-            'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'reduce_scatter'},
+            'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'},
+            'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'},
         },
-        'attention': {'forward': None, 'backward': {'kind': 'reduce_scatter', 'participants': 'cp', 'interconnect': 'cp', 'suffix': 'reduce_scatter'}},
+        'attention': {'forward': None, 'backward': {'kind': 'reduce_scatter', 'participants': 'cp', 'interconnect': 'cp'}},
         'output_proj': {
-            'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'cp', 'suffix': 'reduce_scatter'},
-            'backward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp', 'suffix': 'all_gather'},
+            'forward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'cp'},
+            'backward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'cp'},
         },
         'qkv_proj': {
-            'forward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'tp', 'suffix': 'all_gather'},
-            'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp', 'suffix': 'reduce_scatter'},
+            'forward': {'kind': 'all_gather', 'participants': 'cp', 'interconnect': 'tp'},
+            'backward': {'kind': 'reduce_scatter', 'participants': 'tp', 'interconnect': 'tp'},
         },
     },
     ParallelismMode.SINGLE: {
         COMM_RULE_DEFAULT_KEY: {'forward': None, 'backward': None},
     },
 }
+# MoE-specific overrides for router, MLP, and layer norm collectives.
+_MOE_ROUTER_RULE = {
+    'forward': {
+        'kind': 'all_to_all',
+        'participants': 'moe',
+        'interconnect': 'tp',
+    },
+    'backward': {
+        'kind': 'reduce_scatter',
+        'participants': 'moe',
+        'interconnect': 'tp',
+    },
+}
+_MOE_MLP_RULE = {
+    'forward': {
+        'kind': 'all_to_all',
+        'participants': 'moe',
+        'interconnect': 'tp',
+    },
+    'backward': {
+        'kind': 'all_to_all',
+        'participants': 'moe',
+        'interconnect': 'tp',
+    },
+}
+_MOE_LAYER_NORM1_RULE = {
+    'backward': {
+        'kind': 'all_to_all',
+        'participants': 'moe',
+        'interconnect': 'tp',
+    },
+}
+MOE_COMMUNICATION_RULES: Dict[
+    ParallelismMode, Dict[str, Dict[str, Optional[Dict[str, str]]]]
+] = {
+    mode: {
+        'router': _MOE_ROUTER_RULE,
+        'MLP': _MOE_MLP_RULE,
+        'layernorm1': _MOE_LAYER_NORM1_RULE,
+    }
+    for mode in ParallelismMode
+}
 class GemmType(Enum):
     ATTENTION_SCORE = "attention_score"
     ATTENTION_OUTPUT = "attention_output"
@@ -116,6 +193,7 @@ class GemmType(Enum):
     FFN1 = "ffn1"
     FFN2 = "ffn2"
     LINEAR_SOFTMAX = "linear_softmax"
+    LAYER_NORM_1 = "layer_norm1"
 
 
 class TimeCalculationLLM(TimeCalculation):
@@ -130,16 +208,23 @@ def __init__(self, hw_config, model_config, mode, output_dir: Optional[str] = No
             mode,
             astra_policy_override=astra_policy,
         )
-        if self.num_workers > 1 and self.use_moe:
-            raise ValueError("MoE is only supported for single-worker runs (system_hierarchy.num_workers must be 1).")
+
           
         self.output_dir = os.path.abspath(output_dir) if output_dir else os.getcwd()
         os.makedirs(self.output_dir, exist_ok=True)
+        overlap_cfg = getattr(getattr(hw_config, "network_layout", None), "overlap_config", None)
+        if overlap_cfg is None:
+            raise ValueError("network.overlap must be provided in the hardware config")
+        self.tp_overlap = float(overlap_cfg.tp_overlap)
+        self.tp_sp_overlap = float(overlap_cfg.tp_sp_overlap)
+        self.cp_overlap = float(overlap_cfg.cp_overlap)
         self._generate_graphs = _env_flag("DEEPFLOW_VISUALIZE_GRAPHS")
         self.persist_astrasim_artifacts = _env_flag("DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS")
         self._debug_memory = _env_flag("DEEPFLOW_DEBUG_MEMORY")
         self._memory_breakdown_debug = None
         self.execution_mode = execution_mode
+        # Use pipeline-style recomputation (explicit recompute nodes in the pipeline graph).
+        self.pipeline_style_recompute = bool(self.full_recomputation)
         inference_cfg = getattr(hw_config, "inference_config", None)
 
         self.all_reduce = "every layer"
@@ -157,12 +242,15 @@ def __init__(self, hw_config, model_config, mode, output_dir: Optional[str] = No
         self.transformer_backward_root: Optional[Any] = None
         self.transformer_analytical_time_forward: Optional[float] = None
         self.transformer_analytical_time_backward: Optional[float] = None
+        self.transformer_analytical_time_backward_combined: Optional[float] = None
         self.transformer_astrasim_time_forward: Optional[float] = None
         self.transformer_astrasim_time_backward: Optional[float] = None
         self.transformer_astrasim_per_rank_forward: Optional[List[float]] = None
         self.transformer_astrasim_per_rank_backward: Optional[List[float]] = None
         self.pipeline_astrasim_time: Optional[float] = None
         self.pipeline_astrasim_per_rank: Optional[List[float]] = None
+        self.pipeline_graph_no_dp: Optional[Graph] = None
+        self.pipeline_root_no_dp: Optional[Any] = None
 
     def _sequence_parallel_degree(self) -> int:
         """Return tensor-parallel degree used for sequence-parallel collectives.
@@ -189,6 +277,31 @@ def get_parallelism_mode(self):
         else:
             return ParallelismMode.SINGLE
 
+    @property
+    def experts_per_gpu(self) -> int:
+        """
+        Return the number of MoE experts assigned to each GPU assuming each dp group has all experts.
+        """
+        if not self.use_moe or self.moe_num_experts <= 1:
+            return 0
+        moe_ranks = max(1, self.tp * self.cp)
+        if moe_ranks > self.moe_num_experts:
+            warnings.warn(
+                "MoE configuration has fewer experts than tp * cp; DeepFlow only supports at least one expert per GPU (no expert slicing).",
+                RuntimeWarning,
+            )
+            raise ValueError("Number of MoE experts must be at least equal to the number of parallel ranks.")
+        if self.moe_num_experts % moe_ranks != 0:
+            raise ValueError("Number of MoE experts must be divisible by the total number of parallel ranks (tp * cp ).")
+        return max(1, math.ceil(self.moe_num_experts / moe_ranks))
+
+    @experts_per_gpu.setter
+    def experts_per_gpu(self, value: int) -> None:
+        # Preserve the base class assignment for debugging, but the getter drives usage.
+        self._experts_per_gpu_base_value = value
+
+
+
     def _param_stats_per_rank(
         self,
         hidden_dim: int,
@@ -297,6 +410,7 @@ def get_tensor_reduction_time(self, total_bytes: int, kind: str, name: str, part
             ll=self.LLTP,
             local_bytes=0,
             debug_label=name or "comm",
+            axis="tp",
         )
         return reduction_time
 
@@ -316,7 +430,7 @@ def flash_attention_kernel_forward(self, batch_size, hidden_dim, seq_len, num_he
         # assuming key and value are preloaded into shared memory before attention computation
         load_kv_bytes = Bc * d * 2 * num_SMs * self.precision.activations #load key and value for one tile from HBM to SRAM
         initial_load_time = self.roofline(0, load_kv_bytes, "flash_attention_initial_load", mem_level=self.num_levels - 1) #assume key and value of one attention head is loaded from HBM to SRAM
-        initial_load_time_per_head = initial_load_time * math.ceil(Tc/ num_SMs)
+
         
         # attention score gemm
         load_q_bytes = Br * d * self.precision.activations #load query for one tile assuming k is already in shared memory
@@ -333,12 +447,13 @@ def flash_attention_kernel_forward(self, batch_size, hidden_dim, seq_len, num_he
         attn_output_time_per_tile = self.getGEMMTime(Br, Bc, d, "attention_output", read_bytes_l2=output_bytes, write_bytes_l2=output_bytes, flashattn_enable=True)[0] 
         attn_output_time = attn_output_time_per_tile * Tc * Tr #attention output gemm time for one head
         
-        load_time = initial_load_time_per_head * batch_size * kv_heads / max(1, self.tp) #load time for key and value are needed once per kv head
+        
         attn_score_time *= batch_size * num_heads / max(1, self.tp) + self.O
         attn_scale_softmax_time *= batch_size * num_heads / max(1, self.tp)
         attn_output_time *= batch_size * num_heads / max(1, self.tp) + self.O
-        
-        attention_forward_gemm_time = load_time + attn_score_time + attn_scale_softmax_time + attn_output_time
+
+
+        attention_forward_gemm_time = initial_load_time + attn_score_time + attn_scale_softmax_time + attn_output_time
         attention_forward_time = attention_forward_gemm_time + attention_forward_reduction_time
 
         # HBM traffic consists of only reading Q, K, V once and writing output once
@@ -347,48 +462,85 @@ def flash_attention_kernel_forward(self, batch_size, hidden_dim, seq_len, num_he
         return attention_forward_time, attention_forward_gemm_time, attention_forward_reduction_time, attention_size_f, attention_mem
     
     
-    def flash_attention_kernel_backward(self) -> Tuple[Dict[str, float], Dict[str, float], Dict[str, float]]:
-        #TODO: should be repleced with a modified version of getgemmtime that considers flash attention tiling
-        """Return time for flash attention kernel."""
-        #FIXME not finished
-        #TODO gqa support
-        batch_size = self.microB
-        seq_len = self.seq_len
-        hidden_dim = self.hidden_dim
-        num_heads = self.num_heads
+    def flash_attention_kernel_backward(self, batch_size, hidden_dim, seq_len, num_heads, kv_heads, num_SMs) -> Tuple[Dict[str, float], Dict[str, float], Dict[str, float]]:
+        """Return time for flash attention backward kernel."""
+        
         shard_seq = math.ceil(seq_len / max(1, self.cp)) 
-        d = hidden_dim // num_heads #gemm shape for one head is (seq_len, d) x (d, seq_len)
-        Bc = self.attention_tile_size
-        Br = min(Bc, d)
-        Tr = math.ceil(shard_seq / Br) #number of row tiles
-        Tc =  math.ceil(shard_seq / Bc) #number of column tiles
-        print(f"Flash Attention: Br={Br}, Tr={Tr}, d={d}, Bc={Bc}, Tc={Tc}")
-        #TODO load 3 tiled gemm inputs and 1 output from HBM to SRAM
+        d = hidden_dim // num_heads # gemm shape for one head is (seq_len, d) x (d, seq_len)
+        Bc = self.attention_tile_size # kv tile size
+        Br = min(Bc, d) # q tile size
+        Tr = math.ceil(shard_seq / Br) #number of q tiles
+        Tc =  math.ceil(shard_seq / Bc) #number of kv tiles
         
-        #TODO disable HBM access
-        attention_forward_reduction_time = 0
-        attention_size_f = 0
+        attention_backward_reduction_time = 0
+        attention_size_b = 0
+        recompute_time = 0
+        # assuming key and value are preloaded into shared memory before attention backward computation
+        load_kv_bytes = Bc * d * 2 * num_SMs * self.precision.activations #load key and value for one tile from HBM to SRAM
+        initial_load_time = self.roofline(0, load_kv_bytes, "flash_attention_initial_load_b", mem_level=self.num_levels - 1) #assume key and value of one attention head is loaded from HBM to SRAM
+
         
-        attn_score_time = self.getGEMMTime(Br, d, Bc, "attention_score_f")[0] * Tc * Tr #attention score gemm time for one head
-        number_flops =  Br * Bc * d * Tc * Tr * batch_size * num_heads
-        print(f"Flash Attention: attention score gemm flops={number_flops/1e9} GFLOPS")
-         
-        # Softmax time
+        # attention score recompute
+        load_q_bytes = Br * d * self.precision.activations #load query for one tile assuming k is already in shared memory
+        attn_score_time_per_tile = self.getGEMMTime(Br, d, Bc, "attention_score_b",read_bytes_l2=load_q_bytes, flashattn_enable=True)[0] #recompute S = QK^T for backward [Br, Bc]
+        attn_score_time = attn_score_time_per_tile * Tc * Tr #attention score gemm time for one head
+        
+        # Softmax recompute time 
         elements = Br * Bc
-        flops = SOFTMAX_FORWARD_FLOPS_PER_ELEMENT * elements  # exponentiation + subtract max + dropout
-        attn_scale_softmax_time = self.roofline(flops, 1, "attention_scale_softmax", mem_level=self.num_levels - 1) * Tc * Tr #use roofline model for softmax time with no memory access, memory access set to 1 because roofline does not accept 0 memory access
-        attn_output_time = self.getGEMMTime(Br, Bc, d, "attention_output")[0] * Tc * Tr #attention output gemm time for one head
-        attn_score_time *= batch_size * num_heads / max(1, self.tp)
+        flops = SOFTMAX_FORWARD_FLOPS_PER_ELEMENT * elements  
+        attn_scale_softmax_time = self.roofline(flops, 1, "attention_scale_softmax_recompute", mem_level=self.num_levels - 1) * Tc * Tr #use roofline model for softmax time with no memory access, memory access set to 1 because roofline does not accept 0 memory access
+        
+        # dV dP
+        load_o_bytes = Br * d * self.precision.activations #load output for one tile assuming v is already in shared memory
+        act_dO_time_per_tile = self.getGEMMTime(Bc, Br, d, "attention_output_b", read_bytes_l2=load_o_bytes, flashattn_enable=True)[0] #dV = P^T * dO  [Bc, d]
+        act_dP_time_per_tile = self.getGEMMTime(Br, d, Bc, "attention_score_b", read_bytes_l2=0, flashattn_enable=True)[0] #dP = dO * V^T [Br, Bc]
+        act_dO_time = act_dO_time_per_tile * Tc * Tr #dV time for one head
+        act_dP_time = act_dP_time_per_tile * Tc * Tr #dP time for one head
+        
+        # compute dS
+        elements = Br * Bc
+        flops = (SOFTMAX_BACKWARD_FLOPS_PER_ELEMENT ) * elements 
+        softmax_time_backward = self.roofline(flops, 1, "attention_scale_softmax_b", mem_level=self.num_levels - 1) * Tc * Tr #use roofline model for softmax time with no memory access, memory access set to 1 because roofline does not accept 0 memory access
+        
+        # dQ
+        act_dQ_time_per_tile = self.getGEMMTime(Br, Bc, d, "attention_score_b", read_bytes_l2=0, flashattn_enable=True)[0] #dQ = dS * K [Br, d]
+        act_dQ_time = act_dQ_time_per_tile * Tc * Tr #dQ time for one head
+        
+        
+        attn_score_time *= batch_size * num_heads / max(1, self.tp) + self.O
         attn_scale_softmax_time *= batch_size * num_heads / max(1, self.tp)
-        attn_output_time *= batch_size * num_heads / max(1, self.tp)
+        act_dO_time *= batch_size * num_heads / max(1, self.tp) + self.O
+        act_dP_time *= batch_size * num_heads / max(1, self.tp) + self.O
+        softmax_time_backward *= batch_size * num_heads / max(1, self.tp)
+        act_dQ_time *= batch_size * num_heads / max(1, self.tp) + self.O
+        if self.full_recomputation:  #attention recompute is already included in full recomputation   
+            recompute_time = 0
+        else:
+            recompute_time = attn_score_time + attn_scale_softmax_time #selective recomputation only recompute attention score and softmax
+        attention_backward_gemm_time = initial_load_time + recompute_time + act_dO_time + act_dP_time + softmax_time_backward + act_dQ_time
+        attention_backward_time = attention_backward_gemm_time + attention_backward_reduction_time
         
-        attention_forward_gemm_time = attn_score_time + attn_scale_softmax_time + attn_output_time
         
-        attention_forward_time = attention_forward_gemm_time + attention_forward_reduction_time
-        attention_size_f = 0
+        attention_size_b = self.precision.grad_communication * batch_size * seq_len * hidden_dim * 2 / self.tp # weight gradient of K V need to be reduce scattered  *2 account for both attn key and value
+        kind = "reduce_scatter"
+        participants = self.cp
+        axis_hint = "cp"
+        if attention_size_b > 0:
+            attention_backward_reduction_time = self.network_model.collective(
+            kind=kind,
+            size_bytes=attention_size_b,
+            participants=participants,
+            ib=self.IBTP,
+            ll=self.LLTP,
+            local_bytes=0,
+            debug_label= "attention_backward_reduction",
+            axis=axis_hint,
+        )
+           
+        
         
     
-        return attention_forward_time, attention_forward_gemm_time, attention_forward_reduction_time, attention_size_f
+        return attention_backward_time, attention_backward_gemm_time, attention_backward_reduction_time, attention_size_b
         
         
     @staticmethod
@@ -432,9 +584,6 @@ def single_gpu_gemm_forward(self, gemm: Tuple[int, ...], name: str, gemm_type: O
         gemm_type = self._normalize_gemm_type(gemm_type)
         if gemm_type in (GemmType.ATTENTION_SCORE, GemmType.ATTENTION_OUTPUT):  # attention gemm
             gemm_time = self.getGEMMTime(m, k, n, name, disable_overhead=True)[0] * batch + self.O
-        elif (gemm_type == GemmType.FFN1 or gemm_type == GemmType.FFN2) and self.use_moe:
-            gemm_time,_,_, mem_accesses = self.getGEMMTime(m, k, n, name)
-            gemm_time *= self.moe_num_experts
         else :
             gemm_time,_,_, mem_accesses = self.getGEMMTime(m, k, n, name)
         return gemm_time, 0, 0, total_flops, mem_accesses
@@ -445,9 +594,6 @@ def single_gpu_gemm_backward(self, gemm: Tuple[int, ...], name: str, gemm_type:
         if gemm_type in (GemmType.ATTENTION_SCORE, GemmType.ATTENTION_OUTPUT):  # attention gemm
             grad_time_act = self.getGEMMTime(m, k, n, name, disable_overhead=True)[0] * batch + self.O
             grad_time_wt = self.getGEMMTime(k, m, n, name, disable_overhead=True)[0] * batch + self.O
-        elif (gemm_type == GemmType.FFN1 or gemm_type == GemmType.FFN2) and self.use_moe:
-            grad_time_act = self.getGEMMTime(m, n, k, name)[0] * self.moe_num_experts
-            grad_time_wt = self.getGEMMTime(k, m, n, name)[0] * self.moe_num_experts
         else :
             grad_time_act = self.getGEMMTime(m, n, k, name)[0]
             grad_time_wt = self.getGEMMTime(k, m, n, name)[0]
@@ -492,6 +638,8 @@ def _tensor_context_hybrid_gemm_forward(self, gemm: Tuple[int, ...], name: str,
         if gemm_type is None:
             raise ValueError("gemm_type is required for tensor-context hybrid forward GEMM")
         
+        axis_hint = None
+
         if gemm_type == GemmType.ATTENTION_SCORE:  # attention gemm
             gemm_time = self.getGEMMTime(shard_m, k, n, name, disable_overhead=True)[0] * batch / max(1, self.tp) + self.O
         elif gemm_type == GemmType.ATTENTION_OUTPUT:  # attention gemm
@@ -501,16 +649,19 @@ def _tensor_context_hybrid_gemm_forward(self, gemm: Tuple[int, ...], name: str,
             total_bytes = self.get_kv_size_bytes()  / self.tp # each tp group holds a tp shard of kv for each cp group
             kind = "all_gather"
             participants = self.cp # all gather K V for each cp group
+            axis_hint = "cp"
         elif gemm_type == GemmType.OUT_PROJ:
             gemm_time = self.getGEMMTime(shard_m, shard_k, n, name)[0]
             total_bytes = math.ceil(self.precision.activations * shard_m * n)
             kind = "reduce_scatter"
             participants = self.tp # reduce scatter output activation for each tp group
+            axis_hint = "tp"
         elif gemm_type == GemmType.FFN2:  # row wise
             gemm_time = self.getGEMMTime(shard_m, shard_k, n, name)[0]
             total_bytes = math.ceil(self.precision.activations * shard_m * n)
             kind = "reduce_scatter"
             participants = self.tp  # reduce scatter output activation for each tp group
+            axis_hint = "tp"
         elif gemm_type == GemmType.FFN1:  # column wise
             gemm_time = self.getGEMMTime(shard_m, k, shard_n, name)[0]
         else:
@@ -525,6 +676,7 @@ def _tensor_context_hybrid_gemm_forward(self, gemm: Tuple[int, ...], name: str,
             ll=self.LLTP,
             local_bytes=0,
             debug_label=name or "comm",
+            axis=axis_hint,
         )
         return gemm_time, reduction_time, total_bytes
     def _tensor_context_hybrid_gemm_backward(self, gemm: Tuple[int, ...], name: str, gemm_type: Optional[GemmType] = None) -> Tuple[float, float]:
@@ -563,12 +715,15 @@ def _tensor_context_hybrid_gemm_backward(self, gemm: Tuple[int, ...], name: str,
         if gemm_type is None:
             raise ValueError("gemm_type is required for tensor-context hybrid backward GEMM")
         
+        axis_hint = None
+
         if gemm_type == GemmType.ATTENTION_SCORE:  # attention gemm
             grad_time_act = self.getGEMMTime(shard_m, n, k, name, disable_overhead=True)[0] * batch / max(1, self.tp) + self.O
             grad_time_wt = self.getGEMMTime(k, shard_m, n, name, disable_overhead=True)[0] * batch / max(1, self.tp) + self.O
             total_bytes = self.precision.grad_communication * k * n * batch * 2 / self.tp # weight gradient of K V need to be reduce scattered  *2 account for both attn key and value
             kind = "reduce_scatter"
             participants = self.cp
+            axis_hint = "cp"
         elif gemm_type == GemmType.ATTENTION_OUTPUT:  # attention gemm
             grad_time_act = self.getGEMMTime(shard_m, n, k, name)[0] * batch / max(1, self.tp)
             grad_time_wt = self.getGEMMTime(k, shard_m, n, name)[0] * batch / max(1, self.tp)
@@ -578,22 +733,25 @@ def _tensor_context_hybrid_gemm_backward(self, gemm: Tuple[int, ...], name: str,
             total_bytes = math.ceil(self.precision.grad_communication * shard_m * k)
             kind = "reduce_scatter"
             participants = self.tp
+            axis_hint = "tp"
         elif gemm_type == GemmType.OUT_PROJ:
             grad_time_act = self.getGEMMTime(shard_m, n, shard_k, name)[0]
             grad_time_wt = self.getGEMMTime(shard_k, shard_m, n, name)[0]
             total_bytes = self.get_kv_size_bytes() / self.tp
             kind = "all_gather"
             participants = self.cp
+            axis_hint = "cp"
         elif gemm_type == GemmType.FFN2:  # row wise
             grad_time_act = self.getGEMMTime(shard_m, n, shard_k, name)[0]
             grad_time_wt = self.getGEMMTime(shard_k, shard_m, n, name)[0]
         elif gemm_type == GemmType.FFN1:  # column wise
-            
+
             grad_time_act = self.getGEMMTime(shard_m, shard_n, k, name)[0]
             grad_time_wt = self.getGEMMTime(k, shard_m, shard_n, name)[0]
             total_bytes = math.ceil(self.precision.grad_communication * shard_m * k)
             kind = "reduce_scatter"
             participants = self.tp
+            axis_hint = "tp"
         elif gemm_type == GemmType.LINEAR_SOFTMAX:
             grad_time_act = self.getGEMMTime(shard_m, n, shard_k, name)[0]
             grad_time_wt = self.getGEMMTime(shard_k, shard_m, n, name)[0]
@@ -612,6 +770,7 @@ def _tensor_context_hybrid_gemm_backward(self, gemm: Tuple[int, ...], name: str,
             ll=self.LLTP,
             local_bytes=0,
             debug_label=name or "comm",
+            axis=axis_hint,
         )
         return gemm_time, reduction_time, total_bytes
     def _tensor_parallelism_gemm_forward(self, gemm: Tuple[int, ...], name: str, gemm_type: Optional[GemmType] = None) -> Tuple[float, float]:
@@ -657,10 +816,6 @@ def _tensor_parallelism_gemm_forward(self, gemm: Tuple[int, ...], name: str, gem
         if gemm_type in (GemmType.ATTENTION_SCORE, GemmType.ATTENTION_OUTPUT):  # attention gemm
             gemm_time,_,_, mem_accesses = self.getGEMMTime(m, k, n, name, disable_overhead=True)
             gemm_time = gemm_time * batch / max(1, self.tp) + self.O
-        elif gemm_type in (GemmType.FFN1, GemmType.FFN2) and self.use_moe:
-            # MoE FFN layers replicate experts across all TP ranks, so no TP sharding here.
-            gemm_time,_,_, mem_accesses = self.getGEMMTime(m, k, n, name)
-            gemm_time *= self.moe_num_experts
         elif gemm_type in (GemmType.QKV, GemmType.FFN1):  # column wise
             shard_n = math.ceil(n / max(1, self.tp))
             gemm_time,_,_, mem_accesses = self.getGEMMTime(m, k, shard_n, name)
@@ -678,7 +833,7 @@ def _tensor_parallelism_gemm_forward(self, gemm: Tuple[int, ...], name: str, gem
             raise ValueError(f"Unsupported gemm type: {gemm_type}")
             
         if size_bytes > 0:
-            total_bytes = size_bytes #we already has the totol bytes for all reduce not bytes per rank
+            total_bytes = size_bytes # we already has the total bytes for all reduce not bytes per rank
             reduction_time = self.get_tensor_reduction_time(total_bytes, kind=comm_kind_fwd, participants=participants, name=name)
 
 
@@ -791,6 +946,7 @@ def _context_parallelism_gemm_forward(self, gemm: Tuple[int, ...], name: str, ge
                 ll=self.LLTP,
                 local_bytes=0,
                 debug_label=name or "comm",
+                axis="cp",
             )
 
         return gemm_time, reduction_time, total_bytes
@@ -818,7 +974,7 @@ def _context_parallelism_gemm_backward(self, gemm: Tuple[int, ...], name: str, g
         if gemm_type == GemmType.ATTENTION_SCORE:
             grad_time_act = self.getGEMMTime(shard_m, n, k, name, disable_overhead=True)[0] * batch + self.O
             grad_time_wt = self.getGEMMTime(k, shard_m, n, name, disable_overhead=True)[0] * batch + self.O
-            total_bytes = self.precision.grad_communication * k * n * batch * 2 # account for both Q and K
+            total_bytes = self.precision.grad_communication * k * n * batch * 2 # account for both K and V
             kind = "reduce_scatter"
         elif gemm_type == GemmType.ATTENTION_OUTPUT:  # attention gemm
             grad_time_act = self.getGEMMTime(shard_m, n, k, name, disable_overhead=True)[0] * batch + self.O
@@ -843,9 +999,56 @@ def _context_parallelism_gemm_backward(self, gemm: Tuple[int, ...], name: str, g
                 ll=self.LLTP,
                 local_bytes=0,
                 debug_label=name or "comm",
+                axis="cp",
             )
         return gemm_time, reduction_time, total_bytes
-    
+    def get_moe_ffn_f(self, gemm, name, gemm_type: Optional[GemmType] = None):
+        batch, m, k, n = self._expand_gemm_descriptor(gemm)
+        reduction_time = 0
+        total_bytes = 0
+        total_flops = 2 * batch * m * k * n * self.experts_per_gpu
+        gemm_time, _, _ ,mem_accesses= self.getGEMMTime(m, k, n, name)
+        gemm_time *= self.experts_per_gpu
+
+        if gemm_type == GemmType.FFN2:
+            per_rank_bytes = math.ceil(self.precision.activations * m * n * self.experts_per_gpu)
+            total_bytes = int(math.ceil(per_rank_bytes * self.cp * self.tp )) 
+            reduction_time = self.network_model.collective(
+                kind="all_to_all",
+                size_bytes=total_bytes,
+                participants=self.cp * self.tp,
+                ib=self.IBTP,
+                ll=self.LLTP,
+                local_bytes=0,
+                debug_label="moe_ffn_f_all_to_all",
+            )
+        return gemm_time, reduction_time, total_bytes, total_flops, mem_accesses
+
+    def get_moe_ffn_b(self, gemm, name, gemm_type):
+        
+        batch, m, k, n = self._expand_gemm_descriptor(gemm)
+        reduction_time = 0
+        total_bytes = 0
+        grad_time_act = self.getGEMMTime(m, n, k, name)[0] * self.experts_per_gpu
+        grad_time_wt = self.getGEMMTime(k, m, n, name)[0] * self.experts_per_gpu
+        gemm_time = grad_time_act + grad_time_wt
+
+        if gemm_type == GemmType.FFN1:
+            per_rank_bytes = math.ceil(self.precision.activations * m * k * self.experts_per_gpu)
+            total_bytes = int(math.ceil(per_rank_bytes * self.cp * self.tp ))
+            reduction_time = self.network_model.collective(
+                kind="all_to_all",
+                size_bytes=total_bytes,
+                participants=self.tp * self.cp,
+                ib=self.IBTP,
+                ll=self.LLTP,
+                local_bytes=0,
+                debug_label="moe_ffn_b_all_to_all",
+                
+            )
+        
+        return gemm_time, reduction_time, total_bytes
+        
 
                 
     def get_embedding_f(self, vocab_size, seq_len, hidden_dim):
@@ -878,13 +1081,16 @@ def get_linear_softmax_f(self, gemm):
         """
         _, effective_m, k, n = self._effective_dims(gemm)
 
-        gemm_time, reduction_time, size_bytes,_,_ = self._tensor_parallelism_gemm_forward(gemm, "linear_softmax_f", gemm_type=GemmType.LINEAR_SOFTMAX)
+        # Previous TP-aware path kept comm time here:
+        # gemm_time, reduction_time, size_bytes, _, _ = self._tensor_parallelism_gemm_forward(
+        #     gemm, "linear_softmax_f", gemm_type=GemmType.LINEAR_SOFTMAX
+        # )
+        # Linear softmax is modeled as running on a single device within each TP group.
+        gemm_time, _, _, _, _ = self.single_gpu_gemm_forward(gemm, "linear_softmax_f", gemm_type=GemmType.LINEAR_SOFTMAX)
 
             
         elements = effective_m * n / (self.tp * self.cp) # each tp-cp group holds a shard of the vocab dimension
         point_flop = elements * SOFTMAX_FORWARD_FLOPS_PER_ELEMENT
-        # TODO:
-        # unsure if should be precision.activations or precision.stats
         point_mem = self.precision.activations * elements * SOFTMAX_FORWARD_MEM_ACCESSES 
         point_time = self.roofline(
             point_flop,
@@ -901,14 +1107,19 @@ def get_linear_softmax_f(self, gemm):
             )
             print("Linear Softmax (f) point_time: {:,}\n".format(point_time))
 
-        return gemm_time + reduction_time + point_time, point_mem
+        return gemm_time + point_time, point_mem
     
     def get_linear_softmax_b(self, gemm):
 
 
         _, effective_m, k, n = self._effective_dims(gemm)
 
-        gemm_time, reduction_time, size_bytes = self._tensor_parallelism_gemm_backward(gemm, "linear_softmax_b", gemm_type=GemmType.LINEAR_SOFTMAX)
+        # Previous TP-aware path kept comm time here:
+        # gemm_time, reduction_time, size_bytes = self._tensor_parallelism_gemm_backward(
+        #     gemm, "linear_softmax_b", gemm_type=GemmType.LINEAR_SOFTMAX
+        # )
+        # Linear softmax is modeled as running on a single device within each TP group.
+        gemm_time, _, _ = self.single_gpu_gemm_backward(gemm, "linear_softmax_b", gemm_type=GemmType.LINEAR_SOFTMAX)
         elements = effective_m * n / (self.tp * self.cp) # each tp-cp group holds a shard of the vocab dimension
         point_flop = elements * SOFTMAX_BACKWARD_FLOPS_PER_ELEMENT
         # TODO:
@@ -930,7 +1141,7 @@ def get_linear_softmax_b(self, gemm):
             )
             print("Linear Softmax (b) point_time: {:,}\n".format(point_time))
 
-        return gemm_time + reduction_time + point_time
+        return gemm_time + point_time
     def get_scale_softmax_f(self, gemm):
         """
         Estimate time for scale + softmax forward.
@@ -1118,6 +1329,7 @@ def get_layernorm_f(self, batch, seq_len, d_model, comm_after=False):
                 ll=self.LLTP,
                 local_bytes=0,
                 debug_label="layernorm_f_all_gather",
+                axis="tp",
             )
         else:
             reduction_time = 0.0
@@ -1127,7 +1339,7 @@ def get_layernorm_f(self, batch, seq_len, d_model, comm_after=False):
     
     
 
-    def get_layernorm_b(self, batch, seq_len, d_model, comm_after=False):
+    def get_layernorm_b(self, batch, seq_len, d_model, type = Optional):
         tp_mode = self.get_parallelism_mode()
         seq_degree = self._sequence_parallel_degree()
         if tp_mode == ParallelismMode.TENSOR_CONTEXT_HYBRID:
@@ -1143,7 +1355,21 @@ def get_layernorm_b(self, batch, seq_len, d_model, comm_after=False):
             name="pointwise-layernorm-b",
             mem_level=self.num_levels - 1,
         ) + 4 * self.O
-        if tp_mode in (ParallelismMode.TENSOR_SEQUENCE, ParallelismMode.TENSOR_CONTEXT_HYBRID):  # communication after layernorm
+        if self.use_moe and type == GemmType.LAYER_NORM_1:  # communication after layernorm
+            per_rank_bytes = self.precision.grad_communication * elements * self.moe_top_k
+            total_bytes = int(math.ceil(per_rank_bytes * self.tp * self.cp))
+            
+            reduction_time = self.network_model.collective(
+                kind="all_to_all",
+                size_bytes=total_bytes,
+                participants=self.tp * self.cp,
+                ib=self.IBTP,
+                ll=self.LLTP,
+                local_bytes=0,
+                debug_label="layernorm_b_all_to_all",
+            )
+
+        elif tp_mode in (ParallelismMode.TENSOR_SEQUENCE, ParallelismMode.TENSOR_CONTEXT_HYBRID) :
             per_rank_bytes = self.precision.grad_communication * elements
             total_bytes = int(math.ceil(per_rank_bytes * self.tp))
             reduction_time = self.network_model.collective(
@@ -1154,6 +1380,7 @@ def get_layernorm_b(self, batch, seq_len, d_model, comm_after=False):
                 ll=self.LLTP,
                 local_bytes=0,
                 debug_label="layernorm_b_all_gather",
+                axis="tp",
             )
 
         else:
@@ -1162,7 +1389,57 @@ def get_layernorm_b(self, batch, seq_len, d_model, comm_after=False):
 
 
         return compute_time, reduction_time, total_bytes
+    def get_router_f(self, gemm_router, gemm_ffn1):
+        #TODO router overhead for moe is very minimal, can be ignored, implemented for completeness
+        _, effective_m, k, n = self._effective_dims(gemm_router)
+        effective_m = effective_m / (self.cp)
+        gemm_time = self.single_gpu_gemm_forward(gemm_router, "router_f")[0]
+        elements = effective_m * n
+        flops = elements * SOFTMAX_FORWARD_FLOPS_PER_ELEMENT
+        mem_bytes = self.precision.activations * elements * SOFTMAX_FORWARD_MEM_ACCESSES  
+
+        compute_time = gemm_time + self.roofline(flops, mem_bytes, name="pointwise-router-f", mem_level=self.num_levels - 1) + self.O
+
+        per_rank_bytes = math.ceil(self.precision.activations * gemm_ffn1[0] * gemm_ffn1[1] * self.experts_per_gpu)
+        total_bytes = int(math.ceil(per_rank_bytes * self.cp * self.tp )) 
+        reduction_time = self.network_model.collective(
+            kind="all_to_all",
+            size_bytes=total_bytes,
+            participants = self.tp * self.cp,
+            ib=self.IBTP,
+            ll=self.LLTP,
+            local_bytes=0,
+            debug_label="router_f_all_to_all",
+        )
+
+
+        return compute_time, reduction_time, total_bytes
+    def get_router_b(self, gemm_router, gemm_ffn1):
+        #TODO router overhead for moe is very minimal, can be ignored, implemented for completeness
+        _, effective_m, k, n = self._effective_dims(gemm_router)
+        elements = effective_m * k / (self.tp * self.cp)
+
+        gemm_time = self.single_gpu_gemm_backward(gemm_router, "router_f")[0]
+
+        flops = elements * SOFTMAX_BACKWARD_FLOPS_PER_ELEMENT
+        mem_bytes = self.precision.activations * elements * SOFTMAX_BACKWARD_MEM_ACCESSES
+
+        compute_time = gemm_time + self.roofline(flops, mem_bytes, name="pointwise-router-f", mem_level=self.num_levels - 1) + self.O
+        bytes_per_rank = math.ceil(self.precision.activations * elements)
+
+        total_bytes = int(math.ceil(bytes_per_rank * self.tp))
+        reduction_time = self.network_model.collective(
+            kind="reduce_scatter",
+            size_bytes=total_bytes,
+            participants=self.tp ,
+            ib=self.IBTP,
+            ll=self.LLTP,
+            local_bytes=0,
+            debug_label="router_b_reduce_scatter",
+        )
+
 
+        return compute_time, reduction_time, total_bytes
     
     def get_embedding_b(self, vocab_size, seq_len, hidden_dim):
         batch = self._effective_transformer_batch()
@@ -1194,12 +1471,7 @@ def get_data_parallel_reduction_sizes(self, d, intermediate_size):
         """Calculate communication sizes for data parallel reductions (no timing)."""
         if not getattr(self, "dp", 1) or self.dp <= 1:
             # No communication needed for dp=1
-            return {
-                'qkv_size': 0,
-                'output_size': 0,
-                'ffn_size': 0,
-                'total_size': 0
-            }
+            return 0
 
         # Calculate sizes only
         qkv_size = math.ceil(self.precision.grad_communication * d * 3 * d)
@@ -1209,187 +1481,69 @@ def get_data_parallel_reduction_sizes(self, d, intermediate_size):
         ffn2_size = math.ceil(self.precision.grad_communication * intermediate_size * d)
         total_size = qkv_size + output_size + ffn1_size + ffn2_size
 
-        return {
-            'qkv_size': qkv_size,
-            'output_size': output_size,
-            'ffn_size': ffn1_size + ffn2_size,
-            'total_size': total_size
-        }
+        return total_size
 
-    def get_data_parallel_local_computation(self, d, intermediate_size):
-        """Calculate local computation times for apply_grad operations."""
-        qkv_local = self.apply_grad(Dim0=d, Dim1=3*d, name="qkv_proj reduction")
-        output_local = self.apply_grad(Dim0=d, Dim1=d, name="output_proj reduction")
+    def get_data_parallel_reduction_llm(self, d, intermediate_size):
+        """Return apply_grad compute time per rank (no communication), honoring ZeRO sharding."""
+        apply_grad_time = self.apply_grad(int(d * 3 * d)) # QKV
+        apply_grad_time += self.apply_grad(int(d * d)) # Output
         ffn1_dim = self._ffn1_output_dim(intermediate_size)
-        ffn_local = (
-            self.apply_grad(Dim0=ffn1_dim, Dim1=d, name="ffn1 reduction")
-            + self.apply_grad(Dim0=intermediate_size, Dim1=d, name="ffn2 reduction")
-        )
+        apply_grad_time += self.apply_grad(int(ffn1_dim * d)) # FFN1
+        apply_grad_time += self.apply_grad(int(intermediate_size * d)) # FFN2
 
-        return {
-            'qkv_local': qkv_local,
-            'output_local': output_local,
-            'ffn_local': ffn_local,
-            'total_local': qkv_local + output_local + ffn_local
-        }
-
-    def get_data_parallel_reduction_llm(self, d, intermediate_size): 
-        # If no data parallelism, still apply gradients locally but no cross-device reduction
-        if self.dp <= 1:
-            apply_grad_time = 0.0
-            apply_grad_time += self.apply_grad(Dim0=d, Dim1=3*d, name="qkv_proj reduction")
-            apply_grad_time += self.apply_grad(Dim0=d, Dim1=d, name="output_proj reduction")
-            ffn1_dim = self._ffn1_output_dim(intermediate_size)
-            apply_grad_time += self.apply_grad(Dim0=ffn1_dim, Dim1=d, name="ffn1 reduction")
-            apply_grad_time += self.apply_grad(Dim0=intermediate_size, Dim1=d, name="ffn2 reduction")
-            if self.debug:
-                print(f"(dp=1) apply_grad_time: {apply_grad_time}")
-            return apply_grad_time
-
-        reduction_sizes = self.get_data_parallel_reduction_sizes(d, intermediate_size)
-        grad_bytes = int(reduction_sizes['total_size'])
-        (
-            total_params_per_rank,
-            _,
-            _,
-            _,
-            _,
-        ) = self._param_stats_per_rank(d, intermediate_size, self.vocab_size)
-        param_bytes = int(math.ceil(total_params_per_rank * self.precision.parameters))
-        grad_shard = self.dp if self.zero_stage >= 2 else 1
-
-        if grad_shard > 1: # ZeRO stages >= 2: gradients are sharded via reduce-scatter and parameters are gathered.
-            apply_grad_time = 0.0
-            apply_grad_time += self.apply_grad(Dim0=d, Dim1=3 * d, name="qkv_proj reduction")
-            apply_grad_time += self.apply_grad(Dim0=d, Dim1=d, name="output_proj reduction")
-            ffn1_dim = self._ffn1_output_dim(intermediate_size)
-            apply_grad_time += self.apply_grad(Dim0=ffn1_dim, Dim1=d, name="ffn1 reduction")
-            apply_grad_time += self.apply_grad(Dim0=intermediate_size, Dim1=d, name="ffn2 reduction")
-            # divide by dp to get the time per device 
+        grad_shard = self.dp if (self.zero_stage >= 2 and self.dp > 1) else 1
+        if grad_shard > 1:
             apply_grad_time /= grad_shard
 
-            reduction_time = self.network_model.collective(
-                kind="reduce_scatter",
-                size_bytes=grad_bytes,
-                participants=int(self.dp),
-                ib=self.IBD,
-                ll=self.LLD,
-                local_bytes=0.0,
-                debug_label="dp_zero_gradients",
-            )
-
-            gather_time = 0.0
-            if self.zero_stage == 2:
-                gather_bytes = int(math.ceil(param_bytes))
-                gather_time += self.network_model.collective(
-                    kind="all_gather",
-                    size_bytes=gather_bytes,
-                    participants=int(self.dp),
-                    ib=self.IBD,
-                    ll=self.LLD,
-                    local_bytes=0.0,
-                    debug_label="dp_zero_param_gather_zero2",
-                )
-            if self.debug:
-                print(f"apply_grad_time (ZeRO-{self.zero_stage}): {apply_grad_time}")
-            # TODO(ZeRO): revisit optimizer FLOPs once ZeRO-3 materialization is finalized.
-            return reduction_time + gather_time + apply_grad_time
-        else: # legacy path for zero stages <= 2.
-            reduction_time = 0.0
-            apply_grad_time = 0.0
-
-            total_bytes = math.ceil(self.precision.grad_communication * d * 3 * d)
-            reduction_time += self.network_model.collective(
-                kind="all_reduce",
-                size_bytes=total_bytes,
-                participants=int(self.dp),
-                ib=self.IBD,
-                ll=self.LLD,
-                local_bytes=0.0,
-                debug_label="qkv_proj reduction",
-            )
-            apply_grad_time += self.apply_grad(Dim0=d, Dim1=3*d, name="qkv_proj reduction")
-
-            total_bytes = math.ceil(self.precision.grad_communication * d * d)
-            reduction_time += self.network_model.collective(
-                kind="all_reduce",
-                size_bytes=total_bytes,
-                participants=int(self.dp),
-                ib=self.IBD,
-                ll=self.LLD,
-                local_bytes=0.0,
-                debug_label="output_proj reduction",
-            )
-            apply_grad_time += self.apply_grad(Dim0=d, Dim1=d, name="output_proj reduction")
-
-            ffn1_dim = self._ffn1_output_dim(intermediate_size)
-            total_bytes = math.ceil(self.precision.grad_communication * ffn1_dim * d)
-            reduction_time += self.network_model.collective(
-                kind="all_reduce",
-                size_bytes=total_bytes,
-                participants=int(self.dp),
-                ib=self.IBD,
-                ll=self.LLD,
-                local_bytes=0.0,
-                debug_label="ffn1 reduction",
-            )
-            apply_grad_time += self.apply_grad(Dim0=ffn1_dim, Dim1=d, name="ffn1 reduction")
-
-            total_bytes = math.ceil(self.precision.grad_communication * intermediate_size * d)
-            reduction_time += self.network_model.collective(
-                kind="all_reduce",
-                size_bytes=total_bytes,
-                participants=int(self.dp),
-                ib=self.IBD,
-                ll=self.LLD,
-                local_bytes=0.0,
-                debug_label="ffn2 reduction",
-            )
-            apply_grad_time += self.apply_grad(Dim0=intermediate_size, Dim1=d, name="ffn2 reduction")
-
-            if self.debug:
-                print(f"apply_grad_time: {apply_grad_time}")
-
-            return reduction_time + apply_grad_time
+        return apply_grad_time
     
     def _combine_mem(self, *args):
             combined = {}
             for d in args:
                 mem_levels = self._mem_levels(d)
                 for k, v in mem_levels.items():
-                    combined[k] = combined.get(k, 0) + v
+                    combined[k] = combined.get(k, 0.0) + (v or 0.0)
             return combined
         
     def _mem_levels(self, arr):
+        # Accept both dict-like and sequence-like inputs since GEMM helpers return lists while
+        # aggregated OperationTiming instances already carry dicts. MappingABC/SequenceABC cover
+        # any object implementing the mapping/sequence protocol (e.g., dict, defaultdict, numpy arrays).
+        if isinstance(arr, MappingABC):
+            return {str(k): float(v) for k, v in arr.items()}
         if isinstance(arr, (int, float)):
-            return {f"L{i}": (arr if i == self.num_levels - 1 else 0) for i in range(self.num_levels)}
-        return {f"L{i}": v for i, v in enumerate(arr or [])}
+            return {f"L{i}": (float(arr) if i == self.num_levels - 1 else 0.0) for i in range(self.num_levels)}
+        if isinstance(arr, SequenceABC) and not isinstance(arr, (str, bytes)):
+            return {f"L{i}": float(v) for i, v in enumerate(arr)}
+        return {}
 
     # TODO TODO:
     # we need a significant refactor here. The comm sizes are ingested in a weird way and never used. Instead we use old precomputed sizes.
     # FIX at some point!
-    def compute_all_gemm_and_node_times(self, batch_size, vocab_size, hidden_dim, seq_len, num_heads, kv_heads, intermediate_size, num_SMs):
-        """Compute latency for all GEMM operations and node breakdown times in one function."""
+    def compute_all_gemm_and_node_times(
+        self,
+        batch_size,
+        vocab_size,
+        hidden_dim,
+        seq_len,
+        num_heads,
+        kv_heads,
+        intermediate_size,
+        num_SMs,
+    ):
+        """Compute latency for all GEMM operations and node breakdown times."""
 
-        # Process GEMM shapes
         gemm_shapes = LLM_util.process_gemm_shapes(
-            self, batch_size, seq_len, hidden_dim, num_heads, kv_heads, intermediate_size, vocab_size
+            self,
+            batch_size,
+            seq_len,
+            hidden_dim,
+            num_heads,
+            kv_heads,
+            intermediate_size,
+            vocab_size,
         )
-        if self.debug:
-            print(
-                "generating gemm shapes for transformer batch size:",
-                batch_size,
-                "seq_len:",
-                seq_len,
-                "hidden_dim:",
-                hidden_dim,
-                "num_heads:",
-                num_heads,
-                "kv_heads:",
-                kv_heads,
-                "intermediate_size:",
-                intermediate_size,
-            )
+
         gemm_qkv_proj = gemm_shapes["qkv_proj"]
         gemm_attention_score = gemm_shapes["attention_score"]
         gemm_attention_output = gemm_shapes["attention_output"]
@@ -1397,290 +1551,417 @@ def compute_all_gemm_and_node_times(self, batch_size, vocab_size, hidden_dim, se
         gemm_ffn1 = gemm_shapes["ffn1"]
         gemm_ffn2 = gemm_shapes["ffn2"]
         gemm_linear = gemm_shapes["linear"]
+        gemm_router = gemm_shapes["router"]
 
-        # Compute GEMM times and organize in structured dict
-        transformer_results = {}
-
-        # assuming no context-parallelism for now and only running inference
+        transformer_timings: Dict[str, OperationTiming] = {}
 
-        # Helpers to handle variable-length returns (supports TP/Single returns with flops/mem and CP/Hybrid without)
-        def _extract_forward(ret):
-            # ret may be (time, red, size, flops, mem) or (time, red, size)
-            if len(ret) >= 5:
-                return ret[0], ret[1], ret[2], ret[3], ret[4]
-            elif len(ret) == 3:
-                return ret[0], ret[1], ret[2], 0, []
-            else:
-                return ret[0], 0, 0, 0, []
+        def _extract_forward(ret: Sequence[Any]) -> Tuple[float, float, float, float, Any]:
+            if len(ret) == 5:
+                time, reduction, size, flops, mem = ret
+                return time, reduction, size, flops, mem
+            if len(ret) == 3:
+                time, reduction, size = ret
+                return time, reduction, size, 0.0, []
+            raise ValueError(f"Unsupported return length: {len(ret)}")
 
-        # QKV Projection GEMM
-        qkv_proj_gemm_f,  qkv_proj_reduction_f, qkv_proj_size_f, qkv_proj_flops_f, qkv_proj_mem_f = _extract_forward(
+        # QKV
+        qkv_proj_gemm_f, qkv_proj_reduction_f, qkv_proj_size_f, qkv_proj_flops_f, qkv_proj_mem_f = _extract_forward(
             self.parallelism_gemm_forward(gemm_qkv_proj, "qkv_projection_f", gemm_type=GemmType.QKV)
         )
-        qkv_proj_gemm_b,  qkv_proj_reduction_b, qkv_proj_size_b = self.parallelism_gemm_backward(gemm_qkv_proj, "qkv_projection_b", gemm_type=GemmType.QKV)
-        qkv_proj_f = qkv_proj_gemm_f + qkv_proj_reduction_f
-        qkv_proj_b = qkv_proj_gemm_b + qkv_proj_reduction_b
-        transformer_results['qkv_proj'] = {
-            'forward': qkv_proj_f, 'backward': qkv_proj_b,
-            'forward_gemm': qkv_proj_gemm_f, 'forward_reduction': qkv_proj_reduction_f,
-            'backward_gemm': qkv_proj_gemm_b, 'backward_reduction': qkv_proj_reduction_b,
-            'comm_size_forward': qkv_proj_size_f, 'comm_size_backward': qkv_proj_size_b,
-            'flops': qkv_proj_flops_f,
-            'memory_accesses': self._mem_levels(qkv_proj_mem_f),
-        }
+        qkv_proj_gemm_b, qkv_proj_reduction_b, qkv_proj_size_b = self.parallelism_gemm_backward(
+            gemm_qkv_proj, "qkv_projection_b", gemm_type=GemmType.QKV
+        )
+        qkv_forward = DirectionTiming(
+            compute_time=qkv_proj_gemm_f,
+            comm_time=qkv_proj_reduction_f,
+            comm_bytes=int(qkv_proj_size_f or 0),
+            flops=qkv_proj_flops_f,
+            memory_accesses=dict(self._mem_levels(qkv_proj_mem_f)),
+        )
+        qkv_backward = DirectionTiming(
+            compute_time=qkv_proj_gemm_b,
+            comm_time=qkv_proj_reduction_b,
+            comm_bytes=int(qkv_proj_size_b or 0),
+        )
+        transformer_timings["qkv_proj"] = OperationTiming("qkv_proj", forward=qkv_forward, backward=qkv_backward)
 
-        if self.flash_attention is False:
-            # Attention Score GEMM
-            attn_score_gemm_f,  attn_score_reduction_f, attn_score_size_f, attn_score_flops_f, attn_score_mem_f = _extract_forward(
+        if not self.flash_attention:
+            attn_score_gemm_f, attn_score_reduction_f, attn_score_size_f, attn_score_flops_f, attn_score_mem_f = _extract_forward(
                 self.parallelism_gemm_forward(gemm_attention_score, "attention_score_f", gemm_type=GemmType.ATTENTION_SCORE)
             )
-            attn_score_gemm_b,  attn_score_reduction_b, attn_score_size_b = self.parallelism_gemm_backward(gemm_attention_score, "attention_score_b", gemm_type=GemmType.ATTENTION_SCORE)
-            attention_score_f = attn_score_gemm_f + attn_score_reduction_f
-            attention_score_b = attn_score_gemm_b + attn_score_reduction_b
-            transformer_results['attention_score'] = {
-                'forward': attention_score_f, 'backward': attention_score_b,
-                'forward_gemm': attn_score_gemm_f, 'forward_reduction': attn_score_reduction_f,
-                'backward_gemm': attn_score_gemm_b, 'backward_reduction': attn_score_reduction_b,
-                'comm_size_forward': attn_score_size_f, 'comm_size_backward': attn_score_size_b,
-                'flops': attn_score_flops_f,
-                'memory_accesses': self._mem_levels(attn_score_mem_f),
-            }
+            attn_score_gemm_b, attn_score_reduction_b, attn_score_size_b = self.parallelism_gemm_backward(
+                gemm_attention_score, "attention_score_b", gemm_type=GemmType.ATTENTION_SCORE
+            )
+            attn_score_forward = DirectionTiming(
+                compute_time=attn_score_gemm_f,
+                comm_time=attn_score_reduction_f,
+                comm_bytes=int(attn_score_size_f or 0),
+                flops=attn_score_flops_f,
+                memory_accesses=dict(self._mem_levels(attn_score_mem_f)),
+            )
+            attn_score_backward = DirectionTiming(
+                compute_time=attn_score_gemm_b,
+                comm_time=attn_score_reduction_b,
+                comm_bytes=int(attn_score_size_b or 0),
+            )
+            transformer_timings["attention_score"] = OperationTiming(
+                "attention_score",
+                forward=attn_score_forward,
+                backward=attn_score_backward,
+            )
+            
+            
             attention_scale_softmax_f = self.get_scale_softmax_f(gemm=gemm_attention_score)
             attention_scale_softmax_b = self.get_scale_softmax_b(gemm=gemm_attention_score)
-            transformer_results['attention_scale_softmax'] = {'forward': attention_scale_softmax_f, 'backward': attention_scale_softmax_b}
+            
+            transformer_timings["attention_scale_softmax"] = OperationTiming(
+                name="attention_scale_softmax",
+                forward=DirectionTiming(
+                    compute_time=attention_scale_softmax_f,
+                    comm_time=0.0,
+                    comm_bytes=0,
+                ),
+                backward=DirectionTiming(
+                    compute_time=attention_scale_softmax_b,
+                    comm_time=0.0,
+                    comm_bytes=0,
+                ),
+            )
 
-            # Attention Output GEMM
-            attn_out_gemm_f,  attn_out_reduction_f, attn_out_size_f, attn_out_flops_f, attn_out_mem_f = _extract_forward(
+            attn_out_gemm_f, attn_out_reduction_f, attn_out_size_f, attn_out_flops_f, attn_out_mem_f = _extract_forward(
                 self.parallelism_gemm_forward(gemm_attention_output, "attention_output_f", gemm_type=GemmType.ATTENTION_OUTPUT)
             )
-            attn_out_gemm_b,  attn_out_reduction_b, attn_out_size_b = self.parallelism_gemm_backward(gemm_attention_output, "attention_output_b", gemm_type=GemmType.ATTENTION_OUTPUT)
-            attention_output_f = attn_out_gemm_f + attn_out_reduction_f
-            attention_output_b = attn_out_gemm_b + attn_out_reduction_b
-            transformer_results['attention_output'] = {
-                'forward': attention_output_f, 'backward': attention_output_b,
-                'forward_gemm': attn_out_gemm_f, 'forward_reduction': attn_out_reduction_f,
-                'backward_gemm': attn_out_gemm_b, 'backward_reduction': attn_out_reduction_b,
-                'comm_size_forward': attn_out_size_f, 'comm_size_backward': attn_out_size_b,
-                'flops': attn_out_flops_f,
-                'memory_accesses': self._mem_levels(attn_out_mem_f),
-            }
-            # Aggregate attention (score + softmax + output)
-            attention_f = attention_score_f + attention_scale_softmax_f + attention_output_f
-            attention_b = attention_score_b + attention_scale_softmax_b + attention_output_b
-            attention_gemm_f = attn_score_gemm_f + attn_out_gemm_f + attention_scale_softmax_f
-            attention_reduction_f = attn_score_reduction_f + attn_out_reduction_f
-            attention_gemm_b = attn_score_gemm_b + attn_out_gemm_b + attention_scale_softmax_b
-            attention_reduction_b = attn_score_reduction_b + attn_out_reduction_b
-            attention_size_f = attn_score_size_f + attn_out_size_f
-            attention_size_b = attn_score_size_b + attn_out_size_b
-            # Aggregate flops and per-level memory accesses (score + output)
-            attention_flops = (attn_score_flops_f or 0) + (attn_out_flops_f or 0)
-            attention_mem = [
-                (a or 0) + (b or 0)
-                for a, b in zip_longest(attn_score_mem_f or [], attn_out_mem_f or [], fillvalue=0)
-            ]
-            transformer_results["attention"] = {
-                'forward': attention_f,
-                'backward': attention_b,
-                "forward_gemm": attention_gemm_f,
-                "forward_reduction": attention_reduction_f,
-                "backward_gemm": attention_gemm_b,
-                "backward_reduction": attention_reduction_b,
-                "comm_size_forward": attention_size_f,
-                "comm_size_backward": attention_size_b,
-                "flops": attention_flops,
-                "memory_accesses": self._mem_levels(attention_mem),
-            }
-        elif self.flash_attention is True:
-            attention_f, attention_gemm_f, attention_reduction_f, attention_size_f, attention_mem = self.flash_attention_kernel_forward(batch_size,hidden_dim, seq_len, num_heads, kv_heads, num_SMs)
-            # TODO WIP.
-            attention_b, attention_gemm_b, attention_reduction_b, attention_size_b = 2 * attention_f, 2 * attention_gemm_f, 2 * attention_reduction_f, 2 * attention_size_f
+            attn_out_gemm_b, attn_out_reduction_b, attn_out_size_b = self.parallelism_gemm_backward(
+                gemm_attention_output, "attention_output_b", gemm_type=GemmType.ATTENTION_OUTPUT
+            )
+            attn_out_forward = DirectionTiming(
+                compute_time=attn_out_gemm_f,
+                comm_time=attn_out_reduction_f,
+                comm_bytes=int(attn_out_size_f or 0),
+                flops=attn_out_flops_f,
+                memory_accesses=dict(self._mem_levels(attn_out_mem_f)),
+            )
+            attn_out_backward = DirectionTiming(
+                compute_time=attn_out_gemm_b,
+                comm_time=attn_out_reduction_b,
+                comm_bytes=int(attn_out_size_b or 0),
+            )
+            transformer_timings["attention_output"] = OperationTiming(
+                "attention_output",
+                forward=attn_out_forward,
+                backward=attn_out_backward,
+            )
 
-            # flops for flash attention is 2* (QK^T + SV)
-            attention_flops = 2 * (batch_size * seq_len * (hidden_dim ** 2) + batch_size * (seq_len ** 2) * hidden_dim)
-            
-            transformer_results['attention'] = {
-                'forward': attention_f, 'backward': attention_b,
-                'forward_gemm': attention_gemm_f, 'forward_reduction': attention_reduction_f,
-                'backward_gemm': attention_gemm_b, 'backward_reduction': attention_reduction_b,
-                'comm_size_forward': attention_size_f, 'comm_size_backward': attention_size_b,
-                'flops': attention_flops,
-                'memory_accesses': self._mem_levels(attention_mem), # HBM only
-            }
+            attention_forward_compute = (
+                attn_score_forward.compute_time
+                + transformer_timings["attention_scale_softmax"].forward.compute_time
+                + attn_out_forward.compute_time
+            )
+            attention_forward_comm = attn_score_forward.comm_time + attn_out_forward.comm_time
+            attention_backward_compute = (
+                attn_score_backward.compute_time
+                + transformer_timings["attention_scale_softmax"].backward.compute_time
+                + attn_out_backward.compute_time
+            )
+            attention_backward_comm = attn_score_backward.comm_time + attn_out_backward.comm_time
+            attention_comm_bytes_f = attn_score_size_f + attn_out_size_f
+            attention_comm_bytes_b = attn_score_size_b + attn_out_size_b
+            attention_mem = self._combine_mem(attn_score_mem_f, attn_out_mem_f)
+            transformer_timings["attention"] = OperationTiming(
+                "attention",
+                forward=DirectionTiming(
+                    compute_time=attention_forward_compute,
+                    comm_time=attention_forward_comm,
+                    comm_bytes=int(attention_comm_bytes_f or 0),
+                    memory_accesses=dict(self._mem_levels(attention_mem)),
+                    flops=(attn_score_flops_f + attn_out_flops_f),
+                ),
+                backward=DirectionTiming(
+                    compute_time=attention_backward_compute,
+                    comm_time=attention_backward_comm,
+                    comm_bytes=int(attention_comm_bytes_b or 0),
+                ),
+            )
         else:
-            raise ValueError("flash_attention should be either True or False")
+            attention_f, attention_gemm_f, attention_reduction_f, attention_size_f, attention_mem = self.flash_attention_kernel_forward(
+                batch_size,
+                hidden_dim,
+                seq_len,
+                num_heads,
+                kv_heads,
+                num_SMs,
+            )
+            attention_b, attention_gemm_b, attention_reduction_b, attention_size_b = self.flash_attention_kernel_backward(
+                batch_size,
+                hidden_dim,
+                seq_len,
+                num_heads,
+                kv_heads,
+                num_SMs,
+            )
+
+            transformer_timings["attention"] = OperationTiming(
+                "attention",
+                forward=DirectionTiming(
+                    compute_time=attention_gemm_f,
+                    comm_time=attention_reduction_f,
+                    comm_bytes=int(attention_size_f or 0),
+                    memory_accesses=dict(self._mem_levels(attention_mem)),
+                ),
+                backward=DirectionTiming(
+                    compute_time=attention_gemm_b,
+                    comm_time=attention_reduction_b,
+                    comm_bytes=int(attention_size_b or 0),
+                ),
+            )
 
-        # Output Projection GEMM
         out_proj_gemm_f, out_proj_reduction_f, out_proj_size_f, out_proj_flops_f, out_proj_mem_f = _extract_forward(
             self.parallelism_gemm_forward(gemm_output_proj, "output_projection_f", gemm_type=GemmType.OUT_PROJ)
         )
-        out_proj_gemm_b,  out_proj_reduction_b, out_proj_size_b = self.parallelism_gemm_backward(gemm_output_proj, "output_projection_b", gemm_type=GemmType.OUT_PROJ)
-        output_proj_f = out_proj_gemm_f + out_proj_reduction_f
-        output_proj_b = out_proj_gemm_b + out_proj_reduction_b
-        transformer_results['output_proj'] = {
-            'forward': output_proj_f, 'backward': output_proj_b,
-            'forward_gemm': out_proj_gemm_f, 'forward_reduction': out_proj_reduction_f,
-            'backward_gemm': out_proj_gemm_b, 'backward_reduction': out_proj_reduction_b,
-            'comm_size_forward': out_proj_size_f, 'comm_size_backward': out_proj_size_b,
-            'flops': out_proj_flops_f,
-            'memory_accesses': self._mem_levels(out_proj_mem_f),
-        }
+        out_proj_gemm_b, out_proj_reduction_b, out_proj_size_b = self.parallelism_gemm_backward(
+            gemm_output_proj, "output_projection_b", gemm_type=GemmType.OUT_PROJ
+        )
+        transformer_timings["output_proj"] = OperationTiming(
+            "output_proj",
+            forward=DirectionTiming(
+                compute_time=out_proj_gemm_f,
+                comm_time=out_proj_reduction_f,
+                comm_bytes=int(out_proj_size_f or 0),
+                flops=out_proj_flops_f,
+                memory_accesses=dict(self._mem_levels(out_proj_mem_f)),
+            ),
+            backward=DirectionTiming(
+                compute_time=out_proj_gemm_b,
+                comm_time=out_proj_reduction_b,
+                comm_bytes=int(out_proj_size_b or 0),
+            ),
+        )
+        if not self.use_moe: 
+            ffn1_gemm_f, ffn1_reduction_f, ffn1_size_f, ffn1_flops_f, ffn1_mem_f = _extract_forward(
+                self.parallelism_gemm_forward(gemm_ffn1, "ffn_f", gemm_type=GemmType.FFN1)
+            )
+            ffn1_gemm_b, ffn1_reduction_b, ffn1_size_b = self.parallelism_gemm_backward(
+                gemm_ffn1, "ffn_b", gemm_type=GemmType.FFN1
+            )
+            ffn2_gemm_f, ffn2_reduction_f, ffn2_size_f, ffn2_flops_f, ffn2_mem_f = _extract_forward(
+                self.parallelism_gemm_forward(gemm_ffn2, "ffn2_f", gemm_type=GemmType.FFN2)
+            )
+            ffn2_gemm_b, ffn2_reduction_b, ffn2_size_b = self.parallelism_gemm_backward(
+                gemm_ffn2, "ffn2_b", gemm_type=GemmType.FFN2
+            )
 
-        # FFN1 GEMM
-        ffn1_gemm_f,  ffn1_reduction_f, ffn1_size_f, ffn1_flops_f, ffn1_mem_f = _extract_forward(
-            self.parallelism_gemm_forward(gemm_ffn1, "ffn_f", gemm_type=GemmType.FFN1)
+        else: # MOE expert parallelism
+            router_gemm_time_f , router_reduction_f, router_size_f, _, _ = _extract_forward(self.get_router_f(gemm_router, gemm_ffn1))
+            router_gemm_time_b , router_reduction_b, router_size_b= self.get_router_b(gemm_router, gemm_ffn1)
+            ffn1_gemm_f, ffn1_reduction_f, ffn1_size_f, ffn1_flops_f, ffn1_mem_f = _extract_forward(
+                self.get_moe_ffn_f(gemm_ffn1, "ffn1_f", gemm_type=GemmType.FFN1)
+            )
+            ffn1_gemm_b, ffn1_reduction_b, ffn1_size_b = self.get_moe_ffn_b(gemm_ffn1, "ffn1_b", gemm_type=GemmType.FFN1)
+            ffn2_gemm_f, ffn2_reduction_f, ffn2_size_f, ffn2_flops_f, ffn2_mem_f = _extract_forward(
+                self.get_moe_ffn_f(gemm_ffn2, "ffn2_f", gemm_type=GemmType.FFN2)
+            )
+            ffn2_gemm_b, ffn2_reduction_b, ffn2_size_b = self.get_moe_ffn_b(gemm_ffn2, "ffn2_b", gemm_type=GemmType.FFN2)
+        if self.use_moe:
+            transformer_timings["router"] = OperationTiming(
+                "router",
+                forward=DirectionTiming(
+                    compute_time=router_gemm_time_f,
+                    comm_time=router_reduction_f,
+                    comm_bytes=int(router_size_f or 0),
+                ),
+                backward=DirectionTiming(
+                    compute_time=router_gemm_time_b,
+                    comm_time=router_reduction_b,
+                    comm_bytes=int(router_size_b or 0),
+                ),
+            )
+        
+        transformer_timings["ffn1"] = OperationTiming(
+            "ffn1",
+            forward=DirectionTiming(
+                compute_time=ffn1_gemm_f,
+                comm_time=ffn1_reduction_f,
+                comm_bytes=int(ffn1_size_f or 0),
+                flops=ffn1_flops_f,
+                memory_accesses=dict(self._mem_levels(ffn1_mem_f)),
+            ),
+            backward=DirectionTiming(
+                compute_time=ffn1_gemm_b,
+                comm_time=ffn1_reduction_b,
+                comm_bytes=int(ffn1_size_b or 0),
+            ),
         )
-        ffn1_gemm_b,  ffn1_reduction_b, ffn1_size_b = self.parallelism_gemm_backward(gemm_ffn1, "ffn_b", gemm_type=GemmType.FFN1)
-        ffn1_f = ffn1_gemm_f + ffn1_reduction_f
-        ffn1_b = ffn1_gemm_b + ffn1_reduction_b
-        transformer_results['ffn1'] = {
-            'forward': ffn1_f, 'backward': ffn1_b,
-            'forward_gemm': ffn1_gemm_f, 'forward_reduction': ffn1_reduction_f,
-            'backward_gemm': ffn1_gemm_b, 'backward_reduction': ffn1_reduction_b,
-            'comm_size_forward': ffn1_size_f, 'comm_size_backward': ffn1_size_b,
-            'flops': ffn1_flops_f,
-            'memory_accesses': self._mem_levels(ffn1_mem_f),
-        }
 
-        # FFN2 GEMM
-        ffn2_gemm_f, ffn2_reduction_f, ffn2_size_f, ffn2_flops_f, ffn2_mem_f = _extract_forward(
-            self.parallelism_gemm_forward(gemm_ffn2, "ffn2_f", gemm_type=GemmType.FFN2)
+
+        transformer_timings["ffn2"] = OperationTiming(
+            "ffn2",
+            forward=DirectionTiming(
+                compute_time=ffn2_gemm_f,
+                comm_time=ffn2_reduction_f,
+                comm_bytes=int(ffn2_size_f or 0),
+                flops=ffn2_flops_f,
+                memory_accesses=dict(self._mem_levels(ffn2_mem_f)),
+            ),
+            backward=DirectionTiming(
+                compute_time=ffn2_gemm_b,
+                comm_time=ffn2_reduction_b,
+                comm_bytes=int(ffn2_size_b or 0),
+            ),
         )
-        ffn2_gemm_b,  ffn2_reduction_b, ffn2_size_b = self.parallelism_gemm_backward(gemm_ffn2, "ffn2_b", gemm_type=GemmType.FFN2)
-        ffn2_f = ffn2_gemm_f + ffn2_reduction_f
-        ffn2_b = ffn2_gemm_b + ffn2_reduction_b
-        transformer_results['ffn2'] = {
-            'forward': ffn2_f, 'backward': ffn2_b,
-            'forward_gemm': ffn2_gemm_f, 'forward_reduction': ffn2_reduction_f,
-            'backward_gemm': ffn2_gemm_b, 'backward_reduction': ffn2_reduction_b,
-            'comm_size_forward': ffn2_size_f , 'comm_size_backward': ffn2_size_b,
-            'flops': ffn2_flops_f,
-            'memory_accesses': self._mem_levels(ffn2_mem_f),
-        }
-        
-            
 
-        # Calculate non-GEMM operations
         embedding_f, embedding_mem = self.get_embedding_f(vocab_size=vocab_size, seq_len=seq_len, hidden_dim=hidden_dim)
         embedding_b = self.get_embedding_b(vocab_size=vocab_size, seq_len=seq_len, hidden_dim=hidden_dim)
-        transformer_results['embedding'] = {'forward': embedding_f, 'backward': embedding_b, 'memory_accesses': self._mem_levels(embedding_mem)}
+        transformer_timings["embedding"] = OperationTiming(
+            "embedding",
+            forward=DirectionTiming(
+                compute_time=embedding_f,
+                comm_time=0.0,
+                comm_bytes=0,
+                memory_accesses=dict(self._mem_levels(embedding_mem)),
+            ),
+            backward=DirectionTiming(
+                compute_time=embedding_b,
+                comm_time=0.0,
+                comm_bytes=0,
+            ),
+        )
 
         residual1_f = self.get_residual_f(tensor_shape=gemm_output_proj)
         residual1_b = self.get_residual_b(tensor_shape=gemm_output_proj)
-        # transformer_results['residual1'] = {'forward': residual1_f, 'backward': residual1_b}
-
-        layernorm1_f, layernorm_reduction_f, LN1_comm_bytes_f= self.get_layernorm_f(batch=batch_size, seq_len=seq_len, d_model=hidden_dim)
-        layernorm1_b, layernorm_reduction_b, LN1_comm_bytes_b= self.get_layernorm_b(batch=batch_size, seq_len=seq_len, d_model=hidden_dim)
-        transformer_results['layernorm1'] = {
-            'forward': layernorm1_f + residual1_f +layernorm_reduction_f,
-            'forward_compute': layernorm1_f + residual1_f, 
-            'forward_reduction':layernorm_reduction_f, 
-            'backward': layernorm1_b + residual1_b + layernorm_reduction_b,
-            "backward_compute":layernorm1_b + residual1_b , 
-            'backward_reduction': layernorm_reduction_b, 
-            'comm_size_forward': LN1_comm_bytes_f, 'comm_size_backward': LN1_comm_bytes_b
-        }
+        layernorm1_f, layernorm1_reduction_f, LN1_comm_bytes_f = self.get_layernorm_f(
+            batch=batch_size,
+            seq_len=seq_len,
+            d_model=hidden_dim,
+        )
+        layernorm1_b, layernorm1_reduction_b, LN1_comm_bytes_b = self.get_layernorm_b(
+            batch=batch_size,
+            seq_len=seq_len,
+            d_model=hidden_dim,
+            type=GemmType.LAYER_NORM_1,
+        )
+        transformer_timings["layernorm1"] = OperationTiming(
+            "layernorm1",
+            forward=DirectionTiming(
+                compute_time=layernorm1_f + residual1_f,
+                comm_time=layernorm1_reduction_f,
+                comm_bytes=int(LN1_comm_bytes_f or 0),
+            ),
+            backward=DirectionTiming(
+                compute_time=layernorm1_b + residual1_b,
+                comm_time=layernorm1_reduction_b,
+                comm_bytes=int(LN1_comm_bytes_b or 0),
+            ),
+        )
 
         residual2_f = self.get_residual_f(tensor_shape=gemm_ffn2)
         residual2_b = self.get_residual_b(tensor_shape=gemm_ffn2)
-        # transformer_results['residual2'] = {'forward': residual2_f, 'backward': residual2_b}
+        layernorm2_f, layernorm2_reduction_f, LN2_comm_bytes_f = self.get_layernorm_f(
+            batch=batch_size,
+            seq_len=seq_len,
+            d_model=hidden_dim,
+        )
+        layernorm2_b, layernorm2_reduction_b, LN2_comm_bytes_b = self.get_layernorm_b(
+            batch=batch_size,
+            seq_len=seq_len,
+            d_model=hidden_dim,
+        )
+        transformer_timings["layernorm2"] = OperationTiming(
+            "layernorm2",
+            forward=DirectionTiming(
+                compute_time=layernorm2_f + residual2_f,
+                comm_time=layernorm2_reduction_f,
+                comm_bytes=int(LN2_comm_bytes_f or 0),
+            ),
+            backward=DirectionTiming(
+                compute_time=layernorm2_b + residual2_b,
+                comm_time=layernorm2_reduction_b,
+                comm_bytes=int(LN2_comm_bytes_b or 0),
+            ),
+        )
 
-        layernorm2_f, layernorm_reduction_f, LN2_comm_bytes_f = self.get_layernorm_f(batch=batch_size, seq_len=seq_len, d_model=hidden_dim)
-        layernorm2_b, layernorm_reduction_b, LN2_comm_bytes_b = self.get_layernorm_b(batch=batch_size, seq_len=seq_len, d_model=hidden_dim)
-        transformer_results['layernorm2'] = {'forward': layernorm2_f + residual2_f + layernorm_reduction_f, "forward_compute": layernorm2_f + residual2_f, 'forward_reduction':layernorm_reduction_f,'backward': layernorm2_b + residual2_b + layernorm_reduction_b, "backward_compute":layernorm2_b + residual2_b, 'backward_reduction': layernorm_reduction_b, 'comm_size_forward': LN2_comm_bytes_f, 'comm_size_backward': LN2_comm_bytes_b}
-        
         if self.model_type == "llama":
             act_f = self.get_swiglu_f(tensor_shape=gemm_ffn1)
             act_b = self.get_swiglu_b(tensor_shape=gemm_ffn1)
         else:
             act_f = self.get_gelu_f(tensor_shape=gemm_ffn1)
             act_b = self.get_gelu_b(tensor_shape=gemm_ffn1)
-        transformer_results['gelu'] = {'forward': act_f, 'backward': act_b}
+        transformer_timings["gelu"] = OperationTiming(
+            "gelu",
+            forward=DirectionTiming(compute_time=act_f, comm_time=0.0, comm_bytes=0),
+            backward=DirectionTiming(compute_time=act_b, comm_time=0.0, comm_bytes=0),
+        )
 
         linear_softmax_f, linear_softmax_mem = self.get_linear_softmax_f(gemm=gemm_linear)
         linear_softmax_b = self.get_linear_softmax_b(gemm=gemm_linear)
-        transformer_results['linear_softmax'] = {
-            'forward': linear_softmax_f, 
-            'backward': linear_softmax_b, 
-            'memory_accesses': self._mem_levels(linear_softmax_mem)}
-
-        # Calculate MHA and FFN times directly from results dict
-        mha_time_f = ( 
-            transformer_results['qkv_proj']['forward'] + transformer_results['attention']['forward'] + transformer_results['output_proj']['forward'] 
-        )
-        
-        
-        mha_time_b = ( 
-            transformer_results['qkv_proj']['backward'] + transformer_results['attention']['backward'] + transformer_results['output_proj']['backward'] 
-        )
-        transformer_results['MHA'] = {
-            'forward': mha_time_f,
-            'backward': mha_time_b,
-            "forward_reduction": qkv_proj_reduction_f + attention_reduction_f + out_proj_reduction_f,
-            "backward_reduction": qkv_proj_reduction_b + attention_reduction_b + out_proj_reduction_b,
-            "comm_size_forward": qkv_proj_size_f + attention_size_f + out_proj_size_f,
-            "comm_size_backward": qkv_proj_size_b + attention_size_b + out_proj_size_b,
-            "flops": qkv_proj_flops_f + transformer_results['attention']['flops'] + out_proj_flops_f,
-            "memory_accesses": self._combine_mem(
-                qkv_proj_mem_f,
-                attention_mem,
-                out_proj_mem_f,
+        transformer_timings["linear_softmax"] = OperationTiming(
+            "linear_softmax",
+            forward=DirectionTiming(
+                compute_time=linear_softmax_f,
+                comm_time=0.0,
+                comm_bytes=0,
+                memory_accesses=dict(self._mem_levels(linear_softmax_mem)),
+            ),
+            backward=DirectionTiming(
+                compute_time=linear_softmax_b,
+                comm_time=0.0,
+                comm_bytes=0,
             ),
-
-        }
-
-        ffn_time_f = transformer_results['ffn1']['forward'] + transformer_results['ffn2']['forward'] + transformer_results['gelu']['forward']
-        ffn_time_b = (
-            transformer_results['ffn1']['backward'] + transformer_results['ffn2']['backward'] + transformer_results['gelu']['backward']
         )
-        transformer_results['MLP'] = {
-            'forward': ffn_time_f,
-            'backward': ffn_time_b,
-            "forward_reduction": ffn2_reduction_f + ffn1_reduction_f,
-            "backward_reduction": ffn1_reduction_b + ffn2_reduction_b,
-            "comm_size_forward": ffn2_size_f + ffn1_size_f,
-            "comm_size_backward": ffn1_size_b + ffn2_size_b,
-            "flops": ffn1_flops_f + ffn2_flops_f,
-            "memory_accesses": self._combine_mem(
-                ffn1_mem_f,
-                ffn2_mem_f,
+
+        mlp_group = OperationGroup(
+            "MLP",
+            operations=(
+                transformer_timings["ffn1"],
+                transformer_timings["gelu"],
+                transformer_timings["ffn2"],
             ),
-        }
-        # Calculate transformer times directly
-        
+        )
+
         transformer_time_f = (
-            transformer_results['MHA']['forward'] + transformer_results['MLP']['forward']  +
-            transformer_results['layernorm1']['forward'] + transformer_results['layernorm2']['forward']
+            transformer_timings["qkv_proj"].total_forward_time()
+            + transformer_timings["attention"].total_forward_time()
+            + transformer_timings["output_proj"].total_forward_time()
+            + mlp_group.forward_total_time()
+            + transformer_timings["layernorm1"].total_forward_time()
+            + transformer_timings["layernorm2"].total_forward_time()
         )
         transformer_time_b = (
-            transformer_results['MHA']['backward'] + transformer_results['MLP']['backward'] +
-            transformer_results['layernorm1']['backward'] + transformer_results['layernorm2']['backward']
+            transformer_timings["qkv_proj"].total_backward_time()
+            + transformer_timings["attention"].total_backward_time()
+            + transformer_timings["output_proj"].total_backward_time()
+            + mlp_group.backward_total_time()
+            + transformer_timings["layernorm1"].total_backward_time()
+            + transformer_timings["layernorm2"].total_backward_time()
         )
-        
+        # Pipeline-style recompute uses explicit recompute nodes, so keep backward time as-is.
+        transformer_time_b_combined = transformer_time_b
 
         node_breakdown = {
-            'transformer_time_f': transformer_time_f,
-            'transformer_time_b': transformer_time_b,
-            'embedding_f': transformer_results['embedding']['forward'],
-            'embedding_b': transformer_results['embedding']['backward'],
-            'linear_softmax_f': transformer_results['linear_softmax']['forward'],
-            'linear_softmax_b': transformer_results['linear_softmax']['backward']
+            "transformer_time_f": transformer_time_f,
+            "transformer_time_b": transformer_time_b,
+            "transformer_time_b_combined": transformer_time_b_combined,
+            "embedding_f": transformer_timings["embedding"].total_forward_time(),
+            "embedding_b": transformer_timings["embedding"].total_backward_time(),
+            "linear_softmax_f": transformer_timings["linear_softmax"].total_forward_time(),
+            "linear_softmax_b": transformer_timings["linear_softmax"].total_backward_time(),
         }
-
         if self._generate_graphs:
-            results_path = os.path.join(self.output_dir, "transformer_results.txt")
+            results_path = os.path.join(self.output_dir, "transformer_timings.yaml")
             with open(results_path, "w", encoding="utf-8") as results_file:
-                json.dump(
+                yaml.dump(
                     {
-                        "transformer_results": transformer_results,
+                        "transformer_results": {
+                            name: timing.to_dict() for name, timing in transformer_timings.items()
+                        },
                         "node_breakdown": node_breakdown,
                     },
                     results_file,
-                    indent=2,
                     sort_keys=True,
                 )
 
-        return transformer_results, node_breakdown
+        return transformer_timings, node_breakdown
+
+
 
 
     def _effective_transformer_batch(self) -> int:
@@ -1707,11 +1988,11 @@ def _build_comm_metadata(
         grad_collective = 'reduce_scatter' if (self.zero_stage >= 2 and self.dp > 1) else 'all_reduce'
         metadata = {
             'transformer': {
-                'size': reduction_sizes['total_size'],
+                'size': reduction_sizes,
                 'type': grad_collective,
                 'participants': self.dp,
                 'interconnect_type': 'dp',
-                'local_comp_time': local_comp['total_local']
+                'local_comp_time': local_comp
             },
             'embedding': {
                 'size': embedding_size,
@@ -1786,77 +2067,6 @@ def _build_comm_metadata(
             }
         return metadata
 
-    def _populate_transformer_comm_metadata(
-        self,
-        entry: Dict[str, Any],
-        metadata: Dict[str, Dict[str, Any]],
-        # gemm_spec: Tuple[int, ...],
-        comm_bytes_fwd: int,
-        comm_bytes_bwd: int,
-    ) -> None:
-        """Attach tensor-parallel collectives for a GEMM to metadata and entry."""
-
-
-        if not self.tp or self.tp <= 1 and self.cp <= 1:
-            return
-
-
-
-        def add_comm(direction: str, suffix: str, kind: str, size_bytes: float, participants: int, interconnect: str) -> None:
-            if participants <= 1:
-                return
-            bytes_int = int(math.ceil(size_bytes))
-            if bytes_int <= 0:
-                return
-            key = f"{entry['name']}_{direction}_{suffix}"
-            # ensure uniqueness when multiple collectives share the same suffix
-            unique_key = key
-            counter = 1
-            while unique_key in metadata:
-                counter += 1
-                unique_key = f"{key}_{counter}"
-
-            metadata[unique_key] = {
-                'size': bytes_int,
-                'type': kind,
-                'participants': int(participants),
-                'interconnect_type': interconnect,
-            }
-            entry[direction]['comm_keys'].append(unique_key)
-
-        parallelism_mode = self.get_parallelism_mode()
-        rules_by_mode = COMMUNICATION_RULES.get(parallelism_mode)
-        if not rules_by_mode:
-            return
-
-        spec = rules_by_mode.get(entry['name'])
-        if spec is None:
-            spec = rules_by_mode.get(COMM_RULE_DEFAULT_KEY)
-        if not spec:
-            return
-
-
-        participants_lookup = {
-            'tp': getattr(self, 'tp', 0),
-            'cp': getattr(self, 'cp', 0),
-            'dp': getattr(self, 'dp', 0),
-            'lp': getattr(self, 'lp', 0),
-        }
-
-        for direction in ('forward', 'backward'):
-            rule = spec.get(direction)
-            if not rule:
-                continue
-            size_bytes = comm_bytes_fwd if direction == 'forward' else comm_bytes_bwd
-            scope = rule.get('participants')
-            participants = participants_lookup.get(scope or '', 0)
-            if participants <= 1:
-                continue
-            kind = rule['kind']
-            suffix = rule.get('suffix', kind)
-            interconnect = rule.get('interconnect', scope or 'tp')
-            add_comm(direction, suffix, kind, size_bytes, participants, interconnect)
-
     def _build_interconnect_params(self) -> Dict[str, Tuple[float, float]]:
         return {
             'dp': (self.IBD, self.LLD),
@@ -1865,11 +2075,12 @@ def _build_interconnect_params(self) -> Dict[str, Tuple[float, float]]:
             'cp': (self.IBTP, self.LLTP),
         }
 
+
     def _prepare_execution_graphs(
         self,
         *,
         node_breakdown: Dict[str, float],
-        transformer_results: Dict[str, Dict[str, Any]],
+        transformer_timings: Dict[str, OperationTiming],
         batch_size: int,
         seq_len: int,
         hidden_dim: int,
@@ -1884,29 +2095,27 @@ def _prepare_execution_graphs(
         zero3_embedding_gather_bytes: float = 0.0,
         zero3_transformer_gather_bytes: float = 0.0,
         zero3_softmax_gather_bytes: float = 0.0,
-    ) -> Tuple[Graph, Any, Optional[Graph], Optional[Any], Optional[Any], Dict[str, Tuple[float, float]]]:
+    ) -> Tuple[
+        Graph,
+        Any,
+        Optional[Graph],
+        Optional[Any],
+        Optional[Any],
+        Dict[str, Tuple[float, float]],
+        Optional[Graph],
+        Optional[Any],
+    ]:
         """Build pipeline/transformer graphs shared across training and inference."""
 
         if not include_pipeline_backward and not include_transformer_backward:
             # Forward-only inference: skip training all-reduce bookkeeping.
-            reduction_sizes = {
-                'qkv_size': 0,
-                'output_size': 0,
-                'ffn_size': 0,
-                'total_size': 0,
-            }
-            local_comp = {
-                'qkv_local': 0.0,
-                'output_local': 0.0,
-                'ffn_local': 0.0,
-                'total_local': 0.0,
-            }
+            reduction_sizes = 0.0
+            local_comp = 0.0
         else:
             reduction_sizes = self.get_data_parallel_reduction_sizes(hidden_dim, intermediate_size)
-            local_comp = self.get_data_parallel_local_computation(hidden_dim, intermediate_size)
+            local_comp = self.get_data_parallel_reduction_llm(hidden_dim, intermediate_size)
 
         # these are used for dp all-reduce/reduce-scatter.
-        # TODO: check the embedding_size below. I think only the first term is needed.
         embedding_size = math.ceil(self.precision.grad_communication * vocab_size * hidden_dim) + math.ceil(self.precision.grad_communication * seq_len * hidden_dim * batch_size)
         softmax_size = math.ceil(self.precision.grad_communication * hidden_dim * vocab_size)
         cross_layer_bytes = self.get_inter_layer_comm_latency_llm(batch_size, hidden_dim, seq_len)[1]
@@ -1928,79 +2137,179 @@ def _prepare_execution_graphs(
         transformer_operation_entries: List[Dict[str, Any]] = []
         transformer_comm_metadata: Dict[str, Dict[str, Any]] = {}
         parallelism_mode = self.get_parallelism_mode()
-        if parallelism_mode in (ParallelismMode.CONTEXT, ParallelismMode.TENSOR_CONTEXT_HYBRID):
-            for key in ("layernorm1", "qkv_proj", "attention", "output_proj", "layernorm2", "MLP"):
-                try:
-                    spec = transformer_results[key]
-                except KeyError:
-                    print(f"Key {key} not found in transformer_results")
-                    print(transformer_results)
-                    exit()
-                fwd_time = spec['forward_gemm'] if 'forward_gemm' in spec else spec['forward_compute'] if 'forward_compute' in spec else spec['forward']
-                bwd_time = spec['backward_gemm'] if 'backward_gemm' in spec else spec['backward_compute'] if 'backward_compute' in spec else spec['backward']
-                fwd_red = spec.get('forward_reduction', 0.0)
-                bwd_red = spec.get('backward_reduction', 0.0)
-                comm_bytes_fwd = spec.get('comm_size_forward', 0)
-                comm_bytes_bwd = spec.get('comm_size_backward', 0)
-
-                entry = {
-                    "name": key,
-                    "forward": {
-                        "duration": fwd_time,
-                        "reduction": fwd_red,
-                        "comm_keys": [],
-                    },
-                    "backward": {
-                        "duration": bwd_time,
-                        "reduction": bwd_red,
-                        "comm_keys": [],
-                    },
-                }
+        def _clone_comm_rules(rules: Dict[str, Dict[str, Optional[Dict[str, str]]]]) -> Dict[str, Dict[str, Optional[Dict[str, str]]]]:
+            cloned: Dict[str, Dict[str, Optional[Dict[str, str]]]] = {}
+            for op_name, directions in rules.items():
+                cloned[op_name] = {}
+                for direction, spec in directions.items():
+                    cloned[op_name][direction] = dict(spec) if spec else None
+            return cloned
+
+        def _resolve_comm_rules(mode: ParallelismMode) -> Dict[str, Dict[str, Optional[Dict[str, str]]]]:
+            base_rules = COMMUNICATION_RULES.get(mode) or {}
+            if not self.use_moe:
+                return base_rules
+            overrides = MOE_COMMUNICATION_RULES.get(mode)
+            if not overrides:
+                return base_rules
+            merged = _clone_comm_rules(base_rules)
+            for op_name, directions in overrides.items():
+                merged.setdefault(op_name, {})
+                for direction, spec in directions.items():
+                    merged[op_name][direction] = dict(spec) if spec else None
+            return merged
+
+        rules_by_mode = _resolve_comm_rules(parallelism_mode)
+        participants_lookup = {
+            'tp': int(getattr(self, 'tp', 0) or 0),
+            'cp': int(getattr(self, 'cp', 0) or 0),
+            'dp': int(getattr(self, 'dp', 0) or 0),
+            'lp': int(getattr(self, 'lp', 0) or 0),
+        }
+        participants_lookup['moe'] = max(1, participants_lookup['tp'] * participants_lookup['cp'])
+        group_members = {
+            "MLP": ("ffn1", "gelu", "ffn2"),
+        }
 
-                self._populate_transformer_comm_metadata(
-                    entry=entry,
-                    metadata=transformer_comm_metadata,
-                    comm_bytes_fwd=comm_bytes_fwd,
-                    comm_bytes_bwd=comm_bytes_bwd,
-                )
+        def _register_specs(specs: Sequence[CommSpec]) -> List[str]:
+            names: List[str] = []
+            for spec in specs:
+                existing = transformer_comm_metadata.get(spec.name)
+                if existing:
+                    if (
+                        existing["size"] != spec.size_bytes
+                        or existing["type"] != spec.kind
+                        or existing["participants"] != spec.participants
+                        or existing["interconnect_type"] != spec.interconnect
+                    ):
+                        raise ValueError(
+                            f"Conflicting CommSpec registration for '{spec.name}' "
+                            f"(existing={existing}, new={spec})"
+                        )
+                else:
+                    transformer_comm_metadata[spec.name] = {
+                        "size": spec.size_bytes,
+                        "type": spec.kind,
+                        "participants": spec.participants,
+                        "interconnect_type": spec.interconnect,
+                    }
+                names.append(spec.name)
+            return names
+
+        def _build_direction_entry(op_name: str, direction: str, timing: Optional[DirectionTiming]) -> Dict[str, Any]:
+            if timing is None:
+                return {"duration": 0.0, "reduction": 0.0, "comm_keys": []}
+            specs = _make_comm_specs(op_name, direction, timing.comm_bytes)
+            return {
+                "duration": timing.compute_time,
+                "reduction": timing.comm_time,
+                "comm_keys": _register_specs(specs),
+            }
 
-                transformer_operation_entries.append(entry)
-            
-        elif parallelism_mode in (ParallelismMode.TENSOR, ParallelismMode.TENSOR_SEQUENCE, ParallelismMode.SINGLE):
-            for key in ("layernorm1", "MHA", "layernorm2", "MLP"):
-                spec = transformer_results[key]
-                fwd_time = spec['forward']
-                bwd_time = spec['backward']
-                fwd_red = spec.get('forward_reduction', 0.0)
-                bwd_red = spec.get('backward_reduction', 0.0)
-                comm_bytes_fwd = spec.get('comm_size_forward', 0)
-                comm_bytes_bwd = spec.get('comm_size_backward', 0)
-
-                entry = {
-                    "name": key,
-                    "forward": {
-                        "duration": fwd_time,
-                        "reduction": fwd_red,
-                        "comm_keys": [],
-                    },
-                    "backward": {
-                        "duration": bwd_time,
-                        "reduction": bwd_red,
-                        "comm_keys": [],
+        def _make_comm_specs(op_name: str, direction: str, total_bytes: float) -> Tuple[CommSpec, ...]:
+            bytes_int = int(math.ceil(float(total_bytes or 0.0)))
+            if bytes_int <= 0:
+                return ()
+            rule_spec = rules_by_mode.get(op_name) or rules_by_mode.get(COMM_RULE_DEFAULT_KEY)
+            if not rule_spec:
+                raise ValueError(
+                    f"Missing communication rule for op '{op_name}' in mode '{parallelism_mode.value}' "
+                    f"with non-zero comm_bytes ({bytes_int})"
+                )
+            rule = rule_spec.get(direction)
+            if not rule:
+                return ()
+            kind = rule.get("kind")
+            scope = rule.get("participants")
+            if not kind or not scope:
+                return ()
+            participants = participants_lookup.get(scope, 0)
+            if participants <= 1:
+                return ()
+            interconnect = rule.get("interconnect", scope)
+            name = f"{op_name}_{direction}_{kind}"
+            return (
+                CommSpec(
+                    name=name,
+                    kind=kind,
+                    size_bytes=bytes_int,
+                    participants=participants,
+                    interconnect=interconnect,
+                    extra={
+                        "scope": scope,
+                        "direction": direction,
+                        "op_name": op_name,
+                        "parallelism_mode": parallelism_mode.value,
                     },
-                }
+                ),
+            )
 
-                self._populate_transformer_comm_metadata(
-                    entry=entry,
-                    metadata=transformer_comm_metadata,
-                    comm_bytes_fwd=comm_bytes_fwd,
-                    comm_bytes_bwd=comm_bytes_bwd,
+        def _build_group_operation(name: str, members: Tuple[str, ...]) -> OperationTiming:
+            ops = tuple(transformer_timings[m] for m in members)
+            group = OperationGroup(name, ops)
+            forward_memory = self._combine_mem(*(op.forward.memory_accesses for op in ops))
+            forward_flops = sum(op.forward.flops for op in ops)
+            forward_dir = DirectionTiming(
+                compute_time=group.forward_compute_time(),
+                comm_time=group.forward_comm_time(),
+                comm_bytes=group.forward_comm_bytes(),
+                flops=forward_flops,
+                memory_accesses=forward_memory,
+            )
+            if include_transformer_backward:
+                backward_memory = self._combine_mem(
+                    *(op.backward.memory_accesses for op in ops if op.backward is not None)
                 )
+                backward_flops = sum((op.backward.flops if op.backward else 0.0) for op in ops)
+                backward_dir = DirectionTiming(
+                    compute_time=group.backward_compute_time(),
+                    comm_time=group.backward_comm_time(),
+                    comm_bytes=group.backward_comm_bytes(),
+                    flops=backward_flops,
+                    memory_accesses=backward_memory,
+                )
+            else:
+                backward_dir = None
+            return OperationTiming(name, forward_dir, backward_dir)
+
+        def _resolve_operation(name: str) -> OperationTiming:
+            if name in group_members:
+                return _build_group_operation(name, group_members[name])
+            return transformer_timings[name]
+
+        if parallelism_mode not in (
+            ParallelismMode.CONTEXT,
+            ParallelismMode.TENSOR_CONTEXT_HYBRID,
+            ParallelismMode.TENSOR,
+            ParallelismMode.TENSOR_SEQUENCE,
+            ParallelismMode.SINGLE,
+        ):
+            raise ValueError(f"Unsupported parallelism mode: {parallelism_mode}")
+        op_names: Sequence[str] = ("layernorm1", "qkv_proj", "attention", "output_proj", "layernorm2", "MLP")
+        op_names = list(op_names)
+        if "optimizer" in transformer_timings:
+            op_names.append("optimizer")
+        if self.use_moe and "MLP" in op_names and "router" not in op_names: # if MOE is enabled, ensure router is included in the graph
+            mlp_index = op_names.index("MLP")
+            op_names.insert(mlp_index, "router")
+
+        for key in op_names:
+            try:
+                op_timing = _resolve_operation(key)
+            except KeyError as exc:
+                raise KeyError(f"Missing transformer timing for operation '{key}'") from exc
+
+            entry = {
+                "name": key,
+                "forward": _build_direction_entry(key, "forward", op_timing.forward),
+                "backward": _build_direction_entry(key, "backward", op_timing.backward if include_transformer_backward else None),
+            }
 
-                transformer_operation_entries.append(entry)
+            if not include_transformer_backward:
+                # Ensure backward section is zeroed when excluded.
+                entry["backward"] = {"duration": 0.0, "reduction": 0.0, "comm_keys": []}
 
-        else:
-            raise ValueError(f"Unsupported parallelism mode: {parallelism_mode}")
+            transformer_operation_entries.append(entry)
         transformer_graph: Optional[Graph] = None
         transformer_forward_root: Optional[Any] = None
         transformer_backward_root: Optional[Any] = None
@@ -2023,7 +2332,24 @@ def _prepare_execution_graphs(
         )
         transformer_forward_root = transformer_graph.construct_transformer_graph(direction="forward")
         if include_transformer_backward:
-            transformer_backward_root = transformer_graph.construct_transformer_graph(direction="backward")
+            bwd_direction = "backward"
+            transformer_backward_root = transformer_graph.construct_transformer_graph(direction=bwd_direction)
+
+        transformer_forward_root = apply_overlap_transforms(
+            transformer_forward_root,
+            parallelism_mode,
+            self.tp_overlap,
+            self.tp_sp_overlap,
+            self.cp_overlap,
+        )
+        if include_transformer_backward:
+            transformer_backward_root = apply_overlap_transforms(
+                transformer_backward_root,
+                parallelism_mode,
+                self.tp_overlap,
+                self.tp_sp_overlap,
+                self.cp_overlap,
+            )
 
         comp_times = {
             "embedding_f": node_breakdown.get('embedding_f', 0.0),
@@ -2032,14 +2358,21 @@ def _prepare_execution_graphs(
             "linear_softmax_b": node_breakdown.get('linear_softmax_b', 0.0) if include_pipeline_backward else 0.0,
             "transformer_f": node_breakdown.get('transformer_time_f', 0.0),
             "transformer_b": node_breakdown.get('transformer_time_b', 0.0) if include_pipeline_backward else 0.0,
+            "optimizer": self.get_data_parallel_reduction_llm(hidden_dim, intermediate_size),
             "cross_layer_f": 0.0,
             "cross_layer_b": 0.0,
         }
+        flattened_mode = self.execution_mode == ExecutionMode.FULL_ASTRASIM_FLATTENED
+        pipeline_style_recompute_flag = bool(getattr(self, "full_recomputation", False))
         misc_metadata = {
             "num_batch": self.mb,
             "num_layer": self.num_layers,
             "all_reduce": "every_layer",
             "dp_zero_stage": self.zero_stage,
+            "full_recomputation": self.full_recomputation,
+            "flattened_mode": flattened_mode,
+            "pipeline_style_recompute": pipeline_style_recompute_flag,
+            "dp_microbatch_mode": getattr(self, "dp_microbatch", "every_mb"),
         }
 
         pipeline_graph_obj = simulate_LLM.Graph(
@@ -2052,12 +2385,39 @@ def _prepare_execution_graphs(
             comm_metadata=comm_metadata,
             misc_metadata=misc_metadata,
         )
-        graph_root = pipeline_graph_obj.construct_fwd_bwd_graph(include_backward=include_pipeline_backward)
+        
+        graph_root = pipeline_graph_obj.construct_fwd_bwd_graph(
+            include_backward=include_pipeline_backward,
+            include_optimizer=True
+        )
+        pipeline_graph_obj_no_dp = None
+        graph_root_no_dp = None
+        need_no_dp_variant = getattr(self, "gradient_accumulation_steps", 1) > 1
+        if need_no_dp_variant:
+            pipeline_graph_obj_no_dp = simulate_LLM.Graph( # if gradient accumulation is used, we need a no-dp variant of the graph for the non-last step, since dp all-reduce is only done on the last step
+                mode="pipeline",
+                dp=1,
+                lp=self.lp,
+                tp=self.tp,
+                cp=self.cp,
+                comp_times=comp_times,
+                comm_metadata=comm_metadata,
+                misc_metadata=misc_metadata,
+            )
+            
+            graph_root_no_dp = pipeline_graph_obj_no_dp.construct_fwd_bwd_graph(
+                include_backward=include_pipeline_backward,
+                include_optimizer=False
+            )
+
+        
         interconnect_params = self._build_interconnect_params()
 
         return (
             pipeline_graph_obj,
             graph_root,
+            pipeline_graph_obj_no_dp,
+            graph_root_no_dp,
             transformer_graph,
             transformer_forward_root,
             transformer_backward_root,
@@ -2120,8 +2480,15 @@ def calc_time_llm(self):
         )
 
         num_SMs = self.hw_config.tech_config.core.num_bundles
-        transformer_results, node_breakdown = self.compute_all_gemm_and_node_times(
-            batch_size, vocab_size, hidden_dim, seq_len, num_heads, kv_heads, intermediate_size, num_SMs
+        transformer_timings, node_breakdown = self.compute_all_gemm_and_node_times(
+            batch_size,
+            vocab_size,
+            hidden_dim,
+            seq_len,
+            num_heads,
+            kv_heads,
+            intermediate_size,
+            num_SMs,
         )
 
         # transformer mem layer considers zero stage internally
@@ -2129,6 +2496,8 @@ def calc_time_llm(self):
             LLM_util.get_transformer_mem_layer(
                 dp=self.dp,
                 tp=self.tp,
+                lp=self.lp,
+                mb=self.mb,
                 batch_size=batch_size,
                 hidden_dim=hidden_dim,
                 seq_len=seq_len / self.cp,
@@ -2137,6 +2506,8 @@ def calc_time_llm(self):
                 precision=self.precision,
                 model_type=self.model_type,
                 zero_stage=self.zero_stage,
+                flash_attention=self.flash_attention,
+                full_recomputation=self.full_recomputation,
             )
         )
 
@@ -2170,33 +2541,18 @@ def calc_time_llm(self):
                     self.zero3_ephemeral_peak_bytes
                 )
 
-    
-
-        if self.debug:
-            print("self.tp:", self.tp)
-            print("Calculating LLM time...")
-            print(
-                "simulating parallelism with dp = {}, lp = {}, total data batch = {}, "
-                "for each dp node, data batch = {}, for each pipeline stage, data batch = {}".format(
-                    self.dp, self.lp, self.batch_size, self.miniB, self.microB
-                )
-            )
-            print("data-parallel zero stage:", self.zero_stage)
-            print("total number of workers: {}".format(self.num_workers))
-            print("number of workers for each data parallelism batch: {}".format(self.num_workers_dp))
-            print("number of workers for each pipeline stage: {}".format(self.num_workers_lp))
-        
-
         (
             pipeline_graph_obj,
             graph_root,
+            pipeline_graph_obj_no_dp,
+            graph_root_no_dp,
             transformer_graph,
             transformer_forward_root,
             transformer_backward_root,
             interconnect_params,
         ) = self._prepare_execution_graphs(
             node_breakdown=node_breakdown,
-            transformer_results=transformer_results,
+            transformer_timings=transformer_timings,
             batch_size=batch_size,
             seq_len=seq_len,
             hidden_dim=hidden_dim,
@@ -2217,14 +2573,37 @@ def calc_time_llm(self):
         self.transformer_forward_root = transformer_forward_root
         self.transformer_backward_root = transformer_backward_root
         self.transformer_analytical_time_forward = node_breakdown['transformer_time_f']
-        self.transformer_analytical_time_backward = node_breakdown['transformer_time_b']
+        # Report backward including recompute overhead when enabled.
+        self.transformer_analytical_time_backward_combined = node_breakdown['transformer_time_b_combined']
+        self.transformer_analytical_time_backward = self.transformer_analytical_time_backward_combined
 
         self.pipeline_graph = pipeline_graph_obj
         self.pipeline_root = graph_root
         self.pipeline_interconnect = interconnect_params
+        self.pipeline_graph_no_dp = pipeline_graph_obj_no_dp
+        self.pipeline_root_no_dp = graph_root_no_dp
 
         mode = self.execution_mode
-        
+        time_fw_bw_no_dp: Optional[float] = None
+        if self.gradient_accumulation_steps > 1:
+            if not (self.pipeline_graph_no_dp and self.pipeline_root_no_dp):
+                raise RuntimeError("Gradient accumulation steps > 1 requires a no-DP pipeline graph")
+            dispatcher_no_dp = LLMExecutionDispatcher(
+                time_calc=self,
+                pipeline_graph=self.pipeline_graph_no_dp,
+                pipeline_root=self.pipeline_root_no_dp,
+                interconnect_params=self.pipeline_interconnect,
+                transformer_graph=self.transformer_graph,
+                transformer_forward_root=self.transformer_forward_root,
+                transformer_backward_root=self.transformer_backward_root,
+                no_data_parallel=True,
+            )
+            try:
+                result_no_dp = dispatcher_no_dp.run(mode)
+            except NotImplementedError as exc:
+                raise NotImplementedError(f"{exc}. Selected execution mode '{mode.value}'.") from exc
+            time_fw_bw_no_dp = result_no_dp.total_time
+
         dispatcher = LLMExecutionDispatcher(
             time_calc=self,
             pipeline_graph=self.pipeline_graph,
@@ -2233,20 +2612,21 @@ def calc_time_llm(self):
             transformer_graph=self.transformer_graph,
             transformer_forward_root=self.transformer_forward_root,
             transformer_backward_root=self.transformer_backward_root,
+            no_data_parallel=False,
         )
-        mode = self.execution_mode
         try:
             result = dispatcher.run(mode)
         except NotImplementedError as exc:
             raise NotImplementedError(f"{exc}. Selected execution mode '{mode.value}'.") from exc
-
         time_fw_bw = result.total_time
 
+
         pipeline_root = result.graph_root
         self.pipeline_graph = dispatcher.pipeline_graph
         self.pipeline_root = pipeline_root
         self.pipeline_interconnect = dispatcher.interconnect_params
-        
+        _, peak_mem = self._simulate_with_memory(graph_root, memory_data, mode="training")
+
 
 
         # debug helper. If set, print analytical transformer time and actual transformer time
@@ -2257,9 +2637,13 @@ def calc_time_llm(self):
                 print(f"Actual transformer forward time: {self.transformer_astrasim_time_forward:.4f}s")
                 print(f"Actual transformer backward time: {self.transformer_astrasim_time_backward:.4f}s")
 
-        self.tot_time = time_fw_bw
-
-        return time_fw_bw
+        if self.gradient_accumulation_steps > 1: 
+            if time_fw_bw_no_dp is None:
+                raise RuntimeError("Missing no-DP timing for gradient accumulation computation")
+            self.tot_time = time_fw_bw_no_dp * (self.gradient_accumulation_steps - 1) + time_fw_bw #total time is (time for no-dp steps) + (time for last step with dp)
+        else:
+            self.tot_time = time_fw_bw
+        return self.tot_time
         
     def _simulate_with_memory( 
         self,
diff --git a/time_calculation_inf.py b/time_calculation_inf.py
index 3f08e605..f2dbd0e9 100644
--- a/time_calculation_inf.py
+++ b/time_calculation_inf.py
@@ -2,12 +2,18 @@
 
 import math
 import os
-from typing import Dict, List, Optional, Tuple
-from time_calculation_LLM import LLMExecutionDispatcher, TimeCalculationLLM, GemmType
+from typing import Dict, List, Optional, Tuple, Mapping, Set
+from time_calculation_LLM import (
+    LLMExecutionDispatcher,
+    TimeCalculationLLM,
+    GemmType,
+    COMMUNICATION_RULES,
+    COMM_RULE_DEFAULT_KEY,
+)
 from simulate_inf import DecodeSample, InferenceConfig, InferenceEngine
 import LLM_util
 import json
-from itertools import zip_longest  # added for element-wise aggregation of memory access lists
+from timing_model import CommSpec, DirectionTiming, OperationTiming, OperationGroup
 
 def convert_prefix(value: float) -> float:
     """Assign SI unit prefixes to numerical values."""
@@ -27,6 +33,8 @@ class TimeCalculationLLMInference(TimeCalculationLLM):
 
     def __init__(self, hw_config, model_config, mode, output_dir: Optional[str] = None):
         super().__init__(hw_config, model_config, mode, output_dir)
+        if self.zero_stage >= 3 and self.dp > 1:
+            raise ValueError("ZeRO-3 data parallelism is not supported for inference runs (dp_zero_stage must be <3 or dp=1).")
         self._raw_model_config = model_config
 
     def _build_decode_transformer_results(
@@ -35,8 +43,9 @@ def _build_decode_transformer_results(
         batch_size: int,
         total_seq_len: int,
         gemm_shapes: Optional[Dict[str, Tuple[int, ...]]] = None,
-    ) -> Tuple[Dict[str, Dict[str, float]], Dict[str, float]]:
-        """Construct transformer spec + node breakdown for a single decode step."""
+    ) -> Tuple[Dict[str, OperationTiming], Dict[str, float]]:
+        """Construct transformer timings + node breakdown for a single decode step."""
+
         head_dim = self.hidden_dim // self.num_heads
 
         token_bytes = LLM_util.kv_cache_token_bytes(
@@ -46,7 +55,6 @@ def _build_decode_transformer_results(
             precision_bytes=self.precision.kv_cache,
         )
         seq_degree = self._sequence_parallel_degree()
-
         comm_kind = "all_reduce" if seq_degree == 1 else "reduce_scatter"
 
         intermediate_size = self.intermediate_size
@@ -70,15 +78,88 @@ def _build_decode_transformer_results(
         gemm_ffn2 = gemm_shapes["ffn2"]
         # FlashAttention is not used during single-token (incremental) decoding.
 
+
+        output_seq_len = 1
+
+        transformer_timings: Dict[str, OperationTiming] = {}
+
+        def _make_forward(
+            op_name: str,
+            compute_time: float,
+            comm_time: float,
+            comm_bytes: float,
+            *,
+            flops: float = 0.0,
+            memory: Optional[Mapping[str, float]] = None,
+        ) -> DirectionTiming:
+            bytes_int = int(math.ceil(float(comm_bytes or 0.0)))
+            return DirectionTiming(
+                compute_time=compute_time,
+                comm_time=comm_time,
+                comm_bytes=bytes_int,
+                flops=flops,
+                memory_accesses=dict(memory) if memory else {},
+            )
+
+        # QKV projection
         qkv_proj_time, qkv_proj_reduction, qkv_proj_size, qkv_proj_flops, qkv_proj_mem = self.parallelism_gemm_forward(
             gemm_qkv_proj, "decode_qkv_proj_f", gemm_type=GemmType.QKV
         )
+        transformer_timings["qkv_proj"] = OperationTiming(
+            "qkv_proj",
+            forward=_make_forward(
+                "qkv_proj",
+                compute_time=qkv_proj_time,
+                comm_time=qkv_proj_reduction,
+                comm_bytes=qkv_proj_size,
+                flops=qkv_proj_flops,
+                memory=self._mem_levels(qkv_proj_mem),
+            ),
+            backward=None,
+        )
+
+        # Attention components
         attention_score_time, attention_score_reduction, attention_score_size, attention_score_flops, attention_score_mem = self.parallelism_gemm_forward(
             gemm_attention_score, "decode_attention_score_f", gemm_type=GemmType.ATTENTION_SCORE
         )
         attention_output_time, attention_output_reduction, attention_output_size, attention_output_flops, attention_output_mem = self.parallelism_gemm_forward(
             gemm_attention_output, "decode_attention_output_f", gemm_type=GemmType.ATTENTION_OUTPUT
         )
+        attention_scale_softmax_f = self.get_scale_softmax_f(gemm_attention_score)
+
+        attention_reduction = attention_score_reduction + attention_output_reduction
+        attention_comm_bytes = attention_score_size + attention_output_size
+        attention_forward_compute = attention_score_time + attention_scale_softmax_f + attention_output_time
+        attention_forward_time = attention_forward_compute + attention_reduction
+        attention_flops = (attention_score_flops or 0.0) + (attention_output_flops or 0.0)
+        attention_mem = self._combine_mem(attention_score_mem, attention_output_mem)
+
+        transformer_timings["attention"] = OperationTiming(
+            "attention",
+            forward=_make_forward(
+                "attention",
+                compute_time=attention_forward_compute,
+                comm_time=attention_reduction,
+                comm_bytes=attention_comm_bytes,
+                flops=attention_flops,
+                memory=attention_mem,
+            ),
+            backward=None,
+        )
+
+        attention_scale_softmax_op = OperationTiming(
+            "attention_scale_softmax",
+            forward=_make_forward(
+                "attention_scale_softmax",
+                compute_time=attention_scale_softmax_f,
+                comm_time=0.0,
+                comm_bytes=0.0,
+            ),
+            backward=None,
+        )
+        transformer_timings["attention_scale_softmax"] = attention_scale_softmax_op
+
+        # Output projection
         out_proj_time, _, out_proj_size, out_proj_flops, out_proj_mem = self.parallelism_gemm_forward(
             gemm_output_proj, "decode_output_projection_f", gemm_type=GemmType.OUT_PROJ
         )
@@ -87,7 +168,20 @@ def _build_decode_transformer_results(
             if out_proj_size
             else 0.0
         )
+        transformer_timings["output_proj"] = OperationTiming(
+            "output_proj",
+            forward=_make_forward(
+                "output_proj",
+                compute_time=out_proj_time,
+                comm_time=out_proj_reduction,
+                comm_bytes=out_proj_size,
+                flops=out_proj_flops,
+                memory=self._mem_levels(out_proj_mem),
+            ),
+            backward=None,
+        )
 
+        # FFN layers
         ffn1_time, ffn1_reduction, ffn1_size, ffn1_flops, ffn1_mem = self.parallelism_gemm_forward(
             gemm_ffn1, "decode_ffn1_f", gemm_type=GemmType.FFN1
         )
@@ -95,13 +189,45 @@ def _build_decode_transformer_results(
             gemm_ffn2, "decode_ffn2_f", gemm_type=GemmType.FFN2
         )
         ffn2_reduction = (
-            self.get_tensor_reduction_time(ffn2_size, comm_kind, "decode_ffn2")
-            if ffn2_size
-            else 0.0
+            self.get_tensor_reduction_time(ffn2_size, comm_kind, "decode_ffn2") if ffn2_size else 0.0
+        )
+        transformer_timings["ffn1"] = OperationTiming(
+            "ffn1",
+            forward=_make_forward(
+                "ffn1",
+                compute_time=ffn1_time,
+                comm_time=ffn1_reduction,
+                comm_bytes=ffn1_size,
+                flops=ffn1_flops,
+                memory=self._mem_levels(ffn1_mem),
+            ),
+            backward=None,
+        )
+        transformer_timings["ffn2"] = OperationTiming(
+            "ffn2",
+            forward=_make_forward(
+                "ffn2",
+                compute_time=ffn2_time,
+                comm_time=ffn2_reduction,
+                comm_bytes=ffn2_size,
+                flops=ffn2_flops,
+                memory=self._mem_levels(ffn2_mem),
+            ),
+            backward=None,
         )
 
-        output_seq_len = 1
+        # GELU/SwiGLU activation
+        if self.model_type == "llama":
+            act_f = self.get_swiglu_f(gemm_ffn1)
+        else:
+            act_f = self.get_gelu_f(gemm_ffn1)
+        transformer_timings["gelu"] = OperationTiming(
+            "gelu",
+            forward=_make_forward("gelu", compute_time=act_f, comm_time=0.0, comm_bytes=0.0),
+            backward=None,
+        )
 
+        # Layer norms
         output_proj_shape = (
             batch_size,
             output_seq_len,
@@ -112,8 +238,17 @@ def _build_decode_transformer_results(
         layernorm1_f, layernorm1_reduction, layernorm1_bytes = self.get_layernorm_f(
             batch=batch_size, seq_len=output_seq_len, d_model=self.hidden_dim
         )
+        transformer_timings["layernorm1"] = OperationTiming(
+            "layernorm1",
+            forward=_make_forward(
+                "layernorm1",
+                compute_time=layernorm1_f + residual1_f,
+                comm_time=layernorm1_reduction,
+                comm_bytes=layernorm1_bytes,
+            ),
+            backward=None,
+        )
 
-        intermediate_size = self.intermediate_size
         ffn2_shape = (
             batch_size,
             output_seq_len,
@@ -124,6 +259,16 @@ def _build_decode_transformer_results(
         layernorm2_f, layernorm2_reduction, layernorm2_bytes = self.get_layernorm_f(
             batch=batch_size, seq_len=output_seq_len, d_model=self.hidden_dim
         )
+        transformer_timings["layernorm2"] = OperationTiming(
+            "layernorm2",
+            forward=_make_forward(
+                "layernorm2",
+                compute_time=layernorm2_f + residual2_f,
+                comm_time=layernorm2_reduction,
+                comm_bytes=layernorm2_bytes,
+            ),
+            backward=None,
+        )
 
         linear_shape = (
             batch_size,
@@ -131,24 +276,31 @@ def _build_decode_transformer_results(
             self.hidden_dim,
             self.vocab_size,
         )
-        linear_softmax_f, _ = self.get_linear_softmax_f(linear_shape)
+        linear_softmax_f, linear_softmax_mem = self.get_linear_softmax_f(linear_shape)
+        transformer_timings["linear_softmax"] = OperationTiming(
+            "linear_softmax",
+            forward=_make_forward(
+                "linear_softmax",
+                compute_time=linear_softmax_f,
+                comm_time=0.0,
+                comm_bytes=0.0,
+                memory=self._mem_levels(linear_softmax_mem),
+            ),
+            backward=None,
+        )
 
-        if self.model_type == "llama":
-            act_f = self.get_swiglu_f(gemm_ffn1)
-        else:
-            act_f = self.get_gelu_f(gemm_ffn1)
+        mlp_group = OperationGroup(
+            "MLP",
+            operations=(
+                transformer_timings["ffn1"],
+                transformer_timings["gelu"],
+                transformer_timings["ffn2"],
+            ),
+        )
 
-        attention_scale_softmax_f = self.get_scale_softmax_f(gemm_attention_score)
-        # Preserve the same per-op structure and collective accounting used in the training path.
-        # Most decode ops (qkv projection, attention score/output, ffn1) still have zero-sized
-        # reductions under tensor parallelism - we pass through the values returned by the shared
-        # helpers so the context/tensor-context graph builder sees the exact same semantics.
+        # Match exact floating-point operation order from original code
         qkv_proj_forward = qkv_proj_time + qkv_proj_reduction
-        attention_reduction = attention_score_reduction + attention_output_reduction
-        attention_comm_bytes = attention_score_size + attention_output_size
-        attention_forward = (
-            attention_score_time + attention_scale_softmax_f + attention_output_time + attention_reduction
-        )
+        attention_forward = attention_score_time + attention_scale_softmax_f + attention_output_time + attention_reduction
         out_proj_forward = out_proj_time + out_proj_reduction
         mha_forward = qkv_proj_forward + attention_forward + out_proj_forward
 
@@ -168,139 +320,17 @@ def _build_decode_transformer_results(
             + layernorm2_reduction
         )
 
-        # Aggregate attention flops and memory accesses from score + output
-        attention_flops = (attention_score_flops or 0) + (attention_output_flops or 0)
-        attention_mem = [
-            (a or 0) + (b or 0)
-            for a, b in zip_longest(attention_score_mem, attention_output_mem, fillvalue=0)
-        ]
-
-        transformer_results = {
-            "qkv_proj": {
-                "forward": qkv_proj_time + qkv_proj_reduction,
-                "backward": 0.0,
-                "forward_gemm": qkv_proj_time,
-                "forward_reduction": qkv_proj_reduction,
-                "backward_gemm": 0.0,
-                "backward_reduction": 0.0,
-                "comm_size_forward": qkv_proj_size,
-                "comm_size_backward": 0,
-                "flops": qkv_proj_flops,
-                "memory_accesses": self._mem_levels(qkv_proj_mem),  # changed to per-level dict
-            },
-            "attention": {
-                "forward": attention_forward,
-                "backward": 0.0,
-                "forward_gemm": attention_score_time + attention_output_time + attention_scale_softmax_f,
-                "forward_reduction": attention_reduction,
-                "backward_gemm": 0.0,
-                "backward_reduction": 0.0,
-                "comm_size_forward": attention_comm_bytes,
-                "comm_size_backward": 0,
-                "flops": attention_flops,
-                "memory_accesses": self._mem_levels(attention_mem),  # changed to per-level dict
-            },
-            "output_proj": {
-                "forward": out_proj_time + out_proj_reduction,
-                "backward": 0.0,
-                "forward_gemm": out_proj_time,
-                "forward_reduction": out_proj_reduction,
-                "backward_gemm": 0.0,
-                "backward_reduction": 0.0,
-                "comm_size_forward": out_proj_size,
-                "comm_size_backward": 0,
-                "flops": out_proj_flops,
-                "memory_accesses": self._mem_levels(out_proj_mem),  # changed to per-level dict
-            },
-            "MHA": {
-                "forward": mha_forward,
-                "backward": 0.0,
-                "forward_reduction": qkv_proj_reduction + attention_reduction + out_proj_reduction,
-                "backward_reduction": 0.0,
-                "comm_size_forward": qkv_proj_size + attention_comm_bytes + out_proj_size,
-                "comm_size_backward": 0,
-                "flops": qkv_proj_flops + attention_flops + out_proj_flops,
-                "memory_accesses": self._combine_mem(
-                    qkv_proj_mem,
-                    attention_mem,
-                    out_proj_mem,
-                ),
-            },
-            "ffn1": {
-                "forward": ffn1_forward,
-                "backward": 0.0,
-                "forward_gemm": ffn1_time,
-                "forward_reduction": ffn1_reduction,
-                "backward_gemm": 0.0,
-                "backward_reduction": 0.0,
-                "comm_size_forward": ffn1_size,
-                "comm_size_backward": 0,
-                "flops": ffn1_flops,
-                "memory_accesses": self._mem_levels(ffn1_mem),  # changed to per-level dict
-            },
-            "ffn2": {
-                "forward": ffn2_forward,
-                "backward": 0.0,
-                "forward_gemm": ffn2_time,
-                "forward_reduction": ffn2_reduction,
-                "backward_gemm": 0.0,
-                "backward_reduction": 0.0,
-                "comm_size_forward": ffn2_size,
-                "comm_size_backward": 0,
-                "flops": ffn2_flops,
-                "memory_accesses": self._mem_levels(ffn2_mem),  # changed to per-level dict
-            },
-            "MLP": {
-                "forward": mlp_forward,
-                "backward": 0.0,
-                "forward_reduction": ffn1_reduction + ffn2_reduction,
-                "backward_reduction": 0.0,
-                "comm_size_forward": ffn1_size + ffn2_size,
-                "comm_size_backward": 0,
-                "flops": ffn1_flops + ffn2_flops,
-                "memory_accesses": self._combine_mem(
-                    ffn1_mem,
-                    ffn2_mem,
-                ),
-            },
-            "gelu": {
-                "forward": act_f,
-                "backward": 0.0,
-                "forward_gemm": act_f,
-                "forward_reduction": 0.0,
-                "backward_gemm": 0.0,
-                "backward_reduction": 0.0,
-                "comm_size_forward": 0,
-                "comm_size_backward": 0,
-            },
-            "layernorm1": {
-                "forward": layernorm1_forward,
-                "backward": 0.0,
-                "forward_reduction": layernorm1_reduction,
-                "backward_reduction": 0.0,
-                "comm_size_forward": layernorm1_bytes,
-                "comm_size_backward": 0,
-            },
-            "layernorm2": {
-                "forward": layernorm2_forward,
-                "backward": 0.0,
-                "forward_reduction": layernorm2_reduction,
-                "backward_reduction": 0.0,
-                "comm_size_forward": layernorm2_bytes,
-                "comm_size_backward": 0,
-            },
-        }
-
         node_breakdown = {
             "transformer_time_f": transformer_forward,
             "transformer_time_b": 0.0,
-            "linear_softmax_f": linear_softmax_f,
+            "linear_softmax_f": transformer_timings["linear_softmax"].total_forward_time(),
             "linear_softmax_b": 0.0,
             "embedding_f": 0.0,
             "embedding_b": 0.0,
         }
 
-        return transformer_results, node_breakdown
+        return transformer_timings, node_breakdown
+
 
     def prepare_decode_graphs(
         self,
@@ -310,21 +340,23 @@ def prepare_decode_graphs(
         gemm_shapes: Optional[Dict[str, Tuple[int, ...]]] = None,
     ):
         intermediate_size = self.intermediate_size
-        transformer_results, node_breakdown = self._build_decode_transformer_results(
+        transformer_timings, node_breakdown = self._build_decode_transformer_results(
             batch_size=batch_size,
             total_seq_len=total_seq_len,
             gemm_shapes=gemm_shapes,
         )
 
         output_act_bytes = gemm_shapes["qkv_proj"][0] * gemm_shapes["qkv_proj"][1] * self.precision_bytes
-        energy = self.calc_energy(transformer_results, output_act_bytes)
+        energy = self.calc_energy(transformer_timings, output_act_bytes)
 
         if self._generate_graphs:
             results_path = os.path.join(self.output_dir, "decode_transformer_results.txt")
             with open(results_path, "w", encoding="utf-8") as results_file:
                 json.dump(
                     {
-                        "transformer_results": transformer_results,
+                        "transformer_results": {
+                            name: timing.to_dict() for name, timing in transformer_timings.items()
+                        },
                         "node_breakdown": node_breakdown,
                     },
                     results_file,
@@ -334,7 +366,7 @@ def prepare_decode_graphs(
 
         return self._prepare_execution_graphs(
             node_breakdown=node_breakdown,
-            transformer_results=transformer_results,
+            transformer_timings=transformer_timings,
             batch_size=batch_size,
             seq_len=1,
             hidden_dim=self.hidden_dim,
@@ -345,24 +377,42 @@ def prepare_decode_graphs(
             gemm_shapes=gemm_shapes,
         ), energy
     
-    def calc_energy(self, transformer_results: Dict[str, Dict[str, float]], cross_layer_comm) -> float:
+    def calc_energy(self, transformer_timings: Dict[str, OperationTiming], cross_layer_comm) -> float:
         """
         Calculate energy consumption based on transformer results.
         """
-        total_flops = 0
-        total_hbm_bytes = 0
-        inter_comm_bytes = 0 # data parallelism?
-        for op_name, op_results in transformer_results.items():
-            if op_name in ['MHA', 'MLP', 'layernorm1', 'layernorm2', 'embedding', 'linear_softmax']:
-                total_flops += op_results.get("flops", 0)
-                total_hbm_bytes += op_results.get("memory_accesses", {}).get("L3", 0)
-                inter_comm_bytes += op_results.get("comm_size_forward", 0)
+        total_flops = 0.0
+        total_hbm_bytes = 0.0
+        inter_comm_bytes = 0.0  # data parallelism?
+
+        aggregate_groups = {
+            "MLP": ("ffn1", "gelu", "ffn2"),
+        }
+        solo_ops = ("layernorm1", "layernorm2", "embedding", "linear_softmax", "qkv_proj", "attention", "output_proj")
+
+        for members in aggregate_groups.values():
+            for name in members:
+                op = transformer_timings.get(name)
+                if op is None:
+                    continue
+                total_flops += op.forward.flops
+                total_hbm_bytes += op.forward.memory_accesses.get("L3", 0.0)
+                inter_comm_bytes += op.forward.comm_bytes
+
+        for name in solo_ops:
+            op = transformer_timings.get(name)
+            if op is None:
+                continue
+            total_flops += op.forward.flops
+            total_hbm_bytes += op.forward.memory_accesses.get("L3", 0.0)
+            inter_comm_bytes += op.forward.comm_bytes
 
         total_comm_bytes = self.num_layers * inter_comm_bytes + (self.lp - 1) * cross_layer_comm
 
         energy_per_flop = self.core.nominal_energy_per_flop
         energy_hbm_byte = self.DRAM.dynamic_energy_per_bit * 8
-        energy_comm_byte = self.network.intra_network.nominal_energy_per_link * 8  # 8 pJ/bit for interconnect
+        # TODO: honor per-dimension interconnect energies; currently assumes all comms use dimension 0.
+        energy_comm_byte = (self.network.energies_per_bit[0] if self.network.energies_per_bit else 0.0) * 8
 
         total_energy = (total_flops * energy_per_flop) + \
             (total_hbm_bytes * energy_hbm_byte) + \
@@ -389,7 +439,7 @@ def calc_time(self) -> float:
 
         self.readjust_type()
         num_SMs = self.hw_config.tech_config.core.num_bundles
-        transformer_results, node_breakdown = self.compute_all_gemm_and_node_times(
+        transformer_timings, node_breakdown = self.compute_all_gemm_and_node_times(
             batch_size,
             vocab_size,
             hidden_dim,
@@ -402,7 +452,7 @@ def calc_time(self) -> float:
 
         output_act_bytes = batch_size * self.seq_len * hidden_dim * self.precision_bytes
 
-        total_energy = self.calc_energy(transformer_results, output_act_bytes)
+        total_energy = self.calc_energy(transformer_timings, output_act_bytes)
 
         head_dim = hidden_dim // num_heads
         token_bytes = LLM_util.kv_cache_token_bytes(
@@ -415,13 +465,15 @@ def calc_time(self) -> float:
         (
             pipeline_graph,
             pipeline_root,
+            _,
+            _,
             transformer_graph,
             transformer_forward_root,
             _,
             interconnect_params,
         ) = self._prepare_execution_graphs(
             node_breakdown=node_breakdown,
-            transformer_results=transformer_results,
+            transformer_timings=transformer_timings,
             batch_size=batch_size,
             seq_len=prefill_len,
             hidden_dim=hidden_dim,
diff --git a/timing_model.py b/timing_model.py
new file mode 100644
index 00000000..b0a1ea28
--- /dev/null
+++ b/timing_model.py
@@ -0,0 +1,238 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, Mapping, Optional, Sequence, Tuple
+
+
+@dataclass(frozen=True)
+class CommSpec:
+    """
+    Metadata describing a single communication collective.
+
+    Attributes:
+        name: Unique identifier used when wiring graph edges.
+        kind: Collective type (e.g. all_reduce, reduce_scatter).
+        size_bytes: Total bytes transferred (int >= 0).
+        participants: Number of ranks involved (int >= 1).
+        interconnect: Logical interconnect label (tp/dp/cp/lp/etc.).
+        extra: Free-form metadata (e.g. debug labels).
+    """
+
+    name: str
+    kind: str
+    size_bytes: int
+    participants: int
+    interconnect: str
+    extra: Mapping[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if not self.name:
+            raise ValueError("CommSpec.name must be non-empty")
+        if not self.kind:
+            raise ValueError(f"CommSpec '{self.name}' missing kind")
+        if self.size_bytes < 0:
+            raise ValueError(f"CommSpec '{self.name}' has negative size_bytes ({self.size_bytes})")
+        if self.participants <= 0:
+            raise ValueError(f"CommSpec '{self.name}' must have participants >= 1 (got {self.participants})")
+        if not self.interconnect:
+            raise ValueError(f"CommSpec '{self.name}' missing interconnect label")
+
+
+@dataclass(frozen=True)
+class DirectionTiming:
+    """
+    Timing information for a single direction (forward/backward) of an operation.
+
+    Attributes:
+        compute_time: Pure computation time (seconds).
+        comm_time: Time attributed to communication collectives (seconds).
+        comm_bytes: Raw communication bytes (int >= 0). CommSpec instances are created
+                    later during graph building based on COMMUNICATION_RULES.
+        flops: Floating-point operation count (optional).
+        memory_accesses: Tuple describing memory accesses per level (optional).
+        notes: Misc metadata for debugging/reporting.
+    """
+
+    compute_time: float
+    comm_time: float = 0.0
+    comm_bytes: int = 0
+    flops: float = 0.0
+    memory_accesses: Mapping[str, float] = field(default_factory=dict)
+    notes: Mapping[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if self.compute_time < 0:
+            raise ValueError(f"DirectionTiming compute_time must be >= 0 (got {self.compute_time})")
+        if self.comm_time < 0:
+            raise ValueError(f"DirectionTiming comm_time must be >= 0 (got {self.comm_time})")
+        if self.comm_bytes < 0:
+            raise ValueError(f"DirectionTiming comm_bytes must be >= 0 (got {self.comm_bytes})")
+        if self.flops < 0:
+            raise ValueError(f"DirectionTiming flops must be >= 0 (got {self.flops})")
+        for level_bytes in self.memory_accesses.values():
+            if level_bytes < 0:
+                raise ValueError("DirectionTiming memory_accesses entries must be >= 0")
+
+    def total_time(self) -> float:
+        """Return compute + comm time."""
+        return self.compute_time + self.comm_time
+
+    def total_comm_bytes(self) -> int:
+        """Return raw communication bytes."""
+        return self.comm_bytes
+
+    def has_comm(self) -> bool:
+        return self.comm_bytes > 0
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Serialize to a dict for backwards compatibility (e.g. logging)."""
+        return {
+            "compute_time": self.compute_time,
+            "comm_time": self.comm_time,
+            "comm_bytes": self.comm_bytes,
+            "flops": self.flops,
+            "memory_accesses": dict(self.memory_accesses),
+            "notes": dict(self.notes),
+        }
+
+
+@dataclass(frozen=True)
+class OperationTiming:
+    """
+    Timing details for an operation (e.g., QKV GEMM, MLP).
+
+    Attributes:
+        name: Logical operation name.
+        forward: Forward-direction timing (optional).
+        backward: Backward-direction timing (optional).
+        metadata: Additional operation-level metadata (e.g. layer index).
+    """
+
+    name: str
+    forward: Optional[DirectionTiming]
+    backward: Optional[DirectionTiming] = None
+    metadata: Mapping[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if not self.name:
+            raise ValueError("OperationTiming.name must be non-empty")
+        if self.forward is None and self.backward is None:
+            raise ValueError(f"OperationTiming '{self.name}' must have at least one direction")
+        if self.forward is not None and not isinstance(self.forward, DirectionTiming):
+            raise TypeError("OperationTiming.forward must be a DirectionTiming instance or None")
+        if self.backward is not None and not isinstance(self.backward, DirectionTiming):
+            raise TypeError("OperationTiming.backward must be a DirectionTiming instance or None")
+
+    def has_backward(self) -> bool:
+        return self.backward is not None
+
+    def total_forward_time(self) -> float:
+        if self.forward is None:
+            raise RuntimeError(f"Operation '{self.name}' does not have forward timing")
+        return self.forward.total_time()
+
+    def total_backward_time(self) -> float:
+        if self.backward is None:
+            raise RuntimeError(f"Operation '{self.name}' does not have backward timing")
+        return self.backward.total_time()
+
+    def total_time(self) -> float:
+        total = self.forward.total_time() if self.forward is not None else 0.0
+        if self.backward is not None:
+            total += self.backward.total_time()
+        return total
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Serialize to dict for reporting or artifact dumps."""
+        data: Dict[str, Any] = {
+            "name": self.name,
+        }
+        if self.forward is not None:
+            data["forward"] = self.forward.to_dict()
+        if self.backward is not None:
+            data["backward"] = self.backward.to_dict()
+        if self.metadata:
+            data["metadata"] = dict(self.metadata)
+        return data
+
+    def validate(self, require_backward: bool = False) -> None:
+        if require_backward and self.backward is None:
+            raise ValueError(f"Operation '{self.name}' is missing backward timing")
+
+
+@dataclass(frozen=True)
+class OperationGroup:
+    """
+    Aggregates a set of OperationTiming instances (e.g. the GEMMs that form MHA).
+    Provides convenience accessors for summary statistics.
+    """
+
+    name: str
+    operations: Tuple[OperationTiming, ...]
+    metadata: Mapping[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if not self.name:
+            raise ValueError("OperationGroup.name must be non-empty")
+        if not self.operations:
+            raise ValueError(f"OperationGroup '{self.name}' must contain at least one operation")
+        for op in self.operations:
+            if not isinstance(op, OperationTiming):
+                raise TypeError("OperationGroup.operations must contain OperationTiming instances")
+
+    def forward_compute_time(self) -> float:
+        return sum(op.forward.compute_time for op in self.operations)
+
+    def forward_comm_time(self) -> float:
+        return sum(op.forward.comm_time for op in self.operations)
+
+    def forward_comm_bytes(self) -> int:
+        return sum(op.forward.comm_bytes for op in self.operations)
+
+    def forward_total_time(self) -> float:
+        return sum(op.total_forward_time() for op in self.operations)
+
+    def backward_compute_time(self) -> float:
+        total = 0.0
+        for op in self.operations:
+            if op.backward is None:
+                raise RuntimeError(
+                    f"OperationGroup '{self.name}' member '{op.name}' is missing backward timing"
+                )
+            total += op.backward.compute_time
+        return total
+
+    def backward_comm_time(self) -> float:
+        total = 0.0
+        for op in self.operations:
+            if op.backward is None:
+                raise RuntimeError(
+                    f"OperationGroup '{self.name}' member '{op.name}' is missing backward timing"
+                )
+            total += op.backward.comm_time
+        return total
+
+    def backward_comm_bytes(self) -> int:
+        total = 0
+        for op in self.operations:
+            if op.backward is None:
+                raise RuntimeError(
+                    f"OperationGroup '{self.name}' member '{op.name}' is missing backward timing"
+                )
+            total += op.backward.comm_bytes
+        return total
+
+    def backward_total_time(self) -> float:
+        total = 0.0
+        for op in self.operations:
+            if op.backward is None:
+                raise RuntimeError(
+                    f"OperationGroup '{self.name}' member '{op.name}' is missing backward timing"
+                )
+            total += op.total_backward_time()
+        return total
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "name": self.name,
+            "operations": [op.to_dict() for op in self.operations],
+            "metadata": dict(self.metadata),
+        }
diff --git a/tools/exec_with_fault.py b/tools/exec_with_fault.py
new file mode 100644
index 00000000..8ee47b2a
--- /dev/null
+++ b/tools/exec_with_fault.py
@@ -0,0 +1,212 @@
+#!/usr/bin/env python3
+"""
+Execute the LLM AstraSim debug harness under multiple faulty-link scenarios.
+
+For each predefined configuration this script patches the
+``configs/hardware-config/a100_80GB_faraz.yaml`` file with the desired
+``network.faulty_links`` list, invokes ``./examples/llm_astra_inference_debug_graphviz.sh``,
+parses the reported training time, and records the results in a summary table.
+
+The original YAML contents are restored once the sweep finishes (even if a run fails).
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import copy
+import os
+import re
+import subprocess
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Iterable, List, Sequence
+
+import yaml
+
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+CONFIG_PATH = PROJECT_ROOT / "configs" / "hardware-config" / "a100_80GB_faraz.yaml"
+RUN_SCRIPT = "./examples/llm_astra_inference_debug_graphviz.sh"
+
+# Faulty-link sweep configurations.
+# Each entry provides a human-readable label and the list of faulty links to inject.
+# Faulty links follow the AstraSim convention: [src, dst, reliability_weight]
+FAULT_SWEEP_CASES: Sequence[dict[str, Any]] = [
+    {"label": "baseline", "faulty_links": []},
+    {"label": "tp_top_left_harsh", "faulty_links": [[0, 1, 0.3]]},
+    {"label": "tp_top_left", "faulty_links": [[0, 1, 0.5]]},
+    {"label": "tp_sequential", "faulty_links": [[0, 1, 0.5], [1, 2, 0.5]]},
+    {"label": "tp_top_left_vert", "faulty_links": [[0, 1, 0.5], [0, 4, 0.5]]},
+    {"label": "tp_sequential_harsh", "faulty_links": [[0, 1, 0.3], [1, 2, 0.3]]},
+    {"label": "tp_top_left_vert_harsh", "faulty_links": [[0, 1, 0.3], [0, 4, 0.3]]},
+]
+
+
+TRAINING_TIME_PATTERN = re.compile(r"Training time for batch:\s*([0-9.]+)s")
+
+
+@dataclass
+class SweepResult:
+    label: str
+    links: Sequence[Sequence[float]]
+    returncode: int
+    training_time: float | None
+    stdout: str
+    stderr: str
+
+
+def _load_yaml(path: Path) -> Any:
+    with path.open("r", encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def _dump_yaml(path: Path, data: Any) -> None:
+    with path.open("w", encoding="utf-8") as handle:
+        yaml.safe_dump(data, handle, sort_keys=False)
+
+
+def _set_faulty_links(doc: Any, links: Sequence[Sequence[float]]) -> None:
+    network_block = doc.get("network")
+    if isinstance(network_block, dict):
+        network_block["faulty_links"] = copy.deepcopy(list(links))
+        return
+    raise ValueError("Hardware config network section must be a mapping with 'faulty_links'.")
+
+
+def _run_script(cwd: Path) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        [RUN_SCRIPT],
+        cwd=str(cwd),
+        check=False,
+        text=True,
+        capture_output=True,
+        shell=False,
+    )
+
+
+def _parse_training_time(output: str) -> float | None:
+    match = TRAINING_TIME_PATTERN.search(output)
+    if not match:
+        return None
+    try:
+        return float(match.group(1))
+    except ValueError:
+        return None
+
+
+def _format_summary(results: Iterable[SweepResult]) -> str:
+    lines: List[str] = []
+    header = f"{'Label':<24} {'Training Time (s)':>18} {'Return Code':>12}"
+    lines.append(header)
+    lines.append("-" * len(header))
+    for result in results:
+        time_str = f"{result.training_time:.2f}" if result.training_time is not None else "n/a"
+        lines.append(f"{result.label:<24} {time_str:>18} {result.returncode:>12}")
+    return "\n".join(lines)
+
+
+def run_sweep(cases: Sequence[dict[str, Any]], *, dry_run: bool = False) -> List[SweepResult]:
+    original_yaml = _load_yaml(CONFIG_PATH)
+    original_links = copy.deepcopy(original_yaml.get("network", {}).get("faulty_links"))
+    results: List[SweepResult] = []
+
+    try:
+        for case in cases:
+            label = case.get("label", "unnamed")
+            links = case.get("faulty_links", [])
+            print(f"[exec_with_fault] Running case '{label}' with links={links}")
+
+            next_yaml = copy.deepcopy(original_yaml)
+            _set_faulty_links(next_yaml, links)
+            if not dry_run:
+                _dump_yaml(CONFIG_PATH, next_yaml)
+
+            completed: subprocess.CompletedProcess[str]
+            if dry_run:
+                completed = subprocess.CompletedProcess(
+                    args=[RUN_SCRIPT], returncode=0, stdout="", stderr=""
+                )
+            else:
+                completed = _run_script(PROJECT_ROOT)
+
+            combined_output = f"{completed.stdout}\n{completed.stderr}"
+            training_time = _parse_training_time(combined_output)
+            if training_time is None:
+                print(
+                    f"[exec_with_fault] Warning: training time not found in output for case '{label}'.",
+                    file=sys.stderr,
+                )
+
+            results.append(
+                SweepResult(
+                    label=label,
+                    links=links,
+                    returncode=completed.returncode,
+                    training_time=training_time,
+                    stdout=completed.stdout,
+                    stderr=completed.stderr,
+                )
+            )
+            if completed.returncode != 0:
+                print(
+                    f"[exec_with_fault] Case '{label}' exited with return code {completed.returncode}.",
+                    file=sys.stderr,
+                )
+    finally:
+        # Restore original faulty links (or remove the key if it was missing)
+        restored_yaml = _load_yaml(CONFIG_PATH)
+        network_block = restored_yaml.get("network")
+        if isinstance(network_block, dict):
+            if original_links is None:
+                network_block.pop("faulty_links", None)
+            else:
+                network_block["faulty_links"] = original_links
+            if not dry_run:
+                _dump_yaml(CONFIG_PATH, restored_yaml)
+
+    return results
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Execute AstraSim runs with predefined faulty links.")
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Do not modify files or run the benchmark; useful for quick validation.",
+    )
+    parser.add_argument(
+        "--cases",
+        nargs="*",
+        help="Optional subset of case labels to run. Defaults to the entire FAULT_SWEEP_CASES list.",
+    )
+    args = parser.parse_args(argv)
+
+    if args.dry_run:
+        print("[exec_with_fault] Dry run enabled: configurations will not be written or executed.")
+
+    selected_cases = []
+    if args.cases:
+        case_set = set(args.cases)
+        for case in FAULT_SWEEP_CASES:
+            if case.get("label") in case_set:
+                selected_cases.append(case)
+        missing = case_set - {case.get("label") for case in selected_cases}
+        if missing:
+            parser.error(f"Unknown case label(s): {', '.join(sorted(missing))}")
+    else:
+        selected_cases = list(FAULT_SWEEP_CASES)
+
+    if not selected_cases:
+        parser.error("No cases selected for execution.")
+
+    results = run_sweep(selected_cases, dry_run=args.dry_run)
+    print("\nSweep summary:\n")
+    print(_format_summary(results))
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/tools/fault_meta_sweep.py b/tools/fault_meta_sweep.py
new file mode 100644
index 00000000..3425643e
--- /dev/null
+++ b/tools/fault_meta_sweep.py
@@ -0,0 +1,189 @@
+#!/usr/bin/env python3
+"""
+Meta sweep orchestrator for DeepFlow fault sensitivity experiments.
+
+Runs a sequence of configurations by invoking ``fault_sweep.py`` with tailored
+arguments, collects artefacts under ``tools/fault_sweep_results/``, and skips
+any configuration whose plot already exists.
+"""
+
+import subprocess
+import sys
+from pathlib import Path
+from typing import List
+
+
+RESULTS_DIR = Path("tools", "fault_sweep_results")
+RESULTS_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def run_config(name: str, extra_args: List[str]) -> None:
+    plot_path = RESULTS_DIR / f"{name}.png"
+    if plot_path.exists():
+        print(f"[meta] Skipping '{name}' (plot already exists).")
+        return
+
+    report_path = RESULTS_DIR / f"{name}.tsv"
+    results_json = RESULTS_DIR / f"{name}_results.json"
+
+    cmd = [
+        sys.executable,
+        "tools/fault_sweep.py",
+        "--plot-output",
+        str(plot_path),
+        "--report-output",
+        str(report_path),
+        "--results-json",
+        str(results_json),
+    ]
+    cmd.extend(extra_args)
+
+    print(f"[meta] Running configuration '{name}'")
+    completed = subprocess.run(cmd, text=True)
+    if completed.returncode != 0:
+        raise RuntimeError(f"Configuration '{name}' failed with exit code {completed.returncode}")
+
+
+def main() -> None:
+    common_args = [
+        "--target-num-gpus",
+        "256",
+        "--fault-iter",
+        "200",
+        "--fault-workers",
+        "86",
+        "--fault-mag",
+        "0.5,0.1",
+        "--num-faults",
+        "1",
+        "--min-tp",
+        "1",
+        "--max-tp",
+        "16",
+        "--min-cp",
+        "1",
+        "--max-cp",
+        "256",
+        "--min-dp",
+        "1",
+        "--max-dp",
+        "256",
+        "--min-lp",
+        "1",
+        "--max-lp",
+        "128",
+        "--seed",
+        "1337",
+    ]
+
+    configs = [
+        (
+            "dense_1fault_alldim_256",
+            common_args
+            + [
+                "--sample-count",
+                "25",
+                "--fault-iter",
+                "50",
+                "--fault-workers",
+                "96",
+            ],
+        ),
+        (
+            "dense_1fault_dim0_256",
+            common_args
+            + [
+                "--sample-count",
+                "48",
+                "--allowed-fault-dims",
+                "0",
+            ],
+        ),
+        (
+            "sparse_1fault_alldim_tpcponly_256",
+            [
+                "--target-num-gpus",
+                "256",
+                "--fault-iter",
+                "100",
+                "--fault-workers",
+                "86",
+                "--fault-mag",
+                "0.5,0.1",
+                "--num-faults",
+                "1",
+                "--sample-count",
+                "10",
+                "--min-tp",
+                "1",
+                "--max-tp",
+                "64",
+                "--min-cp",
+                "1",
+                "--max-cp",
+                "128",
+                "--min-dp",
+                "1",
+                "--max-dp",
+                "1",
+                "--min-lp",
+                "1",
+                "--max-lp",
+                "1",
+                "--seed",
+                "1337",
+            ],
+        ),
+        (
+            "sparse_1fault_alldim_dplponly_256",
+            [
+                "--target-num-gpus",
+                "256",
+                "--fault-iter",
+                "200",
+                "--fault-workers",
+                "86",
+                "--fault-mag",
+                "0.5,0.1",
+                "--num-faults",
+                "1",
+                "--sample-count",
+                "20",
+                "--min-tp",
+                "1",
+                "--max-tp",
+                "1",
+                "--min-cp",
+                "1",
+                "--max-cp",
+                "1",
+                "--min-dp",
+                "1",
+                "--max-dp",
+                "256",
+                "--min-lp",
+                "1",
+                "--max-lp",
+                "128",
+                "--seed",
+                "1337",
+            ],
+        ),
+        (
+            "dense_1fault_alldim_256_harsh",
+            common_args
+            + [
+                "--sample-count",
+                "48",
+                "--fault-mag",
+                "0.3,0.2",
+            ],
+        ),
+    ]
+
+    for name, args in configs:
+        run_config(name, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/fault_sweep.py b/tools/fault_sweep.py
new file mode 100644
index 00000000..c85d3e57
--- /dev/null
+++ b/tools/fault_sweep.py
@@ -0,0 +1,1209 @@
+#!/usr/bin/env python3
+"""
+Fault sensitivity sweep for DeepFlow parallelism configurations.
+
+Generates a representative set of parallelism tuples that map to a fixed GPU
+budget, evaluates baseline runtimes, and then injects random soft faults across
+network dimensions to gauge runtime variability.
+"""
+
+import argparse
+import atexit
+import datetime
+import copy
+import json
+import math
+import os
+import random
+import signal
+import shlex
+import sys
+import threading
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+
+import matplotlib.pyplot as plt
+import numpy as np
+import yaml
+from tqdm import tqdm
+
+try:  # Optional styling dependency
+    import seaborn as sns  # type: ignore
+except ImportError:  # pragma: no cover
+    sns = None
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import config  # noqa: E402
+
+from parallelism_sweep import (  # noqa: E402
+    HARDWARE_CONFIG_PATH,
+    MODEL_CONFIG_PATH,
+    ASTRA_CACHE_MODE,
+    determine_model_mode,
+    evaluate_parallelism,
+    read_yaml,
+    set_astrasim_cache_mode,
+    tp_cp_product_is_power_of_two_square,
+)
+
+
+class FlowSeq(list):
+    """List subclass that forces YAML flow-style serialization."""
+
+
+def _represent_flow_seq(dumper, data):
+    return dumper.represent_sequence("tag:yaml.org,2002:seq", data, flow_style=True)
+
+
+yaml.SafeDumper.add_representer(FlowSeq, _represent_flow_seq)
+
+
+def _quantize_weight(value: float, decimals: int = 2) -> float:
+    return round(float(value), decimals)
+
+# -----------------------------------------------------------------------------
+# Fault sweep configuration
+# -----------------------------------------------------------------------------
+
+TARGET_NUM_GPUS = 64
+SAMPLE_COUNT = 10
+FAULT_ITER = 25
+FAULT_WORKERS = 100
+FAULT_MAG = [0.5, 0.1]  # May also be a (mean, std) tuple
+NUM_FAULTS = [1]
+
+MIN_ALLOWED_TP = 1
+MAX_ALLOWED_TP = 128
+MIN_ALLOWED_CP = 1
+MAX_ALLOWED_CP = 128
+MIN_ALLOWED_DP = 1
+MAX_ALLOWED_DP = 128
+MIN_ALLOWED_LP = 1
+MAX_ALLOWED_LP = 128
+ALLOWED_FAULT_DIMS: Optional[List[int]] = None
+
+PLOT_OUTPUT_PATH = "tools/fault_sweep.png"
+REPORT_OUTPUT_PATH = "tools/fault_sweep.tsv"
+FAULT_OUTPUT_DIR = os.path.join(os.getcwd(), "astra_cache_faults")
+RESULTS_JSON_PATH = Path("tools/results.json")
+
+RANDOM_SEED = 1337
+RESULTS_WRITE_LOCK = threading.Lock()
+_ACTIVE_EXECUTOR: Optional[ProcessPoolExecutor] = None
+_OLD_SIGINT_HANDLER = None
+_OLD_SIGTERM_HANDLER = None
+
+# When True, discard configurations whose tp*cp product is not a square power of two.
+ENFORCE_SQUARE_TP_CP = True
+
+# Debug dumping controls
+DEBUG_MODE = False
+DEBUG_OUTPUT_BASE = Path("dbg_hw_conf")
+DEBUG_RUN_DIR: Optional[Path] = None
+DEBUG_SAVED_PATHS: List[Path] = []
+_DEBUG_FILE_INDEX = 0
+
+# Failure dump controls
+FAILURE_OUTPUT_BASE = DEBUG_OUTPUT_BASE / "failures"
+FAILURE_DUMP_DIR: Optional[Path] = None
+FAILURE_SAVED_PATHS: List[Path] = []
+_FAILURE_FILE_INDEX = 0
+
+
+def _parse_int_list(value: Optional[str]) -> List[int]:
+    if not value:
+        return []
+    parts = [item.strip() for item in value.split(',')]
+    result: List[int] = []
+    for part in parts:
+        if not part:
+            continue
+        result.append(int(part))
+    return result
+
+
+def _parse_float_tuple(value: Optional[str]) -> Tuple[float, float]:
+    if not value:
+        return (0.0, 0.0)
+    parts = [item.strip() for item in value.split(',') if item.strip()]
+    if len(parts) == 1:
+        val = float(parts[0])
+        return (val, 0.0)
+    if len(parts) >= 2:
+        return (float(parts[0]), float(parts[1]))
+    return (0.0, 0.0)
+
+# Network topologies in AstraSim that understand per-link fault annotations.
+# The neighbor builders below encode each topology's link structure so that
+# fault injection only targets physically valid edges.
+TOPOLOGIES_WITH_FAULTY_LINK_SUPPORT = {
+    "ring",
+    "torus2d",
+    "mesh",
+    "mesh2d",
+    "kingmesh2d",
+}
+
+def _build_ring_neighbors(node_count: int) -> List[List[int]]:
+    if node_count <= 0:
+        return []
+    if node_count == 1:
+        return [[]]
+    neighbors: List[List[int]] = []
+    for idx in range(node_count):
+        left = (idx - 1) % node_count
+        right = (idx + 1) % node_count
+        entries: List[int] = []
+        for candidate in (left, right):
+            if candidate != idx and candidate not in entries:
+                entries.append(candidate)
+        neighbors.append(entries)
+    return neighbors
+
+
+def _build_mesh_neighbors(node_count: int) -> List[List[int]]:
+    if node_count <= 0:
+        return []
+    neighbors: List[List[int]] = []
+    for idx in range(node_count):
+        entries: List[int] = []
+        if idx - 1 >= 0:
+            entries.append(idx - 1)
+        if idx + 1 < node_count:
+            entries.append(idx + 1)
+        neighbors.append(entries)
+    return neighbors
+
+
+def _build_mesh2d_neighbors(node_count: int) -> List[List[int]]:
+    if node_count <= 0:
+        return []
+    dim = int(math.isqrt(node_count))
+    if dim * dim != node_count:
+        raise ValueError(f"Mesh2D topology requires perfect square node count, got {node_count}.")
+    neighbors: List[List[int]] = [[] for _ in range(node_count)]
+    for idx in range(node_count):
+        row, col = divmod(idx, dim)
+        entries: List[int] = []
+        if col - 1 >= 0:
+            entries.append(row * dim + (col - 1))
+        if col + 1 < dim:
+            entries.append(row * dim + (col + 1))
+        if row - 1 >= 0:
+            entries.append((row - 1) * dim + col)
+        if row + 1 < dim:
+            entries.append((row + 1) * dim + col)
+        neighbors[idx] = entries
+    return neighbors
+
+
+def _build_torus2d_neighbors(node_count: int) -> List[List[int]]:
+    if node_count <= 0:
+        return []
+    dim = int(math.isqrt(node_count))
+    if dim * dim != node_count:
+        raise ValueError(f"Torus2D topology requires perfect square node count, got {node_count}.")
+    neighbors: List[List[int]] = [[] for _ in range(node_count)]
+    if dim == 1:
+        return neighbors
+    for idx in range(node_count):
+        row, col = divmod(idx, dim)
+        left = row * dim + ((col - 1) % dim)
+        right = row * dim + ((col + 1) % dim)
+        up = ((row - 1) % dim) * dim + col
+        down = ((row + 1) % dim) * dim + col
+        entries = []
+        for candidate in (left, right, up, down):
+            if candidate != idx and candidate not in entries:
+                entries.append(candidate)
+        neighbors[idx] = entries
+    return neighbors
+
+
+def _build_kingmesh2d_neighbors(node_count: int) -> List[List[int]]:
+    if node_count <= 0:
+        return []
+    dim = int(math.isqrt(node_count))
+    if dim * dim != node_count:
+        raise ValueError(f"KingMesh2D topology requires perfect square node count, got {node_count}.")
+    neighbors: List[List[int]] = [[] for _ in range(node_count)]
+    for idx in range(node_count):
+        row, col = divmod(idx, dim)
+        entries: List[int] = []
+        for dy in (-1, 0, 1):
+            for dx in (-1, 0, 1):
+                if dx == 0 and dy == 0:
+                    continue
+                nr = row + dy
+                nc = col + dx
+                if 0 <= nr < dim and 0 <= nc < dim:
+                    candidate = nr * dim + nc
+                    if candidate != idx:
+                        entries.append(candidate)
+        neighbors[idx] = entries
+    return neighbors
+
+
+_NEIGHBOR_BUILDERS = {
+    "ring": _build_ring_neighbors,
+    "mesh": _build_mesh_neighbors,
+    "mesh2d": _build_mesh2d_neighbors,
+    "torus2d": _build_torus2d_neighbors,
+    "kingmesh2d": _build_kingmesh2d_neighbors,
+}
+
+
+_NEIGHBOR_CACHE: Dict[Tuple[str, int], Tuple[Tuple[int, ...], ...]] = {}
+
+
+def _neighbors_for_topology(topo_type: str, node_count: int) -> Tuple[Tuple[int, ...], ...]:
+    key = (topo_type, int(node_count))
+    cached = _NEIGHBOR_CACHE.get(key)
+    if cached is not None:
+        return cached
+    builder = _NEIGHBOR_BUILDERS.get(topo_type)
+    if builder is None:
+        raise ValueError(f"No neighbor builder registered for topology '{topo_type}'.")
+    neighbor_lists = builder(int(node_count))
+    normalized = tuple(tuple(neigh) for neigh in neighbor_lists)
+    _NEIGHBOR_CACHE[key] = normalized
+    return normalized
+
+
+def _candidate_fault_dimensions(settings: Dict[str, int]) -> List[Tuple[int, int]]:
+    candidates: List[Tuple[int, int]] = []
+    dim0_nodes = int(settings.get("tp", 1)) * int(settings.get("cp", 1))
+    if dim0_nodes >= 2 and (ALLOWED_FAULT_DIMS is None or 0 in ALLOWED_FAULT_DIMS):
+        candidates.append((0, dim0_nodes))
+    dim1_nodes = int(settings.get("lp", 1)) * int(settings.get("dp", 1))
+    if dim1_nodes >= 2 and (ALLOWED_FAULT_DIMS is None or 1 in ALLOWED_FAULT_DIMS):
+        candidates.append((1, dim1_nodes))
+    return candidates
+
+
+def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="DeepFlow fault sensitivity sweep")
+    parser.add_argument("--target-num-gpus", type=int, default=TARGET_NUM_GPUS)
+    parser.add_argument("--sample-count", type=int, default=SAMPLE_COUNT)
+    parser.add_argument("--fault-iter", type=int, default=FAULT_ITER)
+    parser.add_argument("--fault-workers", type=int, default=FAULT_WORKERS)
+    parser.add_argument("--fault-mag", type=str, default=f"{FAULT_MAG[0]},{FAULT_MAG[1] if isinstance(FAULT_MAG, (tuple, list)) else 0.0}", help="Mean[,std] for fault magnitude")
+    parser.add_argument("--num-faults", type=str, default=",".join(str(n) for n in NUM_FAULTS), help="Comma-separated list of number of faults per evaluation")
+    parser.add_argument("--min-tp", type=int, default=MIN_ALLOWED_TP)
+    parser.add_argument("--max-tp", type=int, default=MAX_ALLOWED_TP)
+    parser.add_argument("--min-cp", type=int, default=MIN_ALLOWED_CP)
+    parser.add_argument("--max-cp", type=int, default=MAX_ALLOWED_CP)
+    parser.add_argument("--min-dp", type=int, default=MIN_ALLOWED_DP)
+    parser.add_argument("--max-dp", type=int, default=MAX_ALLOWED_DP)
+    parser.add_argument("--min-lp", type=int, default=MIN_ALLOWED_LP)
+    parser.add_argument("--max-lp", type=int, default=MAX_ALLOWED_LP)
+    parser.add_argument("--allowed-fault-dims", type=str, default=None, help="Comma-separated list of network dimension indices eligible for faults")
+    parser.add_argument("--plot-output", type=str, default=PLOT_OUTPUT_PATH)
+    parser.add_argument("--report-output", type=str, default=REPORT_OUTPUT_PATH)
+    parser.add_argument("--results-json", type=str, default=str(RESULTS_JSON_PATH))
+    parser.add_argument("--seed", type=int, default=RANDOM_SEED)
+    parser.add_argument(
+        "--enforce-square-tp-cp",
+        action="store_true",
+        help="Require tp*cp to be a square power of two when sampling parallelism tuples.",
+    )
+    parser.add_argument(
+        "--debug-mode",
+        action="store_true",
+        help="Enable debug dumps of hardware configs and exit after generation.",
+    )
+    return parser.parse_args(argv)
+
+
+# -----------------------------------------------------------------------------
+# Parallelism sampling helpers
+# -----------------------------------------------------------------------------
+
+def _divisors(n: int) -> List[int]:
+    divs = set()
+    for i in range(1, int(math.sqrt(n)) + 1):
+        if n % i == 0:
+            divs.add(i)
+            divs.add(n // i)
+    return sorted(divs)
+
+
+def _enumerate_parallelism_vectors(num_gpus: int) -> List[Tuple[int, int, int, int]]:
+    vectors: List[Tuple[int, int, int, int]] = []
+    for tp in _divisors(num_gpus):
+        rem_tp = num_gpus // tp
+        for cp in _divisors(rem_tp):
+            rem_cp = rem_tp // cp
+            for dp in _divisors(rem_cp):
+                lp = rem_cp // dp
+                vectors.append((tp, cp, dp, lp))
+    return vectors
+
+
+def _log_distance(a: Sequence[int], b: Sequence[int]) -> float:
+    total = 0.0
+    for x, y in zip(a, b):
+        total += abs(math.log2(max(x, 1)) - math.log2(max(y, 1)))
+    return total
+
+
+def _select_representative_vectors(candidates: List[Tuple[int, int, int, int]], count: int) -> List[Tuple[int, int, int, int]]:
+    if not candidates:
+        return []
+    selected: List[Tuple[int, int, int, int]] = []
+
+    # Seed with extreme configurations
+    extremes = [
+        max(candidates, key=lambda v: v[0]),
+        max(candidates, key=lambda v: v[1]),
+        max(candidates, key=lambda v: v[2]),
+        max(candidates, key=lambda v: v[3]),
+    ]
+    for cand in extremes:
+        if cand not in selected:
+            selected.append(cand)
+        if len(selected) >= count:
+            return selected[:count]
+
+    # Farthest-point sampling to promote coverage
+    remaining = [v for v in candidates if v not in selected]
+    while remaining and len(selected) < count:
+        best_vec = None
+        best_score = -1.0
+        for vec in remaining:
+            min_dist = min(_log_distance(vec, sel) for sel in selected)
+            if min_dist > best_score:
+                best_score = min_dist
+                best_vec = vec
+        if best_vec is None:
+            break
+        selected.append(best_vec)
+        remaining.remove(best_vec)
+
+    # If still short, append random leftovers
+    while remaining and len(selected) < count:
+        candidate = remaining.pop(random.randrange(len(remaining)))
+        selected.append(candidate)
+
+    return selected[:count]
+
+
+def generate_parallelism_samples(num_gpus: int, sample_count: int) -> List[Dict[str, int]]:
+    tuples = _enumerate_parallelism_vectors(num_gpus)
+    filtered: List[Tuple[int, int, int, int]] = []
+    for tp, cp, dp, lp in tuples:
+        if not (MIN_ALLOWED_TP <= tp <= MAX_ALLOWED_TP):
+            continue
+        if not (MIN_ALLOWED_CP <= cp <= MAX_ALLOWED_CP):
+            continue
+        if not (MIN_ALLOWED_DP <= dp <= MAX_ALLOWED_DP):
+            continue
+        if not (MIN_ALLOWED_LP <= lp <= MAX_ALLOWED_LP):
+            continue
+        if ENFORCE_SQUARE_TP_CP and not tp_cp_product_is_power_of_two_square(tp, cp):
+            continue
+        filtered.append((tp, cp, dp, lp))
+    tuples = filtered
+    chosen = _select_representative_vectors(tuples, sample_count)
+    samples: List[Dict[str, int]] = []
+    for tp, cp, dp, lp in chosen:
+        settings = {
+            "tp": tp,
+            "cp": cp,
+            "dp": dp,
+            "lp": lp,
+            "mb": max(lp, 1),
+            "tp_sp": True,
+        }
+        samples.append(settings)
+    return samples
+
+
+# -----------------------------------------------------------------------------
+# Fault injection helpers
+# -----------------------------------------------------------------------------
+
+def sample_fault_magnitude() -> float:
+    if isinstance(FAULT_MAG, (tuple, list)) and len(FAULT_MAG) == 2:
+        mean, std = FAULT_MAG
+        value = random.gauss(mean, std)
+    else:
+        value = float(FAULT_MAG)
+    value = max(0.0, min(0.95, value))
+    value = _quantize_weight(value)
+    return max(0.1, value) # min of 10% bw, max of 95% bw
+
+
+def make_fault_mutator(fault_specs: Sequence[Tuple[int, float, int]]):
+    fault_specs = [
+        (int(dim_idx), float(weight), int(node_count))
+        for dim_idx, weight, node_count in fault_specs
+    ]
+
+    def _mutator(hw_dict: Dict[str, object]) -> None:
+        network = hw_dict.get("network")
+        if not isinstance(network, dict):
+            return
+        dimensions = network.get("dimensions")
+        if not isinstance(dimensions, list):
+            return
+        fault_entries: List[FlowSeq] = []
+        for dimension_index, weight, node_count in fault_specs:
+            if not (0 <= dimension_index < len(dimensions)):
+                continue
+            dim = dimensions[dimension_index]
+            if not isinstance(dim, dict):
+                continue
+            topology = dim.get("topology")
+            topo_type = ""
+            if isinstance(topology, dict):
+                topo_raw = topology.get("type", "")
+                if isinstance(topo_raw, str):
+                    topo_type = topo_raw.strip().lower()
+            supports_faulty_links = topo_type in TOPOLOGIES_WITH_FAULTY_LINK_SUPPORT
+            if not supports_faulty_links:
+                label = dim.get("label", f"dim{dimension_index}")
+                raise ValueError(
+                    f"Topology '{topo_type or 'unknown'}' on network dimension '{label}' "
+                    "does not support faulty_links injection."
+                )
+            if node_count < 2:
+                continue
+            try:
+                neighbor_map = _neighbors_for_topology(topo_type, node_count)
+            except ValueError as exc:
+                raise ValueError(
+                    f"Unable to determine neighbors for topology '{topo_type or 'unknown'}' "
+                    f"with node_count={node_count}: {exc}"
+                ) from exc
+            valid_sources = [idx for idx, neigh in enumerate(neighbor_map) if neigh]
+            if not valid_sources:
+                continue
+            src = random.choice(valid_sources)
+            dst_candidates = neighbor_map[src]
+            dst = random.choice(dst_candidates)
+            weight_clamped = float(max(0.0, min(1.0, weight)))
+            weight_clamped = _quantize_weight(weight_clamped)
+            dim.pop("faulty_links", None)
+            fault_entries.append(FlowSeq([int(src), int(dst), weight_clamped]))
+
+        if fault_entries:
+            network["faulty_links"] = FlowSeq(fault_entries)
+        else:
+            network.pop("faulty_links", None)
+
+    return _mutator
+
+
+# -----------------------------------------------------------------------------
+# Multiprocessing helpers
+# -----------------------------------------------------------------------------
+
+_WORKER_HW_DICT = None
+_WORKER_MODEL_CONFIG = None
+_WORKER_MODE = None
+
+
+def _parallelism_worker_init(base_hw_dict, model_config_path, mode):
+    global _WORKER_HW_DICT, _WORKER_MODEL_CONFIG, _WORKER_MODE
+    set_astrasim_cache_mode(ASTRA_CACHE_MODE)
+    _WORKER_HW_DICT = base_hw_dict
+    _WORKER_MODEL_CONFIG = config.parse_config(model_config_path, config_type=mode)
+    _WORKER_MODE = mode
+
+
+def _execute_parallelism_task(base_hw_dict, model_config_obj, mode, task: Dict[str, object]) -> Dict[str, object]:
+    kind = task.get("kind")
+    settings = dict(task.get("settings", {}))
+    key = task.get("key")
+    settings_key = tuple(task.get("settings_key", ()))
+    task_id = task.get("task_id")
+    try:
+        mutator = None
+        if kind == "fault":
+            raw_faults = task.get("faults", [])
+            faults = [tuple(spec) for spec in raw_faults]
+            mutator = make_fault_mutator(faults)
+        else:
+            faults = []
+        metrics = evaluate_parallelism(
+            base_hw_dict,
+            model_config_obj,
+            mode,
+            settings,
+            hw_mutator=mutator,
+        )
+        result: Dict[str, object] = {
+            "status": "ok",
+            "kind": kind,
+            "key": key,
+            "settings": settings,
+            "settings_key": settings_key,
+            "runtime": float(metrics["runtime"]),
+            "hw_yaml": metrics.get("hw_yaml"),
+            "task_id": task_id,
+        }
+        if kind == "fault":
+            result.update(
+                {
+                    "faults": faults,
+                    "num_faults": int(task.get("num_faults", len(faults))),
+                }
+            )
+        else:
+            result.update({"faults": [], "num_faults": 0})
+        return result
+    except Exception as exc:
+        return {
+            "status": "error",
+            "kind": kind,
+            "key": key,
+            "settings": settings,
+            "settings_key": settings_key,
+            "task_id": task_id,
+            "error": str(exc),
+        }
+
+
+def _parallelism_worker_task(task: Dict[str, object]) -> Dict[str, object]:
+    if _WORKER_HW_DICT is None or _WORKER_MODEL_CONFIG is None or _WORKER_MODE is None:
+        raise RuntimeError("Worker initialisation missing before task execution.")
+    return _execute_parallelism_task(_WORKER_HW_DICT, _WORKER_MODEL_CONFIG, _WORKER_MODE, task)
+
+
+def _parallelism_key(settings: Dict[str, object]) -> Tuple[Tuple[str, object], ...]:
+    return tuple(sorted(settings.items()))
+
+
+def _serialize_settings_key(key: Tuple[Tuple[str, object], ...]) -> List[List[object]]:
+    return [[k, v] for k, v in key]
+
+
+def _deserialize_settings_key(raw: Iterable[Iterable[object]]) -> Tuple[Tuple[str, object], ...]:
+    return tuple((str(k), v) for k, v in raw)
+
+
+def _normalise_faults_raw(faults: Iterable[Iterable[object]]) -> List[List[object]]:
+    normalised: List[List[object]] = []
+    for dim, weight, nodes in faults:
+        normalised.append([int(dim), _quantize_weight(float(weight)), int(nodes)])
+    return normalised
+
+
+def _task_identifier(
+    kind: str,
+    settings_key: Tuple[Tuple[str, object], ...],
+    *,
+    faults: Sequence[Tuple[int, float, int]] | None = None,
+    num_faults: int | None = None,
+) -> str:
+    payload: Dict[str, object] = {
+        "kind": kind,
+        "settings_key": _serialize_settings_key(settings_key),
+    }
+    if kind == "fault":
+        effective_faults = faults or ()
+        payload["num_faults"] = int(num_faults if num_faults is not None else len(effective_faults))
+        payload["faults"] = _normalise_faults_raw(effective_faults)
+    return json.dumps(payload, sort_keys=True)
+
+
+def _load_existing_results(path: Path) -> Dict[str, Dict[str, object]]:
+    if not path.exists():
+        return {}
+    try:
+        with path.open("r", encoding="utf-8") as handle:
+            data = json.load(handle)
+        if isinstance(data, dict):
+            return data
+        print(f"Warning: results file {path} did not contain a JSON object. Ignoring existing data.", file=sys.stderr)
+    except Exception as exc:
+        print(f"Warning: failed to load results file {path}: {exc}", file=sys.stderr)
+    return {}
+
+
+def _record_partial_result(results_store: Dict[str, Dict[str, object]], entry: Dict[str, object]) -> None:
+    task_id = entry.get("task_id")
+    if not task_id:
+        return
+    results_store[task_id] = {
+        "kind": entry.get("kind"),
+        "settings_key": _serialize_settings_key(tuple(entry.get("settings_key", ()))),
+        "runtime": entry.get("runtime"),
+    }
+    if entry.get("kind") == "fault":
+        results_store[task_id]["num_faults"] = int(entry.get("num_faults", 0))
+        results_store[task_id]["faults"] = _normalise_faults_raw(entry.get("faults", []))
+    tmp_path = RESULTS_JSON_PATH.with_suffix(".tmp")
+    tmp_path.parent.mkdir(parents=True, exist_ok=True)
+    with RESULTS_WRITE_LOCK:
+        with tmp_path.open("w", encoding="utf-8") as handle:
+            json.dump(results_store, handle, indent=2)
+        tmp_path.replace(RESULTS_JSON_PATH)
+
+
+def _dump_debug_hw_config(result: Dict[str, object]) -> Optional[Path]:
+    if not DEBUG_MODE:
+        return None
+    yaml_text = result.get("hw_yaml")
+    if not isinstance(yaml_text, str) or not yaml_text.strip():
+        return None
+    if DEBUG_RUN_DIR is None:
+        return None
+    global _DEBUG_FILE_INDEX
+    _DEBUG_FILE_INDEX += 1
+    settings = result.get("settings") or {}
+    kind = str(result.get("kind") or "task")
+    tp = settings.get("tp")
+    cp = settings.get("cp")
+    dp = settings.get("dp")
+    lp = settings.get("lp")
+    label = f"{kind}_{_DEBUG_FILE_INDEX:03d}"
+    if all(isinstance(val, (int, float)) for val in (tp, cp, dp, lp)):
+        label += f"_tp{int(tp)}_cp{int(cp)}_dp{int(dp)}_lp{int(lp)}"
+    num_faults = result.get("num_faults")
+    if kind == "fault" and isinstance(num_faults, int):
+        label += f"_{num_faults}faults"
+    file_path = DEBUG_RUN_DIR / f"{label}.yaml"
+    file_path.parent.mkdir(parents=True, exist_ok=True)
+    with file_path.open("w", encoding="utf-8") as handle:
+        handle.write(yaml_text)
+    DEBUG_SAVED_PATHS.append(file_path)
+    return file_path
+
+
+def _generate_hw_yaml_from_task(base_hw_dict: Dict[str, object], task: Dict[str, object]) -> Optional[str]:
+    if base_hw_dict is None:
+        return None
+    try:
+        hw_copy: Dict[str, object] = copy.deepcopy(base_hw_dict)
+    except Exception as exc:
+        print(f"Warning: failed to clone base hardware config: {exc}", file=sys.stderr)
+        return None
+
+    settings = dict(task.get("settings", {}) or {})
+    parallel_block = hw_copy.setdefault("parallelism", {})
+    if isinstance(parallel_block, dict):
+        for key, value in settings.items():
+            parallel_block[key] = value
+
+    mutator = None
+    faults: Sequence[Tuple[int, float, int]] = ()
+    if task.get("kind") == "fault":
+        raw_faults = task.get("faults") or []
+        try:
+            faults = [tuple(spec) for spec in raw_faults]  # type: ignore[assignment]
+        except Exception:
+            faults = ()
+        mutator = make_fault_mutator(faults)
+
+    try:
+        if mutator is not None:
+            mutator(hw_copy)
+    except Exception as exc:
+        print(f"Warning: failed to apply fault mutator for debug dump: {exc}", file=sys.stderr)
+
+    try:
+        return yaml.safe_dump(hw_copy, default_flow_style=False, sort_keys=False)
+    except Exception as exc:
+        print(f"Warning: failed to serialize debug hardware config: {exc}", file=sys.stderr)
+        return None
+
+
+def _ensure_failure_dump_dir() -> Path:
+    global FAILURE_DUMP_DIR
+    if FAILURE_DUMP_DIR is None:
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        FAILURE_DUMP_DIR = FAILURE_OUTPUT_BASE / timestamp
+        FAILURE_DUMP_DIR.mkdir(parents=True, exist_ok=True)
+        print(f"Saving failing hardware configs to {FAILURE_DUMP_DIR}")
+    return FAILURE_DUMP_DIR
+
+
+def _failure_label_from_settings(kind: str, settings: Dict[str, object], num_faults: Optional[int]) -> str:
+    label = kind or "task"
+    tp = settings.get("tp")
+    cp = settings.get("cp")
+    dp = settings.get("dp")
+    lp = settings.get("lp")
+    if all(isinstance(val, (int, float)) for val in (tp, cp, dp, lp)):
+        label += f"_tp{int(tp)}_cp{int(cp)}_dp{int(dp)}_lp{int(lp)}"
+    if kind == "fault" and isinstance(num_faults, int):
+        label += f"_{num_faults}faults"
+    return label
+
+
+def _dump_failure_hw_config(
+    base_hw_dict: Dict[str, object],
+    result: Dict[str, object],
+    error_message: Optional[str] = None,
+) -> Optional[Path]:
+    task_like = {
+        "kind": result.get("kind"),
+        "settings": result.get("settings"),
+        "faults": result.get("faults"),
+        "num_faults": result.get("num_faults"),
+    }
+    yaml_text = _generate_hw_yaml_from_task(base_hw_dict, task_like)
+    failure_dir = _ensure_failure_dump_dir()
+    global _FAILURE_FILE_INDEX
+    _FAILURE_FILE_INDEX += 1
+    settings = dict(result.get("settings", {}) or {})
+    kind = str(result.get("kind") or "task")
+    num_faults = result.get("num_faults")
+    label = _failure_label_from_settings(kind, settings, num_faults if isinstance(num_faults, int) else None)
+    base_filename = f"failure_{_FAILURE_FILE_INDEX:03d}_{label}"
+    yaml_path = failure_dir / f"{base_filename}.yaml"
+    command_line = None
+    if isinstance(yaml_text, str) and yaml_text.strip():
+        with yaml_path.open("w", encoding="utf-8") as handle:
+            handle.write(yaml_text)
+        FAILURE_SAVED_PATHS.append(yaml_path)
+        print(f"Stored failing hardware config at {yaml_path}")
+        cli_hw_path = os.path.relpath(yaml_path, os.getcwd())
+        command_line = (
+            f"uv run run_perf.py --hardware_config {shlex.quote(cli_hw_path)} "
+            f"--model_config {shlex.quote(MODEL_CONFIG_PATH)}"
+        )
+    else:
+        yaml_path = None
+    if error_message or command_line:
+        log_path = failure_dir / f"{base_filename}_error.txt"
+        with log_path.open("w", encoding="utf-8") as handle:
+            if error_message:
+                handle.write(str(error_message).strip() + "\n")
+            if command_line:
+                handle.write("\n# Reproduce failing configuration\n")
+                handle.write(command_line + "\n")
+        FAILURE_SAVED_PATHS.append(log_path)
+    return yaml_path
+
+
+def _terminate_active_executor() -> None:
+    global _ACTIVE_EXECUTOR
+    executor = _ACTIVE_EXECUTOR
+    if executor is None:
+        return
+    _ACTIVE_EXECUTOR = None
+    try:
+        executor.shutdown(wait=False, cancel_futures=True)
+    except Exception:
+        pass
+    processes = getattr(executor, "_processes", None)
+    if isinstance(processes, dict):
+        for proc in processes.values():
+            try:
+                if proc is not None and proc.is_alive():
+                    proc.kill()
+            except Exception:
+                continue
+
+
+def _signal_handler(signum, frame):  # type: ignore[override]
+    _terminate_active_executor()
+    if signum == signal.SIGINT:
+        raise KeyboardInterrupt()
+    sys.exit(128 + signum)
+
+
+def _install_signal_handlers() -> None:
+    global _OLD_SIGINT_HANDLER, _OLD_SIGTERM_HANDLER
+    _OLD_SIGINT_HANDLER = signal.getsignal(signal.SIGINT)
+    _OLD_SIGTERM_HANDLER = signal.getsignal(signal.SIGTERM)
+    signal.signal(signal.SIGINT, _signal_handler)
+    signal.signal(signal.SIGTERM, _signal_handler)
+    atexit.register(_terminate_active_executor)
+
+
+# -----------------------------------------------------------------------------
+# Reporting and plotting
+# -----------------------------------------------------------------------------
+
+def summarise_fault_runs(runtimes: Iterable[float]) -> Tuple[float, float, float]:
+    values = [float(v) for v in runtimes if math.isfinite(v)]
+    if not values:
+        return float("nan"), float("nan"), float("nan")
+    return min(values), max(values), float(sum(values) / len(values))
+
+
+def write_fault_report(records: List[Dict[str, object]], path: str) -> None:
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    header = [
+        "tp",
+        "cp",
+        "dp",
+        "lp",
+        "baseline_runtime",
+        "fault_min",
+        "fault_max",
+        "fault_mean",
+    ]
+    with open(path, "w") as handle:
+        handle.write("\t".join(header) + "\n")
+        for record in records:
+            row = [
+                str(record["parallelism"]["tp"]),
+                str(record["parallelism"]["cp"]),
+                str(record["parallelism"]["dp"]),
+                str(record["parallelism"]["lp"]),
+                f"{record['baseline_runtime']:.6f}",
+                f"{record['fault_min']:.6f}" if math.isfinite(record["fault_min"]) else "nan",
+                f"{record['fault_max']:.6f}" if math.isfinite(record["fault_max"]) else "nan",
+                f"{record['fault_mean']:.6f}" if math.isfinite(record["fault_mean"]) else "nan",
+            ]
+            handle.write("\t".join(row) + "\n")
+
+
+def plot_fault_sensitivity(records: List[Dict[str, object]], output_path: str, num_gpus: int) -> None:
+    if not records:
+        print("No records to plot.", file=sys.stderr)
+        return
+
+    if sns is not None:  # pragma: no branch
+        sns.set_theme(style="whitegrid")
+
+    xs = np.arange(len(records))
+    baseline = np.array([rec["baseline_runtime"] for rec in records], dtype=float)
+    fault_mean = np.array([rec["fault_mean"] for rec in records], dtype=float)
+    fault_min = np.array([rec["fault_min"] for rec in records], dtype=float)
+    fault_max = np.array([rec["fault_max"] for rec in records], dtype=float)
+
+    config_labels = [rec["label"] for rec in records]
+
+    plt.figure(figsize=(12, 6))
+    plt.scatter(xs, baseline, color="#1f78b4", marker="o", s=80, label="Baseline")
+
+    fault_lower = []
+    fault_upper = []
+    for mean_val, min_val, max_val in zip(fault_mean, fault_min, fault_max):
+        if math.isfinite(mean_val) and math.isfinite(min_val):
+            fault_lower.append(max(mean_val - min_val, 0.0))
+        else:
+            fault_lower.append(0.0)
+        if math.isfinite(mean_val) and math.isfinite(max_val):
+            fault_upper.append(max(max_val - mean_val, 0.0))
+        else:
+            fault_upper.append(0.0)
+
+    plt.errorbar(
+        xs,
+        fault_mean,
+        yerr=[fault_lower, fault_upper],
+        fmt="s",
+        color="#d7301f",
+        ecolor="#fb6a4a",
+        elinewidth=2,
+        capsize=6,
+        capthick=1.6,
+        markersize=7,
+        label="Fault mean ± range",
+    )
+
+    # Candlestick-style shading for min/max bounds
+    for x, fmin, fmax in zip(xs, fault_min, fault_max):
+        if math.isfinite(fmin) and math.isfinite(fmax):
+            plt.fill_between([x - 0.22, x + 0.22], [fmin, fmin], [fmax, fmax], color="#fdd0a2", alpha=0.4)
+
+    plt.xticks(xs, config_labels, rotation=45, ha="right")
+    plt.ylabel("Runtime (s)")
+    plt.title(f"Fault Sensitivity Across Parallelism Configurations (Num GPUs = {num_gpus})")
+    plt.grid(alpha=0.3, axis="y")
+    plt.legend()
+    plt.tight_layout()
+
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    plt.savefig(output_path, dpi=200)
+    plt.close()
+    print(f"Saved fault sensitivity plot to {output_path}")
+
+
+# -----------------------------------------------------------------------------
+# Main execution
+# -----------------------------------------------------------------------------
+
+def main(cli_args: Optional[Sequence[str]] = None) -> int:
+    global TARGET_NUM_GPUS, SAMPLE_COUNT, FAULT_ITER, FAULT_WORKERS, FAULT_MAG
+    global NUM_FAULTS, MIN_ALLOWED_TP, MAX_ALLOWED_TP, MIN_ALLOWED_CP, MAX_ALLOWED_CP
+    global MIN_ALLOWED_DP, MAX_ALLOWED_DP, MIN_ALLOWED_LP, MAX_ALLOWED_LP
+    global ALLOWED_FAULT_DIMS, PLOT_OUTPUT_PATH, REPORT_OUTPUT_PATH, RESULTS_JSON_PATH
+    global RANDOM_SEED, ENFORCE_SQUARE_TP_CP, DEBUG_MODE, DEBUG_RUN_DIR, DEBUG_SAVED_PATHS
+    global _DEBUG_FILE_INDEX, FAILURE_DUMP_DIR, FAILURE_SAVED_PATHS, _FAILURE_FILE_INDEX
+
+    args = parse_args(cli_args)
+
+    TARGET_NUM_GPUS = args.target_num_gpus
+    SAMPLE_COUNT = args.sample_count
+    FAULT_ITER = args.fault_iter
+    FAULT_WORKERS = args.fault_workers
+
+    mag_mean, mag_std = _parse_float_tuple(args.fault_mag)
+    FAULT_MAG = (mag_mean, mag_std)
+
+    parsed_num_faults = _parse_int_list(args.num_faults)
+    NUM_FAULTS = parsed_num_faults or [1]
+
+    MIN_ALLOWED_TP = max(1, args.min_tp)
+    MAX_ALLOWED_TP = max(MIN_ALLOWED_TP, args.max_tp)
+    MIN_ALLOWED_CP = max(1, args.min_cp)
+    MAX_ALLOWED_CP = max(MIN_ALLOWED_CP, args.max_cp)
+    MIN_ALLOWED_DP = max(1, args.min_dp)
+    MAX_ALLOWED_DP = max(MIN_ALLOWED_DP, args.max_dp)
+    MIN_ALLOWED_LP = max(1, args.min_lp)
+    MAX_ALLOWED_LP = max(MIN_ALLOWED_LP, args.max_lp)
+
+    allowed_dims_list = _parse_int_list(args.allowed_fault_dims) if args.allowed_fault_dims else []
+    ALLOWED_FAULT_DIMS = allowed_dims_list if allowed_dims_list else None
+
+    PLOT_OUTPUT_PATH = args.plot_output
+    REPORT_OUTPUT_PATH = args.report_output
+    RESULTS_JSON_PATH = Path(args.results_json)
+
+    RANDOM_SEED = args.seed
+    random.seed(RANDOM_SEED)
+    ENFORCE_SQUARE_TP_CP = ENFORCE_SQUARE_TP_CP or bool(args.enforce_square_tp_cp)
+    DEBUG_MODE = bool(args.debug_mode)
+    _DEBUG_FILE_INDEX = 0
+    _FAILURE_FILE_INDEX = 0
+    FAILURE_SAVED_PATHS.clear()
+    FAILURE_DUMP_DIR = None
+    if DEBUG_MODE:
+        FAULT_WORKERS = 1
+        DEBUG_SAVED_PATHS.clear()
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        DEBUG_RUN_DIR = DEBUG_OUTPUT_BASE / timestamp
+        DEBUG_RUN_DIR.mkdir(parents=True, exist_ok=True)
+        print(f"Debug mode enabled. Hardware configs will be written to {DEBUG_RUN_DIR}")
+    else:
+        DEBUG_RUN_DIR = None
+        DEBUG_SAVED_PATHS.clear()
+
+    set_astrasim_cache_mode(ASTRA_CACHE_MODE)
+    _install_signal_handlers()
+
+    base_hw_dict = read_yaml(HARDWARE_CONFIG_PATH)
+    mode = determine_model_mode(MODEL_CONFIG_PATH)
+    model_config_obj = config.parse_config(MODEL_CONFIG_PATH, config_type=mode)
+
+    existing_results = _load_existing_results(RESULTS_JSON_PATH)
+    if existing_results:
+        print(f"Loaded {len(existing_results)} cached result entries from {RESULTS_JSON_PATH}")
+    results_store: Dict[str, Dict[str, object]] = dict(existing_results)
+
+    parallelism_samples = generate_parallelism_samples(TARGET_NUM_GPUS, SAMPLE_COUNT)
+    if not parallelism_samples:
+        print("Unable to generate parallelism samples.", file=sys.stderr)
+        return
+
+    records: List[Dict[str, object]] = []
+    records_by_key: Dict[Tuple[Tuple[str, object], ...], Dict[str, object]] = {}
+    task_queue: List[Dict[str, object]] = []
+    total_fault_tasks = 0
+
+    for settings in parallelism_samples:
+        settings_copy = dict(settings)
+        key = _parallelism_key(settings_copy)
+        settings_key = key
+
+        candidates_for_config = _candidate_fault_dimensions(settings_copy)
+        if not candidates_for_config:
+            continue
+
+        records_by_key[key] = {
+            "parallelism": settings_copy,
+            "baseline_runtime": None,
+            "fault_runtimes": [],
+            "fault_details": [],
+        }
+
+        baseline_task_id = _task_identifier("baseline", settings_key)
+        baseline_cached = existing_results.get(baseline_task_id)
+        if baseline_cached is not None:
+            runtime = baseline_cached.get("runtime")
+            if runtime is not None:
+                records_by_key[key]["baseline_runtime"] = float(runtime)
+        else:
+            task_queue.append(
+                {
+                    "kind": "baseline",
+                    "settings": settings_copy,
+                    "key": key,
+                    "settings_key": settings_key,
+                    "task_id": baseline_task_id,
+                }
+            )
+
+        for num_faults in NUM_FAULTS:
+            repeats = max(1, int(num_faults))
+            for _ in range(FAULT_ITER):
+                fault_specs: List[Tuple[int, float, int]] = []
+                for _ in range(repeats):
+                    fault_value = sample_fault_magnitude()
+                    dim_choice, node_count = random.choice(candidates_for_config)
+                    fault_specs.append((dim_choice, fault_value, node_count))
+
+                fault_task_id = _task_identifier(
+                    "fault",
+                    settings_key,
+                    faults=fault_specs,
+                    num_faults=repeats,
+                )
+                cached_fault = existing_results.get(fault_task_id)
+                if cached_fault is not None:
+                    runtime = cached_fault.get("runtime")
+                    if runtime is not None:
+                        records_by_key[key]["fault_runtimes"].append(float(runtime))
+                        cached_faults = _normalise_faults_raw(cached_fault.get("faults", []))
+                        records_by_key[key]["fault_details"].append(
+                            {
+                                "num_faults": cached_fault.get("num_faults"),
+                                "faults": cached_faults,
+                            }
+                        )
+                    continue
+                total_fault_tasks += 1
+                task_queue.append(
+                    {
+                        "kind": "fault",
+                        "settings": settings_copy,
+                        "key": key,
+                        "settings_key": settings_key,
+                        "task_id": fault_task_id,
+                        "faults": fault_specs,
+                        "num_faults": repeats,
+                    }
+                )
+
+    if not task_queue:
+        print("No new evaluation tasks; reusing cached results.")
+    else:
+        print(
+            f"Prepared {len(task_queue)} evaluation task(s) across "
+            f"{len(parallelism_samples)} configuration(s) "
+            f"(fault tasks: {total_fault_tasks})."
+        )
+
+    if DEBUG_MODE:
+        if not task_queue:
+            print("Debug mode enabled but no evaluation tasks were generated.")
+        else:
+            for task in task_queue:
+                yaml_text = _generate_hw_yaml_from_task(base_hw_dict, task)
+                debug_entry = {
+                    "kind": task.get("kind"),
+                    "settings": task.get("settings"),
+                    "num_faults": task.get("num_faults"),
+                    "faults": task.get("faults"),
+                    "hw_yaml": yaml_text,
+                }
+                _dump_debug_hw_config(debug_entry)
+        if DEBUG_RUN_DIR:
+            if DEBUG_SAVED_PATHS:
+                print(f"Debug mode: saved {len(DEBUG_SAVED_PATHS)} hardware config(s) to {DEBUG_RUN_DIR}")
+            else:
+                print(f"Debug mode enabled but no hardware configs were generated. Directory: {DEBUG_RUN_DIR}")
+        else:
+            print("Debug mode enabled but no output directory was initialised.")
+        return 0
+
+    def handle_task_result(result: Dict[str, object]) -> None:
+        if not isinstance(result, dict):
+            return
+        key = result.get("key")
+        rec = records_by_key.get(key)
+        if rec is None:
+            return
+        status = result.get("status")
+        if status != "ok":
+            error_msg = result.get("error")
+            print(
+                f"Warning: task failed for settings {result.get('settings')}: {error_msg}",
+                file=sys.stderr,
+            )
+            _dump_failure_hw_config(base_hw_dict, result, error_msg)
+            return
+        kind = result.get("kind")
+        runtime = float(result.get("runtime", float("nan")))
+        if kind == "baseline":
+            rec["baseline_runtime"] = runtime
+        else:
+            rec["fault_runtimes"].append(runtime)
+            rec["fault_details"].append(
+                {
+                    "num_faults": result.get("num_faults"),
+                    "faults": _normalise_faults_raw(result.get("faults", [])),
+                }
+            )
+        _record_partial_result(results_store, result)
+        _dump_debug_hw_config(result)
+
+    if task_queue:
+        if FAULT_WORKERS and FAULT_WORKERS > 1:
+            available_cpus = max(1, os.cpu_count() or 1)
+            worker_count = min(max(1, FAULT_WORKERS), len(task_queue), available_cpus)
+            print(f"Launching ProcessPoolExecutor with {worker_count} worker(s).")
+            executor = ProcessPoolExecutor(
+                max_workers=worker_count,
+                initializer=_parallelism_worker_init,
+                initargs=(base_hw_dict, MODEL_CONFIG_PATH, mode),
+            )
+            try:
+                global _ACTIVE_EXECUTOR
+                _ACTIVE_EXECUTOR = executor
+                futures = [executor.submit(_parallelism_worker_task, task) for task in task_queue]
+                for fut in tqdm(as_completed(futures), total=len(futures), desc="Evaluations", unit="task"):
+                    try:
+                        result = fut.result()
+                    except Exception as exc:  # pragma: no cover - defensive
+                        print(f"Warning: parallel task raised an exception: {exc}", file=sys.stderr)
+                        continue
+                    handle_task_result(result)
+            finally:
+                _terminate_active_executor()
+        else:
+            print("Running tasks sequentially (FAULT_WORKERS <= 1).")
+            for task in tqdm(task_queue, desc="Evaluations", unit="task"):
+                try:
+                    result = _execute_parallelism_task(base_hw_dict, model_config_obj, mode, task)
+                except Exception as exc:  # pragma: no cover - defensive
+                    print(f"Warning: sequential task raised an exception: {exc}", file=sys.stderr)
+                    continue
+                handle_task_result(result)
+
+    for settings in parallelism_samples:
+        key = _parallelism_key(settings)
+        state = records_by_key.get(key)
+        if not state:
+            continue
+        baseline_runtime = state.get("baseline_runtime")
+        if baseline_runtime is None:
+            print(f"Warning: missing baseline runtime for settings {settings}", file=sys.stderr)
+            continue
+        fault_min, fault_max, fault_mean = summarise_fault_runs(state.get("fault_runtimes", []))
+        record = {
+            "parallelism": state["parallelism"],
+            "baseline_runtime": float(baseline_runtime),
+            "fault_min": fault_min,
+            "fault_max": fault_max,
+            "fault_mean": fault_mean,
+            "label": f"tp{settings['tp']}-cp{settings['cp']}-dp{settings['dp']}-lp{settings['lp']}",
+        }
+        records.append(record)
+
+    if not records:
+        print("No successful sweep results collected.", file=sys.stderr)
+        return 0
+
+    write_fault_report(records, REPORT_OUTPUT_PATH)
+    print(f"Wrote fault sweep report to {REPORT_OUTPUT_PATH}")
+    plot_fault_sensitivity(records, PLOT_OUTPUT_PATH, TARGET_NUM_GPUS)
+    if FAILURE_SAVED_PATHS and FAILURE_DUMP_DIR:
+        print(
+            f"Captured {len(FAILURE_SAVED_PATHS)} failing hardware config artefact(s) in {FAILURE_DUMP_DIR}"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/tools/gmap_sweep.py b/tools/gmap_sweep.py
new file mode 100644
index 00000000..e0a506ab
--- /dev/null
+++ b/tools/gmap_sweep.py
@@ -0,0 +1,582 @@
+#!/usr/bin/env python3
+"""
+GMap sweep utility.
+
+This tool explores tensor/context/pipeline parallelism combinations that live in
+the first network dimension (TP/CP/LP) for a square, power-of-two GPU count.
+Each configuration is evaluated by running DeepFlow with the
+DEEPFLOW_GMAP_ONLY=1 environment variable so execution stops immediately after
+SCOTCH emits the mapping artefacts.  The script parses the resulting
+`first_dim_comm.metrics` file (and the `[GMapDebug]` printouts) to recover the
+before/after CommExpan values, computes the percentage delta, and produces both
+ a tab-separated report and a scatter plot coloured by % change.
+
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import io
+import math
+import os
+import re
+import shutil
+import sys
+import tempfile
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from contextlib import redirect_stdout, redirect_stderr
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt  # noqa: E402
+import numpy as np  # noqa: E402
+import yaml  # noqa: E402
+from tqdm import tqdm  # noqa: E402
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import config  # noqa: E402
+from time_calculation_LLM import TimeCalculationLLM  # noqa: E402
+from parallelism_sweep import set_astrasim_cache_mode  # type: ignore
+
+# -----------------------------------------------------------------------------#
+# Global configuration                                                         #
+# -----------------------------------------------------------------------------#
+
+# Total GPU count in the first dimension (must be a square, power-of-two).
+NUM_GPUS = 256
+NUM_WORKERS = 95
+ASTRA_CACHE_MODE = "NO_CACHE"
+TOPOLOGY_VARIANTS = ("Mesh2D", "Torus2D")
+# TOPOLOGY_VARIANTS = ("Mesh2D",)
+# Hard bounds (inclusive) on each parallelism axis.
+TP_BOUNDS = (1, 128)
+CP_BOUNDS = (1, 128)
+LP_BOUNDS = (1, 128)
+
+# -----------------------------------------------------------------------------#
+# Helpers lifted / adapted from existing sweep tools                           #
+# -----------------------------------------------------------------------------#
+
+
+def read_yaml(path: str) -> Dict[str, object]:
+    with open(path, "r") as handle:
+        return yaml.safe_load(handle)
+
+
+def determine_model_mode(model_config_path: str) -> str:
+    model_dict = read_yaml(model_config_path)
+    model_param = model_dict.get("model_param") or {}
+    mode = model_param.get("mode")
+    if not mode:
+        raise ValueError(f"model_param.mode must be defined in {model_config_path}")
+    return str(mode)
+
+
+def make_temp_hw_config(
+    base_hw_dict: Dict[str, object],
+    parallel_settings: Dict[str, int],
+    hw_mutator=None,
+):
+    updated = copy.deepcopy(base_hw_dict)
+    parallel_block = updated.setdefault("parallelism", {})
+    for key, value in parallel_settings.items():
+        parallel_block[key] = int(value)
+
+    if hw_mutator:
+        hw_mutator(updated)
+
+    tmp_file = tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False)
+    try:
+        yaml.safe_dump(updated, tmp_file, default_flow_style=False, sort_keys=False)
+        tmp_file.flush()
+        tmp_file.close()
+        hw_config = config.parse_config(tmp_file.name, config_type="hardware")
+        return hw_config
+    finally:
+        try:
+            tmp_file.close()
+        except Exception:
+            pass
+        try:
+            os.unlink(tmp_file.name)
+        except Exception:
+            pass
+
+
+# -----------------------------------------------------------------------------#
+# Core sweep logic                                                             #
+# -----------------------------------------------------------------------------#
+
+
+def is_power_of_two(value: int) -> bool:
+    return value > 0 and (value & (value - 1) == 0)
+
+
+def is_perfect_square(value: int) -> bool:
+    root = math.isqrt(value)
+    return root * root == value
+
+
+def enumerate_parallelisms(num_gpus: int) -> List[Dict[str, int]]:
+    limit = max(1, num_gpus // 2)
+    def _powers_of_two_between(lo: int, hi: int) -> List[int]:
+        lo = max(1, lo)
+        if hi < lo:
+            return []
+        vals: List[int] = []
+        val = 1
+        while val < lo:
+            val <<= 1
+        while val <= hi:
+            vals.append(val)
+            val <<= 1
+        return vals
+
+    tp_min, tp_max = TP_BOUNDS
+    cp_min, cp_max = CP_BOUNDS
+    lp_min, lp_max = LP_BOUNDS
+    tp_values = _powers_of_two_between(tp_min, min(tp_max, limit))
+    cp_values = _powers_of_two_between(cp_min, min(cp_max, limit))
+    lp_values = _powers_of_two_between(lp_min, min(lp_max, limit))
+    combos: List[Dict[str, int]] = []
+    for tp in tp_values:
+        for cp in cp_values:
+            for lp in lp_values:
+                if tp * cp * lp == num_gpus:
+                    combos.append({"tp": tp, "cp": cp, "lp": lp})
+    combos.sort(key=lambda item: (item["tp"], item["cp"], item["lp"]))
+    return combos
+
+
+def update_first_dimension_size(hw_dict: Dict[str, object], num_gpus: int, topology: str) -> None:
+    network = hw_dict.get("network")
+    if not isinstance(network, dict):
+        return
+    dimensions = network.get("dimensions")
+    if isinstance(dimensions, list) and dimensions:
+        first = dimensions[0]
+        if isinstance(first, dict):
+            first["size"] = int(num_gpus)
+            topology_block = first.setdefault("topology", {})
+            topology_block["type"] = topology
+
+
+def parse_metrics_file(path: str) -> Tuple[Optional[float], Optional[float]]:
+    if not os.path.exists(path):
+        return None, None
+    section: Optional[str] = None
+    before = after = None
+    with open(path, "r") as handle:
+        for raw in handle:
+            line = raw.strip()
+            if not line:
+                continue
+            if line.startswith("Before"):
+                section = "before"
+                continue
+            if line.startswith("After"):
+                section = "after"
+                continue
+            if "CommExpan=" in line:
+                try:
+                    value = float(line.split("CommExpan=")[1].split("\t", 1)[0])
+                except (ValueError, IndexError):
+                    value = None
+                if value is None:
+                    continue
+                if section == "before" and before is None:
+                    before = value
+                elif section == "after" and after is None:
+                    after = value
+    return before, after
+
+
+_COMMEXPAN_REGEX = re.compile(
+    r"CommExpan before\s+([0-9eE+\-\.]+)\s+after\s+([0-9eE+\-\.]+)"
+)
+
+
+def parse_stdout_for_commexpan(stdout_text: str) -> Tuple[Optional[float], Optional[float]]:
+    before = after = None
+    for line in stdout_text.splitlines():
+        if "CommExpan" not in line:
+            continue
+        match = _COMMEXPAN_REGEX.search(line)
+        if not match:
+            continue
+        try:
+            before = float(match.group(1))
+            after = float(match.group(2))
+        except ValueError:
+            continue
+    return before, after
+
+
+def compute_delta(before: Optional[float], after: Optional[float]) -> Optional[float]:
+    if before is None or after is None or before == 0:
+        return None
+    return (after - before) / before * 100.0
+
+
+def _format_delta(delta: Optional[float]) -> str:
+    if delta is None or delta != delta:
+        return "nan"
+    return f"{delta:+.2f}"
+
+
+def _format_float(value: Optional[float]) -> str:
+    if value is None or value != value:
+        return "nan"
+    return f"{value:.6f}"
+
+
+def run_single_configuration(
+    base_hw_dict: Dict[str, object],
+    model_config,
+    mode: str,
+    parallel_settings: Dict[str, int],
+    num_gpus: int,
+    keep_artifacts: bool,
+    topology: str,
+    capture_output: bool = True,
+) -> Dict[str, object]:
+    hw_config = make_temp_hw_config(
+        base_hw_dict,
+        parallel_settings,
+        hw_mutator=lambda cfg: update_first_dimension_size(cfg, num_gpus, topology),
+    )
+    if keep_artifacts:
+        base_dir = os.path.join("tools", "gmap_artifacts")
+        os.makedirs(base_dir, exist_ok=True)
+        dir_name = "gmap_{topology}_tp{tp}_cp{cp}_lp{lp}".format(
+            topology=topology.lower(),
+            tp=parallel_settings["tp"],
+            cp=parallel_settings["cp"],
+            lp=parallel_settings["lp"],
+        )
+        temp_dir = os.path.join(base_dir, dir_name)
+        shutil.rmtree(temp_dir, ignore_errors=True)
+        os.makedirs(temp_dir, exist_ok=True)
+    else:
+        temp_dir = tempfile.mkdtemp(prefix="gmap_sweep_")
+    prev_flag = os.environ.get("DEEPFLOW_GMAP_ONLY")
+    prev_persist = os.environ.get("DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS")    
+    os.environ["DEEPFLOW_GMAP_ONLY"] = "1"
+    if keep_artifacts:
+        os.environ["DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS"] = "1"
+    else:
+        os.environ.pop("DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS", None)
+    stdout_buffer = io.StringIO()
+    try:
+        calculator = TimeCalculationLLM(hw_config, model_config, mode, output_dir=temp_dir)
+        if capture_output:
+            with redirect_stdout(stdout_buffer), redirect_stderr(stdout_buffer):
+                try:
+                    calculator.calc_time_llm()
+                except SystemExit as exc:
+                    if exc.code not in (0, None):
+                        raise
+        else:
+            try:
+                calculator.calc_time_llm()
+            except SystemExit as exc:
+                if exc.code not in (0, None):
+                    raise
+    finally:
+        if prev_flag is None:
+            os.environ.pop("DEEPFLOW_GMAP_ONLY", None)
+        else:
+            os.environ["DEEPFLOW_GMAP_ONLY"] = prev_flag
+        if prev_persist is None:
+            os.environ.pop("DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS", None)
+        else:
+            os.environ["DEEPFLOW_PERSIST_ASTRASIM_ARTIFACTS"] = prev_persist
+    stdout_text = stdout_buffer.getvalue() if capture_output else ""
+    metrics_path = os.path.join(temp_dir, "first_dim_comm.metrics")
+    before, after = parse_metrics_file(metrics_path)
+    if before is None or after is None:
+        alt_before, alt_after = parse_stdout_for_commexpan(stdout_text)
+        if before is None:
+            before = alt_before
+        if after is None:
+            after = alt_after
+    delta_pct = compute_delta(before, after)
+    applied = bool(
+        before is not None and after is not None and before > 0 and after < before
+    )
+    artifact_dir = temp_dir if keep_artifacts else None
+    if not keep_artifacts:
+        shutil.rmtree(temp_dir, ignore_errors=True)
+    else:
+        print(f"Saved artifacts to {temp_dir}")
+    return {
+        "parallelism": dict(parallel_settings),
+        "before": before,
+        "after": after,
+        "delta_pct": delta_pct,
+        "applied": applied,
+        "stdout": stdout_text,
+        "artifact_dir": artifact_dir,
+        "topology": topology,
+    }
+
+
+# -----------------------------------------------------------------------------#
+# Multiprocessing glue                                                         #
+# -----------------------------------------------------------------------------#
+
+_GLOBAL_HW_DICT = None
+_GLOBAL_MODEL_CONFIG = None
+_GLOBAL_MODE = None
+_GLOBAL_NUM_GPUS = None
+_GLOBAL_KEEP_ARTIFACTS = False
+
+
+def _worker_init(hw_dict, model_config_path, mode, num_gpus, keep_artifacts):
+    global _GLOBAL_HW_DICT, _GLOBAL_MODEL_CONFIG, _GLOBAL_MODE, _GLOBAL_NUM_GPUS, _GLOBAL_KEEP_ARTIFACTS
+    _GLOBAL_HW_DICT = hw_dict
+    _GLOBAL_MODE = mode
+    _GLOBAL_NUM_GPUS = num_gpus
+    _GLOBAL_KEEP_ARTIFACTS = bool(keep_artifacts)
+    _GLOBAL_MODEL_CONFIG = config.parse_config(model_config_path, config_type=mode)
+
+
+def _worker_task(parallel_items: Tuple[Tuple[str, int], ...]) -> Dict[str, object]:
+    parallel_settings: Dict[str, int] = {}
+    topology: Optional[str] = None
+    for key, value in parallel_items:
+        if key == "__topology":
+            topology = str(value)
+        else:
+            parallel_settings[key] = int(value)
+    if topology is None:
+        topology = TOPOLOGY_VARIANTS[0]
+    try:
+        result = run_single_configuration(
+            _GLOBAL_HW_DICT,
+            _GLOBAL_MODEL_CONFIG,
+            _GLOBAL_MODE,
+            parallel_settings,
+            _GLOBAL_NUM_GPUS,
+            _GLOBAL_KEEP_ARTIFACTS,
+            topology,
+        )
+        return {"status": "ok", "result": result}
+    except Exception as exc:
+        return {
+            "status": "error",
+            "error": str(exc),
+            "parallelism": parallel_settings,
+            "topology": topology,
+        }
+
+
+# -----------------------------------------------------------------------------#
+# Plotting / reporting                                                         #
+# -----------------------------------------------------------------------------#
+
+
+def plot_results(results: Sequence[Dict[str, object]], num_gpus: int, topology: str, path: str) -> None:
+    if not results:
+        return
+    xs = np.arange(len(results))
+    ys = np.array(
+        [
+            entry["delta_pct"] if entry["delta_pct"] == entry["delta_pct"] else float("nan")
+            for entry in results
+        ],
+        dtype=float,
+    )
+    labels = [
+        "tp={tp}, cp={cp}, lp={lp}".format(
+            tp=entry["parallelism"]["tp"],
+            cp=entry["parallelism"]["cp"],
+            lp=entry["parallelism"]["lp"],
+        )
+        for entry in results
+    ]
+    applied_flags = [bool(entry.get("applied")) for entry in results]
+    applied_x = [x for x, flag in zip(xs, applied_flags) if flag]
+    applied_y = [y for y, flag in zip(ys, applied_flags) if flag]
+    skipped_x = [x for x, flag in zip(xs, applied_flags) if not flag]
+    skipped_y = [y for y, flag in zip(ys, applied_flags) if not flag]
+
+    plt.figure(figsize=(12, 6))
+    if applied_x:
+        plt.scatter(applied_x, applied_y, c="tab:green", s=80, label="Applied", edgecolors="black", linewidths=0.5)
+    if skipped_x:
+        plt.scatter(skipped_x, skipped_y, c="tab:red", s=80, label="Skipped", edgecolors="black", linewidths=0.5)
+    plt.axhline(0.0, color="gray", linestyle="--", linewidth=1)
+    plt.ylabel("Δ CommExpan (%) (negative is better)")
+    plt.xlabel("Parallelism configuration")
+    plt.title(f"GMap sweep ({topology}, num_gpus={num_gpus})")
+    plt.grid(alpha=0.3, axis="y")
+    plt.xticks(xs, labels, rotation=45, ha="right")
+    plt.legend()
+    plt.tight_layout()
+    plt.savefig(path, dpi=200)
+    plt.close()
+    print(f"Saved scatter plot to {path}")
+
+
+def write_report(results: Sequence[Dict[str, object]], path: str, topology: str) -> None:
+    header = [
+        "tp",
+        "cp",
+        "lp",
+        "before_commexpan",
+        "after_commexpan",
+        "delta_percent",
+        "applied",
+        "artifact_dir",
+        "topology",
+    ]
+    with open(path, "w") as handle:
+        handle.write("\t".join(header) + "\n")
+        for entry in results:
+            parallelism = entry["parallelism"]
+            row = [
+                str(parallelism["tp"]),
+                str(parallelism["cp"]),
+                str(parallelism["lp"]),
+                _format_float(entry.get("before")),
+                _format_float(entry.get("after")),
+                _format_delta(entry.get("delta_pct")),
+                "yes" if entry.get("applied") else "no",
+                entry.get("artifact_dir") or "",
+                topology,
+            ]
+            handle.write("\t".join(row) + "\n")
+    print(f"Wrote report to {path}")
+
+
+# -----------------------------------------------------------------------------#
+# CLI                                                                          #
+# -----------------------------------------------------------------------------#
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Sweep TP/CP/LP combinations for GMap cost deltas.")
+    parser.add_argument("--hardware-config", default="configs/hardware-config/gmap_mesh2d_demo.yaml", help="Base hardware config path.")
+    parser.add_argument("--model-config", default="configs/model-config/Llama3.1-405B.yaml", help="Model config path.")
+    parser.add_argument("--max-workers", type=int, default=None, help="Maximum concurrent worker processes.")
+    parser.add_argument("--plot-output", default="tools/gmap_sweep.png", help="Path for the scatter plot.")
+    parser.add_argument("--report-output", default="tools/gmap_sweep.tsv", help="Path for the TSV report.")
+    parser.add_argument("--keep-artifacts", action="store_true", help="Preserve per-run artefacts (first_dim_comm.*) on disk.")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    set_astrasim_cache_mode(ASTRA_CACHE_MODE)
+    if args.max_workers is None or args.max_workers <= 0:
+        args.max_workers = max(1, (os.cpu_count() or 2) - 1)
+    num_gpus = int(NUM_GPUS)
+    if not is_power_of_two(num_gpus):
+        raise ValueError("num_gpus must be a power of two.")
+    if not is_perfect_square(num_gpus):
+        raise ValueError("num_gpus must be a perfect square to satisfy the Mesh2D constraint.")
+
+    base_hw_dict = read_yaml(args.hardware_config)
+    mode = determine_model_mode(args.model_config)
+    cases = enumerate_parallelisms(num_gpus)
+    if not cases:
+        print("No TP/CP/LP combinations satisfy the constraints.")
+        return
+    print(f"Evaluating {len(cases)} configuration(s) for num_gpus={num_gpus}")
+
+    def _suffix_path(base: str, suffix: str) -> str:
+        root, ext = os.path.splitext(base)
+        return f"{root}_{suffix.lower()}{ext or ''}"
+
+    task_specs: List[Dict[str, object]] = []
+    for topology in TOPOLOGY_VARIANTS:
+        for combo in cases:
+            entry = dict(combo)
+            entry["__topology"] = topology
+            task_specs.append(entry)
+
+    results_by_topology: Dict[str, List[Dict[str, object]]] = {topo: [] for topo in TOPOLOGY_VARIANTS}
+    errors_by_topology: Dict[str, List[str]] = {topo: [] for topo in TOPOLOGY_VARIANTS}
+
+    if len(task_specs) == 1:
+        spec = task_specs[0]
+        topology = spec["__topology"]
+        combo = {k: v for k, v in spec.items() if k != "__topology"}
+        model_config_obj = config.parse_config(args.model_config, config_type=mode)
+        try:
+            result = run_single_configuration(
+                base_hw_dict,
+                model_config_obj,
+                mode,
+                combo,
+                num_gpus,
+                args.keep_artifacts,
+                topology,
+                capture_output=False,
+            )
+            results_by_topology[topology].append(result)
+        except Exception as exc:
+            errors_by_topology[topology].append(f"{combo}: {exc}")
+    else:
+        args.max_workers = NUM_WORKERS
+        max_workers = args.max_workers if args.max_workers and args.max_workers > 0 else 1
+        worker_count = min(max_workers, max(1, (os.cpu_count() or 1) - 1), len(task_specs))
+        task_items = [tuple(sorted(spec.items())) for spec in task_specs]
+        with ProcessPoolExecutor(
+            max_workers=worker_count,
+            initializer=_worker_init,
+            initargs=(base_hw_dict, args.model_config, mode, num_gpus, args.keep_artifacts),
+        ) as executor:
+            futures = {executor.submit(_worker_task, items): items for items in task_items}
+            with tqdm(total=len(task_items), desc="Evaluating", unit="config") as progress:
+                for future in as_completed(futures):
+                    progress.update(1)
+                    outcome = future.result()
+                    topo = outcome.get("result", {}).get("topology") if outcome.get("status") == "ok" else outcome.get("topology")
+                    if topo is None:
+                        topo = TOPOLOGY_VARIANTS[0]
+                    if outcome["status"] == "ok":
+                        results_by_topology[topo].append(outcome["result"])
+                    else:
+                        errors_by_topology[topo].append(f"{outcome['parallelism']}: {outcome.get('error')}")
+
+    for topology in TOPOLOGY_VARIANTS:
+        print(f"\n=== Topology: {topology} ===")
+        results = results_by_topology[topology]
+        errors = errors_by_topology[topology]
+        if errors:
+            print("Encountered errors:")
+            for msg in errors:
+                print("  -", msg)
+        if not results:
+            print("No successful evaluations for this topology; skipping.")
+            continue
+        results.sort(key=lambda entry: (entry["parallelism"]["tp"], entry["parallelism"]["cp"], entry["parallelism"]["lp"]))
+        best = min(
+            (entry for entry in results if entry.get("delta_pct") == entry.get("delta_pct")),
+            key=lambda e: e.get("delta_pct"),
+            default=None,
+        )
+        if best:
+            best_delta = best.get("delta_pct")
+            print(
+                "Best improvement: tp={tp}, cp={cp}, lp={lp}, Δ={delta}".format(
+                    tp=best["parallelism"]["tp"],
+                    cp=best["parallelism"]["cp"],
+                    lp=best["parallelism"]["lp"],
+                    delta=_format_delta(best_delta),
+                )
+            )
+        report_path = _suffix_path(args.report_output, topology)
+        plot_path = _suffix_path(args.plot_output, topology)
+        write_report(results, report_path, topology)
+        plot_results(results, num_gpus, topology, plot_path)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/mod_field.py b/tools/mod_field.py
new file mode 100644
index 00000000..6f8105d9
--- /dev/null
+++ b/tools/mod_field.py
@@ -0,0 +1,236 @@
+#!/usr/bin/env python3
+"""
+Modify or insert a scalar field in YAML configs (PyYAML only, minimal text edit).
+
+- If the dotted path exists, overwrite its scalar value (preserving inline comment).
+- If it does not exist, create it (with optional comment) under existing ancestor mappings.
+- Skips files where an ancestor is not a mapping.
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Any, Dict, List, Mapping, Optional, Tuple
+
+import yaml
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Modify/insert a YAML field.")
+    parser.add_argument("--scope", choices=["hardware", "software"], required=True, help="which config family to update")
+    parser.add_argument("--path", required=True, help="dotted path to set (e.g., network.overlap.tp_overlap)")
+    parser.add_argument("--value", required=True, help='YAML literal for the value (e.g., 0.0, true, "text")')
+    parser.add_argument("--comment", default="", help="comment to place above the inserted field (add-only)")
+    parser.add_argument("--dry-run", action="store_true", help="print planned changes without writing files")
+    return parser.parse_args()
+
+
+def parse_value(raw: str) -> Any:
+    try:
+        return yaml.safe_load(raw)
+    except Exception:
+        return raw
+
+
+def find_roots(scope: str) -> List[Path]:
+    repo_root = Path(__file__).resolve().parents[1]
+    if scope == "hardware":
+        roots = [
+            repo_root / "configs" / "hardware-config",
+            repo_root / "validation_scripts" / "validation_configs" / "hardware-config",
+        ]
+    else:
+        roots = [
+            repo_root / "configs" / "model-config",
+            repo_root / "validation_scripts" / "validation_configs" / "model-config",
+        ]
+    return [p for p in roots if p.exists()]
+
+
+def key_from_line(line: str) -> Optional[Tuple[int, str]]:
+    if ":" not in line:
+        return None
+    stripped = line.split("#", 1)[0].rstrip()
+    if ":" not in stripped:
+        return None
+    prefix, _ = stripped.split(":", 1)
+    if not prefix.strip():
+        return None
+    indent = len(prefix) - len(prefix.lstrip(" "))
+    key = prefix.strip()
+    return indent, key
+
+
+def locate_block(lines: List[str], path_parts: List[str]) -> Optional[Tuple[int, int]]:
+    idx = 0
+    indent = 0
+    last_idx = None
+    last_indent = None
+    for part in path_parts:
+        found = False
+        for i in range(idx, len(lines)):
+            res = key_from_line(lines[i])
+            if res is None:
+                continue
+            ind, key = res
+            if ind < indent:
+                return None
+            if ind == indent and key == part:
+                last_idx, last_indent = i, ind
+                idx = i + 1
+                indent = ind + 2
+                found = True
+                break
+        if not found:
+            return None
+    if last_idx is None or last_indent is None:
+        return None
+    return last_idx, last_indent
+
+
+def find_insertion_index(lines: List[str], parent_line_idx: int, parent_indent: int) -> int:
+    for i in range(parent_line_idx + 1, len(lines)):
+        res = key_from_line(lines[i])
+        if res is None:
+            continue
+        ind, _ = res
+        if ind <= parent_indent:
+            return i
+    return len(lines)
+
+
+def render_scalar(value: Any) -> str:
+    dumped = yaml.safe_dump(value, default_flow_style=True)
+    # Take only the first line to avoid document end marker ('...')
+    first_line = dumped.split('\n', 1)[0]
+    return first_line
+
+
+def ancestors_are_mappings(data: Dict[str, Any], path_parts: List[str]) -> bool:
+    cursor = data
+    for part in path_parts[:-1]:
+        if not isinstance(cursor, dict):
+            return False
+        if part not in cursor:
+            return True
+        cursor = cursor[part]
+    return isinstance(cursor, dict)
+
+
+def insert_block(lines: List[str], path_parts: List[str], value: Any, comment: str) -> List[str]:
+    # Find deepest existing ancestor
+    deepest_idx = -1
+    deepest_indent = -2
+    missing_at = 0
+    for i in range(len(path_parts)):
+        loc = locate_block(lines, path_parts[: i + 1])
+        if loc is None:
+            missing_at = i
+            break
+        deepest_idx, deepest_indent = loc
+        missing_at = i + 1
+
+    # All parts exist
+    if missing_at >= len(path_parts):
+        return lines
+
+    missing_parts = path_parts[missing_at:]
+    insertion_idx = find_insertion_index(lines, deepest_idx, deepest_indent) if deepest_idx >= 0 else len(lines)
+    indent = deepest_indent + 2
+
+    snippet_lines: List[str] = []
+    for depth, key in enumerate(missing_parts):
+        pad = " " * (indent + depth * 2)
+        is_leaf = depth == len(missing_parts) - 1
+        if is_leaf:
+            if comment:
+                snippet_lines.append(f"{pad}# {comment}")
+            snippet_lines.append(f"{pad}{key}: {render_scalar(value)}")
+        else:
+            snippet_lines.append(f"{pad}{key}:")
+
+    return lines[:insertion_idx] + snippet_lines + lines[insertion_idx:]
+
+
+def update_leaf_value(lines: List[str], path_parts: List[str], value: Any) -> List[str]:
+    loc = locate_block(lines, path_parts)
+    if loc is None:
+        return lines
+    leaf_idx, leaf_indent = loc
+    key = path_parts[-1]
+    line = lines[leaf_idx]
+    comment = ""
+    if "#" in line:
+        _, comment = line.split("#", 1)
+    new_line = f"{' ' * leaf_indent}{key}: {render_scalar(value)}"
+    if comment:
+        new_line += f"  #{comment.lstrip()}"
+    lines[leaf_idx] = new_line
+    return lines
+
+
+def set_field(path: Path, dotted_path: str, value: Any, comment: str = "", dry_run: bool = False) -> str:
+    path_parts = dotted_path.split(".")
+    try:
+        text = path.read_text()
+    except Exception as exc:
+        return f"[SKIP] {path}: read error: {exc}"
+
+    lines = [ln for ln in text.splitlines() if ln.strip() != "..."]
+
+    try:
+        data = yaml.safe_load("\n".join(lines)) or {}
+    except Exception as exc:
+        return f"[SKIP] {path}: yaml parse error: {exc}"
+
+    if not isinstance(data, dict):
+        return f"[SKIP] {path}: top-level is not a mapping"
+
+    if not ancestors_are_mappings(data, path_parts):
+        return f"[SKIP] {path}: ancestor path is missing or not a mapping; not inserting"
+
+    exists = True
+    cursor: Any = data
+    for part in path_parts:
+        if isinstance(cursor, dict) and part in cursor:
+            cursor = cursor[part]
+        else:
+            exists = False
+            break
+
+    if exists:
+        new_lines = update_leaf_value(lines, path_parts, value)
+    else:
+        new_lines = insert_block(lines, path_parts, value, comment)
+
+    new_text = "\n".join(new_lines) + ("\n" if text.endswith("\n") else "")
+    if dry_run:
+        return f"[DRY-RUN] {path}"
+    try:
+        path.write_text(new_text)
+    except Exception as exc:
+        return f"[SKIP] {path}: write error: {exc}"
+    return f"[UPDATED] {path}"
+
+
+def main() -> None:
+    args = parse_args()
+    value = parse_value(args.value)
+    roots = find_roots(args.scope)
+    path_parts = args.path
+    modified = 0
+    skipped = 0
+    for root in roots:
+        for yaml_path in sorted(root.rglob("*.yaml")):
+            result = set_field(yaml_path, path_parts, value, args.comment, args.dry_run)
+            if result.startswith("[UPDATED]") or result.startswith("[DRY-RUN]"):
+                modified += 1
+            else:
+                skipped += 1
+            print(result)
+    print(f"Done. Modified: {modified}, skipped: {skipped}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/overlap_finder.py b/tools/overlap_finder.py
new file mode 100644
index 00000000..19d0659b
--- /dev/null
+++ b/tools/overlap_finder.py
@@ -0,0 +1,243 @@
+#!/usr/bin/env python3
+"""
+Lightweight overlap tuner using golden-section search (no SciPy).
+
+Global knobs (edit below):
+  - RUN_INFERENCE / RUN_TRAINING: enable which suites to optimize
+  - INF_DEVICES / TRAIN_DEVICES: which devices to run from validation CSVs
+  - GOLD_TOL / GOLD_MAX_ITER: search tolerance and iteration cap
+
+Current behavior:
+  - Inference: optimize network.overlap.tp_overlap using validation_scripts/nvidia_inf.py
+  - Training:  optimize network.overlap.tp_sp_overlap using validation_scripts/nvidia_train.py
+
+Extensible: add more evaluators that return an average pct error for a target overlap key.
+"""
+
+from __future__ import annotations
+
+import io
+import math
+import sys
+from contextlib import redirect_stdout
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple
+
+# Ensure project root on path
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from validation_scripts import nvidia_inf, nvidia_train  # type: ignore
+from validation_scripts.validation_helpers import ValidationSpec, run_validation_suite  # type: ignore
+from tools import mod_field
+
+# ======= Global config =======
+RUN_INFERENCE = True
+RUN_TRAINING = True
+# Set to True to write best overlaps back to hardware config files (rounded to 2 decimals).
+APPLY_BEST = True
+
+INF_DEVICES: Sequence[str] = ("A100",)
+TRAIN_DEVICES: Sequence[str] = ("A100_korthi",)
+
+GOLD_TOL = 0.01  # stop when interval width < GOLD_TOL
+GOLD_MAX_ITER = 20
+
+# =============================
+
+
+def _merge_dict(target: Dict[str, Any], updates: Mapping[str, Any]) -> Dict[str, Any]:
+    for key, val in updates.items():
+        if isinstance(val, Mapping):
+            current = target.get(key, {})
+            if not isinstance(current, dict):
+                current = {}
+            target[key] = _merge_dict(dict(current), val)
+        else:
+            target[key] = val
+    return target
+
+
+def _with_overlap(spec: ValidationSpec, overlap_key: str, overlap_value: float) -> ValidationSpec:
+    base_overrides: Dict[str, Any] = {}
+    if spec.hardware_overrides and isinstance(spec.hardware_overrides, Mapping):
+        base_overrides = _merge_dict({}, spec.hardware_overrides)
+    base_overrides = _merge_dict(base_overrides, {"network": {"overlap": {overlap_key: float(overlap_value)}}})
+    return ValidationSpec(
+        label=spec.label,
+        model_overrides=spec.model_overrides,
+        hardware_overrides=base_overrides,
+        metadata=spec.metadata,
+        model_config_path=spec.model_config_path,
+        hardware_config_path=spec.hardware_config_path,
+        order=spec.order,
+    )
+
+
+def _avg_pct_error(errors: Iterable[float]) -> float:
+    vals = [e for e in errors if not math.isnan(e)]
+    if not vals:
+        return float("inf")
+    return sum(vals) / len(vals)
+
+
+def _all_hardware_configs() -> List[Path]:
+    roots = [
+        PROJECT_ROOT / "configs" / "hardware-config",
+        PROJECT_ROOT / "validation_scripts" / "validation_configs" / "hardware-config",
+    ]
+    paths: List[Path] = []
+    for root in roots:
+        if root.exists():
+            paths.extend(sorted(root.glob("*.yaml")))
+    return paths
+
+
+def eval_inference(tp_overlap: float) -> float:
+    """Average pct error across INF_DEVICES for given tp_overlap."""
+    all_errors: List[float] = []
+    for device in INF_DEVICES:
+        specs, actual_lookup, base_model_path, base_hw_path = nvidia_inf.build_specs_for_device(
+            device, network_ignored=True
+        )
+        specs = [_with_overlap(spec, "tp_overlap", tp_overlap) for spec in specs]
+        with io.StringIO() as buf, redirect_stdout(buf):
+            results = run_validation_suite(
+                specs,
+                base_model_config_path=base_model_path,
+                base_hardware_config_path=base_hw_path,
+                result_parser=nvidia_inf.parse_inference_time,
+                run_perf_path=nvidia_inf.RUN_PERF,
+            )
+        rows = nvidia_inf.compute_pct_errors(results, actual_lookup)
+        all_errors.extend(float(row["pct_error"]) for row in rows)
+    return _avg_pct_error(all_errors)
+
+
+def eval_training(tp_sp_overlap: float) -> float:
+    """Average pct error across TRAIN_DEVICES for given tp_sp_overlap."""
+    all_errors: List[float] = []
+    for device in TRAIN_DEVICES:
+        specs, actual_lookup, base_model_path, base_hw_path = nvidia_train.build_specs_for_device(device)
+        specs = [_with_overlap(spec, "tp_sp_overlap", tp_sp_overlap) for spec in specs]
+        with io.StringIO() as buf, redirect_stdout(buf):
+            results = run_validation_suite(
+                specs,
+                base_model_config_path=base_model_path,
+                base_hardware_config_path=base_hw_path,
+                result_parser=nvidia_train.parse_training_time,
+                run_perf_path=nvidia_train.RUN_PERF,
+            )
+        rows = nvidia_train.compute_pct_errors(results, actual_lookup)
+        all_errors.extend(float(row["pct_error"]) for row in rows)
+    return _avg_pct_error(all_errors)
+
+
+@dataclass
+class SearchResult:
+    best_x: float
+    best_y: float
+    evaluations: List[Tuple[float, float]]
+
+
+def golden_section_search(
+    fn: Callable[[float], float],
+    lo: float = 0.0,
+    hi: float = 1.0,
+    tol: float = GOLD_TOL,
+    max_iter: int = GOLD_MAX_ITER,
+    log_fn: Optional[Callable[[int, int, float, float], None]] = None,
+) -> SearchResult:
+    phi = (1 + 5 ** 0.5) / 2
+    inv_phi = 1 / phi
+    inv_phi_sq = inv_phi ** 2
+
+    a, b = lo, hi
+    h = b - a
+    if h <= tol:
+        mid = (a + b) / 2
+        val = fn(mid)
+        return SearchResult(best_x=mid, best_y=val, evaluations=[(mid, val)])
+
+    n = int(math.ceil(math.log(tol / h) / math.log(inv_phi)))
+    c = a + inv_phi_sq * h
+    d = a + inv_phi * h
+    yc = fn(c)
+    yd = fn(d)
+    if log_fn:
+        log_fn(1, max_iter, c, yc)
+        log_fn(2, max_iter, d, yd)
+    evals: List[Tuple[float, float]] = [(c, yc), (d, yd)]
+
+    iter_idx = 2
+    for _ in range(min(n, max_iter)):
+        if yc < yd:
+            b = d
+            d, yd = c, yc
+            h = inv_phi * h
+            c = a + inv_phi_sq * h
+            yc = fn(c)
+            evals.append((c, yc))
+        else:
+            a = c
+            c, yc = d, yd
+            h = inv_phi * h
+            d = a + inv_phi * h
+            yd = fn(d)
+            evals.append((d, yd))
+        iter_idx += 1
+        if log_fn:
+            log_fn(iter_idx, max_iter, evals[-1][0], evals[-1][1])
+        if h < tol:
+            break
+
+    best_x, best_y = min(evals, key=lambda t: t[1])
+    return SearchResult(best_x=best_x, best_y=best_y, evaluations=evals)
+
+
+def main() -> None:
+    if RUN_INFERENCE:
+        print("=== Optimizing tp_overlap for inference ===")
+        res = golden_section_search(
+            eval_inference,
+            log_fn=lambda i, m, x, y: print(f"[{i}/{m}] Tried {x:.5f} --> {y:.2f}"),
+        )
+        print(f"Best tp_overlap={res.best_x:.4f}, avg pct error={res.best_y:.4f}")
+        for x, y in res.evaluations:
+            print(f"  tried {x:.4f} -> {y:.4f}")
+        err0 = eval_inference(0.0)
+        err1 = eval_inference(1.0)
+        print(f"  error at 0.0: {err0:.4f}")
+        print(f"  error at 1.0: {err1:.4f}")
+        if APPLY_BEST:
+            best_val = round(res.best_x, 2)
+            print(f"Applying tp_overlap={best_val:.2f} to all hardware configs")
+            for cfg_path in _all_hardware_configs():
+                msg = mod_field.set_field(cfg_path, "network.overlap.tp_overlap", best_val, dry_run=False)
+                print(f"  {msg}")
+
+    if RUN_TRAINING:
+        print("\n=== Optimizing tp_sp_overlap for training ===")
+        res = golden_section_search(
+            eval_training,
+            log_fn=lambda i, m, x, y: print(f"[{i}/{m}] Tried {x:.5f} --> {y:.2f}"),
+        )
+        print(f"Best tp_sp_overlap={res.best_x:.4f}, avg pct error={res.best_y:.4f}")
+        for x, y in res.evaluations:
+            print(f"  tried {x:.4f} -> {y:.4f}")
+        err0 = eval_training(0.0)
+        err1 = eval_training(1.0)
+        print(f"  error at 0.0: {err0:.4f}")
+        print(f"  error at 1.0: {err1:.4f}")
+        if APPLY_BEST:
+            best_val = round(res.best_x, 2)
+            print(f"Applying tp_sp_overlap={best_val:.2f} to all hardware configs")
+            for cfg_path in _all_hardware_configs():
+                msg = mod_field.set_field(cfg_path, "network.overlap.tp_sp_overlap", best_val, dry_run=False)
+                print(f"  {msg}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/parallelism_sweep.py b/tools/parallelism_sweep.py
new file mode 100644
index 00000000..0637f6a1
--- /dev/null
+++ b/tools/parallelism_sweep.py
@@ -0,0 +1,1018 @@
+#!/usr/bin/env python3
+"""
+Parallelism sweep utility for DeepFlow LLM configurations.
+
+Update the global configuration section below to point at the desired hardware
+and model config files and to tailor the parallelism search space. The tool
+enumerates every combination, filters those whose total GPU count falls inside
+the configured bounds, evaluates runtime with DeepFlow, and plots a scatter
+chart of accuracy (by default, 1 / runtime) versus GPU count with horizontal
+jitter to avoid overlap.
+"""
+
+from __future__ import print_function
+
+import argparse
+import ast
+import copy
+import itertools
+import math
+import os
+import random
+import shutil
+import sys
+import tempfile
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from contextlib import redirect_stdout, redirect_stderr
+from typing import Dict, Iterable, List, Tuple
+
+import matplotlib.pyplot as plt
+import yaml
+from tqdm import tqdm
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import config
+from time_calculation_LLM import TimeCalculationLLM
+from LLM_util import process_gemm_shapes
+
+import seaborn as sns
+import numpy as np
+import pandas as pd
+import math
+import numpy as np
+from matplotlib.patches import Polygon
+
+# -----------------------------------------------------------------------------
+# Global configuration knobs (no CLI)
+# -----------------------------------------------------------------------------
+
+# Paths to the baseline configuration files
+HARDWARE_CONFIG_PATH = "configs/hardware-config/a100_80GB_faraz.yaml"
+MODEL_CONFIG_PATH = "configs/model-config/Llama3.1-405B.yaml"
+
+# Parallelism values to sweep (dense grid). Edit to suit your search space.
+# Keys should match the entries under the YAML "parallelism" section.
+PARALLELISM_SWEEP = {
+    "tp": [2**i for i in range(0, 9)],
+    "cp": [2**i for i in range(0, 9)],
+    "dp": [2**i for i in range(0, 9)],
+    "lp": [2**i for i in range(0, 9)],
+}
+
+# Optional knobs that still live inside the parallelism section but do not
+# affect GPU counts. Leave empty if you do not want to explore them.
+OTHER_PARALLELISM_OPTIONS = {
+    "tp_sp": [True],
+}
+
+# GPU count filter: only evaluate combinations whose TP*CP*DP*LP fall inside
+# this inclusive range.
+GPU_COUNT_MIN = 64
+GPU_COUNT_MAX = 1024
+
+# When True, discard configurations whose tp*cp product is not a square power of two.
+ENFORCE_SQUARE_TP_CP = False
+
+# Select the y-axis metric for the scatter plot and report. Accepted values:
+#   "runtime"     -> plot raw runtime in seconds (lower is better)
+#   "performance" -> plot 1 / runtime (higher is better)
+PLOT_METRIC = "runtime"
+
+# Random seed for reproducible jitter in the scatter plot.
+PLOT_JITTER_SEED = 1234
+# Maximum absolute horizontal jitter (in GPU units).
+PLOT_JITTER_WIDTH = 0.175
+
+# Default output artefacts
+PLOT_OUTPUT_PATH = "tools/parallelism_sweep.png"
+PLOT_MFU_OUTPUT_PATH = "tools/parallelism_sweep_mfu.png"
+REPORT_OUTPUT_PATH = "tools/parallelism_sweep.tsv"
+
+# AstraSim cache handling within DeepFlow (mirrors run_perf default options).
+ASTRA_CACHE_MODE = "NO_CACHE"  # Options: NO_CACHE, CACHE_READONLY, CACHE_READWRITE
+
+# Plotting behaviour toggles
+MEM_AWARE_FILTER = False  # When True, skip memory-violating configurations in plots.
+
+# Maximum number of parallel worker processes (set <= available CPUs - 1). Set to 1 to disable multiprocessing.
+MAX_WORKERS = 100
+
+
+# -----------------------------------------------------------------------------
+# Helper utilities
+# -----------------------------------------------------------------------------
+
+def add_rgb_ternary_legend(ax,
+                           corner_labels=("R = log2(tp+cp)",
+                                          "G = log2(lp)",
+                                          "B = log2(dp)"),
+                           gamma=0.85,
+                           inset_xywh=(0.72, 0.58, 0.24, 0.24),  # Axes fraction: x, y, w, h
+                           n=120):
+    """
+    Draw a tiny triangular legend that shows how the RGB mix works.
+    - Uses barycentric blending inside an equilateral triangle.
+    - 'gamma' should match your color gamma used in the scatter.
+    - inset_xywh is (x0,y0,w,h) in axes-relative coords.
+    """
+    # Inset axis
+    iax = ax.inset_axes(inset_xywh, transform=ax.transAxes)
+    iax.set_aspect("equal")
+    iax.set_axis_off()
+
+    # Equilateral triangle vertices (R, G, B corners)
+    A = np.array([0.0, 0.0])                        # R corner
+    B = np.array([1.0, 0.0])                        # G corner
+    C = np.array([0.5, math.sqrt(3)/2.0])           # B corner
+
+    # Sample interior with barycentric coords r,g,b (r+g+b=1)
+    xs, ys, cols = [], [], []
+    for i in range(n + 1):
+        for j in range(n + 1 - i):
+            r = i / n
+            g = j / n
+            b = 1.0 - r - g
+            p = r * A + g * B + b * C
+            # apply same gamma tweak used in your plot colors
+            rr = r ** gamma
+            gg = g ** gamma
+            bb = b ** gamma
+            xs.append(p[0]); ys.append(p[1]); cols.append((rr, gg, bb))
+
+    # iax.scatter(xs, ys, s=2, c=cols, edgecolors="none")
+    iax.add_patch(Polygon([A, B, C], facecolor="none", edgecolor="black", linewidth=1))
+
+    # Corner labels
+    iax.text(A[0]-0.06, A[1]-0.04, corner_labels[0], ha="right", va="top", fontsize=8)
+    iax.text(B[0]+0.06, B[1]-0.04, corner_labels[1], ha="left",  va="top", fontsize=8)
+    iax.text(C[0],      C[1]+0.06, corner_labels[2], ha="center", va="bottom", fontsize=8)
+
+    # Optional tick marks (25/50/75%) along edges (comment out if you want it cleaner)
+    for t in (0.25, 0.5, 0.75):
+        # Edge AB (vary r vs g, b=0)
+        p = (1-t) * A + t * B
+        iax.plot([p[0]], [p[1]], marker="|", color="black", ms=6)
+        # Edge BC (vary g vs b, r=0)
+        p = (1-t) * B + t * C
+        iax.plot([p[0]], [p[1]], marker="_", color="black", ms=6)
+        # Edge CA (vary b vs r, g=0)
+        p = (1-t) * C + t * A
+        iax.plot([p[0]], [p[1]], marker="|", color="black", ms=6)
+
+    # Tiny caption (how channels were normalized)
+    iax.text(0.5, -0.18, "Each channel min–max normalized\nthen gamma-adjusted",
+             ha="center", va="top", fontsize=7, transform=iax.transAxes)
+
+
+def read_yaml(path):
+    with open(path, "r") as handle:
+        return yaml.safe_load(handle)
+
+
+def determine_model_mode(model_config_path):
+    model_dict = read_yaml(model_config_path)
+    model_param = model_dict.get("model_param") or {}
+    mode = model_param.get("mode")
+    if not mode:
+        raise ValueError("model_param.mode must be defined in {}".format(model_config_path))
+    return mode
+
+
+def cartesian_product(option_map):
+    """Yield dictionaries for every combination inside option_map."""
+    if not option_map:
+        yield {}
+        return
+    keys = sorted(option_map.keys())
+    value_lists = [option_map[key] for key in keys]
+    for values in itertools.product(*value_lists):
+        yield dict(zip(keys, values))
+
+
+def total_gpu_count(parallel_cfg):
+    total = 1
+    for axis in ("tp", "cp", "dp", "lp"):
+        value = int(parallel_cfg.get(axis, 1) or 1)
+        total *= max(1, value)
+    return total
+
+
+def tp_cp_product_is_power_of_two_square(tp_value, cp_value):
+    """Return True when tp*cp is a square whose root is also a power of two."""
+    if tp_value == 1 and cp_value == 1:
+        return False # special case
+    try:
+        tp_int = int(tp_value)
+        cp_int = int(cp_value)
+    except (TypeError, ValueError):
+        return False
+    if tp_int <= 0 or cp_int <= 0:
+        return False
+    product = tp_int * cp_int
+    root = math.isqrt(product)
+    if root * root != product:
+        return False
+    return (root & (root - 1)) == 0
+
+
+def make_temp_hw_config(base_hw_dict, parallel_settings, hw_mutator=None):
+    """Return (parsed HW config, YAML string) for the given override."""
+    updated = copy.deepcopy(base_hw_dict)
+    parallel_block = updated.setdefault("parallelism", {})
+    for key, value in parallel_settings.items():
+        parallel_block[key] = value
+
+    if hw_mutator:
+        hw_mutator(updated)
+    try:
+        debug_yaml = yaml.safe_dump(updated, default_flow_style=False)
+    except Exception:
+        debug_yaml = None
+
+    tmp_file = tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False)
+    try:
+        yaml.safe_dump(updated, tmp_file, default_flow_style=False, sort_keys=False)
+        tmp_file.flush()
+        tmp_file.close()
+        hw_config = config.parse_config(tmp_file.name, config_type="hardware")
+        return hw_config, debug_yaml
+    finally:
+        try:
+            tmp_file.close()
+        except Exception:
+            pass
+        try:
+            os.unlink(tmp_file.name)
+        except Exception:
+            pass
+
+def _rgb_from_parallelism(entries, gamma: float = 0.85):
+    """
+    Build per-point RGBA colors using:
+      R = log2(tp+cp), G = log2(lp), B = log2(dp)
+    Each channel is min-max normalized over the dataset, then gamma-adjusted.
+    """
+    tps = np.array([max(1, int(e["parallelism"].get("tp", 1))) for e in entries], dtype=float)
+    cps = np.array([max(1, int(e["parallelism"].get("cp", 1))) for e in entries], dtype=float)
+    dps = np.array([max(1, int(e["parallelism"].get("dp", 1))) for e in entries], dtype=float)
+    lps = np.array([max(1, int(e["parallelism"].get("lp", 1))) for e in entries], dtype=float)
+
+    r_raw = np.log2(tps + cps)
+    g_raw = np.log2(lps)
+    b_raw = np.log2(dps)
+
+    def _norm(x):
+        xmin, xmax = float(np.min(x)), float(np.max(x))
+        if not np.isfinite(xmin) or not np.isfinite(xmax) or xmax <= xmin:
+            return np.zeros_like(x)
+        y = (x - xmin) / (xmax - xmin)
+        if gamma and gamma != 1.0:
+            y = np.power(y, gamma)
+        return y
+
+    r = _norm(r_raw)
+    g = _norm(g_raw)
+    b = _norm(b_raw)
+
+    # RGBA with slight transparency to reduce overdraw
+    return np.stack([r, g, b, np.full_like(r, 0.9)], axis=1)
+
+
+
+def gpu_peak_flops(hw_config) -> float:
+    """Return the theoretical peak FLOPs/s for a single GPU."""
+    core = hw_config.tech_config.core
+    bundles = core.num_bundles or 1
+    per_cycle_flops = float(core.nominal_flop_rate_per_mcu)
+    mcu_per_bundle = float(core.num_mcu_per_bundle)
+    frequency = float(core.operating_frequency or core.nominal_frequency or 0.0)
+    util = float(getattr(core, "util", 1.0) or 1.0)
+    peak = per_cycle_flops * mcu_per_bundle * float(bundles)
+    if frequency > 0:
+        peak *= frequency
+    return peak * util if peak > 0 else 0.0
+
+
+def compute_total_flops(calculator: TimeCalculationLLM) -> float:
+    """Estimate total FLOPs executed for one global batch (forward + backward)."""
+    global_batch = getattr(calculator, "batch_size", None)
+    if global_batch is None or global_batch <= 0:
+        global_batch = calculator._effective_transformer_batch()
+    if not global_batch or global_batch <= 0:
+        return float("nan")
+
+    vocab_size = calculator.vocab_size
+    hidden_dim = calculator.hidden_dim
+    seq_len = calculator.seq_len
+    num_heads = calculator.num_heads
+    kv_heads = calculator.kv_heads
+    intermediate_size = calculator.intermediate_size
+
+    gemm_shapes = process_gemm_shapes(
+        calculator,
+        global_batch,
+        seq_len,
+        hidden_dim,
+        num_heads,
+        kv_heads,
+        intermediate_size,
+        vocab_size,
+    )
+
+    def _gemm_flops(shape) -> float:
+        if shape is None:
+            return 0.0
+        try:
+            dims = list(shape)
+        except TypeError:
+            return 0.0
+        if len(dims) == 4:
+            b, m, k, n = dims
+            return 2.0 * float(b) * float(m) * float(k) * float(n)
+        if len(dims) == 3:
+            m, k, n = dims
+            return 2.0 * float(m) * float(k) * float(n)
+        return 0.0
+
+    def _forward_backward(shape) -> float:
+        forward = _gemm_flops(shape)
+        if forward <= 0.0:
+            return 0.0
+        backward = 2.0 * forward
+        return forward + backward
+
+    per_layer_flops = 0.0
+    per_layer_flops += _forward_backward(gemm_shapes.get("qkv_proj"))
+    per_layer_flops += _forward_backward(gemm_shapes.get("attention_score"))
+    per_layer_flops += _forward_backward(gemm_shapes.get("attention_output"))
+    per_layer_flops += _forward_backward(gemm_shapes.get("output_proj"))
+    per_layer_flops += _forward_backward(gemm_shapes.get("ffn1"))
+    per_layer_flops += _forward_backward(gemm_shapes.get("ffn2"))
+
+    total_flops = per_layer_flops * float(calculator.num_layers)
+
+    total_flops += _forward_backward(gemm_shapes.get("linear"))
+
+    # Embedding forward/backward is comparatively small; keep placeholder for future refinement.
+
+    return float(total_flops)
+
+
+def evaluate_parallelism(hw_dict, model_config_obj, mode, parallel_settings, hw_mutator=None):
+    hw_config, debug_yaml = make_temp_hw_config(hw_dict, parallel_settings, hw_mutator=hw_mutator)
+    # if debug_yaml:
+    #     print("=== DEBUG HW CONFIG ===")
+    #     print(debug_yaml)
+    #     try:
+    #         with open("debug.yaml", "w") as debug_handle:
+    #             debug_handle.write(debug_yaml)
+    #     except Exception as exc:
+            # print(f"Warning: failed to write debug.yaml: {exc}", file=sys.stderr)
+    temp_dir = tempfile.mkdtemp(prefix="parallelism_sweep_")
+    try:
+        calculator = TimeCalculationLLM(hw_config, model_config_obj, mode, output_dir=temp_dir)
+        with open(os.devnull, "w") as devnull, redirect_stdout(devnull), redirect_stderr(devnull):
+            runtime = calculator.calc_time_llm()
+            total_flops = compute_total_flops(calculator)
+        performance = (1.0 / runtime) if runtime and runtime > 0.0 else float("nan")
+        peak_flops = gpu_peak_flops(hw_config)
+        num_gpus = total_gpu_count(parallel_settings)
+        achieved_flops = (total_flops / runtime) if runtime > 0 else float("nan")
+        denom = peak_flops * num_gpus if peak_flops and num_gpus else float("nan")
+        mfu = (achieved_flops / denom) if denom and denom > 0 else float("nan")
+        return {
+            "runtime": runtime,
+            "performance": performance,
+            "total_flops": total_flops,
+            "peak_flops": peak_flops,
+            "mfu": mfu,
+            "achieved_flops": achieved_flops,
+            "memory_exceeded": getattr(calculator, "memory_capacity_exceeded", False),
+            "memory_violation_gb": getattr(calculator, "memory_capacity_violation_gb", 0.0),
+            "hw_yaml": debug_yaml,
+        }
+    finally:
+        shutil.rmtree(temp_dir, ignore_errors=True)
+
+
+def set_astrasim_cache_mode(mode_str):
+    mapping = {
+        "NO_CACHE": "NO_CACHE",
+        "CACHE_READONLY": "CACHE_READONLY",
+        "CACHE_READWRITE": "CACHE_READWRITE",
+    }
+    env_value = mapping.get(str(mode_str).strip().upper(), "NO_CACHE")
+    os.environ["DEEPFLOW_ASTRA_CACHE_MODE"] = env_value
+
+
+def write_report(results, path):
+    header = [
+        "num_gpus",
+        "runtime_s",
+        "performance_1_over_s",
+        "total_flops",
+        "achieved_flops_per_s",
+        "peak_flops_per_gpu",
+        "mfu",
+        "memory_exceeded",
+        "memory_violation_gb",
+        "parallelism",
+    ]
+    with open(path, "w") as handle:
+        handle.write("\t".join(header) + "\n")
+        for entry in results:
+            row = [
+                str(entry["num_gpus"]),
+                "{:.6f}".format(entry["runtime"]),
+                (
+                    "{:.6f}".format(entry["performance"])
+                    if entry["performance"] == entry["performance"]
+                    else "nan"
+                ),
+                "{:.6e}".format(entry["total_flops"]),
+                "{:.6e}".format(entry["achieved_flops"]) if entry["achieved_flops"] == entry["achieved_flops"] else "nan",
+                "{:.6e}".format(entry["peak_flops"]),
+                "{:.6f}".format(entry["mfu"]) if entry["mfu"] == entry["mfu"] else "nan",
+                str(entry["memory_exceeded"]),
+                "{:.6f}".format(entry["memory_violation_gb"]),
+                repr(entry["parallelism"]),
+            ]
+            handle.write("\t".join(row) + "\n")
+
+
+def _parse_float(value: str) -> float:
+    try:
+        return float(value)
+    except (TypeError, ValueError):
+        return float("nan")
+
+
+def load_results_from_report(path: str) -> List[Dict[str, object]]:
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"No report found at {path}")
+
+    results: List[Dict[str, object]] = []
+    with open(path, "r") as handle:
+        header = handle.readline()
+        if not header:
+            return results
+        for line in handle:
+            line = line.strip()
+            if not line:
+                continue
+            parts = line.split("\t")
+            if len(parts) < 10:
+                print(f"Warning: skipping malformed row in report: {line}", file=sys.stderr)
+                continue
+            try:
+                parallelism = ast.literal_eval(parts[9])
+            except (SyntaxError, ValueError) as exc:
+                print(f"Warning: failed to parse parallelism '{parts[9]}': {exc}", file=sys.stderr)
+                continue
+            entry = {
+                "num_gpus": int(parts[0]),
+                "runtime": _parse_float(parts[1]),
+                "performance": _parse_float(parts[2]),
+                "total_flops": _parse_float(parts[3]),
+                "achieved_flops": _parse_float(parts[4]),
+                "peak_flops": _parse_float(parts[5]),
+                "mfu": _parse_float(parts[6]),
+                "memory_exceeded": str(parts[7]).strip().lower() == "true",
+                "memory_violation_gb": _parse_float(parts[8]),
+                "parallelism": parallelism,
+            }
+            results.append(entry)
+    return results
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="DeepFlow parallelism sweep utility.")
+    parser.add_argument(
+        "--plot_only",
+        action="store_true",
+        help="Skip evaluation and only regenerate plots from the existing TSV report.",
+    )
+    parser.add_argument(
+        "--enforce-square-tp-cp",
+        action="store_true",
+        help="Require tp*cp to be a square power of two when enumerating configurations.",
+    )
+    return parser.parse_args()
+
+
+def jitter_positions(gpu_counts, jitter_width):
+    jittered = []
+    for count in gpu_counts:
+        offset = random.uniform(1 - jitter_width, 1 + jitter_width)
+        jittered.append(max(count * offset, 1e-3))
+    return jittered
+
+
+# def plot_results_legacy(results, output_path):
+#     if not results:
+#         print("No successful configurations to plot.", file=sys.stderr)
+#         return
+
+#     random.seed(PLOT_JITTER_SEED)
+#     xs = jitter_positions([item["num_gpus"] for item in results], PLOT_JITTER_WIDTH)
+#     metric_key = "runtime" if PLOT_METRIC.lower() == "runtime" else "performance"
+#     ys = [item[metric_key] for item in results]
+
+#     plt.figure(figsize=(10, 6))
+#     plt.scatter(xs, ys, s=60, alpha=0.7, edgecolors="none")
+
+#     best = min(results, key=lambda item: item["runtime"])
+#     best_x = best["num_gpus"]
+#     best_y = best[metric_key]
+#     plt.scatter([best_x], [best_y], s=180, marker="*", c="red", label="Best runtime")
+
+#     plt.xlabel("Number of GPUs")
+#     if metric_key == "runtime":
+#         plt.ylabel("Runtime (s)")
+#     else:
+#         plt.ylabel("Performance (1 / s)")
+#     plt.xscale("log")
+#     if metric_key == "runtime":
+#         plt.yscale("log")
+#     plt.title("Parallelism sweep")
+#     plt.grid(alpha=0.3)
+#     xticks = sorted(set(item["num_gpus"] for item in results))
+#     plt.xticks(xticks, [str(int(x)) for x in xticks])
+#     plt.legend(loc="best")
+#     plt.tight_layout()
+#     plt.savefig(output_path, dpi=200)
+#     plt.close()
+#     print("Saved scatter plot to {}".format(output_path))
+
+def _gpu_exp(n: int) -> float:
+    """Return log2(num_gpus). Assumes powers of two; will return floats otherwise."""
+    return math.log2(float(n))
+
+
+def plot_results(results, output_path):
+    if not results:
+        print("No successful configurations to plot.", file=sys.stderr)
+        return
+
+    plot_entries = results
+    if MEM_AWARE_FILTER:
+        plot_entries = [item for item in results if not item.get("memory_exceeded", False)]
+        if not plot_entries:
+            print("Warning: all configurations violate memory limits; skipping plot.", file=sys.stderr)
+            return
+
+    metric_key = "runtime" if PLOT_METRIC.lower() == "runtime" else "performance"
+
+    # Build tidy frame (include GPU exponent)
+    rows = []
+    for i, item in enumerate(plot_entries):
+        p = item["parallelism"]
+        ng = int(item["num_gpus"])
+        rows.append({
+            "row_id": i,
+            "num_gpus": ng,
+            "gpu_exp": _gpu_exp(ng),
+            "tp": int(p.get("tp", 1)),
+            "cp": int(p.get("cp", 1)),
+            "dp": int(p.get("dp", 1)),
+            "lp": int(p.get("lp", 1)),
+            metric_key: float(item[metric_key]),
+        })
+    df = pd.DataFrame(rows)
+
+    # Order categories by exponent (gives even spacing like log2)
+    order = sorted(df["gpu_exp"].unique())
+    df["gpu_exp_cat"] = pd.Categorical(df["gpu_exp"], categories=order, ordered=True)
+
+    # --- Per-point RGB colors: R=log2(tp+cp), G=log2(lp), B=log2(dp) ---
+    tps = df["tp"].to_numpy(dtype=float)
+    cps = df["cp"].to_numpy(dtype=float)
+    dps = df["dp"].to_numpy(dtype=float)
+    lps = df["lp"].to_numpy(dtype=float)
+
+    r_raw = np.log2(tps + cps)
+    g_raw = np.log2(lps)
+    b_raw = np.log2(dps)
+
+    def _norm(x, gamma=0.85):
+        x = x.astype(float)
+        xmin, xmax = float(np.min(x)), float(np.max(x))
+        if not np.isfinite(xmin) or not np.isfinite(xmax) or xmax <= xmin:
+            y = np.zeros_like(x)
+        else:
+            y = (x - xmin) / (xmax - xmin)
+        return np.power(y, gamma) if gamma and gamma != 1.0 else y
+
+    r = _norm(r_raw)
+    g = _norm(g_raw)
+    b = _norm(b_raw)
+    a = np.full_like(r, 0.9)
+
+    # Palette: unique color per row via hue="row_id"
+    palette = {rid: (float(r[i]), float(g[i]), float(b[i]), float(a[i]))
+               for i, rid in enumerate(df["row_id"])}
+
+    plt.figure(figsize=(10, 6))
+    ax = sns.swarmplot(
+        data=df,
+        x="gpu_exp_cat",
+        y=metric_key,
+        hue="row_id",
+        palette=palette,   # per-point RGBA
+        size=3.75,
+        linewidth=0,
+        dodge=False,
+        legend=False
+    )
+    # add_rgb_ternary_legend(
+    #     ax,
+    #     corner_labels=("R = log2(tp+cp)", "G = log2(lp)", "B = log2(dp)"),
+    #     gamma=0.85,                      # match your _rgb_from_parallelism gamma
+    #     inset_xywh=(0.70, 0.50, 0.26, 0.26)  # tweak position/size to taste
+    # )
+
+
+    # Best runtime star — place at the matching exponent category index
+    best = min(plot_entries, key=lambda it: it["runtime"])
+    best_exp = _gpu_exp(int(best["num_gpus"]))
+    best_idx = order.index(best_exp)
+    best_y = float(best[metric_key])
+    ax.scatter([best_idx], [best_y], s=180, marker="*", c="red", zorder=5, label="Best runtime")
+
+    # Labels & scales
+    ax.set_xlabel("Number of GPUs (log2 categories)")
+    if metric_key == "runtime":
+        ax.set_ylabel("Runtime (s)")
+        ax.set_yscale("log")
+    else:
+        ax.set_ylabel("Performance (1 / s)")
+
+    # Tick labels show the REAL GPU count (2**exp) so axis *looks* log2
+    ax.set_xticks(range(len(order)))
+    ax.set_xticklabels([str(int(round(2 ** e))) if float(e).is_integer() else f"2^{e:.2f}"
+                        for e in order])
+
+    ax.set_title("Parallelism sweep – beeswarm on log2(GPUs) categories")
+    ax.grid(alpha=0.3, axis="y")
+    ax.legend(loc="best")
+
+    # Reminder of color encoding
+    ax.text(0.99, 0.01,
+            "Color: R=log2(tp+cp), G=log2(lp), B=log2(dp)",
+            transform=ax.transAxes, ha="right", va="bottom", fontsize=9, alpha=0.8)
+
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=200)
+    plt.close()
+    print(f"Saved swarm plot to {output_path}")
+
+
+# def plot_results_categorical(results, output_path):
+#     if not results:
+#         print("No successful configurations to plot.", file=sys.stderr)
+#         return
+
+#     metric_key = "runtime" if PLOT_METRIC.lower() == "runtime" else "performance"
+
+#     # Build a tidy DataFrame for seaborn (categorical x keeps each GPU count together)
+#     df_rows = []
+#     for item in results:
+#         p = item["parallelism"]
+#         df_rows.append({
+#             "num_gpus": int(item["num_gpus"]),
+#             metric_key: float(item[metric_key]),
+#             "tp": int(p.get("tp", 1)),
+#             "cp": int(p.get("cp", 1)),
+#             "dp": int(p.get("dp", 1)),
+#             "lp": int(p.get("lp", 1)),
+#         })
+#     df = pd.DataFrame(df_rows)
+
+#     # Order GPU categories left->right by numeric value
+#     order = sorted(df["num_gpus"].unique())
+#     df["num_gpus_cat"] = pd.Categorical(df["num_gpus"], categories=order, ordered=True)
+
+#     # --- Beeswarm placement ---
+#     plt.figure(figsize=(10, 6))
+#     ax = sns.swarmplot(
+#         data=df,
+#         x="num_gpus_cat",
+#         y=metric_key,
+#         size=5,
+#         color="k",         # temporary (we'll recolor with our RGBs after layout)
+#         linewidth=0,
+#         alpha=0.0          # invisible; we only want the positions it computes
+#     )
+
+#     # Grab the laid-out positions, remove seaborn's collection, and redraw with our colors
+#     if not ax.collections:
+#         print("Warning: swarmplot produced no collections.", file=sys.stderr)
+#         return
+#     coll = ax.collections[0]
+#     offsets = coll.get_offsets()        # Nx2 array (x_index, y_value) in data coords (x is categorical index)
+#     coll.remove()
+
+#     # Compute per-point RGBA colors from parallelism
+#     colors = _rgb_from_parallelism(results, gamma=0.85)
+
+#     # Plot the recolored points at the computed positions
+#     ax.scatter(
+#         offsets[:, 0], offsets[:, 1],
+#         s=60,
+#         c=colors,
+#         edgecolors="none",
+#         zorder=3
+#     )
+
+#     # Global best marker (use category index for x)
+#     best = min(results, key=lambda item: item["runtime"])
+#     best_x_idx = order.index(int(best["num_gpus"]))  # categorical index
+#     best_y = float(best[metric_key])
+#     ax.scatter([best_x_idx], [best_y], s=180, marker="*", c="red", zorder=5, label="Best runtime")
+
+#     # Axes/labels
+#     ax.set_xlabel("Number of GPUs")
+#     if metric_key == "runtime":
+#         ax.set_ylabel("Runtime (s)")
+#         ax.set_yscale("log")
+#     else:
+#         ax.set_ylabel("Performance (1 / s)")
+
+#     # Pretty ticks for categorical axis
+#     ax.set_xticks(range(len(order)))
+#     ax.set_xticklabels([str(x) for x in order])
+
+#     ax.set_title("Parallelism sweep (beeswarm with log-RGB coloring)")
+#     ax.grid(alpha=0.3, axis="y")
+#     ax.legend(loc="best")
+
+#     # Small caption to recall color encoding
+#     ax.text(0.99, 0.01,
+#             "Color channels: R=log2(tp+cp), G=log2(lp), B=log2(dp)",
+#             transform=ax.transAxes, ha="right", va="bottom", fontsize=9, alpha=0.8)
+
+#     plt.tight_layout()
+#     plt.savefig(output_path, dpi=200)
+#     plt.close()
+#     print(f"Saved swarm plot to {output_path}")
+
+
+# def plot_mfu(results, output_path):
+#     if not results:
+#         return
+#     random.seed(PLOT_JITTER_SEED)
+#     xs = jitter_positions([item["num_gpus"] for item in results], PLOT_JITTER_WIDTH)
+#     ys = [item["mfu"] for item in results]
+#     plt.figure(figsize=(10, 6))
+#     plt.scatter(xs, ys, s=60, alpha=0.7, edgecolors="none")
+#     valid = [item for item in results if item["mfu"] == item["mfu"]]
+#     if valid:
+#         best = max(valid, key=lambda item: item["mfu"])
+#         plt.scatter([best["num_gpus"]], [best["mfu"]], s=180, marker="*", c="green", label="Max MFU")
+#     plt.xlabel("Number of GPUs")
+#     plt.ylabel("MFU")
+#     plt.xscale("log")
+#     plt.ylim(0.0, 1.05)
+#     plt.title("Parallelism sweep - MFU")
+#     plt.grid(alpha=0.3)
+#     xticks = sorted(set(item["num_gpus"] for item in results))
+#     plt.xticks(xticks, [str(int(x)) for x in xticks])
+#     if valid:
+#         plt.legend(loc="best")
+#     plt.tight_layout()
+#     plt.savefig(output_path, dpi=200)
+#     plt.close()
+#     print("Saved MFU scatter plot to {}".format(output_path))
+
+
+_GLOBAL_MODEL_CONFIG = None
+_GLOBAL_MODE = None
+_GLOBAL_HW_DICT = None
+
+
+def _worker_init(hw_dict, model_config_path, mode):
+    global _GLOBAL_MODEL_CONFIG, _GLOBAL_MODE, _GLOBAL_HW_DICT
+    _GLOBAL_HW_DICT = hw_dict
+    _GLOBAL_MODE = mode
+    _GLOBAL_MODEL_CONFIG = config.parse_config(model_config_path, config_type=mode)
+
+
+def _worker_task(parallel_items: Tuple[Tuple[str, object], ...]):
+    parallel_settings = {k: v for k, v in parallel_items}
+    try:
+        metrics = evaluate_parallelism(
+            _GLOBAL_HW_DICT,
+            _GLOBAL_MODEL_CONFIG,
+            _GLOBAL_MODE,
+            parallel_settings,
+        )
+        return {
+            "status": "ok",
+            "parallelism": parallel_settings,
+            "metrics": metrics,
+        }
+    except Exception as exc:
+        return {
+            "status": "error",
+            "parallelism": parallel_settings,
+            "error": str(exc),
+        }
+
+
+def _build_tasks(
+    gpu_choices: Iterable[Dict[str, int]],
+    other_choices: Iterable[Dict[str, object]],
+) -> List[Tuple[Tuple[str, object], ...]]:
+    tasks: List[Tuple[Tuple[str, object], ...]] = []
+    for gpu_choice in gpu_choices:
+        for other_choice in other_choices:
+            settings: Dict[str, object] = {}
+            settings.update(gpu_choice)
+            settings.update(other_choice)
+            settings["mb"] = settings.get("lp", 1)
+            tasks.append(tuple(sorted(settings.items())))
+    return tasks
+
+
+def main():
+    args = parse_args()
+    set_astrasim_cache_mode(ASTRA_CACHE_MODE)
+
+    enforce_square = ENFORCE_SQUARE_TP_CP or args.enforce_square_tp_cp
+
+    results: List[Dict[str, object]] = []
+
+    if args.plot_only:
+        try:
+            results = load_results_from_report(REPORT_OUTPUT_PATH)
+        except FileNotFoundError as exc:
+            print(str(exc))
+            return
+
+        if not results:
+            print(f"No entries found in {REPORT_OUTPUT_PATH}; nothing to plot.")
+            return
+
+        best = min(results, key=lambda item: item["runtime"])
+        print(f"Loaded {len(results)} configuration(s) from {REPORT_OUTPUT_PATH}")
+        print("\nBest configuration (lowest runtime):")
+        print("  Parallelism: {}".format(best["parallelism"]))
+        print("  GPUs: {}".format(best["num_gpus"]))
+        print("  Runtime: {:.4f} s".format(best["runtime"]))
+        print("  Performance (1/s): {:.4f}".format(best["performance"]))
+        print("  Total FLOPs: {:.3e}".format(best["total_flops"]))
+        print("  MFU: {:.3f}".format(best["mfu"]))
+        if best["memory_exceeded"]:
+            print("  Memory capacity exceeded by {:.3f} GB".format(best["memory_violation_gb"]))
+    else:
+        base_hw_dict = read_yaml(HARDWARE_CONFIG_PATH)
+        mode = determine_model_mode(MODEL_CONFIG_PATH)
+
+        gpu_axes = list(PARALLELISM_SWEEP.keys())
+        other_axes = list(OTHER_PARALLELISM_OPTIONS.keys())
+        gpu_combos = list(cartesian_product(PARALLELISM_SWEEP))
+        other_combos = list(cartesian_product(OTHER_PARALLELISM_OPTIONS))
+        task_items = _build_tasks(gpu_combos, other_combos)
+        print("Enumerating {} parallelism combinations: {}".format(len(task_items), ", ".join(gpu_axes)))
+
+        skipped_out_of_range = 0
+        skipped_errors = 0
+        error_messages: List[str] = []
+        evaluated = 0
+
+        filtered_tasks: List[Tuple[Tuple[str, object], ...]] = []
+        skipped_square_constraint = 0
+        for items in task_items:
+            settings = dict(items)
+            num_gpus = total_gpu_count(settings)
+            if enforce_square and not tp_cp_product_is_power_of_two_square(settings.get("tp"), settings.get("cp")):
+                skipped_square_constraint += 1
+                continue
+            if GPU_COUNT_MIN <= num_gpus <= GPU_COUNT_MAX:
+                filtered_tasks.append(items)
+            else:
+                skipped_out_of_range += 1
+
+        if not filtered_tasks:
+            print("No configurations within GPU count bounds.")
+            return
+        if enforce_square and skipped_square_constraint:
+            print(f"Skipped {skipped_square_constraint} configuration(s) due to tp*cp square constraint.")
+
+        available_cpus = max(1, os.cpu_count() or 1)
+        if MAX_WORKERS is None or MAX_WORKERS <= 0:
+            worker_limit = max(1, available_cpus - 1)
+        else:
+            worker_limit = min(MAX_WORKERS, max(1, available_cpus - 1))
+        worker_count = max(1, worker_limit)
+        print(f"Using {worker_count} worker process(es) (out of {available_cpus} CPUs).")
+
+        if worker_count > 1 and len(filtered_tasks) > 1:
+            with ProcessPoolExecutor(
+                max_workers=worker_count,
+                initializer=_worker_init,
+                initargs=(base_hw_dict, MODEL_CONFIG_PATH, mode),
+            ) as executor:
+                futures = {executor.submit(_worker_task, items): items for items in filtered_tasks}
+                with tqdm(total=len(filtered_tasks), desc="Evaluating", unit="config") as progress:
+                    for future in as_completed(futures):
+                        progress.update(1)
+                        result = future.result()
+                        settings = result["parallelism"]
+                        num_gpus = total_gpu_count(settings)
+                        if result["status"] != "ok":
+                            skipped_errors += 1
+                            msg = result.get("error") or "unknown error"
+                            error_messages.append(f"{settings}: {msg}")
+                            continue
+                        metrics = result["metrics"]
+                        evaluated += 1
+                        entry = {
+                            "parallelism": settings,
+                            "num_gpus": num_gpus,
+                            "runtime": metrics["runtime"],
+                            "performance": metrics["performance"],
+                            "total_flops": metrics["total_flops"],
+                            "achieved_flops": metrics["achieved_flops"],
+                            "peak_flops": metrics["peak_flops"],
+                            "mfu": metrics["mfu"],
+                            "memory_exceeded": metrics["memory_exceeded"],
+                            "memory_violation_gb": metrics["memory_violation_gb"],
+                        }
+                        results.append(entry)
+        else:
+            model_config_obj = config.parse_config(MODEL_CONFIG_PATH, config_type=mode)
+            with tqdm(total=len(filtered_tasks), desc="Evaluating", unit="config") as progress:
+                for items in filtered_tasks:
+                    settings = dict(items)
+                    num_gpus = total_gpu_count(settings)
+                    progress.update(1)
+                    try:
+                        metrics = evaluate_parallelism(base_hw_dict, model_config_obj, mode, settings)
+                    except Exception as exc:
+                        skipped_errors += 1
+                        error_messages.append(f"{settings}: {exc}")
+                        continue
+
+                    evaluated += 1
+                    entry = {
+                        "parallelism": settings,
+                        "num_gpus": num_gpus,
+                        "runtime": metrics["runtime"],
+                        "performance": metrics["performance"],
+                        "total_flops": metrics["total_flops"],
+                        "achieved_flops": metrics["achieved_flops"],
+                        "peak_flops": metrics["peak_flops"],
+                        "mfu": metrics["mfu"],
+                        "memory_exceeded": metrics["memory_exceeded"],
+                        "memory_violation_gb": metrics["memory_violation_gb"],
+                    }
+                    results.append(entry)
+
+        total_skipped = skipped_out_of_range + skipped_errors
+        if not results:
+            print(
+                "No valid configurations evaluated ({} skipped: {} out-of-range, {} errors).".format(
+                    total_skipped, skipped_out_of_range, skipped_errors
+                )
+            )
+            if error_messages:
+                print("Encountered errors for configurations:")
+                for msg in error_messages:
+                    print(f"  {msg}")
+            return
+
+        best = min(results, key=lambda item: item["runtime"])
+        print(
+            f"Evaluated {evaluated} configuration(s); skipped {total_skipped} "
+            f"(out_of_range={skipped_out_of_range}, errors={skipped_errors})."
+        )
+        if error_messages:
+            print("Some configurations failed:")
+            for msg in error_messages:
+                print(f"  {msg}")
+        print("\nBest configuration (lowest runtime):")
+        print("  Parallelism: {}".format(best["parallelism"]))
+        print("  GPUs: {}".format(best["num_gpus"]))
+        print("  Runtime: {:.4f} s".format(best["runtime"]))
+        print("  Performance (1/s): {:.4f}".format(best["performance"]))
+        print("  Total FLOPs: {:.3e}".format(best["total_flops"]))
+        print("  MFU: {:.3f}".format(best["mfu"]))
+        if best["memory_exceeded"]:
+            print("  Memory capacity exceeded by {:.3f} GB".format(best["memory_violation_gb"]))
+
+        try:
+            write_report(results, REPORT_OUTPUT_PATH)
+            print("Wrote detailed report to {}".format(REPORT_OUTPUT_PATH))
+        except Exception as exc:
+            print("Warning: failed to write report: {}".format(exc), file=sys.stderr)
+
+    if not results:
+        return
+
+    plot_results(results, PLOT_OUTPUT_PATH)
+    # plot_mfu(results, PLOT_MFU_OUTPUT_PATH)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/util_finder.py b/tools/util_finder.py
new file mode 100644
index 00000000..8f965f24
--- /dev/null
+++ b/tools/util_finder.py
@@ -0,0 +1,349 @@
+#!/usr/bin/env python3
+"""
+Utility tuner (core util, HBM bandwidth util) for training validation.
+
+Modeled after overlap_finder.py but searches in two dimensions using a
+coordinate-wise golden-section search. Objective is average absolute percent
+error across Korthi+Selene training validation suites (nvidia_train.py).
+
+Global knobs (edit below):
+  - DEVICES: which training devices to validate (defaults: Korthi + Selene)
+  - APPLY_BEST: write the best utils back to all hardware configs via mod_field
+  - RANGES/TOL/ITER: search bounds and golden search controls
+  - COORD_STEPS: number of coordinate-descent passes (keep small; each eval is slow)
+"""
+
+from __future__ import annotations
+
+import io
+import math
+import sys
+from contextlib import redirect_stdout
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from validation_scripts import nvidia_train  # type: ignore
+from validation_scripts.validation_helpers import ValidationSpec, run_validation_suite  # type: ignore
+from tools import mod_field
+
+# ======= Global config =======
+DEVICES: Sequence[str] = ("A100_korthi", "A100_selene")
+APPLY_BEST = False
+
+# Search bounds for utilities (floats in (0,1]). Set to None to skip tuning that axis.
+CORE_UTIL_RANGE: Optional[Tuple[float, float]] = None
+HBM_UTIL_RANGE: Optional[Tuple[float, float]] = (0.25, 0.55)
+GOLD_TOL = 0.02       # stop when interval width < GOLD_TOL
+GOLD_MAX_ITER = 6     # cap per-axis golden iterations
+COORD_STEPS = 2       # how many coordinate-descent passes (keep small)
+# =============================
+
+# Exclude specific heavy/slow validation cases by metadata match.
+EXCLUDE_TESTS: Tuple[Dict[str, Any], ...] = (
+    # Skip 3k-GPU Selene GPT-1T case (dp=6, tp=8, pp=64, mb=512, batch=3072).
+    {"device": "A100_selene", "model": "GPT 1T", "pp": 64},
+)
+
+
+def _merge_dict(target: Dict[str, Any], updates: Mapping[str, Any]) -> Dict[str, Any]:
+    for key, val in updates.items():
+        if isinstance(val, Mapping):
+            current = target.get(key, {})
+            if not isinstance(current, dict):
+                current = {}
+            target[key] = _merge_dict(dict(current), val)
+        else:
+            target[key] = val
+    return target
+
+
+def _with_utils(spec: ValidationSpec, core_util: Optional[float], hbm_util: Optional[float]) -> ValidationSpec:
+    base_overrides: Dict[str, Any] = {}
+    if spec.hardware_overrides and isinstance(spec.hardware_overrides, Mapping):
+        base_overrides = _merge_dict({}, spec.hardware_overrides)
+    util_overrides: Dict[str, Any] = {"tech_param": {}}
+    if core_util is not None:
+        util_overrides["tech_param"]["core"] = {"util": float(core_util)}
+    if hbm_util is not None:
+        util_overrides["tech_param"]["DRAM"] = {"util": float(hbm_util)}
+    if not util_overrides["tech_param"]:
+        return ValidationSpec(
+            label=spec.label,
+            model_overrides=spec.model_overrides,
+            hardware_overrides=base_overrides,
+            metadata=spec.metadata,
+            model_config_path=spec.model_config_path,
+            hardware_config_path=spec.hardware_config_path,
+            order=spec.order,
+        )
+    base_overrides = _merge_dict(base_overrides, util_overrides)
+    return ValidationSpec(
+        label=spec.label,
+        model_overrides=spec.model_overrides,
+        hardware_overrides=base_overrides,
+        metadata=spec.metadata,
+        model_config_path=spec.model_config_path,
+        hardware_config_path=spec.hardware_config_path,
+        order=spec.order,
+    )
+
+
+def _avg_pct_error(errors: Iterable[float]) -> float:
+    vals = [e for e in errors if not math.isnan(e)]
+    if not vals:
+        return float("inf")
+    return sum(vals) / len(vals)
+
+
+def _all_hardware_configs() -> List[Path]:
+    roots = [
+        PROJECT_ROOT / "configs" / "hardware-config",
+        PROJECT_ROOT / "validation_scripts" / "validation_configs" / "hardware-config",
+    ]
+    paths: List[Path] = []
+    for root in roots:
+        if root.exists():
+            paths.extend(sorted(root.glob("*.yaml")))
+    return paths
+
+
+def _should_exclude(meta: Mapping[str, Any]) -> bool:
+    for rule in EXCLUDE_TESTS:
+        matches = True
+        for key, val in rule.items():
+            if str(meta.get(key)) != str(val):
+                matches = False
+                break
+        if matches:
+            return True
+    return False
+
+
+def _filter_specs(
+    specs: Sequence[ValidationSpec],
+    actual_lookup: Mapping[Tuple[str, int, int, int, int, int, int, bool, str], float],
+) -> Tuple[List[ValidationSpec], Dict[Tuple[str, int, int, int, int, int, int, bool, str], float]]:
+    kept: List[ValidationSpec] = []
+    kept_lookup: Dict[Tuple[str, int, int, int, int, int, int, bool, str], float] = {}
+    for spec in specs:
+        meta = spec.metadata or {}
+        if _should_exclude(meta):
+            continue
+        key = (
+            meta.get("model"),
+            int(meta.get("batch")),
+            int(meta.get("mb")),
+            int(meta.get("dp")),
+            int(meta.get("tp")),
+            int(meta.get("pp")),
+            int(meta.get("cp")),
+            bool(meta.get("tp_sp")),
+            str(meta.get("recomputation")),
+        )
+        if key in actual_lookup:
+            kept_lookup[key] = actual_lookup[key]
+        kept.append(spec)
+    return kept, kept_lookup
+
+
+@dataclass
+class SearchResult:
+    best_x: float
+    best_y: float
+    evaluations: List[Tuple[float, float]]
+
+
+def golden_section_search(
+    fn: Callable[[float], float],
+    lo: float,
+    hi: float,
+    tol: float = GOLD_TOL,
+    max_iter: int = GOLD_MAX_ITER,
+    log_fn: Optional[Callable[[int, int, float, float], None]] = None,
+) -> SearchResult:
+    phi = (1 + 5 ** 0.5) / 2
+    inv_phi = 1 / phi
+    inv_phi_sq = inv_phi ** 2
+
+    a, b = lo, hi
+    h = b - a
+    if h <= tol:
+        mid = (a + b) / 2
+        val = fn(mid)
+        return SearchResult(best_x=mid, best_y=val, evaluations=[(mid, val)])
+
+    n = int(math.ceil(math.log(tol / h) / math.log(inv_phi)))
+    c = a + inv_phi_sq * h
+    d = a + inv_phi * h
+    yc = fn(c)
+    yd = fn(d)
+    if log_fn:
+        log_fn(1, max_iter, c, yc)
+        log_fn(2, max_iter, d, yd)
+    evals: List[Tuple[float, float]] = [(c, yc), (d, yd)]
+
+    iter_idx = 2
+    for _ in range(min(n, max_iter)):
+        if yc < yd:
+            b = d
+            d, yd = c, yc
+            h = inv_phi * h
+            c = a + inv_phi_sq * h
+            yc = fn(c)
+            evals.append((c, yc))
+        else:
+            a = c
+            c, yc = d, yd
+            h = inv_phi * h
+            d = a + inv_phi * h
+            yd = fn(d)
+            evals.append((d, yd))
+        iter_idx += 1
+        if log_fn:
+            log_fn(iter_idx, max_iter, evals[-1][0], evals[-1][1])
+        if h < tol:
+            break
+
+    best_x, best_y = min(evals, key=lambda t: t[1])
+    return SearchResult(best_x=best_x, best_y=best_y, evaluations=evals)
+
+
+class UtilTuner:
+    def __init__(self) -> None:
+        self._eval_cache: Dict[Tuple[float, float], float] = {}
+
+    def eval_pair(self, core_util: float, hbm_util: float) -> float:
+        key = (round(core_util, 6), round(hbm_util, 6))
+        if key in self._eval_cache:
+            print(f"[CACHE HIT] core.util={core_util:.4f}, DRAM.util={hbm_util:.4f} -> {self._eval_cache[key]:.3f}")
+            return self._eval_cache[key]
+
+        all_errors: List[float] = []
+        for device in DEVICES:
+            specs, actual_lookup, base_model_path, base_hw_path = nvidia_train.build_specs_for_device(device)
+            specs, actual_lookup = _filter_specs(specs, actual_lookup)
+            if not specs:
+                print(f"[SKIP DEVICE] {device}: no specs after filtering.")
+                continue
+            specs = [_with_utils(spec, core_util, hbm_util) for spec in specs]
+            print(f"[EVAL] device={device}, specs={len(specs)}, core.util={core_util:.4f}, DRAM.util={hbm_util:.4f}")
+            with io.StringIO() as buf, redirect_stdout(buf):
+                results = run_validation_suite(
+                    specs,
+                    base_model_config_path=base_model_path,
+                    base_hardware_config_path=base_hw_path,
+                    result_parser=nvidia_train.parse_training_time,
+                    run_perf_path=nvidia_train.RUN_PERF,
+                )
+            rows = nvidia_train.compute_pct_errors(results, actual_lookup)
+            all_errors.extend(float(row["pct_error"]) for row in rows)
+            print(f"[RESULTS] device={device}, pct_errors={ [float(row['pct_error']) for row in rows] }")
+
+        avg_error = _avg_pct_error(all_errors)
+        self._eval_cache[key] = avg_error
+        print(f"[EVAL DONE] core.util={core_util:.4f}, DRAM.util={hbm_util:.4f}, avg_abs_pct_err={avg_error:.3f}")
+        return avg_error
+
+    def optimize(self) -> Tuple[Optional[float], Optional[float], float]:
+        core_util = None if CORE_UTIL_RANGE is None else (CORE_UTIL_RANGE[0] + CORE_UTIL_RANGE[1]) / 2
+        hbm_util = None if HBM_UTIL_RANGE is None else (HBM_UTIL_RANGE[0] + HBM_UTIL_RANGE[1]) / 2
+
+        if CORE_UTIL_RANGE is None and HBM_UTIL_RANGE is None:
+            raise ValueError("At least one of CORE_UTIL_RANGE or HBM_UTIL_RANGE must be set.")
+
+        # If only one axis is active, do a single golden search on that axis.
+        if CORE_UTIL_RANGE is None and HBM_UTIL_RANGE is not None:
+            h_lo, h_hi = HBM_UTIL_RANGE
+            print("Tuning only DRAM.util")
+            res_hbm = golden_section_search(
+                lambda y: self.eval_pair(core_util or 1.0, y),
+                lo=h_lo,
+                hi=h_hi,
+                tol=GOLD_TOL,
+                max_iter=GOLD_MAX_ITER,
+                log_fn=lambda i, m, y, v: print(f"  [{i}/{m}] hbm_util={y:.4f} -> {v:.3f}"),
+            )
+            best_val = res_hbm.best_y
+            return core_util, res_hbm.best_x, best_val
+
+        if HBM_UTIL_RANGE is None and CORE_UTIL_RANGE is not None:
+            c_lo, c_hi = CORE_UTIL_RANGE
+            print("Tuning only core.util")
+            res_core = golden_section_search(
+                lambda x: self.eval_pair(x, hbm_util or 1.0),
+                lo=c_lo,
+                hi=c_hi,
+                tol=GOLD_TOL,
+                max_iter=GOLD_MAX_ITER,
+                log_fn=lambda i, m, x, y: print(f"  [{i}/{m}] core_util={x:.4f} -> {y:.3f}"),
+            )
+            best_val = res_core.best_y
+            return res_core.best_x, hbm_util, best_val
+
+        # Both axes active: coordinate descent.
+        core_lo, core_hi = CORE_UTIL_RANGE  # type: ignore[assignment]
+        hbm_lo, hbm_hi = HBM_UTIL_RANGE    # type: ignore[assignment]
+        core = float(core_util)
+        hbm = float(hbm_util)
+        best_val = self.eval_pair(core, hbm)
+
+        for step in range(COORD_STEPS):
+            print(f"\n--- Coordinate step {step + 1}/{COORD_STEPS} ---")
+            # Optimize core util with HBM fixed
+            print(f"Optimizing core util (hbm_util={hbm:.4f})")
+            res_core = golden_section_search(
+                lambda x: self.eval_pair(x, hbm),
+                lo=core_lo,
+                hi=core_hi,
+                tol=GOLD_TOL,
+                max_iter=GOLD_MAX_ITER,
+                log_fn=lambda i, m, x, y: print(f"  [{i}/{m}] core_util={x:.4f} -> {y:.3f}"),
+            )
+            core = res_core.best_x
+            best_val = res_core.best_y
+
+            # Optimize HBM util with core fixed
+            print(f"Optimizing HBM util (core_util={core:.4f})")
+            res_hbm = golden_section_search(
+                lambda y: self.eval_pair(core, y),
+                lo=hbm_lo,
+                hi=hbm_hi,
+                tol=GOLD_TOL,
+                max_iter=GOLD_MAX_ITER,
+                log_fn=lambda i, m, y, v: print(f"  [{i}/{m}] hbm_util={y:.4f} -> {v:.3f}"),
+            )
+            hbm = res_hbm.best_x
+            best_val = min(best_val, res_hbm.best_y)
+
+        best_val = self.eval_pair(core, hbm)
+        return core, hbm, best_val
+
+
+def main() -> None:
+    tuner = UtilTuner()
+    best_core, best_hbm, best_err = tuner.optimize()
+    print("\n=== Best utils ===")
+    print(f"  core.util = {best_core if best_core is None else f'{best_core:.4f}'}")
+    print(f"  DRAM.util = {best_hbm if best_hbm is None else f'{best_hbm:.4f}'}")
+    print(f"  avg abs pct error = {best_err:.3f}")
+
+    if APPLY_BEST:
+        core_val = None if best_core is None else round(best_core, 2)
+        hbm_val = None if best_hbm is None else round(best_hbm, 2)
+        print("\nApplying tuned utils to all hardware configs")
+        for cfg_path in _all_hardware_configs():
+            if core_val is not None:
+                msg1 = mod_field.set_field(cfg_path, "tech_param.core.util", core_val, dry_run=False)
+                print(f"  {msg1}")
+            if hbm_val is not None:
+                msg2 = mod_field.set_field(cfg_path, "tech_param.DRAM.util", hbm_val, dry_run=False)
+                print(f"  {msg2}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/util.py b/util.py
index 3afcf49e..b3eef075 100644
--- a/util.py
+++ b/util.py
@@ -2,6 +2,8 @@
 import sys
 import os
 import config
+import threading
+from typing import Iterable, List, Optional, Tuple
 
 core=1
 DRAM=1
@@ -12,6 +14,195 @@
 
 DF_CACHE_DIR = "./cache"
 
+_LOG_MESSAGES: List[Tuple[Optional[str], str]] = []
+_LOG_LOCK = threading.Lock()
+
+_SECTION_ORDER = [
+    ("network", "TOPOLOGY/NETWORK"),
+    ("faults", "FAULTY LINKS"),
+    ("results", "RESULTS"),
+]
+
+_REPO_ROOT = os.path.abspath(os.environ.get("DEEPFLOW_REPO_ROOT", os.getcwd()))
+
+
+def log_message(message: str, category: Optional[str] = None) -> None:
+    """Append ``message`` to the shared log queue."""
+    if message is None:
+        return
+    text = str(message)
+    if not text:
+        return
+    cat_normalized = str(category).strip().lower() if category else None
+    with _LOG_LOCK:
+        _LOG_MESSAGES.append((cat_normalized, text))
+
+
+def extend_log(lines: Iterable[str], category: Optional[str] = None) -> None:
+    """Append multiple lines to the log queue."""
+    for line in lines:
+        log_message(line, category=category)
+
+
+def drain_log_messages() -> List[Tuple[Optional[str], str]]:
+    """Return and clear all queued log messages (category, message)."""
+    with _LOG_LOCK:
+        if not _LOG_MESSAGES:
+            return []
+        drained = list(_LOG_MESSAGES)
+        _LOG_MESSAGES.clear()
+        return drained
+
+
+def flush_log_queue() -> None:
+    """Print and clear the queued log messages grouped by category."""
+    entries = drain_log_messages()
+    print("\n")
+    if not entries:
+        return
+    categorized = {cat: [] for cat, _ in _SECTION_ORDER}
+    uncategorized: List[str] = []
+    for category, message in entries:
+        if category in categorized:
+            categorized[category].append(message)
+        else:
+            uncategorized.append(message)
+
+    sections_printed = False
+    section_border = "=" * 60
+    for cat, title in _SECTION_ORDER:
+        lines = categorized.get(cat) or []
+        if not lines:
+            continue
+        sections_printed = True
+        print(f"{section_border}\n{title}\n{section_border}")
+        for line in lines:
+            print(line)
+    if sections_printed:
+        print(section_border)
+
+    for line in uncategorized:
+        print(line)
+
+
+def relpath_display(path: str) -> str:
+    """Return ``path`` relative to the repo root when possible."""
+    if not path:
+        return ""
+    abs_path = os.path.abspath(path)
+    try:
+        rel = os.path.relpath(abs_path, start=_REPO_ROOT)
+    except Exception:
+        return abs_path
+    if rel.startswith(".."):
+        return abs_path
+    return rel
+
+
+def _collect_parallelism_values(hw_config):
+    sch_config = getattr(hw_config, "sch_config", None)
+    if sch_config is None:
+        return {}
+
+    values = {}
+    for name in getattr(sch_config, "_fields", []):
+        values[str(name).lower()] = getattr(sch_config, name)
+    return values
+
+
+def _format_parallelism_terms(dim, parallelism_values):
+    terms = []
+    for axis in getattr(dim, "parallelisms", ()):
+        if axis not in parallelism_values:
+            continue
+        factor = parallelism_values[axis]
+        try:
+            factor_value = int(factor)
+        except (TypeError, ValueError):
+            factor_value = factor
+        terms.append(f"{axis} {factor_value}")
+    return terms
+
+
+def network_topology_summary_training(hw_config):
+    parallelism_values = _collect_parallelism_values(hw_config)
+    dimensions = list(getattr(hw_config.network_layout, "dimensions", ()))
+    ordered_axes = ["tp", "cp", "lp", "dp"]
+    formatted_terms = []
+    for axis in ordered_axes:
+        value = parallelism_values.get(axis)
+        if value is None:
+            continue
+        formatted_terms.append(f"{axis}:{value}")
+    formatted_parallelisms = ", ".join(formatted_terms) if formatted_terms else "none"
+    lines = [
+        f"Parallelisms: {formatted_parallelisms}",
+        f"Network Topology [dims={len(dimensions)}]",
+    ]
+    aggregate = 1
+
+    for dim in dimensions:
+        terms = _format_parallelism_terms(dim, parallelism_values)
+        axis_repr = " × ".join(terms) if terms else "(none)"
+        size_value = getattr(dim, "size", 1)
+        try:
+            size_int = int(size_value)
+        except (TypeError, ValueError):
+            size_int = None
+            size_display = size_value
+        else:
+            aggregate *= size_int
+            size_display = size_int
+        lines.append(
+            f"  • {dim.id} {dim.label} : {axis_repr} ⇒ size {size_display}"
+        )
+
+    lines.append(f"  ⇒  total {aggregate} devices")
+    return lines
+
+
+def network_topology_summary_inference(hw_config):
+    parallelism_values = _collect_parallelism_values(hw_config)
+    all_dimensions = list(getattr(hw_config.network_layout, "dimensions", ()))
+    filtered_dimensions = [
+        dim
+        for dim in all_dimensions
+        if any(axis != "dp" for axis in getattr(dim, "parallelisms", ())) or not dim.parallelisms
+    ]
+    lines = [f"Network Topology [dims={len(filtered_dimensions)}]"]
+
+    aggregate_per_replica = 1
+    for dim in filtered_dimensions:
+        terms = _format_parallelism_terms(dim, parallelism_values)
+        axis_repr = " × ".join(terms) if terms else "(none)"
+        size_value = getattr(dim, "size", 1)
+        try:
+            size_int = int(size_value)
+        except (TypeError, ValueError):
+            size_int = None
+            size_display = size_value
+        else:
+            aggregate_per_replica *= size_int
+            size_display = size_int
+        lines.append(
+            f"  • {dim.id} {dim.label} : {axis_repr} ⇒ size {size_display}"
+        )
+
+    dp_factor = parallelism_values.get("dp", 1)
+    try:
+        dp_replicas = max(1, int(dp_factor))
+    except (TypeError, ValueError):
+        dp_replicas = dp_factor if dp_factor else 1
+    if isinstance(aggregate_per_replica, (int, float)) and isinstance(dp_replicas, (int, float)):
+        total_aggregate = aggregate_per_replica * dp_replicas
+    else:
+        total_aggregate = aggregate_per_replica
+    lines.append(f"  replicas (dp): {dp_replicas}")
+    lines.append(
+        f"  => aggregate = {aggregate_per_replica} GPUs per replica ({total_aggregate} total)"
+    )
+    return lines
+
 def printError(message):
   sys.exit(message)
 
diff --git a/uv.lock b/uv.lock
index ae382dd3..2142aa36 100644
--- a/uv.lock
+++ b/uv.lock
@@ -130,13 +130,14 @@ dependencies = [
     { name = "pandas" },
     { name = "protobuf" },
     { name = "pydot" },
+    { name = "pytest" },
     { name = "python-dateutil" },
     { name = "pytz" },
     { name = "pyyaml" },
-    { name = "ruamel-yaml" },
-    { name = "ruamel-yaml-clib" },
+    { name = "scotchpy64" },
     { name = "seaborn" },
     { name = "six" },
+    { name = "tqdm" },
     { name = "tzdata" },
 ]
 
@@ -149,13 +150,14 @@ requires-dist = [
     { name = "pandas", specifier = "==2.3.1" },
     { name = "protobuf", specifier = "==5.28.2" },
     { name = "pydot", specifier = ">=3.0.0" },
+    { name = "pytest", specifier = ">=9.0.0" },
     { name = "python-dateutil", specifier = "==2.9.0.post0" },
     { name = "pytz", specifier = "==2025.2" },
     { name = "pyyaml", specifier = "==6.0.2" },
-    { name = "ruamel-yaml", specifier = "==0.17.40" },
-    { name = "ruamel-yaml-clib", specifier = "==0.2.8" },
+    { name = "scotchpy64", specifier = ">=1.0.1" },
     { name = "seaborn", specifier = ">=0.13.2" },
     { name = "six", specifier = "==1.17.0" },
+    { name = "tqdm", specifier = ">=4.67.1" },
     { name = "tzdata", specifier = "==2025.2" },
 ]
 
@@ -217,6 +219,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/91/4c/e0ce1ef95d4000ebc1c11801f9b944fa5910ecc15b5e351865763d8657f8/graphviz-0.21-py3-none-any.whl", hash = "sha256:54f33de9f4f911d7e84e4191749cac8cc5653f815b06738c54db9a15ab8b1e42", size = 47300 },
 ]
 
+[[package]]
+name = "iniconfig"
+version = "2.3.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484 },
+]
+
 [[package]]
 name = "joblib"
 version = "1.5.2"
@@ -598,6 +609,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/95/7e/f896623c3c635a90537ac093c6a618ebe1a90d87206e42309cb5d98a1b9e/pillow-12.0.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:b290fd8aa38422444d4b50d579de197557f182ef1068b75f5aa8558638b8d0a5", size = 6997850 },
 ]
 
+[[package]]
+name = "pluggy"
+version = "1.6.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538 },
+]
+
 [[package]]
 name = "protobuf"
 version = "5.28.2"
@@ -624,6 +644,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/7e/32/a7125fb28c4261a627f999d5fb4afff25b523800faed2c30979949d6facd/pydot-4.0.1-py3-none-any.whl", hash = "sha256:869c0efadd2708c0be1f916eb669f3d664ca684bc57ffb7ecc08e70d5e93fee6", size = 37087 },
 ]
 
+[[package]]
+name = "pygments"
+version = "2.19.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/b0/77/a5b8c569bf593b0140bde72ea885a803b82086995367bf2037de0159d924/pygments-2.19.2.tar.gz", hash = "sha256:636cb2477cec7f8952536970bc533bc43743542f70392ae026374600add5b887", size = 4968631 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b", size = 1225217 },
+]
+
 [[package]]
 name = "pyparsing"
 version = "3.2.5"
@@ -633,6 +662,22 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/10/5e/1aa9a93198c6b64513c9d7752de7422c06402de6600a8767da1524f9570b/pyparsing-3.2.5-py3-none-any.whl", hash = "sha256:e38a4f02064cf41fe6593d328d0512495ad1f3d8a91c4f73fc401b3079a59a5e", size = 113890 },
 ]
 
+[[package]]
+name = "pytest"
+version = "9.0.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+    { name = "iniconfig" },
+    { name = "packaging" },
+    { name = "pluggy" },
+    { name = "pygments" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/da/1d/eb34f286b164c5e431a810a38697409cca1112cee04b287bb56ac486730b/pytest-9.0.0.tar.gz", hash = "sha256:8f44522eafe4137b0f35c9ce3072931a788a21ee40a2ed279e817d3cc16ed21e", size = 1562764 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/72/99/cafef234114a3b6d9f3aaed0723b437c40c57bdb7b3e4c3a575bc4890052/pytest-9.0.0-py3-none-any.whl", hash = "sha256:e5ccdf10b0bac554970ee88fc1a4ad0ee5d221f8ef22321f9b7e4584e19d7f96", size = 373364 },
+]
+
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"
@@ -690,39 +735,15 @@ wheels = [
 ]
 
 [[package]]
-name = "ruamel-yaml"
-version = "0.17.40"
+name = "scotchpy64"
+version = "1.0.1"
 source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "ruamel-yaml-clib", marker = "python_full_version < '3.13' and platform_python_implementation == 'CPython'" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/d1/d6/eb2833ccba5ea36f8f4de4bcfa0d1a91eb618f832d430b70e3086821f251/ruamel.yaml-0.17.40.tar.gz", hash = "sha256:6024b986f06765d482b5b07e086cc4b4cd05dd22ddcbc758fa23d54873cf313d", size = 137672 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/35/79/5e2cffa1c77432f11cd93a5351f30732c997a239d3a3090856a72d6d8ba7/ruamel.yaml-0.17.40-py3-none-any.whl", hash = "sha256:b16b6c3816dff0a93dca12acf5e70afd089fa5acb80604afd1ffa8b465b7722c", size = 113666 },
-]
-
-[[package]]
-name = "ruamel-yaml-clib"
-version = "0.2.8"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/46/ab/bab9eb1566cd16f060b54055dd39cf6a34bfa0240c53a7218c43e974295b/ruamel.yaml.clib-0.2.8.tar.gz", hash = "sha256:beb2e0404003de9a4cab9753a8805a8fe9320ee6673136ed7f04255fe60bb512", size = 213824 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/b1/15/971b385c098e8d0d170893f5ba558452bb7b776a0c90658b8f4dd0e3382b/ruamel.yaml.clib-0.2.8-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:bef08cd86169d9eafb3ccb0a39edb11d8e25f3dae2b28f5c52fd997521133069", size = 148870 },
-    { url = "https://files.pythonhosted.org/packages/01/b0/4ddef56e9f703d7909febc3a421d709a3482cda25826816ec595b73e3847/ruamel.yaml.clib-0.2.8-cp311-cp311-macosx_13_0_arm64.whl", hash = "sha256:b16420e621d26fdfa949a8b4b47ade8810c56002f5389970db4ddda51dbff248", size = 134475 },
-    { url = "https://files.pythonhosted.org/packages/a4/f7/22d6b620ed895a05d40802d8281eff924dc6190f682d933d4efff60db3b5/ruamel.yaml.clib-0.2.8-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_24_x86_64.whl", hash = "sha256:25c515e350e5b739842fc3228d662413ef28f295791af5e5110b543cf0b57d9b", size = 544020 },
-    { url = "https://files.pythonhosted.org/packages/7c/e4/0d19d65e340f93df1c47f323d95fa4b256bb28320290f5fddef90837853a/ruamel.yaml.clib-0.2.8-cp311-cp311-manylinux_2_24_aarch64.whl", hash = "sha256:1707814f0d9791df063f8c19bb51b0d1278b8e9a2353abbb676c2f685dee6afe", size = 642643 },
-    { url = "https://files.pythonhosted.org/packages/c9/ff/f781eb5e2ae011e586d5426e2086a011cf1e0f59704a6cad1387975c5a62/ruamel.yaml.clib-0.2.8-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:46d378daaac94f454b3a0e3d8d78cafd78a026b1d71443f4966c696b48a6d899", size = 695832 },
-    { url = "https://files.pythonhosted.org/packages/e3/41/f62e67ac651358b8f0d60cfb12ab2daf99b1b69eeaa188d0cec809d943a6/ruamel.yaml.clib-0.2.8-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:09b055c05697b38ecacb7ac50bdab2240bfca1a0c4872b0fd309bb07dc9aa3a9", size = 730923 },
-    { url = "https://files.pythonhosted.org/packages/9f/f0/19ab8acbf983cd1b37f47d27ceb8b10a738d60d36316a54bad57e0d73fbb/ruamel.yaml.clib-0.2.8-cp311-cp311-win32.whl", hash = "sha256:53a300ed9cea38cf5a2a9b069058137c2ca1ce658a874b79baceb8f892f915a7", size = 99999 },
-    { url = "https://files.pythonhosted.org/packages/ec/54/d8a795997921d87224c65d44499ca595a833093fb215b133f920c1062956/ruamel.yaml.clib-0.2.8-cp311-cp311-win_amd64.whl", hash = "sha256:c2a72e9109ea74e511e29032f3b670835f8a59bbdc9ce692c5b4ed91ccf1eedb", size = 118008 },
-    { url = "https://files.pythonhosted.org/packages/7a/a2/eb5e9d088cb9d15c24d956944c09dca0a89108ad6e2e913c099ef36e3f0d/ruamel.yaml.clib-0.2.8-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:ebc06178e8821efc9692ea7544aa5644217358490145629914d8020042c24aa1", size = 144636 },
-    { url = "https://files.pythonhosted.org/packages/66/98/8de4f22bbfd9135deb3422e96d450c4bc0a57d38c25976119307d2efe0aa/ruamel.yaml.clib-0.2.8-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:edaef1c1200c4b4cb914583150dcaa3bc30e592e907c01117c08b13a07255ec2", size = 135684 },
-    { url = "https://files.pythonhosted.org/packages/30/d3/5fe978cd01a61c12efd24d65fa68c6f28f28c8073a06cf11db3a854390ca/ruamel.yaml.clib-0.2.8-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d176b57452ab5b7028ac47e7b3cf644bcfdc8cacfecf7e71759f7f51a59e5c92", size = 734571 },
-    { url = "https://files.pythonhosted.org/packages/55/b3/e2531a050758b717c969cbf76c103b75d8a01e11af931b94ba656117fbe9/ruamel.yaml.clib-0.2.8-cp312-cp312-manylinux_2_24_aarch64.whl", hash = "sha256:1dc67314e7e1086c9fdf2680b7b6c2be1c0d8e3a8279f2e993ca2a7545fecf62", size = 643946 },
-    { url = "https://files.pythonhosted.org/packages/0d/aa/06db7ca0995b513538402e11280282c615b5ae5f09eb820460d35fb69715/ruamel.yaml.clib-0.2.8-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:3213ece08ea033eb159ac52ae052a4899b56ecc124bb80020d9bbceeb50258e9", size = 692169 },
-    { url = "https://files.pythonhosted.org/packages/27/38/4cf4d482b84ecdf51efae6635cc5483a83cf5ca9d9c13e205a750e251696/ruamel.yaml.clib-0.2.8-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:aab7fd643f71d7946f2ee58cc88c9b7bfc97debd71dcc93e03e2d174628e7e2d", size = 740325 },
-    { url = "https://files.pythonhosted.org/packages/6f/67/c62c6eea53a4feb042727a3d6c18f50dc99683c2b199c06bd2a9e3db8e22/ruamel.yaml.clib-0.2.8-cp312-cp312-win32.whl", hash = "sha256:5c365d91c88390c8d0a8545df0b5857172824b1c604e867161e6b3d59a827eaa", size = 98639 },
-    { url = "https://files.pythonhosted.org/packages/10/d2/52a3d810d0b5b3720725c0504a27b3fced7b6f310fe928f7019d79387bc1/ruamel.yaml.clib-0.2.8-cp312-cp312-win_amd64.whl", hash = "sha256:1758ce7d8e1a29d23de54a16ae867abd370f01b5a69e1a3ba75223eaa3ca1a1b", size = 115305 },
+    { url = "https://files.pythonhosted.org/packages/bd/b3/5e171536989e4fa154b22a480f256dd5fada8d9ce47614c77b747adb250c/scotchpy64-1.0.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:78b17354dafe8d99471ff01fc54ce2d26fca2061ea27db95c8f8bb73813fc579", size = 2123702 },
+    { url = "https://files.pythonhosted.org/packages/5f/be/3b344a1beba0abd011d0193f151d510a286787d3910c3ae3d7940dd70931/scotchpy64-1.0.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:87705f9c4105740306d5048c7c44a51e62b2f03c9a95fa7610d6a3386989e8c5", size = 2123702 },
+    { url = "https://files.pythonhosted.org/packages/3b/ea/12491425a0148b8b0ccbe3694a24d2b498a99c21084c8cec09e3e94dd78f/scotchpy64-1.0.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e2ea7f50b3a7797e0cd71537b2c072114dd9bb6ca48dc5f826d555370eb3150a", size = 2123703 },
+    { url = "https://files.pythonhosted.org/packages/82/cd/b2f1f81a7adbec6e4c70cea0b518f89c5b2650e2fdc669aed380bcbb2bb4/scotchpy64-1.0.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3cd52831289d592cca3ca9ec4093467855f3c77d0af0966ba182fe49ceadc4d2", size = 2123703 },
+    { url = "https://files.pythonhosted.org/packages/e4/f0/25613cd7f193b23dd5f886555b04012816fcebe0589c39d51e75f9fcb687/scotchpy64-1.0.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e0c8c77a8aa5cc75d59b40cfb47f02814c5bdfb605fc82deaa57cac4f486f125", size = 2123708 },
 ]
 
 [[package]]
@@ -748,6 +769,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050 },
 ]
 
+[[package]]
+name = "tqdm"
+version = "4.67.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a8/4b/29b4ef32e036bb34e4ab51796dd745cdba7ed47ad142a9f4a1eb8e0c744d/tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2", size = 169737 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540 },
+]
+
 [[package]]
 name = "tzdata"
 version = "2025.2"
diff --git a/validate_gemm.py b/validate_gemm.py
new file mode 100644
index 00000000..08862490
--- /dev/null
+++ b/validate_gemm.py
@@ -0,0 +1,460 @@
+import os
+import csv
+import itertools
+import argparse
+import logging
+from collections import OrderedDict
+import re
+from validation_scripts.helpers import (
+    run_cutlass_profiler,
+    CUTLASS_ENV_VAR,
+    run_ncu_profiler,
+    load_metrics_spec,
+    parse_int_list,            
+    read_cutlass_results,      
+    ensure_cutlass_profiler,   
+    get_metric_labels,         
+)
+
+from run_perf import import_config
+
+logger = logging.getLogger(__name__)
+if not logger.handlers:
+    h = logging.StreamHandler()
+    h.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s"))
+    logger.addHandler(h)
+logger.setLevel(logging.INFO)
+
+
+def parse_llm_shapes_file(filepath):
+    """
+    Parse llm_shapes.txt into a list of dicts, each containing model name and OrderedDict of operations.
+    Each operation maps to a tuple (m,k,n) or (b,m,k,n).
+    
+    Returns: List[dict] where each dict has 'model' (str) and 'shapes' (OrderedDict[str, tuple])
+    """
+    models = []
+    with open(filepath, 'r') as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            # Expected format: "ModelName: OrderedDict([('op1', (dims)), ...])"
+            match = re.match(r"^([^:]+):\s*OrderedDict\((.+)\)$", line)
+            if not match:
+                logger.warning(f"Skipping malformed line: {line}")
+                continue
+            model_name = match.group(1).strip()
+            shapes_str = match.group(2).strip()
+            
+            # Parse the OrderedDict content: [('key', (tuple)), ...]
+            shapes = OrderedDict()
+            # Remove outer brackets
+            shapes_str = shapes_str.strip('[]')
+            # Split by '), (' pattern to separate entries
+            entries = re.findall(r"\('([^']+)',\s*\(([^)]+)\)\)", shapes_str)
+            for op_name, dims_str in entries:
+                dims = tuple(map(int, dims_str.split(',')))
+                shapes[op_name] = dims
+            
+            models.append({'model': model_name, 'shapes': shapes})
+            logger.info(f"Parsed model {model_name} with {len(shapes)} operations")
+    
+    return models
+
+
+def extract_gemm_shapes(models_data):
+    """
+    Extract GEMM shapes from parsed models data.
+    Returns: List[dict] with keys: 'model', 'op_name', 'm', 'k', 'n', 'batch' (optional)
+    """
+    gemm_shapes = []
+    for model_info in models_data:
+        model_name = model_info['model']
+        for op_name, dims in model_info['shapes'].items():
+            if len(dims) == 3:
+                m, k, n = dims
+                gemm_shapes.append({
+                    'model': model_name,
+                    'op_name': op_name,
+                    'm': m, 'k': k, 'n': n
+                })
+            elif len(dims) == 4:
+                b, m, k, n = dims
+                # For batched GEMMs, we profile with the actual m dimension
+                gemm_shapes.append({
+                    'model': model_name,
+                    'op_name': op_name,
+                    'm': m, 'k': k, 'n': n,
+                    'batch': b
+                })
+            else:
+                logger.warning(f"Skipping {model_name}:{op_name} with unsupported dims: {dims}")
+    
+    return gemm_shapes
+
+
+def _safe_err(ncu_val, pred_val):
+    if ncu_val is None or pred_val is None:
+        return "", ""
+    try:
+        abs_err = pred_val - ncu_val
+        pct_err = abs_err / ncu_val * 100 if ncu_val != 0 else ""
+        return f"{abs_err:<.3f}", f"{pct_err:<.2f}"
+    except Exception:
+        return "", ""
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Search CUTLASS kernels across GEMM and CTA sizes."
+    )
+    # parser.add_argument("--cutlass-path", required=True, help="Path to CUTLASS profiler binary")
+    parser.add_argument("--m-list", default="256,1024,4096")
+    parser.add_argument("--n-list", default="256,1024,4096")
+    parser.add_argument("--k-list", default="256,1024,4096")
+    parser.add_argument("--shapes-file", default="validation_scripts/llm_shapes.txt", 
+                        help="Path to file containing LLM GEMM shapes (e.g., llm_shapes.txt)")
+    parser.add_argument("--cta-m", default="64,128,256")
+    parser.add_argument("--cta-n", default="64,128,256")
+    parser.add_argument("--cta-k", default="32,64")
+    parser.add_argument("--cutlass-output", default="validation_scripts/kernel_search_results.csv")
+    parser.add_argument("--ncu-output", default="validation_scripts/cta_sweep/fastest.csv")
+    parser.add_argument("--ncu-metrics", default="validation_scripts/metrics.yaml")
+    parser.add_argument("--log-level", default="INFO")
+    parser.add_argument("--cutlass-path", default=None, help="Path to cutlass_profiler binary; overrides env")
+    parser.add_argument("--deepflow-compare", action="store_true",
+                        help="Enable DeepFlow model comparison CSV output")
+    parser.add_argument("--deepflow-output", default="validation_scripts/deepflow_comparison.csv",
+                        help="Output CSV for DeepFlow vs NCU errors")
+    parser.add_argument("--hw-config", default="configs/hardware-config/a100_80GB_no_parallelism.yaml",
+                        help="Path to A100 hardware config")
+    parser.add_argument("--sw-config", default="configs/model-config/GEMM.yaml",
+                        help="Path to GEMM software config")
+    args = parser.parse_args()
+
+    logger.setLevel(getattr(logging, args.log_level.upper(), logging.INFO))
+
+    # Ensure CUTLASS profiler path (no hardcoded default)
+    ensure_cutlass_profiler(args.cutlass_path)  # added
+    logger.info(f"{CUTLASS_ENV_VAR}={os.environ.get(CUTLASS_ENV_VAR, '')}")
+
+    # Load NCU metric spec (labels, units) from file or fallback
+    if args.ncu_metrics and os.path.isfile(args.ncu_metrics):
+        try:
+            ncu_metric_spec = load_metrics_spec(args.ncu_metrics)
+            logger.info(f"Loaded NCU metrics spec from {args.ncu_metrics} ({len(ncu_metric_spec)} metrics)")
+        except Exception as e:
+            logger.warning(f"Failed to load metrics spec from {args.ncu_metrics}: {e}.")
+            exit(1)
+    else:
+        logger.info("NCU metrics spec file not found.")
+        exit(1)
+
+    # Determine GEMM shapes source
+    if args.shapes_file:
+        if not os.path.isfile(args.shapes_file):
+            logger.error(f"Shapes file not found: {args.shapes_file}")
+            exit(1)
+        logger.info(f"Loading GEMM shapes from {args.shapes_file}")
+        models_data = parse_llm_shapes_file(args.shapes_file)
+        gemm_configs = extract_gemm_shapes(models_data)
+        logger.info(f"Extracted {len(gemm_configs)} GEMM shapes from {len(models_data)} models")
+    else:
+        # Use m/n/k lists
+        m_vals = parse_int_list(args.m_list)   
+        n_vals = parse_int_list(args.n_list)   
+        k_vals = parse_int_list(args.k_list)
+        gemm_configs = [
+            {'m': m, 'k': k, 'n': n, 'model': 'synthetic', 'op_name': f'm{m}_k{k}_n{n}'}
+            for m in m_vals for k in k_vals for n in n_vals
+        ]
+        logger.info(f"Using {len(gemm_configs)} synthetic GEMM shapes from m/k/n lists")
+
+    cta_m_vals = parse_int_list(args.cta_m)
+    cta_n_vals = parse_int_list(args.cta_n)
+    cta_k_vals = parse_int_list(args.cta_k)
+
+    cta_tuples = list(itertools.product(cta_m_vals, cta_n_vals, cta_k_vals))
+    logger.info(f"CTA tuples ({len(cta_tuples)}): {cta_tuples}")
+
+    out_path = args.cutlass_output
+    ncu_out_path = args.ncu_output
+    ncu_exists = os.path.isfile(ncu_out_path)          
+    perform_ncu = not ncu_exists                       
+    ncu_write_header = not ncu_exists
+
+    # DeepFlow TimeCalculation initialization
+    if args.deepflow_compare:
+        try:
+            deepflow_tc = import_config(args.hw_config, args.sw_config, "GEMM")
+            logger.info("DeepFlow TimeCalculation initialized.")
+        except Exception as e:
+            logger.warning(f"Failed to initialize DeepFlow TimeCalculation: {e}")
+            deepflow_tc = None
+    else:
+        deepflow_tc = None
+
+    # Accumulate NCU results for DeepFlow comparison
+    ncu_rows = []  # each: dict(m,k,n, metrics_dict)
+
+    # If NCU output already exists, load it and skip further NCU profiling
+    if ncu_exists:
+        logger.info(f"{ncu_out_path} exists. Skipping NCU profiling and loading existing metrics.")
+        try:
+            with open(ncu_out_path, "r", newline="") as f_in:
+                reader = csv.reader(f_in)
+                header = next(reader, None)
+                if header is None:
+                    logger.warning(f"{ncu_out_path} is empty. No NCU metrics to load.")
+                else:
+                    # First 7 columns are fixed: m,k,n,cta_m,cta_k,cta_n,stages
+                    metric_labels = header[7:]
+                    for row in reader:
+                        if len(row) < 7:
+                            continue
+                        try:
+                            m = int(row[0])
+                            n = int(row[1])
+                            k = int(row[2])
+                            metrics = {
+                                "cta_m": int(row[3]),
+                                "cta_n": int(row[4]),
+                                "cta_k": int(row[5]),
+                                "stages": row[6],
+                            }
+                            # Parse remaining metrics as float when possible
+                            for lab, val in zip(metric_labels, row[7:]):
+                                try:
+                                    metrics[lab] = float(val)
+                                except ValueError:
+                                    metrics[lab] = val
+                            # Provide canonical aliases if present
+                            if "latency_ms" not in metrics and "latency" in metrics:
+                                metrics["latency_ms"] = metrics["latency"]
+                            ncu_rows.append({"m": m, "k": k, "n": n, "metrics": metrics})
+                        except Exception as e:
+                            logger.debug(f"Skipping malformed NCU row: {row} ({e})")
+            logger.info(f"Loaded {len(ncu_rows)} NCU metric rows from {ncu_out_path}")
+        except Exception as e:
+            logger.warning(f"Failed to read existing NCU metrics file {ncu_out_path}: {e}")
+
+    # If CUTLASS results already exist, reuse them; only run NCU if perform_ncu
+    if os.path.isfile(out_path): # broken
+        logger.info(f"{out_path} exists. Skipping CUTLASS profiling and using existing results.")
+        existing = read_cutlass_results(out_path)
+        if not existing:
+            logger.warning(f"No valid entries found in {out_path}. Nothing to do.")
+            # If DeepFlow requested and we have NCU rows, still do comparison
+            if deepflow_tc and ncu_rows:
+                pass  # fall-through to comparison block below
+            else:
+                return
+        else:
+            if perform_ncu:
+                with open(ncu_out_path, "a", newline="") as f_ncu:
+                    writer_ncu = csv.writer(f_ncu)
+                    if ncu_write_header:
+                        metric_labels = get_metric_labels(ncu_metric_spec)
+                        writer_ncu.writerow(["m", "k", "n", "cta_m", "cta_k", "cta_n", "stages"] + metric_labels)
+                    for m, k, n, best_kernel in existing:
+                        try:
+                            logger.info(f"Running NCU for m={m} k={k} n={n} kernel={best_kernel.name}")
+                            ncu_metrics = run_ncu_profiler(
+                                m=m, k=k, n=n, kernel_name=best_kernel.name,
+                                metric_spec=ncu_metric_spec
+                            )
+                            ordered_values = [ncu_metrics.get(spec.label, "") for spec in ncu_metric_spec.values()]
+                            writer_ncu.writerow([m, k, n, best_kernel.cta_m, best_kernel.cta_k,
+                                                 best_kernel.cta_n, best_kernel.stages] + ordered_values)
+                            ncu_rows.append({"m": m, "k": k, "n": n, "metrics": ncu_metrics})
+                        except Exception as e:
+                            logger.warning(f"NCU profiling failed for m={m} k={k} n={n} kernel={best_kernel.name}: {e}")
+                logger.info(f"NCU metrics written to {ncu_out_path}")
+            else:
+                logger.info("NCU profiling skipped (existing NCU metrics file present).")
+        # DeepFlow comparison with existing (and/or loaded) NCU rows
+        if deepflow_tc and ncu_rows:
+            try:
+                write_header_df = not os.path.isfile(args.deepflow_output)
+                # if deepflow_output already exists, clear it and rewrite
+                if not write_header_df:
+                    os.remove(args.deepflow_output)
+                    write_header_df = True
+                with open(args.deepflow_output, "a", newline="") as f_df:
+                    w_df = csv.writer(f_df)
+                    if write_header_df:
+                        w_df.writerow([
+                            "m","k","n","cta_m","cta_k","cta_n","l1_m","l1_k","l1_n",
+                            "ncu_latency_us","df_latency_us","latency_abs_err","latency_pct_err",
+                            "ncu_l2_read","df_l2_read","l2_read_pct_err",
+                            "ncu_l2_write","df_l2_write","l2_write_pct_err",
+                            "ncu_l2_total","df_l2_total","l2_total_pct_err",
+                            "ncu_dram_read","df_dram_read","dram_read_pct_err",
+                            "ncu_dram_write","df_dram_write","dram_write_pct_err",
+                            "ncu_dram_total","df_dram_total","dram_total_pct_err"
+                        ])
+                    for row in ncu_rows:
+                        m,k,n = row["m"], row["k"], row["n"]
+                        metrics = row["metrics"]
+                        ncu_latency = metrics.get("latency_ms")
+                        ncu_latency *= 1e3  # ms → us
+                        ncu_l2_read = metrics.get("l2_read")
+                        ncu_l2_write = metrics.get("l2_write")
+                        ncu_l2_total = ncu_l2_read + ncu_l2_write
+                        ncu_dram_read = metrics.get("dram_read")
+                        ncu_dram_write = metrics.get("dram_write")
+                        ncu_dram_total = ncu_dram_read + ncu_dram_write
+                        ncu_cta_m = metrics.get("cta_m")
+                        ncu_cta_n = metrics.get("cta_n")
+                        ncu_cta_k = metrics.get("cta_k")
+                        try:
+                            df_latency, _, df_tile_dims, _, best_rw_access = deepflow_tc.getGEMMTime(m, k, n, name=f"m{m}_k{k}_n{n}")
+                        except Exception as e:
+                            logger.warning(f"DeepFlow getGEMMTime failed for {m},{k},{n}: {e}")
+                            continue
+                        # logger.info(f"({m},{k},{n}): {df_tile_dims}")
+                        df_latency *= 1e6  # s → us
+                        l1_m, l1_k, l1_n = df_tile_dims[1]
+                        df_l2_read = best_rw_access.reads[2]
+                        df_l2_write = best_rw_access.writes[2]
+                        df_l2_total = df_l2_read + df_l2_write
+                        df_dram_read = best_rw_access.reads[3]
+                        df_dram_write = best_rw_access.writes[3]
+                        df_dram_total = df_dram_read + df_dram_write
+                        lat_abs, lat_pct = _safe_err(ncu_latency, df_latency)
+                        _, l2_read_pct = _safe_err(ncu_l2_read, df_l2_read)
+                        _, l2_write_pct = _safe_err(ncu_l2_write, df_l2_write)
+                        _, l2_total_pct = _safe_err(ncu_l2_total, df_l2_total)
+                        _, dram_read_pct = _safe_err(ncu_dram_read, df_dram_read)
+                        _, dram_write_pct = _safe_err(ncu_dram_write, df_dram_write)
+                        _, dram_total_pct = _safe_err(ncu_dram_total, df_dram_total)
+                        w_df.writerow([
+                            m,k,n,ncu_cta_m,ncu_cta_k,ncu_cta_n,l1_m,l1_k,l1_n,
+                            f"{ncu_latency:.3f}", f"{df_latency:.3f}", lat_abs, lat_pct,
+                            ncu_l2_read, df_l2_read, l2_read_pct,
+                            ncu_l2_write, df_l2_write, l2_write_pct,
+                            ncu_l2_total, df_l2_total, l2_total_pct,
+                            ncu_dram_read, df_dram_read, dram_read_pct,
+                            ncu_dram_write, df_dram_write, dram_write_pct,
+                            ncu_dram_total, df_dram_total, dram_total_pct
+                        ])
+                logger.info(f"DeepFlow comparison written to {args.deepflow_output}")
+            except Exception as e:
+                logger.warning(f"Failed DeepFlow comparison output: {e}")
+        return
+
+    # Otherwise (CUTLASS results absent), run CUTLASS search; only run NCU if perform_ncu
+    write_header = not os.path.isfile(out_path)
+    with open(out_path, "a", newline="") as f_kernel:
+        writer_kernel = csv.writer(f_kernel)
+        if write_header:
+            writer_kernel.writerow([
+                "model","op_name","m","k","n","best_kernel_name","cta_m","cta_k","cta_n",
+                "stages","inst_m","inst_n","inst_k","best_time_ms",
+            ])
+        total_jobs = len(gemm_configs)
+        job_idx = 0
+        for config in gemm_configs:
+            m, k, n = config['m'], config['k'], config['n']
+            model_name = config.get('model', 'unknown')
+            op_name = config.get('op_name', 'unknown')
+            job_idx += 1
+            logger.info(f"[{job_idx}/{total_jobs}] Profiling {model_name}:{op_name} GEMM m={m} k={k} n={n}")
+            best_kernel, best_time = run_cutlass_profiler(m=m, k=k, n=n, cta_tuples=cta_tuples)
+            if best_kernel is None:
+                logger.warning(f"No kernel found for {model_name}:{op_name} m={m} k={k} n={n}")
+                continue
+            writer_kernel.writerow([
+                model_name,op_name,m,k,n,best_kernel.name,best_kernel.cta_m,best_kernel.cta_k,
+                best_kernel.cta_n,best_kernel.stages,best_kernel.inst_m,
+                best_kernel.inst_n,best_kernel.inst_k,f"{best_time:.6f}",
+            ])
+            logger.info(f"Best kernel: {best_kernel.name} CTA=({best_kernel.cta_m},{best_kernel.cta_k},{best_kernel.cta_n}) time={best_time:.4f} ms")
+            if perform_ncu:
+                # Open (append) NCU file lazily only if needed
+                if ncu_write_header:
+                    with open(ncu_out_path, "a", newline="") as f_ncu:
+                        writer_ncu = csv.writer(f_ncu)
+                        metric_labels = get_metric_labels(ncu_metric_spec)
+                        writer_ncu.writerow(["m","k","n","cta_m","cta_k","cta_n","stages"] + metric_labels)
+                    ncu_write_header = False
+                try:
+                    ncu_metrics = run_ncu_profiler(
+                        m=m, k=k, n=n, kernel_name=best_kernel.name,
+                        metric_spec=ncu_metric_spec
+                    )
+                    ordered_values = [ncu_metrics.get(spec.label, "") for spec in ncu_metric_spec.values()]
+                    with open(ncu_out_path, "a", newline="") as f_ncu:
+                        writer_ncu = csv.writer(f_ncu)
+                        writer_ncu.writerow([m, k, n, best_kernel.cta_m, best_kernel.cta_k,
+                                             best_kernel.cta_n, best_kernel.stages] + ordered_values)
+                    ncu_rows.append({"m": m, "k": k, "n": n, "metrics": ncu_metrics})
+                    logger.info(f"NCU metrics collected for kernel {best_kernel.name}: latency={ncu_metrics.get('latency_ms', 'N/A')}")
+                except Exception as e:
+                    logger.warning(f"NCU profiling failed for kernel {best_kernel.name}: {e}")
+
+    logger.info(f"Kernel results written to {out_path}")
+    if perform_ncu:
+        logger.info(f"NCU metrics written to {ncu_out_path}")
+    else:
+        logger.info(f"Skipped writing NCU metrics (existing file retained: {ncu_out_path})")
+
+    # DeepFlow comparison after new profiling (only if NCU metrics available or loaded)
+    if deepflow_tc and ncu_rows:
+        try:
+            write_header_df = not os.path.isfile(args.deepflow_output)
+            with open(args.deepflow_output, "a", newline="") as f_df:
+                w_df = csv.writer(f_df)
+                if write_header_df:
+                    w_df.writerow([
+                        "m","k","n","cta_m","cta_k","cta_n","l1_m","l1_k","l1_n",
+                        "ncu_latency_ms","df_latency_ms","latency_abs_err","latency_pct_err",
+                        "ncu_l2_read","df_l2_read","l2_read_pct_err",
+                        "ncu_l2_write","df_l2_write","l2_write_pct_err",
+                        "ncu_dram_read","df_dram_read","dram_read_pct_err",
+                        "ncu_dram_write","df_dram_write","dram_write_pct_err"
+                    ])
+                for row in ncu_rows:
+                    m,k,n = row["m"], row["k"], row["n"]
+                    metrics = row["metrics"]
+                    ncu_latency = metrics.get("latency_ms")
+                    ncu_l2_read = metrics.get("l2_read")
+                    ncu_l2_write = metrics.get("l2_write")
+                    ncu_dram_read = metrics.get("dram_read")
+                    ncu_dram_write = metrics.get("dram_write")
+                    ncu_cta_m = metrics.get("cta_m")
+                    ncu_cta_n = metrics.get("cta_n")
+                    ncu_cta_k = metrics.get("cta_k")
+                    try:
+                        df_latency, _, df_tile_dims, _, best_rw_access = deepflow_tc.getGEMMTime(m, k, n)
+                    except Exception as e:
+                        logger.warning(f"DeepFlow getGEMMTime failed for {m},{k},{n}: {e}")
+                        continue
+                    l1_m, l1_k, l1_n = df_tile_dims[1]
+                    df_l2_read = best_rw_access.reads[2]
+                    df_l2_write = best_rw_access.writes[2]
+                    df_dram_read = best_rw_access.reads[3]
+                    df_dram_write = best_rw_access.writes[3]
+                    lat_abs, lat_pct = _safe_err(ncu_latency, df_latency * 1e3)
+                    _, l2_read_pct = _safe_err(ncu_l2_read, df_l2_read)
+                    _, l2_write_pct = _safe_err(ncu_l2_write, df_l2_write)
+                    _, dram_read_pct = _safe_err(ncu_dram_read, df_dram_read)
+                    _, dram_write_pct = _safe_err(ncu_dram_write, df_dram_write)
+                    w_df.writerow([
+                        m,k,n,ncu_cta_m,ncu_cta_k,ncu_cta_n,l1_m,l1_k,l1_n,
+                        ncu_latency, df_latency, lat_abs, lat_pct,
+                        ncu_l2_read, df_l2_read, l2_read_pct,
+                        ncu_l2_write, df_l2_write, l2_write_pct,
+                        ncu_dram_read, df_dram_read, dram_read_pct,
+                        ncu_dram_write, df_dram_write, dram_write_pct
+                    ])
+            logger.info(f"DeepFlow comparison written to {args.deepflow_output}")
+        except Exception as e:
+            logger.warning(f"Failed DeepFlow comparison output: {e}")
+
+
+if __name__ == "__main__":
+    os.environ[CUTLASS_ENV_VAR] = (
+        "/u1/ee/nanoproj/users/yaoe888/cutlass/build/tools/profiler/cutlass_profiler"
+    )
+    main()
diff --git a/validation_scripts/__init__.py b/validation_scripts/__init__.py
new file mode 100644
index 00000000..36426389
--- /dev/null
+++ b/validation_scripts/__init__.py
@@ -0,0 +1 @@
+"""Validation helpers for DeepFlow."""
diff --git a/validation_scripts/imec_data/A100_inf.csv b/validation_scripts/imec_data/A100_inf.csv
new file mode 100644
index 00000000..99b20f95
--- /dev/null
+++ b/validation_scripts/imec_data/A100_inf.csv
@@ -0,0 +1,12 @@
+device,model,TP,actual
+A100,Llama 2-7B,1,2.19
+A100,Llama 2-7B,2,1.544
+A100,Llama 2-7B,4,1.28
+A100,Llama 2-7B,8,1.187
+A100,Llama 2-13B,1,3.884
+A100,Llama 2-13B,2,2.499
+A100,Llama 2-13B,4,1.894
+A100,Llama 2-13B,8,1.693
+A100,Llama 2-70B,2,10.5
+A100,Llama 2-70B,4,6.403
+A100,Llama 2-70B,8,4.735
\ No newline at end of file
diff --git a/validation_scripts/imec_data/A100_train.csv b/validation_scripts/imec_data/A100_train.csv
new file mode 100644
index 00000000..f8428204
--- /dev/null
+++ b/validation_scripts/imec_data/A100_train.csv
@@ -0,0 +1,12 @@
+device,model,batch,mb,dp,tp,pp,cp,tp_sp,recomputation,actual
+A100_korthi,GPT 22B,4,1,1,8,1,1,False,full,1.4
+A100_korthi,GPT 175B,64,64,1,8,8,1,False,full,16.9
+A100_korthi,GPT 530B,280,280,1,8,35,1,False,full,46.8
+A100_korthi,GPT 1T,512,512,1,8,64,1,False,full,87.9
+A100_korthi,GPT 22B,4,1,1,8,1,1,True,selective,1.1
+A100_korthi,GPT 175B,64,64,1,8,8,1,True,selective,12.9
+A100_korthi,GPT 530B,280,280,1,8,35,1,True,selective,35.5
+A100_korthi,GPT 1T,512,512,1,8,64,1,True,selective,69.1
+A100_selene,GPT 310B,2160,144,15,8,16,1,False,full,34.1
+A100_selene,GPT 530B,2520,280,9,8,35,1,False,full,51.2
+A100_selene,GPT 1T,3072,512,6,8,64,1,False,full,100.7
diff --git a/validation_scripts/imec_data/H100_inf.csv b/validation_scripts/imec_data/H100_inf.csv
new file mode 100644
index 00000000..9db0e2c0
--- /dev/null
+++ b/validation_scripts/imec_data/H100_inf.csv
@@ -0,0 +1,12 @@
+device,model,TP,actual
+H100,Llama 2-7B,1,1.44
+H100,Llama 2-7B,2,1.143
+H100,Llama 2-7B,4,0.924
+H100,Llama 2-7B,8,0.828
+H100,Llama 2-13B,1,2.396
+H100,Llama 2-13B,2,1.717
+H100,Llama 2-13B,4,1.431
+H100,Llama 2-13B,8,1.201
+H100,Llama 2-70B,2,6.267
+H100,Llama 2-70B,4,4.1226
+H100,Llama 2-70B,8,3.202
\ No newline at end of file
diff --git a/validation_scripts/imec_data/uci_train.csv b/validation_scripts/imec_data/uci_train.csv
new file mode 100644
index 00000000..7abf0763
--- /dev/null
+++ b/validation_scripts/imec_data/uci_train.csv
@@ -0,0 +1,15 @@
+variant,TP,PP,CP,DP,Status,Error,Avg_Step_Time_s,Total_Throughput_tokens_s,Throughput_per_GPU_tokens_s
+DDP,1,1,1,4,SUCCESS,,36.9333,14195.52,3548.88
+DDP,4,1,1,1,SUCCESS,,70.3667,7450.8,1862.7
+DDP,1,4,1,1,SUCCESS,,40.5667,12924.11,3231.03
+DDP,1,1,4,1,SUCCESS,,107.0,4899.89,1224.97
+DDP,2,2,1,1,SUCCESS,,42.9722,12200.63,3050.16
+DDP,2,1,2,1,SUCCESS,,63.2944,8283.32,2070.83
+DDP,1,2,2,1,SUCCESS,,43.9556,11927.68,2981.92
+DDP,2,1,1,2,SUCCESS,,41.7444,12559.47,3139.87
+DDP,1,2,1,2,SUCCESS,,38.2194,13717.83,3429.46
+DDP,1,1,2,2,SUCCESS,,43.025,12185.66,3046.41
+FSDP,1,1,1,4,SUCCESS,,56.2639,9318.37,2329.59
+FSDP,2,1,1,2,SUCCESS,,74.1639,7069.32,1767.33
+FSDP,1,2,1,2,INSUFFICIENT_DATA,AssertionError: FSDP does not support pipeline parallelism,,,
+FSDP,1,1,2,2,SUCCESS,,104.2222,5030.48,1257.62
diff --git a/validation_scripts/koyeb.py b/validation_scripts/koyeb.py
index 107a8446..d1f3a000 100644
--- a/validation_scripts/koyeb.py
+++ b/validation_scripts/koyeb.py
@@ -1,16 +1,17 @@
 #!/usr/bin/env python3
 import os
-import sys
-import re
-import tempfile
-import subprocess
-import yaml
 import math
-from typing import Dict, Tuple, List
+from typing import Dict, Tuple, List, Any
 import seaborn as sns
 
 import matplotlib.pyplot as plt
 
+from .validation_helpers import (
+  ValidationSpec,
+  parse_decode_time,
+  run_validation_suite,
+)
+
 
 # Llama-3.1-8B Instruct - A100 SXM (from Koyeb matrix, using actual Input/Output tokens)
 # https://www.koyeb.com/docs/hardware/gpu-benchmarks
@@ -70,64 +71,9 @@
 
 PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 RUN_PERF = os.path.join(PROJECT_ROOT, "run_perf.py")
-HARDWARE_CONFIG = os.path.join(PROJECT_ROOT, "configs/hardware-config/a100_80GB_no_parallelism.yaml")
-BASE_MODEL_CONFIG = os.path.join(PROJECT_ROOT, "configs/model-config/Llama3.1-8B_inf.yaml")
-
-def _load_yaml(path: str) -> dict:
-  with open(path, "r") as f:
-    return yaml.safe_load(f)
-
-def _write_yaml(path: str, data: dict) -> None:
-  with open(path, "w") as f:
-    yaml.safe_dump(data, f, sort_keys=False)
-
-
-def _update_model_config(base_cfg: dict, *, prefill_len: int, decode_len: int, batch_size: int) -> dict:
-  cfg = dict(base_cfg)  # shallow copy is fine; we reassign primitives
-  model_param = dict(cfg.get("model_param", {}))
-  # seq_len = prefill + decode
-  model_param["seq_len"] = int(prefill_len + decode_len)
-  model_param["decode_len"] = int(decode_len)
-  model_param["batch_size"] = int(batch_size)
-  cfg["model_param"] = model_param
-  return cfg
-
-
-def _run_deepflow_with_temp_model(prefill_len: int, decode_len: int, batch_size: int) -> Tuple[str, float]:
-  base_cfg = _load_yaml(BASE_MODEL_CONFIG)
-  updated_cfg = _update_model_config(base_cfg, prefill_len=prefill_len, decode_len=decode_len, batch_size=batch_size)
-  # if .tmp does not exist, make it
-  if not os.path.exists("./tmp"):
-    os.makedirs("./tmp")
-  fd, tmp_path = tempfile.mkstemp(prefix="koyeb_model_", suffix=".yaml", dir="./tmp")
-  os.close(fd)
-
-  try:
-    _write_yaml(tmp_path, updated_cfg)
-    cmd = [
-      sys.executable,
-      RUN_PERF,
-      "--hardware_config", HARDWARE_CONFIG,
-      "--model_config", tmp_path,
-    ]
-
-    proc = subprocess.run(cmd, cwd=PROJECT_ROOT, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, check=False)
-    output = proc.stdout
-
-    # Parse decode time from line: [prefill] time: Xs, [decode] time: Ys, [total] time: Zs
-    decode_time = None
-    m = re.search(r"\[prefill\]\s*time:\s*[0-9.]+s,\s*\[decode\]\s*time:\s*([0-9.]+)s,\s*\[total\]\s*time:\s*[0-9.]+s", output)
-    if m:
-      decode_time = float(m.group(1))
-
-    if decode_time is None:
-      raise RuntimeError(f"Failed to parse decode time from DeepFlow output: {output}")
-    return output, decode_time
-  finally:
-    try:
-      os.remove(tmp_path)
-    except Exception:
-      pass
+VALIDATION_CONFIG_ROOT = os.path.join(PROJECT_ROOT, "validation_scripts", "validation_configs")
+HARDWARE_CONFIG = os.path.join(VALIDATION_CONFIG_ROOT, "hardware-config/a100_80GB_no_parallelism.yaml")
+BASE_MODEL_CONFIG = os.path.join(VALIDATION_CONFIG_ROOT, "model-config/Llama3.1-8B_inf.yaml")
 
 
 def _collect_points_in_order(data: Dict[str, Dict[Tuple[int, int, int], Tuple[float, float]]]) -> List[Tuple[str, Tuple[int, int, int], Tuple[float, float]]]:
@@ -142,83 +88,173 @@ def _collect_points_in_order(data: Dict[str, Dict[Tuple[int, int, int], Tuple[fl
   return ordered
 
 
-def run_all_and_plot() -> None:
+def run_all_and_plot(*, enable_plot: bool = True, verbose: bool = True) -> Dict[str, Any]:
   experiments = _collect_points_in_order(KOYEB_DATA)
   if not experiments:
-    print("No experiments to run.")
-    return
-
+    if verbose:
+      print("No experiments to run.")
+    return {
+      "labels": [],
+      "expected_tokps": [],
+      "predicted_tokps": [],
+      "percent_errors": [],
+      "avg_error": float("nan"),
+      "avg_abs_error": float("nan"),
+      "plot_path": None,
+    }
+
+  log = print if verbose else (lambda *args, **kwargs: None)
   labels: List[str] = []
   expected_tokps: List[float] = []
   predicted_tokps: List[float] = []
   percent_errors: List[float] = []
+  specs: List[ValidationSpec] = []
 
-  for size, (prefill_len, decode_len, batch_size), (expected_decode_time_s, expected_tokps_agg) in experiments:
+  for idx, (size, (prefill_len, decode_len, batch_size), (expected_decode_time_s, expected_tokps_agg)) in enumerate(experiments):
     label = f"{size} bs={batch_size}"
-    print(f"\n=== Running {label} (prefill={prefill_len}, decode={decode_len}, batch={batch_size}) ===")
-    try:
-      console, decode_time_s = _run_deepflow_with_temp_model(prefill_len, decode_len, batch_size)
-      # Aggregate tokens per second using overall definition: total tokens / decode_time
-      total_tokens_agg = int(decode_len) * int(batch_size)  # dp=1 for no-parallelism hardware config
-      predicted_tokps_agg = float(total_tokens_agg) / float(decode_time_s) if decode_time_s > 0 else 0.0
-
-      # Calculate percentage error
-      if expected_tokps_agg > 0:
-        pct_error = (predicted_tokps_agg - expected_tokps_agg) / expected_tokps_agg * 100.0
-      else:
-        pct_error = float("nan")
-
-      print(f"Predicted agg tok/s = {predicted_tokps_agg:.2f} (decode_time={decode_time_s:.2f}s)")
-      print(f"Expected agg tok/s  = {expected_tokps_agg:.2f}")
-      print(f"Error = {pct_error:+.2f}%")
-    except Exception as e:
-      print(f"ERROR running DeepFlow for {label}: {e}")
+    metadata = {
+      "size": size,
+      "prefill_len": int(prefill_len),
+      "decode_len": int(decode_len),
+      "batch_size": int(batch_size),
+      "expected_tokps": float(expected_tokps_agg),
+      "expected_decode_time_s": float(expected_decode_time_s),
+    }
+    model_overrides = {
+      "model_param": {
+        "seq_len": int(prefill_len + decode_len),
+        "decode_len": int(decode_len),
+        "global_batch_size": int(batch_size),
+      },
+    }
+    specs.append(
+      ValidationSpec(
+        label=label,
+        model_overrides=model_overrides,
+        hardware_overrides=None,
+        metadata=metadata,
+        order=idx,
+      )
+    )
+
+  validation_results = run_validation_suite(
+    specs,
+    base_model_config_path=BASE_MODEL_CONFIG,
+    base_hardware_config_path=HARDWARE_CONFIG,
+    result_parser=parse_decode_time,
+    run_perf_path=RUN_PERF,
+  )
+
+  for result in validation_results:
+    meta = result.spec.metadata
+    label = result.spec.label
+    prefill_len = int(meta["prefill_len"])
+    decode_len = int(meta["decode_len"])
+    batch_size = int(meta["batch_size"])
+    expected_tokps_agg = float(meta["expected_tokps"])
+    labels.append(label)
+    expected_tokps.append(expected_tokps_agg)
+
+    if result.success:
+      decode_time_s = float(result.metrics.get("decode_time_s", float("nan")))
+    else:
+      decode_time_s = float("nan")
+
+    total_tokens_agg = int(decode_len) * int(batch_size)
+    if decode_time_s and decode_time_s > 0:
+      predicted_tokps_agg = float(total_tokens_agg) / float(decode_time_s)
+    else:
       predicted_tokps_agg = float("nan")
+
+    if expected_tokps_agg > 0 and not math.isnan(predicted_tokps_agg):
+      pct_error = (predicted_tokps_agg - expected_tokps_agg) / expected_tokps_agg * 100.0
+    else:
       pct_error = float("nan")
 
-    labels.append(label)
-    expected_tokps.append(expected_tokps_agg)
     predicted_tokps.append(predicted_tokps_agg)
     percent_errors.append(pct_error)
 
+    block_lines = [
+      f"\n=== Result {label} (prefill={prefill_len}, decode={decode_len}, batch={batch_size}) ===",
+    ]
+    if result.success and not math.isnan(predicted_tokps_agg):
+      block_lines.append(f"Predicted agg tok/s = {predicted_tokps_agg:.2f} (decode_time={decode_time_s:.2f}s)")
+      block_lines.append(f"Expected agg tok/s  = {expected_tokps_agg:.2f}")
+      block_lines.append(f"Error = {pct_error:+.2f}%")
+    else:
+      err_detail = result.error or "DeepFlow run failed"
+      block_lines.append(f"ERROR: {err_detail}")
+      if result.raw_output:
+        block_lines.append(result.raw_output.strip())
+    log("\n".join(block_lines))
+
   # Print summary statistics
-  print("\n" + "="*70)
-  print("SUMMARY")
-  print("="*70)
+  log("\n" + "="*70)
+  log("SUMMARY")
+  log("="*70)
   valid_errors = [e for e in percent_errors if not math.isnan(e)]
   if valid_errors:
     avg_error = sum(valid_errors) / len(valid_errors)
     avg_abs_error = sum(abs(e) for e in valid_errors) / len(valid_errors)
-    print(f"Average error:          {avg_error:+.2f}%")
-    print(f"Average absolute error: {avg_abs_error:.2f}%")
+    log(f"Average error:          {avg_error:+.2f}%")
+    log(f"Average absolute error: {avg_abs_error:.2f}%")
   else:
-    print("No valid error measurements.")
-  print("="*70)
-
-  # Plot
-  fig = plt.figure(figsize=(12, 6))
-  x = list(range(len(labels)))
-  plt.plot(x, predicted_tokps, marker="o", linestyle="-", color="#1f77b4", label="DeepFlow predicted tok/s (aggregate)")
-  plt.plot(x, expected_tokps, marker="o", linestyle="--", color="#ff7f0e", alpha=0.7, label="Expected tok/s (Koyeb)")
-  plt.xticks(x, labels, rotation=30, ha="right")
-  plt.ylabel("Tokens per second (aggregate)")
-  plt.title("Koyeb vs DeepFlow (single A100 SXM, Llama3.1-8B)")
-  plt.grid(True, linestyle=":", alpha=0.4)
-  plt.legend()
-
-  # Add note about large bs=32 datapoint
-  note_text = ( "X-axis is prefill length + batch size following Koyebs format (small = 512x512, medium = 1024x1024, large = 4096x1024)\n"
-                "For large bs=32 datapoint, we use PCIE data point, as SXM data is very noisy \n"
-               "on the official website (it's 2x slower than PCIE). You can change this in the code.")
-  fig.text(0.5, 0.02, note_text, ha='center', fontsize=9, style='italic', color='#555555', wrap=True)
-
-  out_dir = os.path.join(PROJECT_ROOT, "output", "validation")
-  os.makedirs(out_dir, exist_ok=True)
-  out_path = os.path.join(out_dir, "koyeb_a100_sxm_no_parallelism.png")
-  plt.tight_layout(rect=[0, 0.06, 1, 1])  # Leave space at bottom for note
-  plt.savefig(out_path, dpi=160)
-  print(f"Saved plot to: {out_path}")
+    avg_error = float("nan")
+    avg_abs_error = float("nan")
+    log("No valid error measurements.")
+  log("="*70)
+
+  plot_path = None
+  if enable_plot:
+    # Plot
+    fig = plt.figure(figsize=(12, 6))
+    x = list(range(len(labels)))
+    plt.plot(x, predicted_tokps, marker="o", linestyle="-", color="#1f77b4", label="DeepFlow predicted tok/s (aggregate)")
+    plt.plot(x, expected_tokps, marker="o", linestyle="--", color="#ff7f0e", alpha=0.7, label="Expected tok/s (Koyeb)")
+    plt.xticks(x, labels, rotation=30, ha="right")
+    plt.ylabel("Tokens per second (aggregate)")
+    plt.title("Koyeb vs DeepFlow (single A100 SXM, Llama3.1-8B)")
+    plt.grid(True, linestyle=":", alpha=0.4)
+    plt.legend()
+
+    # Add note about large bs=32 datapoint
+    note_text = ( "X-axis is prefill length + batch size following Koyebs format (small = 512x512, medium = 1024x1024, large = 4096x1024)\n"
+                  "For large bs=32 datapoint, we use PCIE data point, as SXM data is very noisy \n"
+                 "on the official website (it's 2x slower than PCIE). You can change this in the code.")
+    fig.text(0.5, 0.02, note_text, ha='center', fontsize=9, style='italic', color='#555555', wrap=True)
+
+    out_dir = os.path.join(PROJECT_ROOT, "output", "validation")
+    os.makedirs(out_dir, exist_ok=True)
+    plot_path = os.path.join(out_dir, "koyeb_a100_sxm_no_parallelism.png")
+    plt.tight_layout(rect=[0, 0.06, 1, 1])  # Leave space at bottom for note
+    plt.savefig(plot_path, dpi=160)
+    log(f"Saved plot to: {plot_path}")
+
+  return {
+    "labels": labels,
+    "expected_tokps": expected_tokps,
+    "predicted_tokps": predicted_tokps,
+    "percent_errors": percent_errors,
+    "avg_error": avg_error,
+    "avg_abs_error": avg_abs_error,
+    "plot_path": plot_path,
+  }
+
+
+def _parse_args():
+  import argparse
+  parser = argparse.ArgumentParser(description="Validate DeepFlow against Koyeb benchmarks.")
+  parser.add_argument("--no-plot", action="store_true", help="Skip generating the matplotlib plot/output file.")
+  parser.add_argument("--quiet", action="store_true", help="Suppress detailed per-experiment logs.")
+  return parser.parse_args()
 
 
 if __name__ == "__main__":
-  run_all_and_plot()
+  args = _parse_args()
+  results = run_all_and_plot(enable_plot=not args.no_plot, verbose=not args.quiet)
+  if args.quiet:
+    avg_abs = results["avg_abs_error"]
+    if math.isnan(avg_abs):
+      print("Average absolute error: NaN (no valid measurements)")
+    else:
+      print(f"Average absolute error: {avg_abs:.2f}%")
diff --git a/validation_scripts/llm_shapes.txt b/validation_scripts/llm_shapes.txt
new file mode 100644
index 00000000..47113a44
--- /dev/null
+++ b/validation_scripts/llm_shapes.txt
@@ -0,0 +1,6 @@
+Llama2-7B: OrderedDict([('qkv_proj', (16384, 4096, 12288)), ('attention_score', (512, 1024, 128, 1024)), ('attention_output', (512, 1024, 1024, 128)), ('output_proj', (16384, 4096, 4096)), ('ffn1', (16384, 4096, 22016)), ('ffn2', (16384, 11008, 4096)), ('linear', (16384, 4096, 32000)), ('router', (16384, 4096, 1))])
+Llama2-13B: OrderedDict([('qkv_proj', (16384, 5120, 15360)), ('attention_score', (640, 1024, 128, 1024)), ('attention_output', (640, 1024, 1024, 128)), ('output_proj', (16384, 5120, 5120)), ('ffn1', (16384, 5120, 27648)), ('ffn2', (16384, 13824, 5120)), ('linear', (16384, 5120, 32000)), ('router', (16384, 5120, 1))])
+Llama2-70B: OrderedDict([('qkv_proj', (16384, 8192, 10240)), ('attention_score', (128, 8192, 128, 1024)), ('attention_output', (128, 8192, 1024, 128)), ('output_proj', (16384, 8192, 8192)), ('ffn1', (16384, 8192, 57344)), ('ffn2', (16384, 28672, 8192)), ('linear', (16384, 8192, 32000)), ('router', (16384, 8192, 1))])
+Llama3.1-8B: OrderedDict([('qkv_proj', (1024, 4096, 6144)), ('attention_score', (8, 4096, 128, 1024)), ('attention_output', (8, 4096, 1024, 128)), ('output_proj', (1024, 4096, 4096)), ('ffn1', (1024, 4096, 28672)), ('ffn2', (1024, 14336, 4096)), ('linear', (1024, 4096, 128256)), ('router', (1024, 4096, 1))])
+Llama3.1-70B: OrderedDict([('qkv_proj', (1024, 8192, 10240)), ('attention_score', (8, 8192, 128, 1024)), ('attention_output', (8, 8192, 1024, 128)), ('output_proj', (1024, 8192, 8192)), ('ffn1', (1024, 8192, 57344)), ('ffn2', (1024, 28672, 8192)), ('linear', (1024, 8192, 128256)), ('router', (1024, 8192, 1))])
+Llama3.1-405B: OrderedDict([('qkv_proj', (1024, 16384, 18432)), ('attention_score', (8, 16384, 128, 1024)), ('attention_output', (8, 16384, 1024, 128)), ('output_proj', (1024, 16384, 16384)), ('ffn1', (1024, 16384, 106496)), ('ffn2', (1024, 53248, 16384)), ('linear', (1024, 16384, 128256)), ('router', (1024, 16384, 1))])
\ No newline at end of file
diff --git a/validation_scripts/ncu_metrics_results.csv b/validation_scripts/ncu_metrics_results.csv
new file mode 100644
index 00000000..9b32e1d3
--- /dev/null
+++ b/validation_scripts/ncu_metrics_results.csv
@@ -0,0 +1,49 @@
+m,k,n,cta_m,cta_k,cta_n,stages,latency_ms,clock_Ghz,dram_read,dram_write,l2_hit,l2_total,l2_read,l2_write,l2_load,l2_store,global_st,global_ldgsts,global_ldgsts_sectors,global_ldgsts_requests,shared_ldsm,shared_ldgsts,shared_ldgsts_cache_bypass,grid_size,block_size,regs_per_thread,waves_per_sm,dynamic_shared_mem_per_block,shared_mem_config_size,theoretical_warps_per_sm,achieved_warps_per_sm,sm_efficiency,memory_efficiency
+16384,4096,12288,128,32,128,3,7.444512,1.040302,2161039638,418370214,87.638415,32071083439,27670116351,422212465,27670116351,422212465,422212465,20860773194,607123552.0,50331648.0,55556404848,8436628955,20707303275,12288,128,164,37.93,50331,171966,12.0,11.906987,96.281251,66.35622
+1024,128,1024,64,32,256,4,0.010848000000000002,1.036226,565575,0,87.893802,10259905,5497558,2199023,5497558,2199023,2199023,5222680,155648.0,10240.0,9895604,947912,4365834,64,128,162,0.3,83886,171966,8.0,3.940938,10.796474,13.058038
+1024,1024,128,64,64,64,5,0.013344,1.038394,2494033,0,60.83675,11745561,8796093,268435,8796093,268435,268435,7971459,237568.0,16384.0,17867063,2005749,2005749,32,128,80,0.15,83886,171966,8.0,3.959352,8.758664,12.668023
+16384,4096,4096,128,32,128,3,2.487136,1.049378,869111059,136524259,85.753672,10816774203,9223372474,140737488,9223372474,140737488,140737488,6953589633,202374768.0,16777216.0,18518801616,2802948270,8075642323,4096,128,164,12.64,50331,171966,12.0,11.749281,95.232402,66.604182
+16384,4096,22016,128,32,128,3,13.175808,1.052499,3763901032,752867233,87.948712,57659793141,49575624234,756463999,49575624234,756463999,756463999,37375514079,1087762656.0,90177536.0,99538558955,15122685075,36510827641,22016,128,164,67.95,50331,171966,12.0,11.945483,96.33759,67.126956
+16384,11008,4096,128,32,256,3,6.545024,1.052425,2060710276,136070872,82.480282,22647972437,18590859356,140737488,18590859356,140737488,140737488,15510385331,451412288.0,33816576.0,49647681974,7001813525,6881838983,2048,256,218,18.96,75497,104857,8.0,7.997923,96.975567,52.639561
+16384,4096,32000,128,32,128,3,19.085312,1.053524,5325413702,1122295355,87.88347,83858419336,72057594037,1125899906,72057594037,1125899906,1125899906,54324947530,1581051296.0,131072000.0,144678138029,21995077278,51393319892,32000,128,164,98.77,50331,171966,12.0,11.955218,96.574459,67.351305
+16384,4096,1,64,32,64,2,0.21443199999999998,1.050273,140777216,2436857,21.688324,272897826,141665872,33554,141666677,16777,33554,0,0.0,0.0,564049465,0,0,256,128,92,0.47,16777,104857,20.0,9.787831,45.110381,70.066126
+16384,5120,15360,128,32,128,3,11.458784,1.052614,3768151976,524792256,85.992715,50864445135,43234555993,527765581,43234555993,527765581,527765581,32561167125,947648800.0,78643200.0,86739328780,13195250856,32929264076,15360,128,164,47.41,50331,171966,12.0,11.93105,96.593947,67.314009
+1024,128,1024,128,64,128,3,0.010944,1.030935,567017,128,87.118525,9091741,4398046,2199023,4398046,2199023,2199023,4398046,131072.0,8192.0,9895604,394264,4179607,64,128,168,0.59,100663,104857,4.0,4.00477,10.756691,13.041437
+1024,1024,128,64,64,64,5,0.01344,1.036274,2494168,0,60.839544,11745628,8796093,268435,8796093,268435,268435,7971459,237568.0,16384.0,17867063,2005749,2005749,32,128,80,0.15,83886,171966,8.0,3.978921,8.713897,12.558307
+16384,5120,5120,128,32,128,3,3.918304,1.039268,1466355521,173123689,84.373827,17126994915,14411519022,175921860,14411519022,175921860,175921860,10853711995,315882928.0,26214400.0,28913109951,4389332022,12245953562,5120,128,164,15.8,50331,171966,12.0,11.807363,95.369729,65.667447
+16384,5120,27648,128,32,128,3,20.562912,1.053423,6631522920,946931303,86.317107,91676702924,77822201432,949978046,77822201432,949978046,949978046,58610010631,1705765152.0,141557760.0,156130791804,23754743633,59409309414,27648,128,164,85.33,50331,171966,12.0,11.955874,96.81497,67.513753
+16384,13824,5120,128,32,128,3,10.335744,1.052809,5727252989,175648324,79.520888,46680338461,38911101252,175921860,38911101252,175921860,175921860,29228419606,850656656.0,70778880.0,77912273338,11891435150,31561409334,5120,128,164,15.8,50331,171966,12.0,11.811011,96.36264,67.170841
+16384,5120,32000,128,32,128,3,23.775616,1.053654,7625098106,1123201593,86.04038,106586295345,90071992547,1125899906,90071992547,1125899906,1125899906,67835691651,1974266528.0,163840000.0,180706935048,27532093282,62675846791,32000,128,164,98.77,50331,171966,12.0,11.958393,96.891624,67.579012
+16384,5120,1,64,32,64,2,0.265024,1.0517,175964004,2437796,22.555699,343609496,177083045,33554,177082038,16777,33554,0,0.0,0.0,704786953,0,0,256,128,92,0.47,16777,104857,20.0,9.752035,45.375592,70.629169
+16384,8192,10240,128,32,128,3,12.22576,1.053123,6179862012,350559659,81.287171,55614438446,46116860227,351843720,46116860227,351843720,351843720,34677784991,1009253536.0,83886080.0,92413864267,14087177050,36503870867,10240,128,164,31.6,50331,171966,12.0,11.906335,96.523044,67.296226
+8192,128,1024,128,32,128,3,0.027232,1.010449,2500342,581828,92.320141,65372489,35184372,17592186,35184372,17592186,17592186,30811021,917952.0,65536.0,79164837,8224275,32008881,512,128,164,1.58,50331,171966,12.0,10.229684,35.284347,43.133191
+8192,1024,128,128,32,128,5,0.027584,1.019056,17902765,107610,44.176805,48397604,35184372,2199023,35184372,2199023,2199023,35184372,1048576.0,65536.0,71468255,8821862,9027819,64,128,220,0.3,83886,171966,8.0,3.993464,34.539874,32.281163
+16384,8192,8192,128,32,256,3,9.919744,1.04257,3498759816,278153624,79.428101,34873624762,27670116351,281474976,27670116351,281474976,281474976,23094437508,672137072.0,50331648.0,73931090984,10412166809,10226037958,4096,256,218,37.93,75497,104857,8.0,7.999085,96.132498,52.467213
+16384,8192,57344,128,32,128,3,68.45232,1.055575,44016446199,2016652501,76.548944,332141729870,258254417487,2017612426,258254417487,2017612426,2017612426,194195640618,5651812560.0,469762048.0,517517640114,79076574548,173258501831,57344,128,164,176.99,50331,171966,12.0,11.985653,96.315677,67.290225
+16384,28672,8192,128,32,256,3,33.594464,1.054792,13954805084,280242455,78.001315,120923102338,96845406155,281474976,96845406155,281474976,281474976,80740515745,2349858816.0,176160768.0,258398531893,36580164519,35932477008,4096,256,218,37.93,75497,104857,8.0,7.999301,98.199608,52.781711
+16384,8192,32000,128,32,128,3,37.99808,1.053792,21499678570,1125221302,79.021507,180300408877,144115188075,1125899906,144115188075,1125899906,1125899906,108368023873,3153912800.0,262144000.0,288793326105,44079942996,104244490197,32000,128,164,98.77,50331,171966,12.0,11.970652,96.988343,67.650502
+16384,8192,1,64,32,64,2,0.418208,1.053273,281527992,2453634,23.717734,555372255,283333221,33554,283333422,16777,33554,0,0.0,0.0,1154046975,0,0,256,128,92,0.47,16777,104857,20.0,9.73105,45.656654,71.296692
+1024,4096,6144,128,32,128,4,0.27532799999999996,1.036044,73249727,10017608,85.279225,994988364,844424930,13194139,844424930,13194139,13194139,638279921,19022048.0,1572864.0,1736137450,229657538,361850071,384,128,160,1.78,67108,138412,8.0,7.311376,81.688114,56.934269
+4096,128,1024,128,32,128,2,0.01856,1.013349,1400293,0,91.381847,32724530,16636555,8796093,16638569,8796093,8796093,0,0.0,0.0,37383395,0,0,256,128,148,0.79,33554,104857,12.0,9.262791,25.811254,42.585261
+4096,1024,128,64,32,128,6,0.019584,1.027625,9097948,128,55.225712,35732685,26388279,1099511,26388279,1099511,1099511,18966575,565248.0,49152.0,53601191,15470472,15470472,64,128,158,0.3,75497,171966,8.0,4.023061,24.121831,25.748215
+1024,4096,4096,128,32,128,4,0.20668799999999998,1.042899,50164010,5554600,84.81866,661345238,562949953,8796093,562949953,8796093,8796093,425522811,12681312.0,1048576.0,1157424966,153003780,261456620,256,128,160,1.19,67108,138412,8.0,7.175642,72.067339,50.544367
+1024,4096,28672,128,32,128,3,1.110944,1.048413,355381297,58233985,85.360376,4830676419,4035224853,61572651,4035224853,61572651,61572651,3042206940,88539024.0,7340032.0,8101975841,1224356663,3539359068,1792,128,164,5.53,50331,171966,12.0,11.455518,93.361916,65.234636
+1024,14336,4096,128,32,128,4,0.683584,1.048398,224542769,9869566,81.75315,2464348081,2017612426,8796093,2017612426,8796093,8796093,2017614574,58720320.0,3670016.0,4039729200,542552180,920171661,256,128,220,1.19,67108,138412,8.0,7.159021,75.865809,53.030863
+1024,4096,128256,128,32,128,3,4.833504,1.049829,1497540205,271802844,86.476888,21399952651,18050426842,275427662,18050426842,275427662,275427662,13608405235,396053200.0,32833536.0,36241873567,5497703094,14409956727,8016,128,164,24.74,50331,171966,12.0,11.866872,95.859057,66.896463
+1024,4096,1,64,32,64,2,0.153536,1.055087,8834881,0,5.093275,13975756,8933531,2097,8933531,1048,2097,0,0.0,0.0,35253091,0,0,16,128,92,0.03,16777,104857,20.0,3.991608,3.919674,6.088024
+1024,8192,10240,128,32,256,3,0.782656,1.047707,228733046,17291404,82.754327,2613189099,2161727907,21990232,2161727907,21990232,21990232,1804251336,52510800.0,3932160.0,5775866650,793798657,781186771,320,256,218,2.96,75497,104857,8.0,8.001965,94.722959,51.697916
+8192,128,1024,128,32,128,3,0.0272,1.021026,2500342,577241,92.176214,65492547,35184372,17592186,35184372,17592186,17592186,30811558,917904.0,65536.0,79164837,8204511,31703301,512,128,164,1.58,50331,171966,12.0,10.226931,34.959922,42.739866
+8192,1024,128,64,32,256,4,0.02576,1.020876,17896457,325451,48.999562,61575603,43980465,2199023,43980465,2199023,2199023,36009005,1073152.0,81920.0,71468255,17085379,43980465,64,128,162,0.3,83886,171966,8.0,3.929513,36.91961,34.962418
+1024,8192,8192,128,32,128,4,0.6597759999999999,1.047157,249002339,17269795,82.040977,2825939230,2305843387,17592186,2305843387,17592186,17592186,2305850903,67109008.0,4194304.0,4620693320,618336861,1076558248,512,128,220,2.37,67108,138412,8.0,7.465593,89.938927,62.866808
+1024,8192,57344,128,32,128,3,4.355616,1.049975,1686570306,122060017,83.414221,19465585280,16140901562,123145302,16140901562,123145302,123145302,12137223243,353238848.0,29360128.0,32344852171,4933365866,12677924192,3584,128,164,11.06,50331,171966,12.0,11.719001,95.109853,66.403134
+1024,28672,8192,128,32,128,4,2.246528,1.049467,985278114,17286303,80.305796,10016486504,8070450781,17592186,8070450781,17592186,17592186,8070457223,234881168.0,14680064.0,16149908109,2227016778,3781289207,512,128,220,2.37,67108,138412,8.0,7.465995,92.245134,64.417251
+1024,8192,128256,128,32,128,3,9.630464,1.04046,4266578635,274132864,81.167765,44408272844,36100854759,275427662,36100854759,275427662,275427662,27146214092,790056208.0,65667072.0,72342728326,11026818505,28633316454,8016,128,164,24.74,50331,171966,12.0,11.869756,97.089292,66.89006
+1024,8192,1,64,32,64,2,0.30160000000000003,1.057017,17639430,13762,3.106613,27316730,17867063,2097,17867063,1048,2097,0,0.0,0.0,70437463,0,0,16,128,92,0.03,16777,104857,20.0,3.991963,3.942794,6.156956
+1024,16384,18432,128,32,128,3,2.85632,1.040205,1199291234,39560138,81.461899,12735532589,10376293094,39582418,10376293094,39582418,39582418,7792364533,226787168.0,18874368.0,20772853066,3162183630,9015129911,1152,128,164,3.56,50331,171966,12.0,11.134014,94.111493,64.905723
+16384,128,1024,128,32,128,3,0.046784,1.021131,4700707,16606088,93.136978,137373253,70368744,35184372,70368744,35184372,35184372,61612379,1836096.0,131072.0,158329674,16498479,61105003,1024,128,164,3.16,50331,171966,12.0,11.292969,40.646908,50.773999
+16384,1024,128,128,32,256,3,0.042912,1.027601,35485556,2425045,36.116753,83941746,52776558,4398046,52776558,4398046,4398046,44530220,1327072.0,98304.0,142936511,18903895,52643934,64,256,152,0.59,75497,104857,8.0,7.990849,44.035483,43.610425
+1024,16384,16384,128,32,256,3,2.553408,1.04966,738906307,32129173,82.773994,8334510739,6917529087,35184372,6917529087,35184372,35184372,5769107446,167903312.0,12582912.0,18464758042,2609435297,2565774807,512,256,218,4.74,75497,104857,8.0,7.998441,92.735751,50.075447
+1024,16384,106496,128,32,256,3,15.763296,1.053701,5023419519,225993528,83.527452,52959598901,44963938533,228698418,44963938533,228698418,228698418,37499200547,1091370976.0,81788928.0,120020930498,16948851483,16659416858,3328,256,218,30.81,75497,104857,8.0,7.999406,97.266869,52.522158
+1024,53248,16384,128,32,256,3,8.191552,1.05317,2459006215,32135616,82.876442,26753356655,22481969803,35184372,22481969803,35184372,35184372,18739474888,545390592.0,40894464.0,59969932569,8494344726,8353434365,512,256,218,4.74,75497,104857,8.0,7.999665,93.634246,50.184621
+1024,16384,128256,128,32,128,3,18.9544,1.05332,9522446080,275487121,79.420551,89680657675,72201709518,275427662,72201709518,275427662,275427662,54221841471,1578060528.0,131334144.0,144544437844,22083531056,55151889796,8016,128,164,24.74,50331,171966,12.0,11.869806,97.454889,67.949664
+1024,16384,1,64,32,64,2,0.5973120000000001,1.061986,35250138,2350420,7.858025,57477903,35734127,2097,35734127,1048,2097,0,0.0,0.0,140806207,0,0,16,128,92,0.03,16777,104857,20.0,4.031643,3.942572,6.173377
diff --git a/validation_scripts/nvidia_inf.py b/validation_scripts/nvidia_inf.py
new file mode 100644
index 00000000..eb73cfd3
--- /dev/null
+++ b/validation_scripts/nvidia_inf.py
@@ -0,0 +1,455 @@
+import argparse
+import math
+import os
+from functools import lru_cache
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+import sys
+import matplotlib.pyplot as plt
+import pandas as pd
+
+try:
+    from .validation_helpers import (
+        ValidationSpec,
+        run_validation_suite,
+        parse_inference_time,
+    )
+except ImportError:
+    sys.path.append(str(Path(__file__).resolve().parents[1]))
+    from validation_scripts.validation_helpers import (  # type: ignore
+        ValidationSpec,
+        run_validation_suite,
+        parse_inference_time,
+    )
+
+# Data used to be in  https://docs.nvidia.com/nemo-framework/user-guide/latest/llms/llama/performance.html in 2022.
+# no longer available now, found via paper. "Performance Modeling and Workload Analysis of Distributed Large Language Model Training and Inference" by J. Kundu et al.
+
+HW_CONFIGS = {
+    "A100": "a100_80GB.yaml",
+    "H100": "H100_SXM5_80GB.yaml",
+}
+
+MODEL_CONFIGS = {
+    "Llama 2-7B": "Llama2-7B_inf.yaml",
+    "Llama 2-13B": "Llama2-13B_inf.yaml",
+    "Llama 2-70B": "Llama2-70B_inf.yaml",
+}
+
+MODEL_DISPLAY = {
+    "Llama 2-7B": "Llama 2-7B",
+    "Llama 2-13B": "Llama 2-13B",
+    "Llama 2-70B": "Llama 2-70B",
+}
+
+PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+RUN_PERF = os.path.join(PROJECT_ROOT, "run_perf.py")
+VALIDATION_CONFIG_ROOT = os.path.join(PROJECT_ROOT, "validation_scripts", "validation_configs")
+HARDWARE_CONFIG_PATH = os.path.join(VALIDATION_CONFIG_ROOT, "hardware-config")
+MODEL_CONFIG_PATH = os.path.join(VALIDATION_CONFIG_ROOT, "model-config")
+DATA_DIR = Path(__file__).parent / "imec_data"
+
+
+def _load_data(csv_path: Path) -> pd.DataFrame:
+    # Skip the leading "// filepath: ..." line by treating '/' as a comment char.
+    df = pd.read_csv(csv_path, comment="/")
+    # Normalize dtypes
+    df["TP"] = df["TP"].astype(int)
+    df["actual"] = pd.to_numeric(df["actual"], errors="coerce")
+    return df.dropna(subset=["device", "model", "TP", "actual"])
+
+
+@lru_cache(maxsize=None)
+def _load_device_data(device: str) -> pd.DataFrame:
+    csv_path = DATA_DIR / f"{device}_inf.csv"
+    if not csv_path.exists():
+        raise FileNotFoundError(f"Validation CSV not found for device {device}: {csv_path}")
+    return _load_data(csv_path)
+
+
+def _iter_tests(df: pd.DataFrame):
+    for _, row in df.iterrows():
+        yield (row["device"], row["model"], row["TP"])
+
+
+def _build_spec(device: str, model: str, tp: int, idx: int, network_ignored: bool) -> Tuple[ValidationSpec, str, str]:
+    label = f"{device} {model} TP={tp}"
+
+    model_overrides = {
+        "model_param": {
+            "global_batch_size": 1,
+            "seq_len": 400,
+            "decode_len": 200
+        }
+    }
+
+    hw_overrides: Dict[str, Dict[str, object]] = {
+        "parallelism": {
+            "dp": 1,
+            "tp": int(tp),
+            "tp_sp": False,
+            "cp": 1,
+            "lp": 1,
+            "mb": 1,
+        }
+    }
+    if network_ignored:
+        override_bw = {
+            "network": {
+                "dimensions": [
+                    {
+                        "id": "dim0",
+                        "topology": {"bandwidth": "100000 GB", "latency": "1e-9" },
+                    }
+                ]
+            },
+        }
+        hw_overrides.update(override_bw)
+
+    spec = ValidationSpec(
+        label=label,
+        model_overrides=model_overrides,
+        hardware_overrides=hw_overrides,
+        model_config_path=os.path.join(MODEL_CONFIG_PATH, MODEL_CONFIGS.get(model)),
+        hardware_config_path=os.path.join(HARDWARE_CONFIG_PATH, HW_CONFIGS.get(device)),
+        metadata={"device": device, "model": model, "tp": int(tp)},
+        order=idx,
+    )
+    hw_config_path = spec.hardware_config_path  # type: ignore[arg-type]
+    model_config_path = spec.model_config_path  # type: ignore[arg-type]
+    return spec, hw_config_path, model_config_path
+
+
+def _lookup_actual(device: str, model: str, tp: int) -> float:
+    data = _load_device_data(device)
+    matches = data.loc[
+        (data["device"] == device) & (data["model"] == model) & (data["TP"] == int(tp)),
+        "actual",
+    ]
+    if matches.empty:
+        raise ValueError(f"No reference inference time for {device} {model} TP={tp}")
+    return float(matches.values[0])
+
+
+def run_single(
+    device: str,
+    model: str,
+    tp: int,
+    *,
+    network_ignored: bool = True,
+    actual_inference_time_s: Optional[float] = None,
+    emit_logs: bool = False,
+) -> Dict[str, object]:
+    spec, hw_config_path, model_config_path = _build_spec(device, model, tp, idx=0, network_ignored=network_ignored)
+    validation_results = run_validation_suite(
+        [spec],
+        base_model_config_path=model_config_path,
+        base_hardware_config_path=hw_config_path,
+        result_parser=parse_inference_time,
+        run_perf_path=RUN_PERF,
+    )
+
+    result = validation_results[0]
+    inference_time_s = float(result.metrics.get("inference_time_s", float("nan"))) if result.success else float("nan")
+    err_detail = None if result.success else (result.error or "DeepFlow run failed")
+
+    if actual_inference_time_s is None:
+        try:
+            actual_inference_time_s = _lookup_actual(device, model, tp)
+        except Exception:
+            actual_inference_time_s = float("nan")
+
+    if math.isnan(inference_time_s) or actual_inference_time_s == 0 or math.isnan(actual_inference_time_s):
+        pct_error = float("nan")
+    else:
+        pct_error = abs(inference_time_s - actual_inference_time_s) / actual_inference_time_s * 100.0
+
+    if emit_logs:
+        block_lines = [
+            f"\n=== Result (device={device}, model={model}, TP={tp}) ===",
+        ]
+        if not math.isnan(inference_time_s):
+            block_lines.append(f"  DeepFlow Inference Time: {inference_time_s:.2f}s")
+            block_lines.append(f"  Actual Inference Time:   {actual_inference_time_s:.2f}s")
+            block_lines.append(f"  Percent Error:          {pct_error:.2f}%")
+        else:
+            block_lines.append(f"  DeepFlow run failed. {err_detail or ''}".rstrip())
+            if result.raw_output:
+                block_lines.append(result.raw_output.strip())
+        print("\n".join(block_lines))
+
+    return {
+        "success": result.success,
+        "inference_time_s": inference_time_s,
+        "actual_inference_time_s": actual_inference_time_s,
+        "pct_error": pct_error,
+        "error": err_detail,
+        "raw_output": result.raw_output,
+    }
+
+
+def build_specs_for_device(
+    device: str,
+    *,
+    network_ignored: bool = True,
+    models: Optional[Iterable[str]] = None,
+) -> Tuple[List[ValidationSpec], Dict[Tuple[str, int], float], str, str]:
+    data = _load_device_data(device)
+    model_filter = set(models) if models is not None else None
+    specs: List[ValidationSpec] = []
+    actual_lookup: Dict[Tuple[str, int], float] = {}
+    base_model_path: Optional[str] = None
+    hw_config_path: Optional[str] = None
+    idx = 0
+    for device_val, model, tp in _iter_tests(data):
+        if model_filter and model not in model_filter:
+            continue
+        spec, hw_path, model_path = _build_spec(device_val, model, tp, idx, network_ignored)
+        specs.append(spec)
+        actual_lookup[(model, int(tp))] = _lookup_actual(device_val, model, int(tp))
+        base_model_path = base_model_path or model_path
+        hw_config_path = hw_config_path or hw_path
+        idx += 1
+
+    if not specs:
+        raise ValueError(f"No validation specs generated for device={device} (models={models}).")
+    return specs, actual_lookup, base_model_path, hw_config_path
+
+
+def compute_pct_errors(results, actual_lookup: Dict[Tuple[str, int], float]):
+    rows: List[Dict[str, object]] = []
+    for res in results:
+        metadata = res.spec.metadata or {}
+        model = metadata.get("model")
+        tp = int(metadata.get("tp")) if "tp" in metadata else None
+        inf_time = float(res.metrics.get("inference_time_s", float("nan"))) if res.success else float("nan")
+        actual = actual_lookup.get((model, tp)) if model is not None and tp is not None else float("nan")
+        if math.isnan(inf_time) or actual is None or actual == 0 or math.isnan(actual):
+            pct_error = float("nan")
+        else:
+            pct_error = abs(inf_time - actual) / actual * 100.0
+        rows.append(
+            {
+                "device": metadata.get("device"),
+                "model": model,
+                "tp": tp,
+                "inference_time_s": inf_time,
+                "actual_inference_time_s": actual,
+                "pct_error": pct_error,
+                "display_model": MODEL_DISPLAY.get(str(model), str(model)),
+                "success": res.success,
+                "error": res.error,
+            }
+        )
+    return rows
+
+
+def plot_device(df: pd.DataFrame, device: str, outdir: Path) -> Path | None:
+    sub = df[df["device"] == device].copy()
+    if sub.empty:
+        return None
+    sub.sort_values(["model", "TP"], inplace=True)
+
+    # Hierarchical grouping: model cluster (primary) -> TP subgroup (secondary with two bars: seconds & actual)
+    # Enforce explicit primary model order
+    desired_order = ["Llama 2-7B", "Llama 2-13B", "Llama 2-70B"]
+    unique_models = list(pd.unique(sub["model"]))
+    models = [m for m in desired_order if m in unique_models] + [
+        m for m in unique_models if m not in desired_order
+    ]
+    display_models = [MODEL_DISPLAY.get(m, m) for m in models]
+
+    bar_width = 0.28  # width per individual bar
+    tp_gap = 0.08  # gap between TP subgroups inside a model cluster
+    model_gap = 0.7  # gap between model clusters
+
+    x_seconds = []
+    x_actual = []
+    seconds_vals = []
+    actual_vals = []
+    tp_midpoints = []
+    tp_labels = []
+    model_centers = []
+    model_bounds = []
+
+    current_x = 0.0
+    for m in models:
+        m_rows = sub[sub["model"] == m]
+        tps = sorted(m_rows["TP"].unique())
+        cluster_start = current_x
+        for tp in tps:
+            row = m_rows[m_rows["TP"] == tp].iloc[0]
+            x_sec = current_x
+            x_act = current_x + bar_width
+            x_seconds.append(x_sec)
+            x_actual.append(x_act)
+            seconds_vals.append(row["seconds"])
+            actual_vals.append(row["actual"])
+            # Midpoint of the TP subgroup (between the two bars)
+            tp_mid = x_sec + bar_width * 0.5
+            tp_midpoints.append(tp_mid)
+            tp_labels.append(f"TP{tp}")
+            current_x += 2 * bar_width + tp_gap
+        cluster_end = current_x - tp_gap  # last subgroup end (exclude trailing tp_gap)
+        model_centers.append((cluster_start + cluster_end) / 2)
+        model_bounds.append((cluster_start, cluster_end))
+        current_x += model_gap  # gap after cluster
+
+    # Dynamic figure width
+    fig_w = max(8.0, 0.12 * len(tp_midpoints) + 2.5)
+    fig, ax = plt.subplots(figsize=(fig_w, 5))
+
+    bars_seconds = ax.bar(
+        x_seconds, seconds_vals, bar_width, label="DeepFlow", color="#1f77b4"
+    )
+    bars_actual = ax.bar(
+        x_actual, actual_vals, bar_width, label="Actual", color="#0bbd37"
+    )
+
+    # Primary ticks: model names at cluster centers
+    ax.set_xticks(model_centers)
+    ax.set_xticklabels(display_models, fontsize=11)
+    ax.set_ylabel("Inference Latency (s)")
+    ax.set_title(
+        f"Validation of Inference Latency on Systems of {device} (network ignored)"
+    )
+
+    # Secondary TP labels ABOVE bars (centered over subgroup)
+    ymin, ymax = ax.get_ylim()
+    pad = 0.02 * (ymax - ymin)
+    subgroup_max = [max(s, a) for s, a in zip(seconds_vals, actual_vals)]
+    new_ymax = max(ymax, max(subgroup_max) + 3 * pad)
+    ax.set_ylim(ymin, new_ymax)
+    for mid, lbl, v in zip(tp_midpoints, tp_labels, subgroup_max):
+        ax.text(mid, v + pad, lbl, ha="center", va="bottom", fontsize=8)
+
+    # Draw light separators between model clusters
+    for start, end in model_bounds:
+        ax.axvspan(
+            start - bar_width * 0.5,
+            end - bar_width * 0.5,
+            facecolor="#000000",
+            alpha=0.02,
+        )
+
+    # Custom legend
+    ax.legend()
+    ax.grid(axis="y", linestyle="--", alpha=0.3)
+    ax.margins(x=0.02, y=0.1)
+
+    outpath = outdir / f"inf_{device}.png"
+    fig.savefig(outpath, dpi=200, bbox_inches="tight")
+    plt.close(fig)
+    return outpath
+
+
+def run(
+    enable_plot: bool = True,
+    network_ignored: bool = True,
+    device: str = "A100",
+    models: Optional[Sequence[str]] = None,
+    emit_logs: bool = True,
+):
+    pct_errors = []
+    data = _load_device_data(device)
+
+    specs, actual_lookup, base_model_path, hw_config_path = build_specs_for_device(
+        device, network_ignored=network_ignored, models=models
+    )
+
+    validation_results = run_validation_suite(
+        specs,
+        base_model_config_path=base_model_path,
+        base_hardware_config_path=hw_config_path,
+        result_parser=parse_inference_time,
+        run_perf_path=RUN_PERF,
+    )
+
+    rows = compute_pct_errors(validation_results, actual_lookup)
+
+    for row in rows:
+        pct_error = float(row["pct_error"])
+        pct_errors.append(pct_error)
+        data.loc[
+            (data["device"] == row["device"]) & (data["model"] == row["model"]) & (data["TP"] == row["tp"]),
+            "seconds",
+        ] = row["inference_time_s"]
+
+        data.loc[
+            (data["device"] == row["device"]) & (data["model"] == row["model"]) & (data["TP"] == row["tp"]),
+            "pct_error",
+        ] = pct_error
+
+        if emit_logs:
+            block_lines = [
+                f"\n=== Result (device={row['device']}, model={row['model']}, TP={row['tp']}) ===",
+            ]
+            if not math.isnan(row["inference_time_s"]):
+                block_lines.append(f"  DeepFlow Inference Time: {float(row['inference_time_s']):.2f}s")
+                block_lines.append(f"  Actual Inference Time:   {float(row['actual_inference_time_s']):.2f}s")
+                block_lines.append(f"  Percent Error:          {pct_error:.2f}%")
+            else:
+                block_lines.append(f"  DeepFlow run failed. {(row.get('error') or '')}".rstrip())
+            print("\n".join(block_lines))
+
+    if enable_plot:
+        out_dir = os.path.join(PROJECT_ROOT, "output", "validation")
+        os.makedirs(out_dir, exist_ok=True)
+        out_dir = Path(out_dir)
+
+        outputs = []
+        for device in ["A100", "H100"]:
+            out = plot_device(data, device, out_dir)
+            if out is not None:
+                outputs.append(out)
+        # Secondary bar plot matching the training style (combined devices).
+        bar_path = _plot_error_bars(
+            rows,
+            path=out_dir / "inf_errors_bar.png",
+            title="Inference validation (combined)",
+        )
+        outputs.append(bar_path)
+        if not outputs:
+            print("No plots generated (no matching device rows).")
+        else:
+            for p in outputs:
+                print(f"Saved: {p}")
+    valid_errors = [e for e in pct_errors if not math.isnan(e)]
+    avg_abs_error = sum(valid_errors) / len(valid_errors) if valid_errors else float("nan")
+    if emit_logs:
+        print("Average absolute percent error across all tests: {:.2f}%".format(avg_abs_error))
+    return {"avg_abs_error": avg_abs_error, "rows": rows}
+
+
+def _plot_error_bars(rows: List[Dict[str, object]], path: Path, title: str) -> Path:
+    # Build labels like "<model> xTP<tp>", drop device prefixes.
+    labels: List[str] = []
+    errors: List[float] = []
+    for row in rows:
+        model = row.get("display_model") or row.get("model")
+        tp = row.get("tp")
+        labels.append(f"{model} xTP{tp}")
+        errors.append(float(row.get("pct_error", float("nan"))))
+
+    fig_w = max(6.0, 0.6 * len(labels))
+    fig, ax = plt.subplots(figsize=(fig_w, 4))
+    bars = ax.bar(range(len(errors)), errors, color="#1f77b4")
+    ax.set_xticks(range(len(labels)))
+    ax.set_xticklabels(labels, rotation=20, ha="right", fontsize=9)
+    ax.set_ylabel("Percent Error")
+    ax.set_title(title)
+    for rect, err in zip(bars, errors):
+        if math.isnan(err):
+            continue
+        height = rect.get_height()
+        ax.text(rect.get_x() + rect.get_width() / 2, height, f"{err:.1f}%", ha="center", va="bottom", fontsize=8)
+    ax.grid(axis="y", linestyle="--", alpha=0.3)
+    fig.tight_layout()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(path, dpi=200)
+    plt.close(fig)
+    return path
+
+
+if __name__ == "__main__":
+    run(enable_plot=True, network_ignored=True)
diff --git a/validation_scripts/nvidia_train.py b/validation_scripts/nvidia_train.py
new file mode 100644
index 00000000..b75d3c73
--- /dev/null
+++ b/validation_scripts/nvidia_train.py
@@ -0,0 +1,380 @@
+import math
+import os
+from functools import lru_cache
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+import sys
+import matplotlib.pyplot as plt
+import pandas as pd
+
+
+try:
+    from .validation_helpers import (
+        ValidationSpec,
+        run_validation_suite,
+        parse_training_time,
+    )
+except ImportError:
+    sys.path.append(str(Path(__file__).resolve().parents[1]))
+    from validation_scripts.validation_helpers import (  # type: ignore
+        ValidationSpec,
+        run_validation_suite,
+        parse_training_time,
+    )
+
+
+PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+RUN_PERF = os.path.join(PROJECT_ROOT, "run_perf.py")
+VALIDATION_CONFIG_ROOT = os.path.join(PROJECT_ROOT, "validation_scripts", "validation_configs")
+HARDWARE_CONFIG_PATH = os.path.join(VALIDATION_CONFIG_ROOT, "hardware-config")
+MODEL_CONFIG_PATH = os.path.join(VALIDATION_CONFIG_ROOT, "model-config")
+DATA_DIR = Path(__file__).parent / "imec_data"
+
+HW_CONFIGS = {
+  "A100": "a100_80GB.yaml",
+  "A100_korthi": "a100_80GB_korthikanti.yaml",
+  "A100_selene": "a100_80GB_selene_sc.yaml",
+}
+
+MODEL_CONFIGS = {
+  "Llama 2-7B": "Llama2-7B_inf.yaml",
+  "GPT 22B": "GPT_22_B.yaml",
+  "GPT 175B": "GPT_175_B.yaml",
+  "GPT 310B": "GPT_310_B.yaml",
+  "GPT 530B": "GPT_530_B.yaml",
+  "GPT 1T": "GPT_1T.yaml",
+}
+
+
+def _load_data(csv_path: Path):
+  # Load training validation CSVs; treat leading '//' as comments for path hints.
+  df = pd.read_csv(csv_path, comment="/")
+  if "mb" not in df.columns:
+    raise ValueError(f"Training validation CSV missing required 'mb' column: {csv_path}")
+  for col in ("batch", "mb", "dp", "tp", "pp", "cp"):
+    if col in df.columns:
+      df[col] = df[col].astype(int)
+  if "tp_sp" in df.columns:
+    df["tp_sp"] = df["tp_sp"].astype(bool)
+  if "recomputation" in df.columns:
+    df["recomputation"] = df["recomputation"].astype(str)
+  return df
+
+
+@lru_cache(maxsize=None)
+def _load_device_data(device: str):
+  candidates = [DATA_DIR / f"{device}_train.csv"]
+  base = device.split("_")[0]
+  if base and base != device:
+    candidates.append(DATA_DIR / f"{base}_train.csv")
+  candidates.append(DATA_DIR / "A100_train.csv")
+
+  for path in candidates:
+    if path.exists():
+      df = _load_data(path)
+      if "device" in df.columns:
+        sub = df[df["device"] == device]
+        if not sub.empty:
+          return sub
+        else:
+          continue
+      return df
+  raise FileNotFoundError(f"Validation CSV not found for device {device}: tried {', '.join(str(p) for p in candidates)}")
+
+
+def _iter_tests(df):
+  for _, row in df.iterrows():
+    yield (
+      row["device"],
+      row["model"],
+      row["batch"],
+      row["mb"],
+      row["dp"],
+      row["tp"],
+      row["pp"],
+      row["cp"],
+      row.get("tp_sp", False),
+      row.get("recomputation", "partial"),
+    )
+
+
+def _build_spec(
+  device: str,
+  model: str,
+  batch: int,
+  mb: int,
+  dp: int,
+  tp: int,
+  pp: int,
+  cp: int,
+  tp_sp: bool,
+  recomputation: str,
+  idx: int,
+) -> Tuple[ValidationSpec, str, str]:
+  label = f"{device} {model} bs={batch} mb={mb} dp={dp} tp={tp} pp={pp} cp={cp} tp_sp={tp_sp}"
+
+  recomputation_mode = str(recomputation).strip().lower()
+  full_recompute = recomputation_mode in {"full", "true", "yes", "on", "1"}
+
+  model_overrides = {
+    "model_param": {
+      "global_batch_size": int(batch),
+    }
+  }
+
+  hw_overrides = {
+    "parallelism": {
+      "dp": int(dp),
+      "tp": int(tp),
+      "tp_sp": bool(tp_sp),
+      "cp": int(cp),
+      "lp": int(pp),
+      "mb": max(1, int(mb)),
+    },
+    "sw_param": {
+      # sw_param.full_recomputation toggles full activation recompute during backward.
+      "full_recomputation": full_recompute,
+    },
+  }
+
+  spec = ValidationSpec(
+    label=label,
+    model_overrides=model_overrides,
+    hardware_overrides=hw_overrides,
+    model_config_path=os.path.join(MODEL_CONFIG_PATH, MODEL_CONFIGS.get(model)),
+    hardware_config_path=os.path.join(HARDWARE_CONFIG_PATH, HW_CONFIGS.get(device)),
+    metadata={
+      "device": device,
+      "model": model,
+      "batch": int(batch),
+      "mb": int(mb),
+      "dp": int(dp),
+      "tp": int(tp),
+      "pp": int(pp),
+      "cp": int(cp),
+      "tp_sp": bool(tp_sp),
+      "recomputation": str(recomputation),
+    },
+    order=idx,
+  )
+  return spec, spec.model_config_path, spec.hardware_config_path  # type: ignore[return-value]
+
+
+def _lookup_actual(device: str, model: str, batch: int, mb: int, dp: int, tp: int, pp: int, cp: int, tp_sp: bool, recomputation: str) -> float:
+  data = _load_device_data(device)
+  matches = data.loc[
+    (data["device"] == device)
+    & (data["model"] == model)
+    & (data["batch"] == int(batch))
+    & (data["mb"] == int(mb))
+    & (data["dp"] == int(dp))
+    & (data["tp"] == int(tp))
+    & (data["pp"] == int(pp))
+    & (data["cp"] == int(cp))
+    & (data.get("tp_sp", False) == bool(tp_sp))
+    & (data.get("recomputation", "").astype(str) == str(recomputation)),
+    "actual",
+  ]
+  if matches.empty:
+    raise ValueError(
+      f"No reference training time for {device} {model} bs={batch} mb={mb} dp={dp} tp={tp} pp={pp} cp={cp} tp_sp={tp_sp} recomputation={recomputation}"
+    )
+  return float(matches.values[0])
+
+
+def compute_pct_errors(results, actual_lookup: Dict[Tuple[str, int, int, int, int, int, int, bool, str], float]):
+  rows: List[Dict[str, object]] = []
+  for res in results:
+    meta = res.spec.metadata or {}
+    model = meta.get("model")
+    batch = int(meta.get("batch")) if "batch" in meta else None
+    mb = int(meta.get("mb")) if "mb" in meta else None
+    dp = int(meta.get("dp")) if "dp" in meta else None
+    tp = int(meta.get("tp")) if "tp" in meta else None
+    pp = int(meta.get("pp")) if "pp" in meta else None
+    cp = int(meta.get("cp")) if "cp" in meta else None
+    tp_sp = bool(meta.get("tp_sp")) if "tp_sp" in meta else False
+    recomputation = str(meta.get("recomputation")) if "recomputation" in meta else "partial"
+    key = (model, batch, mb, dp, tp, pp, cp, tp_sp, recomputation)
+    train_time = float(res.metrics.get("training_time_s", float("nan"))) if res.success else float("nan")
+    actual = actual_lookup.get(key, float("nan"))
+    if math.isnan(train_time) or actual == 0 or math.isnan(actual):
+      signed_pct_error = float("nan")
+      pct_error = float("nan")
+    else:
+      signed_pct_error = (train_time - actual) / actual * 100.0
+      pct_error = abs(signed_pct_error)
+    rows.append(
+      {
+        "device": meta.get("device"),
+        "model": model,
+        "batch": batch,
+        "mb": mb,
+        "dp": dp,
+        "tp": tp,
+        "pp": pp,
+        "cp": cp,
+        "tp_sp": tp_sp,
+        "recomputation": recomputation,
+        "training_time_s": train_time,
+        "actual_inference_time_s": actual,
+        "signed_pct_error": signed_pct_error,
+        "pct_error": pct_error,
+        "success": res.success,
+        "error": res.error,
+        "raw_output": res.raw_output,
+      }
+    )
+  return rows
+
+
+def build_specs_for_device(
+  device: str,
+  *,
+  models: Optional[Iterable[str]] = None,
+  tp_sp_only: Optional[bool] = None,
+) -> Tuple[List[ValidationSpec], Dict[Tuple[str, int, int, int, int, int, int, bool, str], float], str, str]:
+  data = _load_device_data(device)
+  model_filter = set(models) if models is not None else None
+  specs: List[ValidationSpec] = []
+  actual_lookup: Dict[Tuple[str, int, int, int, int, int, int, bool, str], float] = {}
+  base_model_path: Optional[str] = None
+  hw_config_path: Optional[str] = None
+  idx = 0
+  for device_val, model, batch, mb, dp, tp, pp, cp, tp_sp, recomputation in _iter_tests(data):
+    if model_filter and model not in model_filter:
+      continue
+    if tp_sp_only is not None and bool(tp_sp) != bool(tp_sp_only):
+      continue
+    spec, model_path, hw_path = _build_spec(
+      device_val, model, batch, mb, dp, tp, pp, cp, tp_sp, recomputation, idx
+    )
+    specs.append(spec)
+    actual_lookup[(model, int(batch), int(mb), int(dp), int(tp), int(pp), int(cp), bool(tp_sp), str(recomputation))] = _lookup_actual(
+      device_val, model, batch, mb, dp, tp, pp, cp, tp_sp, recomputation
+    )
+    base_model_path = base_model_path or model_path
+    hw_config_path = hw_config_path or hw_path
+    idx += 1
+  if not specs:
+    raise ValueError(f"No validation specs generated for device={device} (models={models}).")
+  return specs, actual_lookup, base_model_path, hw_config_path
+
+
+def run(
+  device: str = "A100",
+  models: Optional[Sequence[str]] = None,
+  tp_sp: Optional[bool] = None,
+  enable_plot: bool = True,
+  show_progress: bool = False,
+  emit_logs: bool = True,
+):
+  specs, actual_lookup, base_model_path, hw_config_path = build_specs_for_device(
+    device, models=models, tp_sp_only=tp_sp
+  )
+
+  validation_results = run_validation_suite(
+    specs,
+    base_model_config_path=base_model_path,
+    base_hardware_config_path=hw_config_path,
+    result_parser=parse_training_time,
+    run_perf_path=RUN_PERF,
+    show_progress=show_progress,
+  )
+
+  rows = compute_pct_errors(validation_results, actual_lookup)
+  pct_errors = [r["pct_error"] for r in rows if not math.isnan(r["pct_error"])]
+
+  if emit_logs:
+    for row in rows:
+      block = [
+        f"\n=== Result (device={row['device']}, model={row['model']}, bs={row['batch']}, mb={row['mb']}, dp={row['dp']}, tp={row['tp']}, pp={row['pp']}, cp={row['cp']}, tp_sp={row['tp_sp']}, recomputation={row['recomputation']}) ==="
+      ]
+      if row["success"] and not math.isnan(row["pct_error"]):
+        block.append(f"  DeepFlow train time: {float(row['training_time_s']):.2f}s")
+        block.append(f"  Actual train time:   {float(row['actual_inference_time_s']):.2f}s")
+        block.append(f"  Percent Error:  {float(row['signed_pct_error']):+.2f}%")
+      else:
+        block.append(f"  DeepFlow run failed. {(row.get('error') or '')}".rstrip())
+        if row.get("raw_output"):
+          block.append("  --- Raw output ---")
+          block.append(str(row["raw_output"]).strip())
+      print("\n".join(block))
+
+  plot_path = None
+  if enable_plot and rows:
+    out_dir = Path(PROJECT_ROOT) / "output" / "validation" / "train"
+    plot_path = _plot_training_rows(
+      rows,
+      title=f"Training validation ({device})",
+      path=out_dir / f"train_{device}.png",
+      include_device=False,
+    )
+
+  avg_abs_error = sum(pct_errors) / len(pct_errors) if pct_errors else float("nan")
+  if emit_logs:
+    print("Average absolute percent error across all training tests: {:.2f}%".format(avg_abs_error))
+    if plot_path:
+      print(f"Saved plot: {plot_path}")
+  return {"avg_abs_error": avg_abs_error, "rows": rows, "plot": plot_path}
+
+
+def _plot_training_rows(rows, title: str, path: Path, include_device: bool = True) -> Path:
+  out_dir = path.parent
+  out_dir.mkdir(parents=True, exist_ok=True)
+  labels: List[str] = []
+  errors: List[float] = []
+  colors: List[str] = []
+  color_map = {"full": "#ff8c00", "partial": "#1f77b4", "selective": "#1f77b4"}
+  for row in rows:
+    if str(row.get("model")) == "GPT 22B":
+      continue
+    gpus = int(row["dp"]) * int(row["tp"]) * int(row["pp"]) * int(row["cp"])
+    prefix = ""
+    labels.append(f"{prefix}{row['model']} x{gpus}")
+    errors.append(row["signed_pct_error"])
+    colors.append(color_map.get(str(row.get("recomputation")).lower(), "#1f77b4"))
+
+  fig_w = max(6.0, 0.6 * len(labels))
+  fig, ax = plt.subplots(figsize=(fig_w, 4))
+  bars = ax.bar(range(len(errors)), errors, color=colors)
+  ax.set_xticks(range(len(labels)))
+  ax.set_xticklabels(labels, rotation=20, ha="right", fontsize=9)
+  ax.set_ylabel("Percent Error")
+  ax.set_title(title)
+  handles = [
+    plt.Rectangle((0, 0), 1, 1, color="#ff8c00"),
+    plt.Rectangle((0, 0), 1, 1, color="#1f77b4"),
+  ]
+  ax.legend(handles, ["recompute=full", "recompute=partial/selective"])
+  for rect, err in zip(bars, errors):
+    if math.isnan(err):
+      continue
+    height = rect.get_height()
+    ax.text(rect.get_x() + rect.get_width() / 2, height, f"{err:.1f}%", ha="center", va="bottom", fontsize=8)
+  ax.grid(axis="y", linestyle="--", alpha=0.3)
+  fig.tight_layout()
+  fig.savefig(path, dpi=200)
+  plt.close(fig)
+  return path
+
+
+if __name__ == "__main__":
+  devices = [
+    "A100_korthi",  # Uncomment to include korthi runs.
+    "A100_selene",
+  ]
+  combined_rows: List[Dict[str, object]] = []
+  for device in devices:
+    print(f"=== Running {device} training validation ===")
+    result = run(device=device, emit_logs=True, show_progress=True)
+    combined_rows.extend(result.get("rows", []))
+
+  if combined_rows:
+    combined_path = _plot_training_rows(
+      combined_rows,
+      title="Training validation (combined)",
+      path=Path(PROJECT_ROOT) / "output" / "validation" / "train" / "train_combined.png",
+      include_device=False,
+    )
+    print(f"Saved combined plot: {combined_path}")
diff --git a/validation_scripts/uci_train.py b/validation_scripts/uci_train.py
new file mode 100644
index 00000000..30067dd2
--- /dev/null
+++ b/validation_scripts/uci_train.py
@@ -0,0 +1,317 @@
+#!/usr/bin/env python3
+"""
+UCI training validation harness combining DDP and FSDP grid-search measurements.
+
+Inspired by nvidia_train.py: builds ValidationSpecs from benchmark CSVs,
+runs DeepFlow, compares training time to measured Avg_Step_Time_s, and
+produces bar plots of percent error.
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import os
+from pathlib import Path
+from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Tuple
+import sys
+import pandas as pd
+import matplotlib.pyplot as plt
+
+try:
+    from .validation_helpers import (
+        ValidationSpec,
+        run_validation_suite,
+        parse_training_time,
+    )
+except ImportError:
+    sys.path.append(str(Path(__file__).resolve().parents[1]))
+    from validation_scripts.validation_helpers import (  # type: ignore
+        ValidationSpec,
+        run_validation_suite,
+        parse_training_time,
+    )
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+RUN_PERF = PROJECT_ROOT / "run_perf.py"
+DEFAULT_HW_CONFIG = PROJECT_ROOT / "validation_scripts" / "validation_configs" / "hardware-config" / "a100_80GB.yaml"
+DEFAULT_MODEL_CONFIG = PROJECT_ROOT / "validation_scripts" / "validation_configs" / "model-config" / "Llama2-7B.yaml"
+
+# Unified benchmark CSV with variant column (DDP/FSDP)
+MERGED_CSV = PROJECT_ROOT / "validation_scripts" / "imec_data" / "uci_train.csv"
+
+
+def _load_grid(csv_path: Path, default_variant: Optional[str] = None) -> pd.DataFrame:
+    df = pd.read_csv(csv_path)
+    if "variant" not in df.columns and default_variant is not None:
+        df.insert(0, "variant", default_variant)
+    # Normalize numeric fields
+    for col in ("TP", "PP", "CP", "DP"):
+        if col in df.columns:
+            df[col] = pd.to_numeric(df[col], errors="coerce").astype("Int64")
+    if "Avg_Step_Time_s" in df.columns:
+        df["Avg_Step_Time_s"] = pd.to_numeric(df["Avg_Step_Time_s"], errors="coerce")
+    # Keep only successes and drop unused columns
+    df = df[df["Status"].str.upper() == "SUCCESS"].copy()
+    for col in ("Status", "Error"):
+        if col in df.columns:
+            df.drop(columns=[col], inplace=True)
+    return df
+
+
+def _iter_tests(df: pd.DataFrame):
+    for _, row in df.iterrows():
+        actual = row.get("Avg_Step_Time_s")
+        if actual is None or pd.isna(actual):
+            continue
+        yield (
+            str(row.get("variant")),
+            int(row["TP"]),
+            int(row["PP"]),
+            int(row["CP"]),
+            int(row["DP"]),
+            float(actual),
+        )
+
+
+def _build_spec(
+    variant: str,
+    tp: int,
+    pp: int,
+    cp: int,
+    dp: int,
+    idx: int,
+    model_config_path: str,
+    hardware_config_path: str,
+) -> ValidationSpec:
+    label = f"{variant} TP={tp} PP={pp} CP={cp} DP={dp}"
+    hw_overrides = {
+        "parallelism": {
+            "dp": int(dp),
+            "tp": int(tp),
+            "tp_sp": False,
+            "cp": int(cp),
+            "lp": int(pp),
+            "mb": int(pp),  # align microbatch count with pipeline stages
+        },
+        "sw_param": {
+            # DDP -> zero_stage 0, FSDP -> zero_stage 3
+            "dp_zero_stage": 0 if variant.upper() == "DDP" else 3,
+        },
+    }
+    return ValidationSpec(
+        label=label,
+        model_overrides={
+            "model_param": {
+                "seq_len": 4096,
+                "run_type": "training",
+            }
+        },
+        hardware_overrides=hw_overrides,
+        model_config_path=model_config_path,
+        hardware_config_path=hardware_config_path,
+        metadata={
+            "variant": variant,
+            "tp": int(tp),
+            "pp": int(pp),
+            "cp": int(cp),
+            "dp": int(dp),
+        },
+        order=idx,
+    )
+
+
+def build_specs(
+    variants: Sequence[str],
+    model_config_path: str,
+    hardware_config_path: str,
+) -> Tuple[List[ValidationSpec], Dict[Tuple[str, int, int, int, int], float]]:
+    specs: List[ValidationSpec] = []
+    actual_lookup: Dict[Tuple[str, int, int, int, int], float] = {}
+    idx = 0
+    df = _load_grid(MERGED_CSV)
+    df = df[df["variant"].str.lower().isin(set(variants))]
+    for variant, tp, pp, cp, dp, actual in _iter_tests(df):
+        spec = _build_spec(variant, tp, pp, cp, dp, idx, model_config_path, hardware_config_path)
+        specs.append(spec)
+        actual_lookup[(variant, tp, pp, cp, dp)] = actual
+        idx += 1
+
+    if not specs:
+        raise ValueError(f"No validation specs generated for variants={variants}")
+    return specs, actual_lookup
+
+
+def compute_pct_errors(results, actual_lookup: Dict[Tuple[str, int, int, int, int], float]):
+    rows: List[Dict[str, object]] = []
+    for res in results:
+        meta = res.spec.metadata or {}
+        variant = str(meta.get("variant"))
+        tp = int(meta.get("tp", 0))
+        pp = int(meta.get("pp", 0))
+        cp = int(meta.get("cp", 0))
+        dp = int(meta.get("dp", 0))
+        inf_time = float(res.metrics.get("training_time_s", float("nan"))) if res.success else float("nan")
+        actual = actual_lookup.get((variant, tp, pp, cp, dp), float("nan"))
+
+        if math.isnan(inf_time) or actual == 0 or math.isnan(actual):
+            signed_pct_error = float("nan")
+            pct_error = float("nan")
+        else:
+            signed_pct_error = (inf_time - actual) / actual * 100.0
+            pct_error = abs(signed_pct_error)
+
+        rows.append(
+            {
+                "variant": variant,
+                "tp": tp,
+                "pp": pp,
+                "cp": cp,
+                "dp": dp,
+                "training_time_s": inf_time,
+                "actual_training_time_s": actual,
+                "signed_pct_error": signed_pct_error,
+                "pct_error": pct_error,
+                "success": res.success,
+                "error": res.error,
+                "raw_output": res.raw_output,
+            }
+        )
+    return rows
+
+
+def _plot_results(rows: List[Dict[str, object]], title: str, path: Path) -> Optional[Path]:
+    if not rows:
+        return None
+    labels: List[str] = []
+    errors: List[float] = []
+    colors: List[str] = []
+    color_map = {"DDP": "#1f77b4", "FSDP": "#ff8c00"}
+
+    for row in rows:
+        variant = str(row.get("variant"))
+        tp = int(row.get("tp", 0))
+        pp = int(row.get("pp", 0))
+        cp = int(row.get("cp", 0))
+        dp = int(row.get("dp", 0))
+        labels.append(f"{variant} TP{tp} PP{pp} CP{cp} DP{dp}")
+        errors.append(float(row.get("signed_pct_error", float("nan"))))
+        colors.append(color_map.get(variant.upper(), "#888888"))
+
+    fig_w = max(6.0, 0.65 * len(labels))
+    fig, ax = plt.subplots(figsize=(fig_w, 4))
+    bars = ax.bar(range(len(errors)), errors, color=colors)
+    ax.set_xticks(range(len(labels)))
+    ax.set_xticklabels(labels, rotation=25, ha="right", fontsize=9)
+    ax.set_ylabel("Percent Error")
+    ax.set_title(title)
+    handles = [
+        plt.Rectangle((0, 0), 1, 1, color="#1f77b4"),
+        plt.Rectangle((0, 0), 1, 1, color="#ff8c00"),
+    ]
+    ax.legend(handles, ["DDP", "FSDP"])
+    for rect, err in zip(bars, errors):
+        if math.isnan(err):
+            continue
+        height = rect.get_height()
+        ax.text(rect.get_x() + rect.get_width() / 2, height, f"{err:.1f}%", ha="center", va="bottom", fontsize=8)
+    ax.grid(axis="y", linestyle="--", alpha=0.3)
+    fig.tight_layout()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(path, dpi=200)
+    plt.close(fig)
+    return path
+
+
+def run(
+    *,
+    hardware_config: Optional[str] = None,
+    model_config: Optional[str] = None,
+    variants: Sequence[str] = ("ddp", "fsdp"),
+    enable_plot: bool = True,
+    show_progress: bool = False,
+    emit_logs: bool = True,
+):
+    variants = tuple(v.lower() for v in variants)
+    hw_cfg = hardware_config or str(DEFAULT_HW_CONFIG)
+    model_cfg = model_config or str(DEFAULT_MODEL_CONFIG)
+    outputs: List[Dict[str, object]] = []
+    for variant in variants:
+        specs, actual_lookup = build_specs((variant,), model_cfg, hw_cfg)
+
+        validation_results = run_validation_suite(
+            specs,
+            base_model_config_path=model_cfg,
+            base_hardware_config_path=hw_cfg,
+            result_parser=parse_training_time,
+            run_perf_path=str(RUN_PERF),
+            show_progress=show_progress,
+        )
+
+        rows = compute_pct_errors(validation_results, actual_lookup)
+        pct_errors = [r["pct_error"] for r in rows if not math.isnan(r["pct_error"])]
+
+        if emit_logs:
+            for row in rows:
+                block = [
+                    f"\n=== Result ({row['variant']}, TP={row['tp']}, PP={row['pp']}, CP={row['cp']}, DP={row['dp']}) ==="
+                ]
+                if row["success"] and not math.isnan(row["pct_error"]):
+                    block.append(f"  DeepFlow train time: {float(row['training_time_s']):.2f}s")
+                    block.append(f"  Actual train time:   {float(row['actual_training_time_s']):.2f}s")
+                    block.append(f"  Percent Error:  {float(row['signed_pct_error']):+.2f}%")
+                else:
+                    block.append(f"  DeepFlow run failed. {(row.get('error') or '')}".rstrip())
+                    if row.get("raw_output"):
+                        block.append("  --- Raw output ---")
+                        block.append(str(row["raw_output"]).strip())
+                print("\n".join(block))
+
+        plot_path = None
+        if enable_plot and rows:
+            out_dir = PROJECT_ROOT / "output" / "validation" / "train"
+            plot_path = _plot_results(rows, f"Training validation ({variant.upper()} benchmarks)", out_dir / f"uci_train_{variant}.png")
+            if emit_logs and plot_path:
+                print(f"Saved plot: {plot_path}")
+
+        avg_abs_error = sum(pct_errors) / len(pct_errors) if pct_errors else float("nan")
+        if emit_logs:
+            print(f"Average absolute percent error for {variant.upper()}: {avg_abs_error:.2f}%")
+        outputs.append({"variant": variant, "avg_abs_error": avg_abs_error, "rows": rows, "plot": plot_path})
+    return outputs
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="UCI training validation (DDP + FSDP gridsearch).")
+    parser.add_argument(
+        "--hardware_config",
+        default=str(DEFAULT_HW_CONFIG),
+        help=f"Path to hardware config YAML (default: {DEFAULT_HW_CONFIG}).",
+    )
+    parser.add_argument(
+        "--model_config",
+        default=str(DEFAULT_MODEL_CONFIG),
+        help=f"Path to model config YAML (default: {DEFAULT_MODEL_CONFIG}).",
+    )
+    parser.add_argument(
+        "--variants",
+        nargs="+",
+        default=["ddp", "fsdp"],
+        choices=["ddp", "fsdp"],
+        help="Which benchmark variants to include (runs each separately).",
+    )
+    parser.add_argument("--no-plot", dest="enable_plot", action="store_false", help="Disable plot generation.")
+    parser.add_argument("--show-progress", action="store_true", help="Show per-run progress.")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = _parse_args()
+    run(
+        hardware_config=args.hardware_config,
+        model_config=args.model_config,
+        variants=args.variants,
+        enable_plot=args.enable_plot,
+        show_progress=args.show_progress,
+    )
diff --git a/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml b/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml
new file mode 100644
index 00000000..f5e72fbc
--- /dev/null
+++ b/validation_scripts/validation_configs/hardware-config/H100_SXM5_80GB.yaml
@@ -0,0 +1,159 @@
+# DeepFlow configuration for NVIDIA H100 PCIe 80 GB
+# specifications from the following sources:
+# [NVIDIA] https://resources.nvidia.com/en-us-hopper-architecture/nvidia-h100-tensor-c
+# [TePU] https://www.techpowerup.com/gpu-specs/h100-sxm5-80-gb.c3900
+# [Semi] https://semianalysis.com/2025/06/23/nvidia-tensor-core-evolution-from-volta-to-blackwell/#4th-generation-tensor-core-warpgroup-level-asynchronous-mma
+
+sw_param:
+  full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation
+  kernel_launch_overhead: 9e-6   # (s)
+  # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3
+  # only used if data parallelism (dp) > 1
+  dp_zero_stage: 1
+  precision:
+    # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32).
+    # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats.
+    # For FP32 or other formats they need to be manually adjusted.
+    # The following are equivalent: (tensor_format: 2.0, tensor_format: fp16, tensor_format: bf16).
+    tensor_format: bf16  # baseline storage width for activations/gradients.
+    # mixed_precision: true | false
+    # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes.
+    # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode.
+    # if true, optimizer states and stats are always FP32.
+    mixed_precision: true
+    # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params.
+    # -> as_tensor_format makes weights/gradients use tensor_format bytes.
+    # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and
+    #    keeps a separate FP32 master copy for weights.
+    # -> fp32_params makes weights *and* gradients use FP32 bytes without a master copy.
+    param_storage_mode: as_tensor_format
+    # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format".
+    kv_cache: as_tensor_format
+  h2d_bandwidth: -1  # Bytes/s; set -1 to disable H2D timing
+  # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb)
+  dp_microbatch: last_mb
+tech_param:
+  core:
+    operating_frequency: 1.98e9     # (Hz) GPU boost clock frequency [TePU]
+    num_bundles: 132     # number of bundles (SMs) [TePU]
+    nominal_flop_rate_per_mcu: 1024     # F16 FLOP / cycle / SM is 4096 [Semi]
+    nominal_energy_per_flop: 1.8e-13
+    num_mcu_per_bundle: 4     # number of tensorcores per SM [TePU]
+    FMA_d1: 16                # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock
+    FMA_d2: 8
+    dataflow: best
+    util: 1.0
+  DRAM:
+    size: 80 GB     # [TePU]
+    bandwidth: 3430 GB     # (/s) [TePU]
+    dynamic_energy_per_bit: 4.5e-12     # (J) [https://passlab.github.io/mchpc/mchpc2019/presentations/MCHPC_Pawlowski_keynote.pdf, slide 36]
+    latency: 100e-9     # (s)
+    util: 1
+  SRAM-L2:
+    size: 50 MB     # [TePU]
+    bandwidth: 8138 GB     # NCU L2 peak traffic
+    dynamic_energy_per_bit: 130e-15
+    latency: 0
+    util: 1
+  SRAM-L1:
+    size: 30096 KB     # max configurable shared memory per SM * number of SMs [NVIDIA]
+    bandwidth: 22392 GB     # NCU L1/TEX peak traffic
+    dynamic_energy_per_bit: 130e-15
+    latency: 0
+    util: 1
+  SRAM-R:
+    size: 33 MB     # register file memory per SM * number of SMs [NVIDIA]
+    bandwidth: 122 TB     # traffic for 2 reads and 1 write for GEMMs computed per clock
+    dynamic_energy_per_bit: 110e-15
+    latency: 0
+    util: 1
+
+memory_hierarchy:
+  l0:   #Register Memory
+    type: SRAM-R
+    scope: mcu
+  l1:   #Shared Memory
+    type: SRAM-L1
+    scope: mcu-bundle
+  l2:   #L2
+    type: SRAM-L2
+    scope: global
+  l3:   #Global Memory
+    type: DRAM
+    scope: global
+
+parallelism:
+  auto: false
+  ### LLM PARAMS ###
+  dp: 1 # data parallelism. Acts as replica count during inference
+  tp: 8 # tensor parallelism
+  tp_sp: false # tensor-sequence parallelism
+  cp: 1 # context parallelism. Not supported for inference.
+  lp: 1 # pipeline parallelism. Number of pipeline stages
+  # mb is not a true parallelism parameter, does not affect num of GPUs.
+  mb: 1 # micro-batches (for best results set to >= lp). 
+
+network:
+  dimensions:
+    - id: dim0
+      label: torus2d_tp
+      size: auto
+      topology:
+        type: Ring
+        bandwidth: 100 GB # We assume 8 25GB/s NVLINKs in the torus. Torus2D is hardcoded bidirectional, so it's actually 4 links, 25GB/s each (because they get replicated). So 100GB/s total. 
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+        optimize_2dmap: false
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [tp, cp]
+    - id: dim1
+      label: inter_torus_infiniband
+      size: auto
+      topology:
+        type: Ring 
+        bandwidth: 25 GB
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [lp,dp]
+    - id: dim2
+      label: inter_torus_infiniband_dp
+      size: auto
+      topology:
+        type: Ring
+        bandwidth: 25 GB
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+      parallelisms: []
+      collective_override: {} # overrides execution_backend.collectives when provided
+
+  overlap:
+    # fraction of TP+SP compute that can overlap TP/SP comm
+    tp_sp_overlap: 0.63
+    # fraction of TP compute that can overlap TP comm
+    tp_overlap: 0.6
+    # fraction of CP comm that can overlap CP compute
+    cp_overlap: 0.0
+execution_backend:
+  model: analytical # analytical | astra
+  astra:
+    backend: astra
+    mode: full_astrasim_flattened  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
+    collectives:
+      # Default collective algorithms. Per-dimension collective_override entries can override these.
+      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d
+      # Note: leave to auto unless you have a specific reason to change it.
+      all_gather: auto
+      all_reduce: auto
+      reduce_scatter: auto
+      all_to_all: auto
+    sys_options:
+      endpoint_delay: 5
+      active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism
+      preferred_dataset_splits: 4 # Should be same as active_chunks_per_dimension
+
+inference:
+  kvcache_type: hbm_only
diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml
new file mode 100644
index 00000000..c7fac8e4
--- /dev/null
+++ b/validation_scripts/validation_configs/hardware-config/a100_80GB.yaml
@@ -0,0 +1,154 @@
+sw_param:
+  full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation
+  kernel_launch_overhead: 9e-6   # (s)
+  # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3
+  # only used if data parallelism (dp) > 1
+  dp_zero_stage: 1
+  precision:
+    # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32).
+    # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats.
+    # For FP32 or other formats they need to be manually adjusted.
+    # The following are equivalent: (tensor_format: 2.0, tensor_format: fp16, tensor_format: bf16).
+    tensor_format: bf16  # baseline storage width for activations/gradients.
+    # mixed_precision: true | false
+    # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes.
+    # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode.
+    # if true, optimizer states and stats are always FP32.
+    mixed_precision: true
+    # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params.
+    # -> as_tensor_format makes weights/gradients use tensor_format bytes.
+    # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and
+    #    keeps a separate FP32 master copy for weights.
+    # -> fp32_params makes weights *and* gradients use FP32 bytes without a master copy.
+    param_storage_mode: as_tensor_format
+    # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format".
+    kv_cache: as_tensor_format
+  h2d_bandwidth: -1  # Bytes/s; set -1 to disable H2D timing
+  # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb)
+  dp_microbatch: last_mb
+tech_param:
+  core:
+    operating_frequency: 1.41e9     # (Hz)
+    num_bundles: 108
+    nominal_flop_rate_per_mcu: 512
+    nominal_energy_per_flop: 1.8e-13 # (J)
+    num_mcu_per_bundle: 4     # number of tensorcores per SM
+    FMA_d1: 8                 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock
+    FMA_d2: 4
+    dataflow: best            # {wst | ast | ost | best}: wst: weight stationary, ast: activation stationary, ost: output stationary, best: best of all 3
+    util: 1.0
+  DRAM:
+    size: 80 GB
+    bandwidth: 1986 GB     # (/s)
+    dynamic_energy_per_bit: 4.5e-12     # (J)
+    latency: 100e-9     # (s)
+    util: 1
+  SRAM-L2:
+    size: 40 MB
+    bandwidth: 7050 GB
+    dynamic_energy_per_bit: 130e-15 # (J)
+    latency: 0
+    util: 1
+  SRAM-L1:
+    size: 20736 KB
+    bandwidth: 18 TB
+    dynamic_energy_per_bit: 130e-15 # (J)
+    latency: 0
+    util: 1
+  SRAM-R:
+    size: 27 MB
+    bandwidth: 122 TB     
+    dynamic_energy_per_bit: 110e-15 # (J)
+    latency: 0
+    util: 1
+
+memory_hierarchy:
+  l0:   #Register Memory
+    type: SRAM-R
+    scope: mcu
+  l1:   #Shared Memory
+    type: SRAM-L1
+    scope: mcu-bundle
+  l2:   #L2
+    type: SRAM-L2
+    scope: global
+  l3:   #Global Memory
+    type: DRAM
+    scope: global
+
+parallelism:
+  dp: 2 # data parallelism. Acts as replica count during inference
+  tp: 2 # tensor parallelism
+  tp_sp: false # tensor-sequence parallelism
+  cp: 2 # context parallelism. Not supported for inference.
+  lp: 4 # pipeline parallelism. Number of pipeline stages
+  # mb is not a true parallelism parameter, does not affect num of GPUs.
+  mb: 1 # micro-batches (for best results set to >= lp). 
+
+network:
+  # Global faulty links apply to the active AstraSim dimension: [[src, dst, weight]]
+  # faulty_links: [[0,1,0.6],[16,17,0.1]]
+  dimensions:
+    # currently only support astrasim integration with all devices in one node.
+    - id: dim0
+      label: torus2d_tp
+      size: auto
+      topology:
+        type: Ring
+        bandwidth: 100 GB # We assume 8 25GB/s NVLINKs in the torus. Torus2D is hardcoded bidirectional, so it's actually 4 links, 25GB/s each (because they get replicated). So 100GB/s total. 
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+        optimize_2dmap: false
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [tp, cp]
+    - id: dim1
+      label: inter_torus_infiniband
+      size: auto
+      topology:
+        type: Ring 
+        bandwidth: 25 GB
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [lp,dp]
+    - id: dim2
+      label: inter_torus_infiniband_dp
+      size: auto
+      topology:
+        type: Ring
+        bandwidth: 25 GB
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+      parallelisms: []
+      collective_override: {} # overrides execution_backend.collectives when provided
+
+  overlap:
+    # fraction of TP+SP compute that can overlap TP/SP comm
+    tp_sp_overlap: 0.63
+    # fraction of TP compute that can overlap TP comm
+    tp_overlap: 0.6
+    # fraction of CP comm that can overlap CP compute
+    cp_overlap: 0.0
+execution_backend:
+  model: astra # analytical | astra
+  astra:
+    backend: astra
+    mode: full_astrasim_hierarchical  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
+    collectives:
+      # Default collective algorithms. Per-dimension collective_override entries can override these.
+      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d
+      # Note: leave to auto unless you have a specific reason to change it.
+      all_gather: auto
+      all_reduce: auto
+      reduce_scatter: auto
+      all_to_all: auto
+    sys_options:
+      endpoint_delay: 5
+      active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism
+      preferred_dataset_splits: 4 # Should be same as active_chunks_per_dimension
+
+inference:
+  kvcache_type: hbm_only # Only supported mode for now.
diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml
new file mode 100644
index 00000000..e0b3b35a
--- /dev/null
+++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_korthikanti.yaml
@@ -0,0 +1,151 @@
+sw_param:
+  full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation
+  kernel_launch_overhead: 9e-6   # (s)
+  # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3
+  # only used if data parallelism (dp) > 1
+  dp_zero_stage: 1
+  precision:
+    # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32).
+    # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats.
+    # For FP32 or other formats they need to be manually adjusted.
+    # The following are equivalent: (tensor_format: 2.0, tensor_format: fp16, tensor_format: bf16).
+    tensor_format: bf16  # baseline storage width for activations/gradients.
+    # mixed_precision: true | false
+    # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes.
+    # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode.
+    # if true, optimizer states and stats are always FP32.
+    mixed_precision: true
+    # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params.
+    # -> as_tensor_format makes weights/gradients use tensor_format bytes.
+    # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and
+    #    keeps a separate FP32 master copy for weights.
+    # -> fp32_params makes weights *and* gradients use FP32 bytes without a master copy.
+    param_storage_mode: tensor_plus_fp32_master
+    # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format".
+    kv_cache: as_tensor_format
+  h2d_bandwidth: -1  # Bytes/s; set -1 to disable H2D timing
+  # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb)
+  dp_microbatch: last_mb
+tech_param:
+  core:
+    operating_frequency: 1.41e9     # (Hz)
+    num_bundles: 108
+    nominal_flop_rate_per_mcu: 512
+    nominal_energy_per_flop: 1.8e-13 # (J)
+    num_mcu_per_bundle: 4     # number of tensorcores per SM
+    FMA_d1: 8                 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock
+    FMA_d2: 4
+    dataflow: best            # {wst | ast | ost | best}: wst: weight stationary, ast: activation stationary, ost: output stationary, best: best of all 3
+    util: 1.0
+  DRAM:
+    size: 80 GB
+    bandwidth: 1986 GB     # (/s)
+    dynamic_energy_per_bit: 4.5e-12     # (J)
+    latency: 100e-9     # (s)
+    util: 0.4
+  SRAM-L2:
+    size: 40 MB
+    bandwidth: 7050 GB
+    dynamic_energy_per_bit: 130e-15 # (J)
+    latency: 0
+    util: 1
+  SRAM-L1:
+    size: 20736 KB
+    bandwidth: 18 TB
+    dynamic_energy_per_bit: 130e-15 # (J)
+    latency: 0
+    util: 1
+  SRAM-R:
+    size: 27 MB
+    bandwidth: 122 TB     
+    dynamic_energy_per_bit: 110e-15 # (J)
+    latency: 0
+    util: 1
+
+memory_hierarchy:
+  l0:   #Register Memory
+    type: SRAM-R
+    scope: mcu
+  l1:   #Shared Memory
+    type: SRAM-L1
+    scope: mcu-bundle
+  l2:   #L2
+    type: SRAM-L2
+    scope: global
+  l3:   #Global Memory
+    type: DRAM
+    scope: global
+
+parallelism:
+  dp: 1 # data parallelism. Acts as replica count during inference
+  tp: 8 # tensor parallelism
+  tp_sp: true # tensor-sequence parallelism
+  cp: 1 # context parallelism. Not supported for inference.
+  lp: 1 # pipeline parallelism. Number of pipeline stages
+  # mb is not a true parallelism parameter, does not affect num of GPUs.
+  mb: 1 # micro-batches (for best results set to >= lp). 
+
+network:
+  overlap:
+    tp_overlap: 0.6
+    tp_sp_overlap: 0.63
+    cp_overlap: 0.8
+  # Global faulty links apply to the active AstraSim dimension: [[src, dst, weight]]
+  # faulty_links: [[0,1,0.6],[16,17,0.1]]
+  dimensions:
+    # currently only support astrasim integration with all devices in one node.
+    - id: dim0
+      label: torus2d_tp
+      size: auto
+      topology:
+        type: FullyConnected
+        bandwidth: 125 GB # based on num nvlinks per GPU * 25 GB/s per link + 50GB/s derate + uni
+        energy_per_bit: 8e-12 # (J/bit)
+        latency: 1e-6 
+        util: 1.0
+        optimize_2dmap: false
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [tp, cp]
+    - id: dim1
+      label: inter_torus_infiniband
+      size: auto
+      topology:
+        type: FullyConnected 
+        bandwidth: 25 GB # 25GB/s per GPU, N GPUs per node. NOT AGGREGATE.
+        latency: 5e-6 
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [lp,dp]
+    - id: dim2
+      label: inter_torus_infiniband_dp
+      size: auto
+      topology:
+        type: Ring
+        bandwidth: 25 GB
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+      parallelisms: []
+      collective_override: {} # overrides execution_backend.collectives when provided
+
+execution_backend:
+  model: astra # analytical | astra
+  astra:
+    backend: astra
+    mode: full_astrasim_hierarchical  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
+    collectives:
+      # Default collective algorithms. Per-dimension collective_override entries can override these.
+      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d
+      # Note: leave to auto unless you have a specific reason to change it.
+      all_gather: auto
+      all_reduce: auto
+      reduce_scatter: auto
+      all_to_all: auto
+    sys_options:
+      endpoint_delay: 5
+      active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism
+      preferred_dataset_splits: 4 # Should be same as active_chunks_per_dimension
+
+inference:
+  kvcache_type: hbm_only # Only supported mode for now.
diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml
new file mode 100644
index 00000000..f0476141
--- /dev/null
+++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_no_parallelism.yaml
@@ -0,0 +1,142 @@
+# EXAMPLE FILE USED FOR TESTING/README ONLY.
+# DO NOT MODIFY.
+
+sw_param:
+  full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation
+  kernel_launch_overhead: 9e-6
+  # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3
+  # only used if data parallelism (dp) > 1
+  dp_zero_stage: 2
+  precision:
+    # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32).
+    # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats.
+    # For FP32 or other formats they need to be manually adjusted.
+    # The following are equivalent: (tensor_format: 2.0, tensor_format: fp16, tensor_format: bf16).
+    tensor_format: bf16  # baseline storage width for activations/gradients.
+    # mixed_precision: true | false
+    # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes.
+    # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode.
+    # if true, optimizer states and stats are always FP32.
+    mixed_precision: false
+    # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params.
+    # -> as_tensor_format makes weights/gradients use tensor_format bytes.
+    # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and
+    #    keeps a separate FP32 master copy for weights.
+    # -> fp32_params makes weights *and* gradients use FP32 bytes without a master copy.
+    param_storage_mode: as_tensor_format
+    # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format".
+    kv_cache: as_tensor_format
+  h2d_bandwidth: -1  # Bytes/s; set -1 to disable H2D timing
+  # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb)
+  dp_microbatch: last_mb
+tech_param:
+  core:
+    operating_frequency: 1.41e9
+    num_bundles: 108
+    nominal_flop_rate_per_mcu: 512
+    nominal_energy_per_flop: 1.8e-13
+    num_mcu_per_bundle: 4
+    FMA_d1: 8
+    FMA_d2: 4
+    dataflow: best
+    util: 1.0
+  DRAM:
+    size: 80 GB
+    bandwidth: 1986 GB
+    dynamic_energy_per_bit: 4.5e-12
+    latency: 100e-9
+    util: 1
+  SRAM-L2:
+    size: 40 MB
+    bandwidth: 7050 GB
+    dynamic_energy_per_bit: 130e-15
+    latency: 0
+    util: 1
+  SRAM-L1:
+    size: 20736 KB
+    bandwidth: 18 TB
+    dynamic_energy_per_bit: 130e-15
+    latency: 0
+    util: 1
+  SRAM-R:
+    size: 27 MB
+    bandwidth: 122 TB
+    dynamic_energy_per_bit: 110e-15
+    latency: 0
+    util: 1
+  network:
+    intra_node:
+      latency: 5e-6
+      util: 1
+      bandwidth: 400 Gb
+    inter_node:
+      latency: 5e-6
+      util: 0.96
+      bandwidth: 400 Gb
+memory_hierarchy:
+  l0:
+    type: SRAM-R
+    scope: mcu
+  l1:
+    type: SRAM-L1
+    scope: mcu-bundle
+  l2:
+    type: SRAM-L2
+    scope: global
+  l3:
+    type: DRAM
+    scope: global
+
+
+parallelism:
+  ### LLM PARAMS ###
+  dp: 1 # acts as replica count during inference
+  lp: 1 # pipeline stages
+  mb: 1 # number of micro-batches
+  tp: 1 # tensor parallelism
+  tp_sp: false # tensor-sequence parallelism
+  cp : 1 # context parallelism
+
+network:
+  dimensions:
+    - id: dim0
+      label: device_fabric
+      size: 1
+      topology:
+        type: Ring
+        bandwidth: 400 Gb
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [tp, cp, lp, dp]
+
+
+
+  overlap:
+    # fraction of TP+SP compute that can overlap TP/SP comm
+    tp_sp_overlap: 0.63
+    # fraction of TP compute that can overlap TP comm
+    tp_overlap: 0.6
+    # fraction of CP comm that can overlap CP compute
+    cp_overlap: 0.0
+execution_backend:
+  model: analytical # analytical | astra
+  astra:
+    backend: astra
+    mode: full_astrasim_flattened  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
+    collectives:
+      # Default collective algorithms. Per-dimension collective_override entries can override these.
+      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d
+      # Note: leave to auto unless you have a specific reason to change it.
+      all_gather: auto
+      all_reduce: auto
+      reduce_scatter: auto
+      all_to_all: auto
+    sys_options:
+      endpoint_delay: 10
+      active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism
+      preferred_dataset_splits: 32 # Should be same as active_chunks_per_dimension
+
+inference:
+  kvcache_type: hbm_only # Only supported mode for now.
diff --git a/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml b/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml
new file mode 100644
index 00000000..3ab41ce0
--- /dev/null
+++ b/validation_scripts/validation_configs/hardware-config/a100_80GB_selene_sc.yaml
@@ -0,0 +1,154 @@
+sw_param:
+  full_recomputation: false # Enable full activation recompute during backward instead of selective recomputation
+  kernel_launch_overhead: 9e-6   # (s)
+  # dp_zero_stage selects ZeRO-style data-parallel sharding: 0=DDP, 1=ZeRO-1, 2=ZeRO-2, 3=ZeRO-3
+  # only used if data parallelism (dp) > 1
+  dp_zero_stage: 1
+  precision:
+    # tensor_format accepts numeric byte counts or dtype strings (fp16, bf16, fp32).
+    # Take care that the number of flops defined below has been carefully tuned for FP16/BF16 formats.
+    # For FP32 or other formats they need to be manually adjusted.
+    # The following are equivalent: (tensor_format: 2.0, tensor_format: fp16, tensor_format: bf16).
+    tensor_format: bf16  # baseline storage width for activations/gradients.
+    # mixed_precision: true | false
+    # -> false forces all activations/weights/gradients/optimizer states/stats to use tensor_format bytes.
+    # -> true enables AMP-style FP32 accumulations/optimizer states depending on param_storage_mode.
+    # if true, optimizer states and stats are always FP32.
+    mixed_precision: true
+    # param_storage_mode: as_tensor_format | tensor_plus_fp32_master | fp32_params.
+    # -> as_tensor_format makes weights/gradients use tensor_format bytes.
+    # -> tensor_plus_fp32_master makes weights/gradients use tensor_format bytes and
+    #    keeps a separate FP32 master copy for weights.
+    # -> fp32_params makes weights *and* gradients use FP32 bytes without a master copy.
+    param_storage_mode: as_tensor_format
+    # kv_cache precision may be set as bytes, dtype string, or "as_tensor_format".
+    kv_cache: as_tensor_format
+  h2d_bandwidth: -1  # Bytes/s; set -1 to disable H2D timing
+  # DP gradient sync frequency: every_mb (default) or last_mb (sync only on final microbatch; ZeRO-3 stays per-mb)
+  dp_microbatch: last_mb
+tech_param:
+  core:
+    operating_frequency: 1.41e9     # (Hz)
+    num_bundles: 108
+    nominal_flop_rate_per_mcu: 512
+    nominal_energy_per_flop: 1.8e-13 # (J)
+    num_mcu_per_bundle: 4     # number of tensorcores per SM
+    FMA_d1: 8                 # tensor core can compute the results for an d1 x d2 x d1 mixed-precision matrix multiplication per clock
+    FMA_d2: 4
+    dataflow: best            # {wst | ast | ost | best}: wst: weight stationary, ast: activation stationary, ost: output stationary, best: best of all 3
+    util: 1.0
+  DRAM:
+    size: 80 GB
+    bandwidth: 1986 GB     # (/s)
+    dynamic_energy_per_bit: 4.5e-12     # (J)
+    latency: 100e-9     # (s)
+    util: 1
+  SRAM-L2:
+    size: 40 MB
+    bandwidth: 7050 GB
+    dynamic_energy_per_bit: 130e-15 # (J)
+    latency: 0
+    util: 1
+  SRAM-L1:
+    size: 20736 KB
+    bandwidth: 18 TB
+    dynamic_energy_per_bit: 130e-15 # (J)
+    latency: 0
+    util: 1
+  SRAM-R:
+    size: 27 MB
+    bandwidth: 122 TB     
+    dynamic_energy_per_bit: 110e-15 # (J)
+    latency: 0
+    util: 1
+
+memory_hierarchy:
+  l0:   #Register Memory
+    type: SRAM-R
+    scope: mcu
+  l1:   #Shared Memory
+    type: SRAM-L1
+    scope: mcu-bundle
+  l2:   #L2
+    type: SRAM-L2
+    scope: global
+  l3:   #Global Memory
+    type: DRAM
+    scope: global
+
+parallelism:
+  dp: 2 # data parallelism. Acts as replica count during inference
+  tp: 2 # tensor parallelism
+  tp_sp: true # tensor-sequence parallelism
+  cp: 2 # context parallelism. Not supported for inference.
+  lp: 4 # pipeline parallelism. Number of pipeline stages
+  # mb is not a true parallelism parameter, does not affect num of GPUs.
+  mb: 4 # micro-batches (for best results set to >= lp). 
+
+network:
+  # Global faulty links apply to the active AstraSim dimension: [[src, dst, weight]]
+  # faulty_links: [[0,1,0.6],[16,17,0.1]]
+  dimensions:
+    # currently only support astrasim integration with all devices in one node.
+    - id: dim0
+      label: torus2d_tp
+      size: auto
+      topology:
+        type: FullyConnected
+        bandwidth: 300 GB # based on num nvlinks per GPU * 25 GB/s per link
+        latency: 1e-6 
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+        optimize_2dmap: false
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [tp, cp]
+    - id: dim1
+      label: inter_torus_infiniband
+      size: auto
+      topology:
+        type: FullyConnected 
+        bandwidth: 25 GB # 25GB/s per GPU, N GPUs per node. NOT AGGREGATE.
+        latency: 5e-6 
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+      collective_override: {} # overrides execution_backend.collectives when provided
+      parallelisms: [lp,dp]
+    - id: dim2
+      label: inter_torus_infiniband_dp
+      size: auto
+      topology:
+        type: Ring
+        bandwidth: 25 GB
+        latency: 5e-6
+        energy_per_bit: 8e-12 # (J/bit)
+        util: 1.0
+      parallelisms: []
+      collective_override: {} # overrides execution_backend.collectives when provided
+
+  overlap:
+    # fraction of TP+SP compute that can overlap TP/SP comm
+    tp_sp_overlap: 0.63
+    # fraction of TP compute that can overlap TP comm
+    tp_overlap: 0.6
+    # fraction of CP comm that can overlap CP compute
+    cp_overlap: 0.0
+execution_backend:
+  model: astra # analytical | astra
+  astra:
+    backend: astra
+    mode: full_astrasim_hierarchical  # hybrid | full_astrasim_hierarchical | full_astrasim_flattened
+    collectives:
+      # Default collective algorithms. Per-dimension collective_override entries can override these.
+      # Allowed values per op: auto | ring | direct | halvingDoubling | doubleBinaryTree | mesh | mesh2d | hypercube | torus2d
+      # Note: leave to auto unless you have a specific reason to change it.
+      all_gather: auto
+      all_reduce: auto
+      reduce_scatter: auto
+      all_to_all: auto
+    sys_options:
+      endpoint_delay: 5
+      active_chunks_per_dimension: 32 # Higher value means significantly slower astrasim runtime, but more accurate especially for tensor parallelism
+      preferred_dataset_splits: 4 # Should be same as active_chunks_per_dimension
+
+inference:
+  kvcache_type: hbm_only # Only supported mode for now.
diff --git a/validation_scripts/validation_configs/model-config/GPT_175_B.yaml b/validation_scripts/validation_configs/model-config/GPT_175_B.yaml
new file mode 100644
index 00000000..ba591aee
--- /dev/null
+++ b/validation_scripts/validation_configs/model-config/GPT_175_B.yaml
@@ -0,0 +1,21 @@
+model_param:
+  mode: "LLM"
+  run_type: "training"
+  tied_embeddings: false
+  model_type: "gpt"
+  global_batch_size: 8
+  gradient_accumulation_steps: 1
+  recomputation: partial
+  seq_len: 2048
+  hidden_dim: 12288
+  attention:
+    attention_type: "mha"
+    num_heads: 96
+    kv_heads: 96
+    use_flashattention: false
+    attention_tile_size: 128
+  num_experts: 1
+  top_k: 1
+  intermediate_size: 49152
+  vocab_size: 51200
+  num_layers: 96
diff --git a/validation_scripts/validation_configs/model-config/GPT_1T.yaml b/validation_scripts/validation_configs/model-config/GPT_1T.yaml
new file mode 100644
index 00000000..5f631bd9
--- /dev/null
+++ b/validation_scripts/validation_configs/model-config/GPT_1T.yaml
@@ -0,0 +1,21 @@
+model_param:
+  mode: "LLM"
+  run_type: "training"
+  tied_embeddings: false
+  model_type: "gpt"
+  global_batch_size: 8
+  gradient_accumulation_steps: 1
+  recomputation: partial
+  seq_len: 2048
+  hidden_dim: 25600
+  attention:
+    attention_type: "mha"
+    num_heads: 160
+    kv_heads: 160
+    use_flashattention: false
+    attention_tile_size: 128
+  num_experts: 1
+  top_k: 1
+  intermediate_size: 102400
+  vocab_size: 51200
+  num_layers: 128
diff --git a/validation_scripts/validation_configs/model-config/GPT_22_B.yaml b/validation_scripts/validation_configs/model-config/GPT_22_B.yaml
new file mode 100644
index 00000000..bc1e43b6
--- /dev/null
+++ b/validation_scripts/validation_configs/model-config/GPT_22_B.yaml
@@ -0,0 +1,21 @@
+model_param:
+  mode: "LLM"
+  run_type: "training"
+  tied_embeddings: false
+  model_type: "gpt"
+  global_batch_size: 4
+  gradient_accumulation_steps: 1
+  recomputation: selective
+  seq_len: 2048
+  hidden_dim: 6144
+  attention:
+    attention_type: "mha"
+    num_heads: 64
+    kv_heads: 64
+    use_flashattention: false
+    attention_tile_size: 128
+  num_experts: 1
+  top_k: 1
+  intermediate_size: 24576
+  vocab_size: 51200
+  num_layers: 48
diff --git a/validation_scripts/validation_configs/model-config/GPT_310_B.yaml b/validation_scripts/validation_configs/model-config/GPT_310_B.yaml
new file mode 100644
index 00000000..3988181a
--- /dev/null
+++ b/validation_scripts/validation_configs/model-config/GPT_310_B.yaml
@@ -0,0 +1,21 @@
+model_param:
+  mode: "LLM"
+  run_type: "training"
+  tied_embeddings: false
+  model_type: "gpt"
+  global_batch_size: 8
+  gradient_accumulation_steps: 1
+  recomputation: partial
+  seq_len: 2048
+  hidden_dim: 16384
+  attention:
+    attention_type: "mha"
+    num_heads: 128
+    kv_heads: 128
+    use_flashattention: false
+    attention_tile_size: 128
+  num_experts: 1
+  top_k: 1
+  intermediate_size: 65536
+  vocab_size: 51200
+  num_layers: 96
diff --git a/validation_scripts/validation_configs/model-config/GPT_530_B.yaml b/validation_scripts/validation_configs/model-config/GPT_530_B.yaml
new file mode 100644
index 00000000..25fcbb10
--- /dev/null
+++ b/validation_scripts/validation_configs/model-config/GPT_530_B.yaml
@@ -0,0 +1,21 @@
+model_param:
+  mode: "LLM"
+  run_type: "training"
+  tied_embeddings: false
+  model_type: "gpt"
+  global_batch_size: 8
+  gradient_accumulation_steps: 1
+  recomputation: partial
+  seq_len: 2048
+  hidden_dim: 20480
+  attention:
+    attention_type: "mha"
+    num_heads: 128
+    kv_heads: 128
+    use_flashattention: false
+    attention_tile_size: 128
+  num_experts: 1
+  top_k: 1
+  intermediate_size: 81920
+  vocab_size: 51200
+  num_layers: 105
diff --git a/validation_scripts/validation_configs/model-config/Llama2-13B_inf.yaml b/validation_scripts/validation_configs/model-config/Llama2-13B_inf.yaml
new file mode 100644
index 00000000..a5105fa0
--- /dev/null
+++ b/validation_scripts/validation_configs/model-config/Llama2-13B_inf.yaml
@@ -0,0 +1,21 @@
+model_param:
+  mode: "LLM"  # Mode of operation
+  run_type: "inference" # Type of run: "training" or "inference"
+  tied_embeddings: true  # Share input and output embedding weights (WIP: affects memory estimation only)
+  model_type: "llama"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
+  global_batch_size: 16  # Sequences per global batch
+  gradient_accumulation_steps: 1  # Number of gradient accumulation micro-steps per optimizer update
+  seq_len: 2048  # Total context length (prefill + decode_len for inference)
+  decode_len: 1024 # Number of tokens to decode during inference
+  hidden_dim: 5120 # Transformer hidden dimension
+  attention: # Attention-related parameters
+    attention_type: "mha"  # Supported attention types: "mha" or "gqa" ("mla" not implemented)
+    num_heads: 40 # Total number of attention heads
+    kv_heads: 8  # Number of key/value heads (used only when attention_type = "gqa")
+    use_flashattention: true  # Enable FlashAttention optimization
+    attention_tile_size: 256  # Tile size used in attention computation
+  num_experts: 1   # Number of experts in MoE layers. Set to 1 to disable MoE.
+  top_k: 1   # Number of active experts per token in MoE
+  intermediate_size: 13824 # Output dimension of the first FFN layer (typically 4×hidden_dim)
+  vocab_size: 32000  # Vocabulary size
+  num_layers: 40 # Number of transformer layers
diff --git a/validation_scripts/validation_configs/model-config/Llama2-70B_inf.yaml b/validation_scripts/validation_configs/model-config/Llama2-70B_inf.yaml
new file mode 100644
index 00000000..24facac3
--- /dev/null
+++ b/validation_scripts/validation_configs/model-config/Llama2-70B_inf.yaml
@@ -0,0 +1,21 @@
+model_param:
+  mode: "LLM"  # Mode of operation
+  run_type: "inference" # Type of run: "training" or "inference"
+  tied_embeddings: true  # Share input and output embedding weights (WIP: affects memory estimation only)
+  model_type: "llama"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
+  global_batch_size: 16  # Sequences per global batch
+  gradient_accumulation_steps: 1  # Number of gradient accumulation micro-steps per optimizer update
+  seq_len: 2048  # Total context length (prefill + decode_len for inference)
+  decode_len: 1024 # Number of tokens to decode during inference
+  hidden_dim: 8192 # Transformer hidden dimension
+  attention: # Attention-related parameters
+    attention_type: "gqa"  # Supported attention types: "mha" or "gqa" ("mla" not implemented)
+    num_heads: 64 # Total number of attention heads
+    kv_heads: 8  # Number of key/value heads (used only when attention_type = "gqa")
+    use_flashattention: true  # Enable FlashAttention optimization
+    attention_tile_size: 256  # Tile size used in attention computation
+  num_experts: 1   # Number of experts in MoE layers. Set to 1 to disable MoE.
+  top_k: 1   # Number of active experts per token in MoE
+  intermediate_size: 28672 # Output dimension of the first FFN layer (typically 4×hidden_dim)
+  vocab_size: 32000  # Vocabulary size
+  num_layers: 80 # Number of transformer layers
diff --git a/validation_scripts/validation_configs/model-config/Llama2-7B.yaml b/validation_scripts/validation_configs/model-config/Llama2-7B.yaml
new file mode 100644
index 00000000..2d13a809
--- /dev/null
+++ b/validation_scripts/validation_configs/model-config/Llama2-7B.yaml
@@ -0,0 +1,20 @@
+model_param:
+  mode: "LLM"  # Mode of operation
+  run_type: "training" # Type of run: "training" or "inference"
+  tied_embeddings: true  # Share input and output embedding weights (WIP: affects memory estimation only)
+  model_type: "llama"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
+  global_batch_size: 128  # Sequences per global batch (validation default)
+  gradient_accumulation_steps: 1  # Number of gradient accumulation micro-steps per optimizer update
+  seq_len: 2048  # Sequence length (tokens per sequence)
+  hidden_dim: 4096 # Transformer hidden dimension
+  attention: # Attention-related parameters
+    attention_type: "mha"  # Supported attention types: "mha" or "gqa" ("mla" not implemented)
+    num_heads: 32 # Total number of attention heads
+    kv_heads: null  # Number of key/value heads (used only when attention_type = "gqa")
+    use_flashattention: true  # Enable FlashAttention optimization
+    attention_tile_size: 128  # Tile size used in attention computation
+  num_experts: 1   # Number of experts in MoE layers. Set to 1 to disable MoE.
+  top_k: 1   # Number of active experts per token in MoE
+  intermediate_size: 11008 # Output dimension of the first FFN layer (typically 4×hidden_dim)
+  vocab_size: 32000  # Vocabulary size
+  num_layers: 32 # Number of transformer layers
diff --git a/validation_scripts/validation_configs/model-config/Llama2-7B_inf.yaml b/validation_scripts/validation_configs/model-config/Llama2-7B_inf.yaml
new file mode 100644
index 00000000..0bed4a63
--- /dev/null
+++ b/validation_scripts/validation_configs/model-config/Llama2-7B_inf.yaml
@@ -0,0 +1,21 @@
+model_param:
+  mode: "LLM"  # Mode of operation
+  run_type: "inference" # Type of run: "training" or "inference"
+  tied_embeddings: true  # Share input and output embedding weights (WIP: affects memory estimation only)
+  model_type: "llama"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
+  global_batch_size: 16  # Sequences per global batch
+  gradient_accumulation_steps: 1  # Number of gradient accumulation micro-steps per optimizer update
+  seq_len: 2048  # Total context length (prefill + decode_len for inference)
+  decode_len: 1024 # Number of tokens to decode during inference
+  hidden_dim: 4096 # Transformer hidden dimension
+  attention: # Attention-related parameters
+    attention_type: "mha"  # Supported attention types: "mha" or "gqa" ("mla" not implemented)
+    num_heads: 32 # Total number of attention heads
+    kv_heads: 8  # Number of key/value heads (used only when attention_type = "gqa")
+    use_flashattention: true  # Enable FlashAttention optimization
+    attention_tile_size: 256  # Tile size used in attention computation
+  num_experts: 1   # Number of experts in MoE layers. Set to 1 to disable MoE.
+  top_k: 1   # Number of active experts per token in MoE
+  intermediate_size: 11008 # Output dimension of the first FFN layer (typically 4×hidden_dim)
+  vocab_size: 32000  # Vocabulary size
+  num_layers: 32 # Number of transformer layers
diff --git a/validation_scripts/validation_configs/model-config/Llama3.1-8B_inf.yaml b/validation_scripts/validation_configs/model-config/Llama3.1-8B_inf.yaml
new file mode 100644
index 00000000..439048b9
--- /dev/null
+++ b/validation_scripts/validation_configs/model-config/Llama3.1-8B_inf.yaml
@@ -0,0 +1,21 @@
+model_param:
+  mode: "LLM"  # Mode of operation
+  run_type: "inference" # Type of run: "training" or "inference"
+  tied_embeddings: true  # Share input and output embedding weights (WIP: affects memory estimation only)
+  model_type: "llama"  # gpt assumes dense-style architecture with GELU; llama assumes llama-style architecture with SwiGLU
+  global_batch_size: 1  # Sequences per global batch
+  gradient_accumulation_steps: 1  # Number of gradient accumulation micro-steps per optimizer update
+  seq_len: 2048  # Total context length (prefill + decode_len for inference)
+  decode_len: 1024 # Number of tokens to decode during inference
+  hidden_dim: 4096 # Transformer hidden dimension
+  attention: # Attention-related parameters
+    attention_type: "gqa"  # Supported attention types: "mha" or "gqa" ("mla" not implemented)
+    num_heads: 32 # Total number of attention heads
+    kv_heads: 8  # Number of key/value heads (used only when attention_type = "gqa")
+    use_flashattention: true  # Enable FlashAttention optimization
+    attention_tile_size: 128  # Tile size used in attention computation
+  num_experts: 1   # Number of experts in MoE layers. Set to 1 to disable MoE.
+  top_k: 1   # Number of active experts per token in MoE
+  intermediate_size: 14336 # Output dimension of the first FFN layer (typically 4×hidden_dim)
+  vocab_size: 128256  # Vocabulary size
+  num_layers: 32 # Number of transformer layers
diff --git a/validation_scripts/validation_helpers.py b/validation_scripts/validation_helpers.py
new file mode 100644
index 00000000..a361cd77
--- /dev/null
+++ b/validation_scripts/validation_helpers.py
@@ -0,0 +1,436 @@
+from __future__ import annotations
+
+import copy
+import os
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+import time
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from dataclasses import dataclass, field, replace
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple
+
+import pytest
+import yaml
+
+PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if PROJECT_ROOT not in sys.path:
+  sys.path.insert(0, PROJECT_ROOT)
+
+from tools.parallelism_sweep import set_astrasim_cache_mode  # type: ignore
+
+VALIDATION_WORKERS_ENV = "DEEPFLOW_VALIDATION_WORKERS"
+DEFAULT_WORKER_COUNT = 8
+
+ResultParser = Callable[[str, "ValidationSpec"], Dict[str, Any]]
+
+
+@dataclass(frozen=True)
+class ValidationSpec:
+  label: str
+  model_overrides: Optional[Mapping[str, Any]] = None
+  hardware_overrides: Optional[Mapping[str, Any]] = None
+  metadata: Dict[str, Any] = field(default_factory=dict)
+  model_config_path: Optional[str] = None
+  hardware_config_path: Optional[str] = None
+  order: int = 0
+
+
+@dataclass
+class ValidationResult:
+  spec: ValidationSpec
+  success: bool
+  metrics: Dict[str, Any]
+  raw_output: str
+  error: Optional[str] = None
+  returncode: Optional[int] = None
+  duration_s: float = 0.0
+  model_config_used: Optional[str] = None
+  hardware_config_used: Optional[str] = None
+
+
+@dataclass
+class _WorkerContext:
+  run_perf_path: str
+  python_executable: str
+  tmp_root: str
+  env_overrides: Dict[str, str]
+  model_config_cache: Dict[str, Dict[str, Any]]
+  hardware_config_cache: Dict[str, Dict[str, Any]]
+  default_model_config_path: str
+  default_hardware_config_path: str
+  result_parser: ResultParser
+  cache_mode: str
+  extra_run_perf_args: Tuple[str, ...]
+
+
+_WORKER_CONTEXT: Optional[_WorkerContext] = None
+
+
+def _worker_init(context: _WorkerContext) -> None:
+  global _WORKER_CONTEXT
+  _WORKER_CONTEXT = context
+  set_astrasim_cache_mode(context.cache_mode)
+
+def _is_list_of_dicts(value: Any) -> bool:
+  if not isinstance(value, list):
+    return False  
+  return all(isinstance(item, dict) for item in value)
+
+def _merge_list_of_dicts(orig: List[Dict[str, Any]], 
+                         overrides: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+  orig_index = {d["id"]: d for d in orig}
+  orig_order = [d["id"] for d in orig]
+  
+  # Apply deep updates (or add new items)
+  for d in overrides:
+    id = d["id"]
+    if id in orig_index:
+      orig_index[id] = _deep_update(copy.deepcopy(orig_index[id]), d)
+    else:
+      orig_index[id] = copy.deepcopy(d)
+
+  # Rebuild list preserving original order, appending new items at the end
+  result: List[Dict[str, Any]] = []
+  for orig_id in orig_order:
+    result.append(orig_index[orig_id])
+  for d in overrides:
+    id = d["id"]
+    if id not in orig_order:
+      result.append(orig_index[id])
+  return result
+
+def _deep_update(target: Dict[str, Any], overrides: Mapping[str, Any]) -> Dict[str, Any]:
+  for key, value in overrides.items():
+    if isinstance(value, Mapping):
+      existing = target.get(key)
+      if not isinstance(existing, dict):
+        existing = {}
+      target[key] = _deep_update(existing, value)
+    elif isinstance(value, list):
+      existing_list = target.get(key)
+      if isinstance(existing_list, list) and _is_list_of_dicts(existing_list) and _is_list_of_dicts(value):
+        target[key] = _merge_list_of_dicts(existing_list, value)
+      else:
+        target[key] = copy.deepcopy(value)
+    else:
+      target[key] = value
+  return target
+
+
+def _load_yaml(path: str) -> Dict[str, Any]:
+  with open(path, "r") as handle:
+    data = yaml.safe_load(handle)
+  if not isinstance(data, dict):
+    raise ValueError(f"{path} did not contain a YAML dictionary.")
+  return data
+
+
+def _write_yaml(path: str, data: Dict[str, Any]) -> None:
+  with open(path, "w") as handle:
+    yaml.safe_dump(data, handle, sort_keys=False)
+
+
+def _prepare_config(
+  tmp_dir: str,
+  base_path: str,
+  base_dict: Dict[str, Any],
+  overrides: Optional[Mapping[str, Any]],
+  filename: str,
+) -> str:
+  if not overrides:
+    return base_path
+  updated = copy.deepcopy(base_dict)
+  _deep_update(updated, overrides)
+  config_path = os.path.join(tmp_dir, filename)
+  _write_yaml(config_path, updated)
+  return config_path
+
+
+def _read_env_worker_setting() -> Tuple[int, str]:
+  env_raw = os.environ.get(VALIDATION_WORKERS_ENV)
+  if env_raw is None:
+    return DEFAULT_WORKER_COUNT, f"{VALIDATION_WORKERS_ENV} unset (default {DEFAULT_WORKER_COUNT})"
+  try:
+    value = int(env_raw)
+  except (TypeError, ValueError):
+    print(f"[validation] WARNING: {VALIDATION_WORKERS_ENV} must be a positive integer (got {env_raw!r}); using default {DEFAULT_WORKER_COUNT}.")
+    return DEFAULT_WORKER_COUNT, f"{VALIDATION_WORKERS_ENV} invalid ({env_raw!r}) -> default {DEFAULT_WORKER_COUNT}"
+  if value <= 0:
+    print(f"[validation] WARNING: {VALIDATION_WORKERS_ENV} must be > 0 (got {env_raw!r}); using default {DEFAULT_WORKER_COUNT}.")
+    return DEFAULT_WORKER_COUNT, f"{VALIDATION_WORKERS_ENV} <= 0 ({env_raw!r}) -> default {DEFAULT_WORKER_COUNT}"
+  return value, f"{VALIDATION_WORKERS_ENV}={value}"
+
+
+def _determine_worker_count(spec_count: int, explicit_max: Optional[int]) -> int:
+  env_workers, env_desc = _read_env_worker_setting()
+  requested = explicit_max if explicit_max is not None else env_workers
+  request_desc = f"max_workers={explicit_max}" if explicit_max is not None else env_desc
+  if requested <= 0:
+    print(f"[validation] WARNING: requested worker budget {requested} is not positive; bumping to 1.")
+    requested = 1
+  cpu_limit = max(1, os.cpu_count() or 1)
+  if explicit_max is None and requested > cpu_limit:
+    print(f"[validation] WARNING: {VALIDATION_WORKERS_ENV} requested {requested} worker(s) but only {cpu_limit} CPU core(s) detected. Capping to {cpu_limit}.")
+    requested = cpu_limit
+  elif explicit_max is not None and requested > cpu_limit:
+    print(f"[validation] WARNING: max_workers requested {requested} worker(s) but only {cpu_limit} CPU core(s) detected. Capping to {cpu_limit}.")
+    requested = cpu_limit
+  worker_count = max(1, min(spec_count, requested))
+  print(f"[validation] Using {worker_count} worker(s) for {spec_count} experiment(s) "
+        f"(requested {request_desc}, cpu_limit={cpu_limit}).")
+  return worker_count
+
+
+def _execute_spec(spec: ValidationSpec) -> ValidationResult:
+  if _WORKER_CONTEXT is None:
+    raise RuntimeError("Worker context is not initialised.")
+  ctx = _WORKER_CONTEXT
+  tmp_dir = tempfile.mkdtemp(prefix="validation_run_", dir=ctx.tmp_root)
+  env = os.environ.copy()
+  env.update(ctx.env_overrides)
+  env.setdefault("ASTRA_CACHE_DIR", os.path.join(tmp_dir, "astra_cache"))
+  start_time = time.time()
+  base_model_path = os.path.abspath(spec.model_config_path or ctx.default_model_config_path)
+  base_hw_path = os.path.abspath(spec.hardware_config_path or ctx.default_hardware_config_path)
+  if base_model_path not in ctx.model_config_cache:
+    raise ValueError(f"Model config {base_model_path} was not preloaded for validation.")
+  if base_hw_path not in ctx.hardware_config_cache:
+    raise ValueError(f"Hardware config {base_hw_path} was not preloaded for validation.")
+  model_config_path = base_model_path
+  hardware_config_path = base_hw_path
+  try:
+    model_config_path = _prepare_config(
+      tmp_dir,
+      base_model_path,
+      ctx.model_config_cache[base_model_path],
+      spec.model_overrides,
+      "model.yaml",
+    )
+    hardware_config_path = _prepare_config(
+      tmp_dir,
+      base_hw_path,
+      ctx.hardware_config_cache[base_hw_path],
+      spec.hardware_overrides,
+      "hardware.yaml",
+    )
+    cmd = [
+      ctx.python_executable,
+      ctx.run_perf_path,
+      "--hardware_config",
+      hardware_config_path,
+      "--model_config",
+      model_config_path,
+    ]
+    if ctx.extra_run_perf_args:
+      cmd.extend(ctx.extra_run_perf_args)
+    proc = subprocess.run(
+      cmd,
+      cwd=tmp_dir,
+      stdout=subprocess.PIPE,
+      stderr=subprocess.STDOUT,
+      text=True,
+      check=False,
+      env=env,
+    )
+    duration = time.time() - start_time
+    output = proc.stdout or ""
+    if proc.returncode != 0:
+      return ValidationResult(
+        spec=spec,
+        success=False,
+        metrics={},
+        raw_output=output,
+        error=f"DeepFlow exited with status {proc.returncode}",
+        returncode=proc.returncode,
+        duration_s=duration,
+        model_config_used=model_config_path,
+        hardware_config_used=hardware_config_path,
+      )
+    try:
+      metrics = ctx.result_parser(output, spec)
+    except Exception as exc:
+      return ValidationResult(
+        spec=spec,
+        success=False,
+        metrics={},
+        raw_output=output,
+        error=f"Failed to parse DeepFlow output: {exc}",
+        returncode=proc.returncode,
+        duration_s=duration,
+        model_config_used=model_config_path,
+        hardware_config_used=hardware_config_path,
+      )
+    return ValidationResult(
+      spec=spec,
+      success=True,
+      metrics=metrics,
+      raw_output=output,
+      error=None,
+      returncode=proc.returncode,
+      duration_s=duration,
+      model_config_used=model_config_path,
+      hardware_config_used=hardware_config_path,
+    )
+  finally:
+    shutil.rmtree(tmp_dir, ignore_errors=True)
+
+
+def run_validation_suite(
+  specs: Sequence[ValidationSpec],
+  *,
+  base_model_config_path: str,
+  base_hardware_config_path: str,
+  result_parser: ResultParser,
+  run_perf_path: Optional[str] = None,
+  python_executable: Optional[str] = None,
+  tmp_root: Optional[str] = None,
+  max_workers: Optional[int] = None,
+  env_overrides: Optional[Dict[str, str]] = None,
+  cache_mode: str = "NO_CACHE",
+  extra_run_perf_args: Optional[Sequence[str]] = None,
+  show_progress: bool = False,
+) -> List[ValidationResult]:
+  if not specs:
+    return []
+  run_perf_path = os.path.abspath(run_perf_path or os.path.join(PROJECT_ROOT, "run_perf.py"))
+  python_executable = python_executable or sys.executable
+  tmp_root = os.path.abspath(tmp_root or os.path.join(PROJECT_ROOT, "tmp", "validation_runs"))
+  os.makedirs(tmp_root, exist_ok=True)
+  set_astrasim_cache_mode(cache_mode)
+  env_map = dict(env_overrides or {})
+  env_map.setdefault("DEEPFLOW_ASTRA_CACHE_MODE", os.environ.get("DEEPFLOW_ASTRA_CACHE_MODE", cache_mode))
+  model_paths = {os.path.abspath(base_model_config_path)}
+  hardware_paths = {os.path.abspath(base_hardware_config_path)}
+  for spec in specs:
+    if spec.model_config_path:
+      model_paths.add(os.path.abspath(spec.model_config_path))
+    if spec.hardware_config_path:
+      hardware_paths.add(os.path.abspath(spec.hardware_config_path))
+  model_cache = {path: _load_yaml(path) for path in model_paths}
+  hardware_cache = {path: _load_yaml(path) for path in hardware_paths}
+  extra_args: Tuple[str, ...] = tuple(extra_run_perf_args or ())
+  assigned_specs = [replace(spec, order=idx) for idx, spec in enumerate(specs)]
+  worker_count = _determine_worker_count(len(assigned_specs), max_workers)
+  context = _WorkerContext(
+    run_perf_path=run_perf_path,
+    python_executable=python_executable,
+    tmp_root=tmp_root,
+    env_overrides=env_map,
+    model_config_cache=model_cache,
+    hardware_config_cache=hardware_cache,
+    default_model_config_path=os.path.abspath(base_model_config_path),
+    default_hardware_config_path=os.path.abspath(base_hardware_config_path),
+    result_parser=result_parser,
+    cache_mode=os.environ.get("DEEPFLOW_ASTRA_CACHE_MODE", cache_mode),
+    extra_run_perf_args=extra_args,
+  )
+  results: List[ValidationResult] = []
+  progress = None
+  if show_progress:
+    try:
+      from tqdm import tqdm  # type: ignore
+      progress = tqdm(total=len(assigned_specs), desc="validation", leave=True)
+    except Exception:
+      progress = None
+
+  with ProcessPoolExecutor(max_workers=worker_count, initializer=_worker_init, initargs=(context,)) as executor:
+    future_map = {executor.submit(_execute_spec, spec): spec for spec in assigned_specs}
+    for future in as_completed(future_map):
+      spec = future_map[future]
+      try:
+        result = future.result()
+      except Exception as exc:
+        result = ValidationResult(
+          spec=spec,
+          success=False,
+          metrics={},
+          raw_output="",
+          error=str(exc),
+          returncode=None,
+          duration_s=0.0,
+          model_config_used=None,
+          hardware_config_used=None,
+        )
+      results.append(result)
+      if progress is not None:
+        progress.update(1)
+  if progress is not None:
+    progress.close()
+  results.sort(key=lambda item: item.spec.order)
+  return results
+
+
+_DECODE_TIME_REGEX = re.compile(
+  r"\[prefill\]\s*time:\s*[0-9.]+s,\s*\[decode\]\s*time:\s*([0-9.]+)s,\s*\[total\]\s*time:\s*[0-9.]+s"
+)
+
+_INF_TIME_REGEX = re.compile(
+  r"LLM inference time:\s*([0-9]+(?:\.[0-9]+)?)s"
+)
+_TRAIN_TIME_REGEX = re.compile(
+  r"Training time for batch:\s*([0-9]+(?:\.[0-9]+)?)s"
+)
+
+def parse_decode_time(output: str, spec: ValidationSpec) -> Dict[str, Any]:
+  match = _DECODE_TIME_REGEX.search(output)
+  if not match:
+    raise ValueError(f"Failed to parse decode time for experiment '{spec.label}'.")
+  return {"decode_time_s": float(match.group(1))}
+
+def parse_inference_time(output: str, spec: ValidationSpec) -> Dict[str, Any]:
+  match = _INF_TIME_REGEX.search(output)
+  if not match:
+    raise ValueError(f"Failed to parse inference time for experiment '{spec.label}'.")
+  return {"inference_time_s": float(match.group(1))}
+
+
+def parse_training_time(output: str, spec: ValidationSpec) -> Dict[str, Any]:
+  match = _TRAIN_TIME_REGEX.search(output)
+  if not match:
+    raise ValueError(f"Failed to parse training time for experiment '{spec.label}'.")
+  return {"training_time_s": float(match.group(1))}
+
+# ---- Pytest integration for quiet validation reporting ----
+_VALIDATION_METRICS: List[Tuple[str, str, float, Optional[float]]] = []
+
+# magic function, needed for pytest to work
+def pytest_configure(config):
+  _VALIDATION_METRICS.clear()
+
+
+@pytest.fixture
+def record_validation(request):
+  def _record(label: str, achieved_pct: float, expected_pct: Optional[float] = None) -> None:
+    #node id is tests/test_IMEC_A100_no_network.py::test_avg_abs_error_below_threshold_network_ignored
+    # I want just the filename under tests/, not the function name
+    filename = request.node.nodeid.split("::")[0]
+    filename = filename.split("/")[-1]
+    _VALIDATION_METRICS.append(
+      (filename, label, achieved_pct, expected_pct)
+    )
+  return _record
+
+
+def pytest_terminal_summary(terminalreporter, exitstatus):
+  if not _VALIDATION_METRICS:
+    return
+  tr = terminalreporter
+  tr.write_sep("=", "Test results")
+  for nodeid, label, achieved, expected in _VALIDATION_METRICS:
+    line = f"{nodeid}: {label} achieved {achieved:.2f}%"
+    if expected is not None:
+      line += f" (expected <= {expected:.2f}%)"
+    tr.write_line(line)
+
+__all__ = [
+  "ValidationSpec",
+  "ValidationResult",
+  "run_validation_suite",
+  "parse_decode_time",
+  "parse_inference_time",
+  "parse_training_time",
+  "record_validation",
+]