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

Preserve global token count in Frontier EP profiles

Browse Source
This commit is contained in:
Gahow Wang
2026-07-15 20:23:21 +08:00
parent 76e0de33ee
commit d64f12eb5b
1 changed files with 75 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

75
runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_ep_feature_contract.patch Normal file
View File

@@ -0,0 +1,75 @@
diff --git a/frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py b/frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py
--- a/frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py
+++ b/frontier/execution_time_predictor/sklearn_moe_execution_time_predictor.py
@@ -424 +424,4 @@
- features = self._build_moe_load_imbalance_features(allocation)
+ features = self._build_moe_load_imbalance_features(
+ allocation,
+ num_tokens=total_routed_tokens // self._router_topk,
+ )
@@ -1365,0 +1369 @@
+ num_tokens: Optional[int] = None,
@@ -1384,3 +1388,6 @@
- approx_num_tokens = max(
- 1, int(round(total_routed_tokens / float(self._router_topk)))
- )
+ if num_tokens is None:
+ num_tokens = max(
+ 1, int(round(total_routed_tokens / float(self._router_topk)))
+ )
+ elif num_tokens <= 0:
+ raise ValueError(f"num_tokens must be positive, got {num_tokens}")
@@ -1389 +1396 @@
- num_tokens=approx_num_tokens,
+ num_tokens=int(num_tokens),
@@ -1475,0 +1483,3 @@
+ profile_num_tokens = int(
+ self._get_effective_moe_total_tokens(batch)
+ )
@@ -1478,0 +1489 @@
+ profile_num_tokens,
@@ -1486 +1497,2 @@
- per_expert_tokens
+ per_expert_tokens,
+ num_tokens=profile_num_tokens,
@@ -1778,4 +1790,4 @@
- # The grouped_gemm predictor expects pre-routing num_tokens metadata, while runtime
- # allocation is post-routing (already expanded by router_topk). Recover the
- # approximate pre-routing token count to keep feature semantics aligned with
- # the profiling dataset contract.
+ # The grouped_gemm predictor keeps global pre-routing num_tokens and
+ # lane-local routed assignments as independent features. Inferring the
+ # former from the latter is valid for EP=1 but undercounts by EP size
+ # for expert-parallel execution.
@@ -1784,2 +1796,7 @@
- approx_num_tokens = max(
- 1, int(round(total_routed_tokens / float(self._router_topk)))
+ profile_num_tokens = (
+ int(self._get_effective_moe_total_tokens(batch))
+ if batch is not None
+ else max(
+ 1,
+ int(round(total_routed_tokens / float(self._router_topk))),
+ )
@@ -1804 +1821 @@
- approx_num_tokens,
+ profile_num_tokens,
@@ -1811 +1828,4 @@
- features = self._build_moe_load_imbalance_features(per_expert_tokens)
+ features = self._build_moe_load_imbalance_features(
+ per_expert_tokens,
+ num_tokens=profile_num_tokens,
+ )
@@ -1851,2 +1871,7 @@
- approx_num_tokens = max(
- 1, int(round(total_routed_tokens / float(self._router_topk)))
+ profile_num_tokens = (
+ int(self._get_effective_moe_total_tokens(batch))
+ if batch is not None
+ else max(
+ 1,
+ int(round(total_routed_tokens / float(self._router_topk))),
+ )
@@ -1854 +1879 @@
- rounded_tokens = self._round_to_valid_key(approx_num_tokens)
+ rounded_tokens = self._round_to_valid_key(profile_num_tokens)