Align Frontier piecewise graph profiles

runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_router.py

@@ -31,6 +31,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--warmup-iters", type=int, default=5)
     parser.add_argument("--repeats", type=int, default=20)
     parser.add_argument("--device", default="cuda:0")
+    parser.add_argument(
+        "--profile-method",
+        choices=("cuda_event", "record_function"),
+        default="cuda_event",
+    )
+    parser.add_argument("--frontier-source", type=Path)
     return parser.parse_args()
 
 
@@ -51,12 +57,34 @@ def stats_ms(samples: list[float]) -> dict[str, float]:
 
 
 def measure_ms(
-    fn: Callable[[], Any], warmup_iters: int, repeats: int
+    fn: Callable[[], Any],
+    warmup_iters: int,
+    repeats: int,
+    *,
+    profile_method: str,
+    trace_root: Path | None = None,
+    operation_name: str | None = None,
+    record_function_tracer: type | None = None,
 ) -> tuple[Any, dict[str, float]]:
     result = None
     for _ in range(warmup_iters):
         result = fn()
     torch.accelerator.synchronize()
+    if profile_method == "record_function":
+        if trace_root is None or operation_name is None or record_function_tracer is None:
+            raise ValueError("record_function profiling requires tracer metadata")
+        tracer = record_function_tracer(str(trace_root))
+        with tracer:
+            for _ in range(repeats):
+                with torch.profiler.record_function(f"vidur_{operation_name}"):
+                    result = fn()
+        stats = tracer.get_operation_time_stats()
+        if operation_name not in stats:
+            raise RuntimeError(f"missing RecordFunctionTracer stats for {operation_name}")
+        return result, {
+            name: float(value) for name, value in stats[operation_name].items()
+        }
+
     samples: list[float] = []
     for _ in range(repeats):
         start = torch.cuda.Event(enable_timing=True)
@@ -76,6 +104,8 @@ def main() -> None:
     source_head = git_head(args.vllm_source)
     if source_head != VLLM_COMMIT:
         raise SystemExit(f"expected vLLM source {VLLM_COMMIT}, got {source_head}")
+    if args.profile_method == "record_function" and args.frontier_source is None:
+        raise SystemExit("--frontier-source is required for --profile-method record_function")
 
     raw_model_config = json.loads(args.model.joinpath("config.json").read_text())
     observed = {
@@ -103,6 +133,15 @@ def main() -> None:
     from vllm.model_executor.layers.fused_moe import fused_topk
     from vllm.model_executor.layers.linear import ReplicatedLinear
 
+    record_function_tracer = None
+    if args.profile_method == "record_function":
+        import sys
+
+        sys.path.insert(0, str(args.frontier_source.resolve()))
+        from frontier.profiling.utils.record_function_tracer import RecordFunctionTracer
+
+        record_function_tracer = RecordFunctionTracer
+
     device = torch.device(args.device)
     torch.accelerator.set_device_index(device)
     torch.manual_seed(20260716)
@@ -113,6 +152,9 @@ def main() -> None:
         skip_tokenizer_init=True,
         generation_config="vllm",
     )
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    if args.profile_method == "record_function":
+        (args.output.parent / "profiler_traces").mkdir(exist_ok=True)
 
     rows: list[dict[str, Any]] = []
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
@@ -141,12 +183,22 @@ def main() -> None:
                     (num_tokens, HIDDEN_DIM), device=device, dtype=torch.bfloat16
                 ).uniform_(-0.1, 0.1)
                 logits, gate_time = measure_ms(
-                    lambda: gate(hidden)[0], args.warmup_iters, args.repeats
+                    lambda: gate(hidden)[0],
+                    args.warmup_iters,
+                    args.repeats,
+                    profile_method=args.profile_method,
+                    trace_root=args.output.parent,
+                    operation_name="moe_gating_linear",
+                    record_function_tracer=record_function_tracer,
                 )
                 topk_result, topk_time = measure_ms(
                     lambda: fused_topk(hidden, logits, TOP_K, renormalize=True),
                     args.warmup_iters,
                     args.repeats,
+                    profile_method=args.profile_method,
+                    trace_root=args.output.parent,
+                    operation_name="moe_gating_routing_topk",
+                    record_function_tracer=record_function_tracer,
                 )
 
                 def gate_and_topk() -> tuple[
@@ -158,7 +210,13 @@ def main() -> None:
                     )
 
                 combined_result, combined_time = measure_ms(
-                    gate_and_topk, args.warmup_iters, args.repeats
+                    gate_and_topk,
+                    args.warmup_iters,
+                    args.repeats,
+                    profile_method=args.profile_method,
+                    trace_root=args.output.parent,
+                    operation_name="moe_gating_linear_and_routing_topk",
+                    record_function_tracer=record_function_tracer,
                 )
                 topk_weights, topk_ids, _ = topk_result
                 combined_weights, combined_ids, _ = combined_result
@@ -207,6 +265,7 @@ def main() -> None:
             "gate_replication": "replicated_across_tp",
             "top_k": TOP_K,
             "norm_topk_prob": True,
+            "profile_method": args.profile_method,
         },
         "measurement_scope": (
             "vLLM ReplicatedLinear gate and fused_topk; measured separately and "
@@ -214,7 +273,6 @@ def main() -> None:
         ),
         "rows": rows,
     }
-    args.output.parent.mkdir(parents=True, exist_ok=True)
     args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")