gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Frontier MoE gating context closure

Browse Source
This commit is contained in:
Gahow Wang
2026-07-15 19:28:39 +08:00
parent c71f379110
commit a54f69352d
1 changed files with 63 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

71
runs/frontier-multicase-sufficiency-v0/best_effort/assemble_profiles.py
View File

@@ -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: {