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

Align Frontier piecewise graph profiles

Browse Source
This commit is contained in:
Gahow Wang
2026-07-17 23:22:42 +08:00
parent 47355a9411
commit bdc357dc6c
10 changed files with 804 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

97
runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_moe.py
View File

@@ -9,7 +9,7 @@ import math
import statistics
import subprocess
from pathlib import Path
from typing import Any
from typing import Any, Callable
import torch
import vllm
@@ -40,6 +40,12 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--repeats", type=int, default=5)
parser.add_argument("--device", default="cuda:0")
parser.add_argument("--check-reference", action="store_true")
parser.add_argument(
"--profile-method",
choices=("cuda_event", "record_function"),
default="cuda_event",
)
parser.add_argument("--frontier-source", type=Path)
return parser.parse_args()
@@ -59,6 +65,34 @@ def stats_ms(samples: list[float]) -> dict[str, float]:
}
def measure_kernel_only_ms(
fn: Callable[[], torch.Tensor],
*,
warmup_iters: int,
repeats: int,
trace_root: Path,
operation_name: str,
record_function_tracer: type,
) -> tuple[torch.Tensor, dict[str, float]]:
"""Use Frontier's KERNEL_ONLY contract, not a CUDA-event relabel."""
result = None
for _ in range(warmup_iters):
result = fn()
torch.accelerator.synchronize()
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}")
if result is None:
raise RuntimeError("kernel-only profiler executed no MoE step")
return result, {name: float(value) for name, value in stats[operation_name].items()}
def routing_inputs(
mode: str, num_tokens: int, device: torch.device
) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
@@ -145,6 +179,8 @@ def main() -> None:
raise SystemExit(
f"model contract mismatch: expected {expected_model}, got {observed_model}"
)
if args.profile_method == "record_function" and args.frontier_source is None:
raise SystemExit("--frontier-source is required for --profile-method record_function")
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
@@ -163,10 +199,22 @@ def main() -> None:
from vllm.utils.math_utils import next_power_of_2
from vllm.v1.worker.workspace import init_workspace_manager
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)
init_workspace_manager(args.device)
args.output.parent.mkdir(parents=True, exist_ok=True)
if args.profile_method == "record_function":
(args.output.parent / "profiler_traces").mkdir(exist_ok=True)
max_num_tokens = next_power_of_2(max(args.num_tokens))
rows: list[dict[str, Any]] = []
@@ -257,8 +305,8 @@ def main() -> None:
routing_mode, num_tokens, device
)
for _ in range(args.warmup_iters):
output = kernel.apply(
def run_kernel() -> torch.Tensor:
return kernel.apply(
hidden_states=hidden,
w1=w13_kernel,
w2=w2_kernel,
@@ -269,27 +317,30 @@ def main() -> None:
expert_map=None,
apply_router_weight_on_input=False,
)
torch.accelerator.synchronize()
samples: list[float] = []
for _ in range(args.repeats):
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
output = kernel.apply(
hidden_states=hidden,
w1=w13_kernel,
w2=w2_kernel,
topk_weights=topk_weights,
topk_ids=topk_ids,
activation=MoEActivation.SILU,
global_num_experts=NUM_EXPERTS,
expert_map=None,
apply_router_weight_on_input=False,
if args.profile_method == "record_function":
output, time_ms = measure_kernel_only_ms(
run_kernel,
warmup_iters=args.warmup_iters,
repeats=args.repeats,
trace_root=args.output.parent,
operation_name="moe_grouped_gemm",
record_function_tracer=record_function_tracer,
)
end.record()
else:
for _ in range(args.warmup_iters):
output = run_kernel()
torch.accelerator.synchronize()
samples.append(float(start.elapsed_time(end)))
samples: list[float] = []
for _ in range(args.repeats):
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
output = run_kernel()
end.record()
torch.accelerator.synchronize()
samples.append(float(start.elapsed_time(end)))
time_ms = stats_ms(samples)
if output.shape != hidden.shape or not torch.isfinite(output).all():
raise SystemExit(
@@ -316,7 +367,7 @@ def main() -> None:
"backend": backend.value,
"intermediate_size_per_partition": INTERMEDIATE_DIM // tp,
"output_is_reduced": kernel.output_is_reduced(),
"time_ms": stats_ms(samples),
"time_ms": time_ms,
"routing_load": load,
}
rows.append(row)
@@ -350,6 +401,7 @@ def main() -> None:
"weight_quantization": "none",
"top_k": TOP_K,
"norm_topk_prob": True,
"profile_method": args.profile_method,
},
"measurement_scope": (
"one TP-local weight shard: vLLM modular MoE prepare+FlashInfer "
@@ -357,7 +409,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")