From a54f69352de209922ed01794d995d3cfd4e72d83 Mon Sep 17 00:00:00 2001
From: Gahow Wang <gahow.wang@gmail.com>
Date: Wed, 15 Jul 2026 19:28:39 +0800
Subject: [PATCH] Add Frontier MoE gating context closure

---
 .../best_effort/assemble_profiles.py          | 71 ++++++++++++++++---
 1 file changed, 63 insertions(+), 8 deletions(-)

diff --git a/runs/frontier-multicase-sufficiency-v0/best_effort/assemble_profiles.py b/runs/frontier-multicase-sufficiency-v0/best_effort/assemble_profiles.py
index a9e4053..a060526 100644
--- a/runs/frontier-multicase-sufficiency-v0/best_effort/assemble_profiles.py
+++ b/runs/frontier-multicase-sufficiency-v0/best_effort/assemble_profiles.py
@@ -28,6 +28,8 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--attention-mixed", type=Path, required=True)
     parser.add_argument("--moe-tp4-ep1", type=Path, required=True)
     parser.add_argument("--moe-tp1-ep8", type=Path, required=True)
+    parser.add_argument("--moe-standalone-tp4-ep1", type=Path, required=True)
+    parser.add_argument("--moe-standalone-tp1-ep8", type=Path, required=True)
     parser.add_argument("--output-root", type=Path, required=True)
     parser.add_argument(
         "--frontier-commit",
@@ -190,7 +192,12 @@ def assemble_attention(standard: pd.DataFrame, mixed: pd.DataFrame) -> pd.DataFr
     return combined.sort_values(sort_columns, kind="stable").reset_index(drop=True)
 
 
-def assemble_moe(tp4_ep1: pd.DataFrame, tp1_ep8: pd.DataFrame) -> pd.DataFrame:
+def assemble_moe(
+    tp4_ep1: pd.DataFrame,
+    tp1_ep8: pd.DataFrame,
+    standalone_tp4_ep1: pd.DataFrame,
+    standalone_tp1_ep8: pd.DataFrame,
+) -> pd.DataFrame:
     timing_columns = [
         "time_stats.moe_gating_linear.median",
         "time_stats.moe_gating_routing_topk.median",
@@ -198,10 +205,33 @@ def assemble_moe(tp4_ep1: pd.DataFrame, tp1_ep8: pd.DataFrame) -> pd.DataFrame:
         "time_stats.moe_grouped_gemm.median",
     ]
     cases = (
-        ("moe-tp4-ep1", tp4_ep1, 4, 1, 128),
-        ("moe-tp1-ep8", tp1_ep8, 1, 8, 16),
+        ("moe-tp4-ep1", tp4_ep1, 4, 1, 128, "prefill_hot"),
+        ("moe-tp1-ep8", tp1_ep8, 1, 8, 16, "prefill_hot"),
+        (
+            "moe-standalone-tp4-ep1",
+            standalone_tp4_ep1,
+            4,
+            1,
+            128,
+            "standalone_legacy",
+        ),
+        (
+            "moe-standalone-tp1-ep8",
+            standalone_tp1_ep8,
+            1,
+            8,
+            16,
+            "standalone_legacy",
+        ),
     )
-    for label, data, expected_tp, expected_ep, expected_local_experts in cases:
+    for (
+        label,
+        data,
+        expected_tp,
+        expected_ep,
+        expected_local_experts,
+        expected_context,
+    ) in cases:
         validate_common(data, label)
         require_columns(
             data,
@@ -228,22 +258,36 @@ def assemble_moe(tp4_ep1: pd.DataFrame, tp1_ep8: pd.DataFrame) -> pd.DataFrame:
         )
         require_exact_values(data, "seed", {0, 1}, label)
         require_exact_values(data, "routing_runtime_path", {"standard_fused_topk"}, label)
-        require_exact_values(data, "gating_runtime_context", {"prefill_hot"}, label)
+        require_exact_values(
+            data, "gating_runtime_context", {expected_context}, label
+        )
         validate_token_grid(
             data,
             ["num_tensor_parallel_workers", "expert_parallel_size"],
             6,
             label,
         )
-        validate_time_columns(data, timing_columns, label)
+        validate_time_columns(
+            data,
+            timing_columns if expected_context == "prefill_hot" else timing_columns[:2],
+            label,
+        )
         if len(data) != 90:
             raise ValueError(f"{label}: expected 90 rows, got {len(data)}")
 
-    combined = pd.concat([tp4_ep1, tp1_ep8], ignore_index=True)
+    # Standalone rows are needed only for Frontier's pure-decode gating models.
+    # Do not duplicate the shuffling/grouped-GEMM observations merely because
+    # the profiler collected them while measuring a second gating context.
+    standalone = pd.concat(
+        [standalone_tp4_ep1, standalone_tp1_ep8], ignore_index=True
+    )
+    standalone[timing_columns[2:]] = float("nan")
+    combined = pd.concat([tp4_ep1, tp1_ep8, standalone], ignore_index=True)
     return combined.sort_values(
         [
             "num_tensor_parallel_workers",
             "expert_parallel_size",
+            "gating_runtime_context",
             "num_tokens",
             "load_distribution",
             "seed",
@@ -261,6 +305,8 @@ def main() -> None:
         "attention_mixed": args.attention_mixed,
         "moe_tp4_ep1": args.moe_tp4_ep1,
         "moe_tp1_ep8": args.moe_tp1_ep8,
+        "moe_standalone_tp4_ep1": args.moe_standalone_tp4_ep1,
+        "moe_standalone_tp1_ep8": args.moe_standalone_tp1_ep8,
     }
     inputs = {name: read_csv(path, name) for name, path in input_paths.items()}
 
@@ -269,7 +315,12 @@ def main() -> None:
         "attention.csv": assemble_attention(
             inputs["attention_standard"], inputs["attention_mixed"]
         ),
-        "moe.csv": assemble_moe(inputs["moe_tp4_ep1"], inputs["moe_tp1_ep8"]),
+        "moe.csv": assemble_moe(
+            inputs["moe_tp4_ep1"],
+            inputs["moe_tp1_ep8"],
+            inputs["moe_standalone_tp4_ep1"],
+            inputs["moe_standalone_tp1_ep8"],
+        ),
     }
 
     output_dir = args.output_root / "compute" / HARDWARE / MODEL
@@ -296,6 +347,10 @@ def main() -> None:
                 {"tensor_parallel_size": 4, "expert_parallel_size": 1},
                 {"tensor_parallel_size": 1, "expert_parallel_size": 8},
             ],
+            "moe_gating_runtime_contexts": [
+                "prefill_hot",
+                "standalone_legacy",
+            ],
         },
         "inputs": {
             name: {