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

132
runs/frontier-qwen30-vllm020-profile-v1/freeze_frontier_profiles.py
View File

@@ -117,8 +117,18 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--attention", type=Path, nargs="+", required=True)
parser.add_argument("--moe", type=Path, required=True)
parser.add_argument("--router", type=Path, required=True)
parser.add_argument("--allreduce", type=Path, nargs=2, required=True)
parser.add_argument("--allreduce", type=Path, nargs=2)
parser.add_argument("--allreduce-frozen", type=Path)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument(
"--measurement-type",
choices=("CUDA_EVENT", "KERNEL_ONLY"),
default="CUDA_EVENT",
)
parser.add_argument(
"--frontier-commit",
default="d9cfeb6d8791fbf2f295dd9744c56a666171776e",
)
return parser.parse_args()
@@ -196,7 +206,7 @@ def write_csv(path: Path, fieldnames: list[str], rows: list[dict[str, Any]]) ->
def freeze_attention(
inputs: list[Path], output: Path
inputs: list[Path], output: Path, *, measurement_type: str
) -> tuple[int, int, list[str]]:
rows: list[dict[str, Any]] = []
mixed_rows: list[dict[str, Any]] = []
@@ -208,6 +218,13 @@ def freeze_attention(
raise ValueError(f"unexpected attention schema in {path}")
if payload["environment"].get("vllm_version") != "0.20.0":
raise ValueError(f"unexpected vLLM version in {path}")
expected_method = (
"record_function" if measurement_type == "KERNEL_ONLY" else "cuda_event"
)
if payload["environment"].get("profile_method", "cuda_event") != expected_method:
raise ValueError(
f"attention profile method mismatch in {path}: expected {expected_method}"
)
for raw in payload["rows"]:
if raw.get("error") is not None:
raise ValueError(f"failed attention row in {path}: {raw['error']}")
@@ -321,7 +338,7 @@ def freeze_attention(
"profiling_precision": "BF16",
"model_arch": "generic",
"quant_signature": "none",
"measurement_type": "CUDA_EVENT",
"measurement_type": measurement_type,
"is_true_mixed_batch": True,
"prefill_seq_lens": json.dumps(prefill_queries),
"prefill_kv_cache_sizes": json.dumps(prefill_contexts),
@@ -418,7 +435,7 @@ def freeze_attention(
"profiling_precision": "BF16",
"model_arch": "generic",
"quant_signature": "none",
"measurement_type": "CUDA_EVENT",
"measurement_type": measurement_type,
"is_true_mixed_batch": False,
"prefill_seq_lens": "",
"prefill_kv_cache_sizes": "",
@@ -498,13 +515,23 @@ def load_features(counts: list[int]) -> dict[str, float]:
}
def freeze_moe(moe_path: Path, router_path: Path, output: Path) -> int:
def freeze_moe(
moe_path: Path, router_path: Path, output: Path, *, measurement_type: str
) -> int:
moe = load_json(moe_path)
router = load_json(router_path)
if moe.get("schema_version") != "qwen30_vllm020_moe_raw.v1":
raise ValueError(f"unexpected MoE schema in {moe_path}")
if router.get("schema_version") != "qwen30_vllm020_router_raw.v1":
raise ValueError(f"unexpected router schema in {router_path}")
expected_method = (
"record_function" if measurement_type == "KERNEL_ONLY" else "cuda_event"
)
for payload, label in ((moe, "moe"), (router, "router")):
if payload["environment"].get("profile_method", "cuda_event") != expected_method:
raise ValueError(
f"{label} profile method mismatch in {payload}: expected {expected_method}"
)
router_by_tokens = {int(row["num_tokens"]): row for row in router["rows"]}
rows: list[dict[str, Any]] = []
seen_pairs: set[tuple[int, int, str]] = set()
@@ -550,7 +577,7 @@ def freeze_moe(moe_path: Path, router_path: Path, output: Path) -> int:
"load_distribution": routing_mode,
"seed": 20260716,
"moe_grouped_gemm_backend": raw["backend"],
"measurement_type": "CUDA_EVENT",
"measurement_type": measurement_type,
"profiling_precision": "BF16",
"model_arch": "generic",
"quant_signature": "none",
@@ -565,9 +592,17 @@ def freeze_moe(moe_path: Path, router_path: Path, output: Path) -> int:
)
rows.append(row)
expected = 3 * 12 * 2
if len(rows) != expected:
raise ValueError(f"expected {expected} MoE rows, got {len(rows)}")
tokens = {int(row["num_tokens"]) for row in router["rows"]}
modes = {str(row["routing_mode"]) for row in moe["rows"]}
expected = {(tp, tokens_value, mode) for tp in (1, 2, 4) for tokens_value in tokens for mode in modes}
actual = {
(int(row["num_tensor_parallel_workers"]), int(row["num_tokens"]), str(row["load_distribution"]))
for row in rows
}
if actual != expected:
raise ValueError(
f"MoE TP/token/routing coverage mismatch: missing={expected - actual}, extra={actual - expected}"
)
moe_fields = [
f"time_stats.{op}.{stat}" for op in MOE_OPS for stat in STAT_NAMES
] + list(MOE_METADATA)
@@ -617,7 +652,13 @@ def freeze_allreduce(inputs: list[Path], output: Path) -> int:
def main() -> None:
args = parse_args()
all_inputs = [args.linear, *args.attention, args.moe, args.router, *args.allreduce]
if args.allreduce is not None and args.allreduce_frozen is not None:
raise SystemExit("provide either --allreduce or --allreduce-frozen, not both")
all_inputs = [args.linear, *args.attention, args.moe, args.router]
if args.allreduce is not None:
all_inputs.extend(args.allreduce)
if args.allreduce_frozen is not None:
all_inputs.append(args.allreduce_frozen)
for path in all_inputs:
if not path.is_file():
raise SystemExit(f"missing input: {path}")
@@ -627,24 +668,39 @@ def main() -> None:
shutil.copyfile(args.linear, linear_output)
with linear_output.open(newline="") as handle:
linear_rows = list(csv.DictReader(handle))
if len(linear_rows) != 36:
raise ValueError(f"expected 36 linear rows, got {len(linear_rows)}")
if not linear_rows:
raise ValueError("linear profile has no rows")
if {row.get("measurement_type") for row in linear_rows} != {args.measurement_type}:
raise ValueError(
f"linear measurement family mismatch: expected {args.measurement_type}"
)
attention_rows, mixed_rows, attention_tps = freeze_attention(
list(args.attention), args.output
list(args.attention), args.output, measurement_type=args.measurement_type
)
moe_rows = freeze_moe(args.moe, args.router, args.output)
allreduce_rows = freeze_allreduce(list(args.allreduce), args.output)
moe_rows = freeze_moe(
args.moe, args.router, args.output, measurement_type=args.measurement_type
)
allreduce_rows = 0
allreduce_source = "not_included"
if args.allreduce is not None:
allreduce_rows = freeze_allreduce(list(args.allreduce), args.output)
allreduce_source = "raw_vllm020_measurements"
elif args.allreduce_frozen is not None:
shutil.copyfile(args.allreduce_frozen, args.output / "allreduce.json")
allreduce_rows = len(load_json(args.allreduce_frozen).get("rows", []))
allreduce_source = "carried_forward_frozen_measurements"
output_files = [
linear_output,
args.output / "attention.csv",
args.output / "attention_true_mixed_fused.csv",
args.output / "moe.csv",
args.output / "allreduce.json",
]
if (args.output / "allreduce.json").is_file():
output_files.append(args.output / "allreduce.json")
batch_composition_augmented = len(args.attention) > 3
long_context_augmented = any(
long_context_augmented = args.measurement_type == "KERNEL_ONLY" or any(
"long-context" in path.name for path in args.attention
)
long_context_coverage: dict[str, Any] = {"included": long_context_augmented}
@@ -659,34 +715,45 @@ def main() -> None:
if int(row["num_tensor_parallel_workers"]) == tp
and row["is_prefill"].lower() == "false"
}
mixed_kv = {
int(float(row["decode_avg_kv_cache_size"]))
for row in frozen_attention
if int(row["num_tensor_parallel_workers"]) == tp
and row.get("is_true_mixed_batch", "").lower() == "true"
}
if not {16384, 32768, 40960}.issubset(decode_kv):
raise ValueError(f"long-context decode coverage mismatch for TP{tp}")
if not {16384, 32768}.issubset(mixed_kv):
raise ValueError(f"long-context mixed coverage mismatch for TP{tp}")
required_decode = (
{128, 1024, 2048, 4096, 8192, 16384, 32768, 40960}
if args.measurement_type == "KERNEL_ONLY"
else {16384, 32768, 40960}
)
if not required_decode.issubset(decode_kv):
raise ValueError(f"decode KV coverage mismatch for TP{tp}")
by_tp[str(tp)] = {
"decode_kv_lengths": sorted(decode_kv),
"true_mixed_decode_avg_kv_lengths": sorted(mixed_kv),
}
if args.measurement_type != "KERNEL_ONLY":
mixed_kv = {
int(float(row["decode_avg_kv_cache_size"]))
for row in frozen_attention
if int(row["num_tensor_parallel_workers"]) == tp
and row.get("is_true_mixed_batch", "").lower() == "true"
}
if not {16384, 32768}.issubset(mixed_kv):
raise ValueError(f"long-context mixed coverage mismatch for TP{tp}")
by_tp[str(tp)]["true_mixed_decode_avg_kv_lengths"] = sorted(mixed_kv)
long_context_coverage["by_tp"] = by_tp
manifest = {
"schema_version": (
"frontier_qwen30_vllm020_frozen_profile.v4"
if long_context_augmented
"frontier_qwen30_vllm020_kernel_only_profile.v1"
if args.measurement_type == "KERNEL_ONLY"
else (
"frontier_qwen30_vllm020_frozen_profile.v4"
if long_context_augmented
else (
"frontier_qwen30_vllm020_frozen_profile.v3"
if batch_composition_augmented
else "frontier_qwen30_vllm020_frozen_profile.v2"
)
)
),
"profile_id": (
"qwen3-30b-a3b-bf16-vllm020-h20-tp1-2-4-"
"fused-mixed-total-conserving"
+ ("-kernel-only-record-function" if args.measurement_type == "KERNEL_ONLY" else "")
+ ("-pure-prefill-batch-composition" if batch_composition_augmented else "")
+ ("-long-context-decode-mixed" if long_context_augmented else "")
),
@@ -696,7 +763,7 @@ def main() -> None:
"dtype": "bfloat16",
"vllm_version": "0.20.0",
"vllm_source_commit": "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1",
"frontier_commit": "d9cfeb6d8791fbf2f295dd9744c56a666171776e",
"frontier_commit": args.frontier_commit,
"tensor_parallel_sizes": [1, 2, 4],
},
"row_counts": {
@@ -740,8 +807,7 @@ def main() -> None:
"expert measurement already includes prepare/finalize so shuffling is zero"
),
"allreduce": (
"Frozen exact runtime measurements; base profile-only comparison keeps the "
"historical Frontier CC backend fixed to isolate compute profile fidelity"
"Frozen exact runtime measurements; source=" + allreduce_source
),
},
"inputs": {str(path.resolve()): sha256(path) for path in all_inputs},

168
runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_flashattn.py
View File

@@ -39,6 +39,12 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--repeats", type=int, default=5)
parser.add_argument("--device", default="cuda:0")
parser.add_argument("--profile-kv-update", 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()
@@ -66,12 +72,17 @@ def main() -> None:
raise SystemExit(f"expected vLLM source {VLLM_COMMIT}, got {source_head}")
if not args.model.joinpath("config.json").is_file():
raise SystemExit(f"missing model config: {args.model / 'config.json'}")
if args.profile_method == "record_function" and args.frontier_source is None:
raise SystemExit("--frontier-source is required for --profile-method record_function")
bench_dir = args.vllm_source / "benchmarks" / "attention_benchmarks"
sys.path.insert(0, str(bench_dir))
import runner # type: ignore[import-not-found] # noqa: PLC0415
from batch_spec import parse_batch_spec # type: ignore[import-not-found] # noqa: PLC0415
from common import BenchmarkConfig # type: ignore[import-not-found] # noqa: PLC0415
from common import ( # type: ignore[import-not-found] # noqa: PLC0415
BenchmarkConfig,
BenchmarkResult,
)
from vllm.config import ( # noqa: PLC0415
CacheConfig,
CompilationConfig,
@@ -90,6 +101,13 @@ def main() -> None:
from vllm.v1.kv_cache_interface import FullAttentionSpec # noqa: PLC0415
from vllm.v1.worker.workspace import init_workspace_manager # noqa: PLC0415
record_function_tracer = None
if args.profile_method == "record_function":
sys.path.insert(0, str(args.frontier_source.resolve()))
from frontier.profiling.utils.record_function_tracer import RecordFunctionTracer
record_function_tracer = RecordFunctionTracer
def create_vllm_config(config: BenchmarkConfig, max_num_blocks: int) -> VllmConfig:
model_config = ModelConfig(
model=str(args.model),
@@ -146,6 +164,9 @@ def main() -> None:
runner._create_vllm_config = create_vllm_config
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)
def profile_kv_cache_update(config: BenchmarkConfig) -> dict[str, float]:
device = torch.device(config.device)
@@ -218,6 +239,137 @@ def main() -> None:
"std_ms": statistics.pstdev(samples),
}
def profile_kernel_only(
config: BenchmarkConfig,
) -> tuple[BenchmarkResult, dict[str, float] | None]:
"""Trace exactly one vLLM FA3 forward/KV-update per annotation.
`RecordFunctionTracer` is Frontier's actual KERNEL_ONLY collector: it
follows CUDA launch correlations and sums kernels under the annotation.
The profiling loop therefore contains no CUDA-event value relabeling.
"""
device = torch.device(config.device)
torch.accelerator.set_device_index(device)
backend_config = runner._get_backend_config(config.backend)
requests = parse_batch_spec(config.batch_spec)
q_lens = [request.q_len for request in requests]
kv_lens = [request.kv_len for request in requests]
total_q = sum(q_lens)
max_kv = max(kv_lens)
max_blocks_per_request = (max_kv + config.block_size - 1) // config.block_size
max_num_blocks = len(requests) * max_blocks_per_request
with runner.log_warnings_and_errors_only():
vllm_config = create_vllm_config(config, max_num_blocks)
dtype = vllm_config.model_config.dtype
with set_current_vllm_config(vllm_config):
backend_class, impl, layer = runner._create_backend_impl(
backend_config, config, device, dtype
)
required_layout = backend_class.get_required_kv_cache_layout()
if required_layout is not None:
set_kv_cache_layout(required_layout)
get_kv_cache_layout.cache_clear()
common_metadata = runner._build_common_attn_metadata(
q_lens, kv_lens, config.block_size, device
)
kv_cache_spec = FullAttentionSpec(
block_size=config.block_size,
num_kv_heads=config.num_kv_heads,
head_size=config.head_dim,
dtype=dtype,
)
builder = runner._create_metadata_builder(
backend_class, kv_cache_spec, vllm_config, device, config.backend
)
attn_metadata = builder.build(
common_prefix_len=0, common_attn_metadata=common_metadata
)
quantize_query = config.kv_cache_dtype.startswith("fp8") and getattr(
impl, "supports_quant_query_input", False
)
q_list, k_list, v_list = runner._create_input_tensors(
config, total_q, device, dtype, quantize_query=quantize_query
)
cache_list = runner._create_kv_cache(
config, max_num_blocks, backend_class, device, dtype
)
output = torch.empty(
total_q,
config.num_q_heads,
config.head_dim,
device=device,
dtype=dtype,
)
def run_core() -> None:
for layer_index in range(config.num_layers):
impl.forward(
layer,
q_list[layer_index],
k_list[layer_index],
v_list[layer_index],
cache_list[layer_index],
attn_metadata,
output=output,
)
for _ in range(config.warmup_iters):
run_core()
torch.accelerator.synchronize()
core_tracer = record_function_tracer(str(args.output.parent))
with core_tracer:
for _ in range(config.repeats):
with torch.profiler.record_function("vidur_attention_core"):
run_core()
core_stats = core_tracer.get_operation_time_stats()
if "attention_core" not in core_stats:
raise RuntimeError("missing KERNEL_ONLY FlashAttention core stats")
core = {
name: float(value) / config.num_layers
for name, value in core_stats["attention_core"].items()
}
kv_stats = None
if args.profile_kv_update:
def run_kv_cache_update() -> None:
for layer_index in range(config.num_layers):
impl.do_kv_cache_update(
layer,
k_list[layer_index],
v_list[layer_index],
cache_list[layer_index],
common_metadata.slot_mapping,
)
for _ in range(config.warmup_iters):
run_kv_cache_update()
torch.accelerator.synchronize()
kv_tracer = record_function_tracer(str(args.output.parent))
with kv_tracer:
for _ in range(config.repeats):
with torch.profiler.record_function("vidur_attn_kv_cache_save"):
run_kv_cache_update()
kv_time_stats = kv_tracer.get_operation_time_stats()
if "attn_kv_cache_save" not in kv_time_stats:
raise RuntimeError("missing KERNEL_ONLY KV-update stats")
kv_stats = {
f"{name}_ms": float(value) / config.num_layers
for name, value in kv_time_stats["attn_kv_cache_save"].items()
}
result = BenchmarkResult(
config=config,
mean_time=core["mean"] / 1000.0,
std_time=core["std"] / 1000.0,
min_time=core["min"] / 1000.0,
max_time=core["max"] / 1000.0,
throughput_tokens_per_sec=(
total_q / (core["mean"] / 1000.0) if core["mean"] > 0 else 0.0
),
)
return result, kv_stats
rows: list[dict[str, object]] = []
for tp in args.tp:
for batch_spec in args.batch_specs:
@@ -237,12 +389,18 @@ def main() -> None:
kv_cache_dtype="auto",
use_cuda_graphs=False,
)
result = runner.run_attention_benchmark(config)
if args.profile_method == "record_function":
result, kv_stats = profile_kernel_only(config)
else:
result = runner.run_attention_benchmark(config)
kv_stats = (
profile_kv_cache_update(config) if args.profile_kv_update else None
)
row = result.to_dict()
row["tensor_parallel_size"] = tp
row["attention_core_excludes_kv_cache_update"] = True
if args.profile_kv_update:
row["kv_cache_update_time"] = profile_kv_cache_update(config)
if kv_stats is not None:
row["kv_cache_update_time"] = kv_stats
rows.append(row)
print(
json.dumps(
@@ -273,10 +431,10 @@ def main() -> None:
"attention_backend": "FLASH_ATTN",
"block_size": 16,
"profile_kv_update": args.profile_kv_update,
"profile_method": args.profile_method,
},
"rows": rows,
}
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(
json.dumps(payload, indent=2, sort_keys=True, default=json_default) + "\n"
)

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")

66
runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_router.py
View File

@@ -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")