diff --git a/transactional_emulator/testbench/emulator_runner.py b/transactional_emulator/testbench/emulator_runner.py index 6df8c88a..cbbf7bd6 100644 --- a/transactional_emulator/testbench/emulator_runner.py +++ b/transactional_emulator/testbench/emulator_runner.py @@ -89,6 +89,7 @@ def run_emulator( stage_profile_out: Path | None = None, run_label: str | None = None, overlap_prefetch_compute: bool | None = None, + dump_cwd: Path | None = None, ) -> dict: """Run the Rust transactional emulator with build artifacts from build_dir. @@ -115,9 +116,15 @@ def run_emulator( --experimental-overlap-prefetch-compute flag. When None, PLENA_EMULATOR_OVERLAP_PREFETCH_COMPUTE controls it. + dump_cwd: optional working directory for emulator dump files. Defaults to + the historical emulator directory. Parallel replay can set this + to build_dir so vram_dump.bin/fpsram_dump.bin are not shared + between concurrent emulator processes. """ emulator_dir = Path(__file__).parent.parent # transactional_emulator/ binary = emulator_dir / "target" / "release" / "transactional_emulator" + dump_dir = Path(dump_cwd) if dump_cwd is not None else emulator_dir + dump_dir.mkdir(parents=True, exist_ok=True) if stage_profile is None: stage_profile = _env_flag("PLENA_EMULATOR_STAGE_PROFILE") @@ -235,7 +242,7 @@ def run_emulator( "started_at_utc": started_at.isoformat(), "build_dir": str(build_dir), "command": cmd, - "cwd": str(emulator_dir), + "cwd": str(dump_dir), "config_path": str(_current_plena_settings_path()), "behavior_config": _current_behavior_config_summary(), "hbm_size_bytes": hbm_size, @@ -261,14 +268,14 @@ def run_emulator( # Avoid copying an HBM dump from a previous debug run when the current run # does not enable DEBUG tracing. - hbm_debug_dump = emulator_dir / "hbm_dump.bin" + hbm_debug_dump = dump_dir / "hbm_dump.bin" if hbm_debug_dump.exists(): hbm_debug_dump.unlink() with log_path.open("w", encoding="utf-8", errors="replace") as log_file: proc = subprocess.Popen( cmd, - cwd=str(emulator_dir), + cwd=str(dump_dir), env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, @@ -320,30 +327,34 @@ def run_emulator( raise RuntimeError(f"Transactional emulator failed (exit code {return_code})") # Copy vram to build dir so subsequent runs don't overwrite it. - vram_src = emulator_dir / "vram_dump.bin" + vram_src = dump_dir / "vram_dump.bin" vram_dst = build_dir / "vram_dump.bin" if vram_src.exists(): import shutil - shutil.copy2(vram_src, vram_dst) - fpsram_src = emulator_dir / "fpsram_dump.bin" + if vram_src.resolve() != vram_dst.resolve(): + shutil.copy2(vram_src, vram_dst) + fpsram_src = dump_dir / "fpsram_dump.bin" fpsram_dst = build_dir / "fpsram_dump.bin" if fpsram_src.exists(): import shutil - shutil.copy2(fpsram_src, fpsram_dst) - intsram_src = emulator_dir / "intsram_dump.bin" + if fpsram_src.resolve() != fpsram_dst.resolve(): + shutil.copy2(fpsram_src, fpsram_dst) + intsram_src = dump_dir / "intsram_dump.bin" intsram_dst = build_dir / "intsram_dump.bin" if intsram_src.exists(): import shutil - shutil.copy2(intsram_src, intsram_dst) - hbm_src = emulator_dir / "hbm_dump.bin" + if intsram_src.resolve() != intsram_dst.resolve(): + shutil.copy2(intsram_src, intsram_dst) + hbm_src = dump_dir / "hbm_dump.bin" hbm_dst = build_dir / "hbm_dump.bin" if hbm_src.exists(): import shutil - shutil.copy2(hbm_src, hbm_dst) + if hbm_src.resolve() != hbm_dst.resolve(): + shutil.copy2(hbm_src, hbm_dst) return metrics diff --git a/transactional_emulator/testbench/window1_p2/__init__.py b/transactional_emulator/testbench/window1_p2/__init__.py new file mode 100644 index 00000000..6b6c217b --- /dev/null +++ b/transactional_emulator/testbench/window1_p2/__init__.py @@ -0,0 +1,2 @@ +"""Window 1 P2 BFCL/GPQA real-route timing pilot sidecars.""" + diff --git a/transactional_emulator/testbench/window1_p2/build_route_traces.py b/transactional_emulator/testbench/window1_p2/build_route_traces.py new file mode 100644 index 00000000..f3b833fd --- /dev/null +++ b/transactional_emulator/testbench/window1_p2/build_route_traces.py @@ -0,0 +1,197 @@ +#!/usr/bin/env python3 +"""Convert P2 true-routing JSONL rows into P1 route trace schema files.""" + +from __future__ import annotations + +import argparse +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from transactional_emulator.testbench.window1_p1.validate_route_trace import validate_trace +from transactional_emulator.testbench.window1_p2.p2_utils import ( + MODEL_CONFIGS, + OUT_ROOT, + ensure_paths, + iter_jsonl, + routing_stats, + stable_id, + write_json, +) + + +def _parse_csv_ints(value: str | None) -> set[int] | None: + if not value: + return None + return {int(part) for part in value.split(",") if part.strip()} + + +def _parse_csv_strings(value: str | None) -> set[str] | None: + if not value: + return None + return {part.strip() for part in value.split(",") if part.strip()} + + +def _trace_id(record: dict[str, Any]) -> str: + phase = str(record["phase"]) + suffix = f"_d{record.get('decode_step')}" if phase == "decode" and record.get("decode_step") is not None else "" + return stable_id( + f"{record['model_key']}_{record['benchmark']}_{record['sample_id']}_" + f"l{record['layer']}_{phase}{suffix}" + ) + + +def _weights_for(record: dict[str, Any], top_k: int, *, allow_uniform_weights: bool) -> tuple[list[list[float]], str]: + weights = record.get("route_weights") + if weights is not None: + return [[float(x) for x in row] for row in weights], "true_router_topk_softmax_weights" + if not allow_uniform_weights: + raise ValueError( + f"routing row {record.get('sample_id')} layer={record.get('layer')} phase={record.get('phase')} " + "has routes but no route_weights; rerun generate_true_routing_with_weights.py or pass " + "--allow-uniform-weights for timing-only shape replay" + ) + rows = int(record["tokens"]) + return [[1.0 / float(top_k) for _ in range(top_k)] for _ in range(rows)], "uniform_weights_timing_only" + + +def trace_from_record( + record: dict[str, Any], + *, + mlen: int, + blen: int, + emu_threads: int, + allow_uniform_weights: bool, +) -> dict[str, Any]: + model_cfg = MODEL_CONFIGS[record["model_key"]] + top_k = int(model_cfg["top_k"]) + num_experts = int(model_cfg["num_experts"]) + indices = [[int(x) for x in row] for row in record["routes"]] + weights, weight_source = _weights_for(record, top_k, allow_uniform_weights=allow_uniform_weights) + stats = routing_stats(indices, num_experts) + source = str(record.get("routing_source", "true_router_logits")) + if weight_source == "uniform_weights_timing_only": + source += "+uniform_weights_timing_only" + trace = { + "schema_version": 1, + "trace_id": _trace_id(record), + "created_by": "transactional_emulator.testbench.window1_p2.build_route_traces", + "created_at_utc": datetime.now(UTC).isoformat(), + "model": { + "name": model_cfg["name"], + "layer_index": int(record["layer"]), + "hidden_size": int(model_cfg["hidden_size"]), + "intermediate_size": int(model_cfg["intermediate_size"]), + "num_experts": num_experts, + "top_k": top_k, + "policy_name": model_cfg["policy_name"], + "activation_policy": model_cfg["activation_policy"], + }, + "workload": { + "benchmark": record["benchmark"], + "sample_id": str(record["sample_id"]), + "phase": record["phase"], + "batch_size": 1, + "seq_len": int(record["tokens"]), + "token_count": int(record["tokens"]), + "input_tokens": int(record.get("input_tokens", record["tokens"])), + "category": record.get("category", ""), + "sample_index": record.get("sample_index"), + "decode_step": record.get("decode_step"), + }, + "routing": { + "source": source, + "weight_source": weight_source, + "topk_indices": indices, + "topk_weights": weights, + **stats, + }, + "logical_bytes": { + "input_hidden_bf16_bytes": int(record["tokens"]) * int(model_cfg["hidden_size"]) * 2, + "routed_hidden_bf16_bytes": int(record["tokens"]) * top_k * int(model_cfg["hidden_size"]) * 2, + "route_weight_bf16_bytes": int(record["tokens"]) * top_k * 2, + "expert_weight_table_note": "Physical HBM bytes come from Rust emulator WithStats; logical model bytes are context only.", + }, + "artifacts": { + "reference_pt": "generated_by_window1_p2_qwen3_trace_replay", + "l1_golden_pt": "generated_by_window1_p2_qwen3_trace_replay", + }, + "replay": { + "harness_module": "transactional_emulator.testbench.window1_p2.qwen3_trace_replay_test", + "stage": "full_vram", + "mlen": mlen, + "blen": blen, + "emu_threads": emu_threads, + }, + "measurement_note": "self-consistent upper bound, absolute accuracy pending RTL (Window 2)", + } + errors = validate_trace(trace, allow_missing_artifacts=True) + if errors: + raise ValueError("Generated invalid trace:\n" + "\n".join(errors)) + return trace + + +def build_traces(args: argparse.Namespace) -> list[Path]: + ensure_paths() + sample_ids = _parse_csv_strings(args.sample_ids) + layers = _parse_csv_ints(args.layers) + phases = _parse_csv_strings(args.phases) + out_dir = args.out_dir or (OUT_ROOT / "route_traces") + out_dir.mkdir(parents=True, exist_ok=True) + written: list[Path] = [] + for record in iter_jsonl(args.input): + if sample_ids is not None and str(record.get("sample_id")) not in sample_ids: + continue + if layers is not None and int(record.get("layer")) not in layers: + continue + if phases is not None and str(record.get("phase")) not in phases: + continue + if "routes" not in record: + raise ValueError(f"routing row has no routes: {record.get('sample_id')} layer={record.get('layer')}") + trace = trace_from_record( + record, + mlen=args.mlen, + blen=args.blen, + emu_threads=args.emu_threads, + allow_uniform_weights=args.allow_uniform_weights, + ) + path = out_dir / f"{trace['trace_id']}.json" + write_json(path, trace) + written.append(path) + if args.limit is not None and len(written) >= args.limit: + break + summary = { + "schema_version": 1, + "input": str(args.input), + "out_dir": str(out_dir), + "written": [str(path) for path in written], + "count": len(written), + } + write_json(args.summary_out or (OUT_ROOT / "route_traces_summary.json"), summary) + return written + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", type=Path, required=True) + parser.add_argument("--out-dir", type=Path) + parser.add_argument("--summary-out", type=Path) + parser.add_argument("--sample-ids") + parser.add_argument("--layers") + parser.add_argument("--phases") + parser.add_argument("--limit", type=int) + parser.add_argument("--mlen", type=int, default=128) + parser.add_argument("--blen", type=int, default=4) + parser.add_argument("--emu-threads", type=int, default=1) + parser.add_argument("--allow-uniform-weights", action="store_true") + args = parser.parse_args() + written = build_traces(args) + print(f"wrote {len(written)} route traces") + for path in written[:10]: + print(path) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) + diff --git a/transactional_emulator/testbench/window1_p2/export_pilot_results.py b/transactional_emulator/testbench/window1_p2/export_pilot_results.py new file mode 100644 index 00000000..296807e9 --- /dev/null +++ b/transactional_emulator/testbench/window1_p2/export_pilot_results.py @@ -0,0 +1,192 @@ +#!/usr/bin/env python3 +"""Export Window 1 P2 pilot trace replay results to CSV/JSON tables.""" + +from __future__ import annotations + +import argparse +import statistics +from pathlib import Path +from typing import Any + +from transactional_emulator.testbench.window1_p1.p1_utils import summarize_run +from transactional_emulator.testbench.window1_p2.p2_utils import OUT_ROOT, load_json, write_csv, write_json + + +MEASUREMENT_NOTE = "self-consistent upper bound, absolute accuracy pending RTL (Window 2)" +ROUTING_STAGE_NAMES = { + "accumulator_init", + "gather", + "expert_weight_address", + "expert_route_weight", + "scatter_combine", +} + + +def _result_dirs(root: Path) -> list[Path]: + return sorted(path for path in root.glob("*") if (path / "qwen3_trace_replay_results.json").exists()) + + +def _sum_stages(profile: dict[str, Any], names: set[str]) -> int: + total = 0 + stages = profile.get("stages", {}) + if not isinstance(stages, dict): + return 0 + for name in names: + stage = stages.get(name, {}) + total += int(stage.get("wall_cycles") or 0) + return total + + +def _stage_lookup(profile: dict[str, Any]) -> dict[str, dict[str, Any]]: + stages = profile.get("stages", {}) + return stages if isinstance(stages, dict) else {} + + +def _mean(values: list[int]) -> float | None: + return float(statistics.mean(values)) if values else None + + +def export(args: argparse.Namespace) -> dict[str, Any]: + root = args.root + runs = [] + stage_rows = [] + tax_rows = [] + bytes_rows = [] + skipped = [] + for build_dir in _result_dirs(root): + result = load_json(build_dir / "qwen3_trace_replay_results.json") + trace = load_json(build_dir / "trace.json") + run_row, stack = summarize_run(result["trace_id"], build_dir) + profile = load_json(Path(run_row["stage_profile_path"])) if run_row.get("stage_profile_path") else {} + routing_tax_cycles = _sum_stages(profile, ROUTING_STAGE_NAMES) + sim_cycles = run_row.get("sim_latency_cycles") + if sim_cycles is None: + skipped.append({"build_dir": str(build_dir), "reason": "missing sim_latency_cycles"}) + continue + routed_fraction = float(routing_tax_cycles) / float(sim_cycles) if int(sim_cycles) else None + common = { + "trace_id": result["trace_id"], + "benchmark": result["benchmark"], + "model": trace["model"]["name"], + "sample_id": result["sample_id"], + "sample_index": trace["workload"].get("sample_index"), + "category": trace["workload"].get("category"), + "layer": result["layer"], + "phase": result["phase"], + "decode_step": trace["workload"].get("decode_step"), + "token_count": result["rows"], + "pair_count": result["pair_count"], + "selected_expert_count": result["selected_expert_count"], + "routing_source": trace["routing"].get("source"), + "route_weight_source": trace["routing"].get("weight_source"), + "issue_model": "in_order_blocking_no_overlap", + "measurement_note": MEASUREMENT_NOTE, + } + runs.append( + { + **common, + "sim_latency_cycles": int(sim_cycles), + "hbm_bytes_read": run_row.get("hbm_bytes_read"), + "hbm_bytes_written": run_row.get("hbm_bytes_written"), + "functional_gate": run_row.get("functional_gate_passed"), + "cycle_accounting_status": run_row.get("cycle_accounting_status"), + "physical_byte_status": run_row.get("physical_byte_status"), + "stage_total_simulation_cycles": run_row.get("stage_total_simulation_cycles"), + "stage_total_wall_cycles": run_row.get("stage_total_wall_cycles"), + "build_dir": str(build_dir), + "stage_profile_path": run_row.get("stage_profile_path"), + "run_stats_path": run_row.get("run_stats_path"), + } + ) + for row in stack: + stage_rows.append({**common, **{k: v for k, v in row.items() if k != "run_id"}}) + tax_rows.append( + { + **common, + "routing_tax_method": "stage_derived_accumulator_gather_address_routeweight_scatter", + "routing_tax_cycles": routing_tax_cycles, + "routing_tax_fraction_of_total": routed_fraction, + "total_cycles": int(sim_cycles), + "note": "Fixed-route trace replay excludes router GEMM/topk; routing tax here is stage-derived movement/dispatch overhead, not device-selected-minus-oracle.", + } + ) + logical = trace.get("logical_bytes", {}) + profile_read = profile.get("total_hbm_bytes_read") + profile_written = profile.get("total_hbm_bytes_written") + bytes_rows.append( + { + **common, + "logical_input_hidden_bf16_bytes": logical.get("input_hidden_bf16_bytes"), + "logical_routed_hidden_bf16_bytes": logical.get("routed_hidden_bf16_bytes"), + "logical_route_weight_bf16_bytes": logical.get("route_weight_bf16_bytes"), + "physical_hbm_bytes_read_run_stats": run_row.get("hbm_bytes_read"), + "physical_hbm_bytes_written_run_stats": run_row.get("hbm_bytes_written"), + "physical_hbm_bytes_read_stage_profile": profile_read, + "physical_hbm_bytes_written_stage_profile": profile_written, + "physical_read_matches": run_row.get("hbm_bytes_read") == profile_read, + "physical_written_matches": run_row.get("hbm_bytes_written") == profile_written, + "physical_source": "rust_emulator WithStats 64B transfer counters", + } + ) + + summary_rows = [] + groups: dict[tuple[str, str], list[dict[str, Any]]] = {} + for row in runs: + groups.setdefault((row["benchmark"], row["phase"]), []).append(row) + for (benchmark, phase), rows in sorted(groups.items()): + cycles = [int(row["sim_latency_cycles"]) for row in rows] + taxes = [ + int(tax["routing_tax_cycles"]) + for tax in tax_rows + for row in rows + if tax["trace_id"] == row["trace_id"] + ] + summary_rows.append( + { + "benchmark": benchmark, + "phase": phase, + "runs": len(rows), + "cycles_min": min(cycles) if cycles else None, + "cycles_mean": _mean(cycles), + "cycles_max": max(cycles) if cycles else None, + "routing_tax_cycles_min": min(taxes) if taxes else None, + "routing_tax_cycles_mean": _mean(taxes), + "routing_tax_cycles_max": max(taxes) if taxes else None, + "measurement_note": MEASUREMENT_NOTE, + } + ) + + out_prefix = args.out_prefix + write_csv(out_prefix.with_name(out_prefix.name + "_timing.csv"), runs) + write_csv(out_prefix.with_name(out_prefix.name + "_stage_stack.csv"), stage_rows) + write_csv(out_prefix.with_name(out_prefix.name + "_routing_tax.csv"), tax_rows) + write_csv(out_prefix.with_name(out_prefix.name + "_bytes_validation.csv"), bytes_rows) + write_csv(out_prefix.with_name(out_prefix.name + "_summary_stats.csv"), summary_rows) + payload = { + "schema_version": 1, + "root": str(root), + "measurement_note": MEASUREMENT_NOTE, + "runs": runs, + "stage_stack": stage_rows, + "routing_tax": tax_rows, + "bytes_validation": bytes_rows, + "summary_stats": summary_rows, + "skipped": skipped, + } + write_json(args.out_json, payload) + return payload + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--root", type=Path, default=OUT_ROOT / "trace_replay") + parser.add_argument("--out-prefix", type=Path, default=OUT_ROOT / "gpqa_bfcl_emulation") + parser.add_argument("--out-json", type=Path, default=OUT_ROOT / "gpqa_bfcl_emulation_timing_summary.json") + args = parser.parse_args() + payload = export(args) + print(f"exported {len(payload['runs'])} runs, {len(payload['stage_stack'])} stage rows") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/transactional_emulator/testbench/window1_p2/generate_true_routing_with_weights.py b/transactional_emulator/testbench/window1_p2/generate_true_routing_with_weights.py new file mode 100644 index 00000000..5f46e86f --- /dev/null +++ b/transactional_emulator/testbench/window1_p2/generate_true_routing_with_weights.py @@ -0,0 +1,655 @@ +#!/usr/bin/env python3 +"""Generate small BFCL/GPQA true-router traces with top-k weights for P2. + +This is intentionally a pilot helper, not a full benchmark runner. It runs the +local HF model on CPU for a small selected subset and writes router top-k expert +ids plus softmax route weights, which P1's route trace schema requires. +""" + +from __future__ import annotations + +import argparse +import time +import traceback +from pathlib import Path +from typing import Any + +import torch +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + +from transactional_emulator.testbench.window1_p2.p2_utils import ( + INPUT_FILES, + MODEL_CONFIGS, + OUT_ROOT, + append_jsonl, + ensure_paths, + iter_jsonl, +) + + +def _input_records(path: Path, model_key: str) -> list[dict[str, Any]]: + return [row for row in iter_jsonl(path) if row.get("model_key") == model_key] + + +def _parse_csv_ints(value: str) -> list[int]: + return [int(part) for part in value.split(",") if part.strip()] + + +def _parse_csv_strings(value: str | None) -> set[str] | None: + if not value: + return None + return {part.strip() for part in value.split(",") if part.strip()} + + +def _route_key(row: dict[str, Any]) -> tuple[str, str, int, int | None]: + decode_step = row.get("decode_step") + return ( + str(row["sample_id"]), + str(row["phase"]), + int(row["layer"]), + int(decode_step) if decode_step is not None else None, + ) + + +def _existing_route_keys(path: Path) -> set[tuple[str, str, int, int | None]]: + if not path.exists(): + return set() + keys = set() + for row in iter_jsonl(path): + try: + keys.add(_route_key(row)) + except (KeyError, TypeError, ValueError): + continue + return keys + + +def _existing_route_keys_from_paths(paths: list[Path]) -> set[tuple[str, str, int, int | None]]: + keys: set[tuple[str, str, int, int | None]] = set() + for path in paths: + keys.update(_existing_route_keys(path)) + return keys + + +def _resume_scan_paths( + *, + output_path: Path, + resume_from: list[Path], + benchmark: str, + model_key: str, +) -> list[Path]: + paths = [output_path, *resume_from] + prefix = f"{benchmark}_{model_key}" + if output_path.parent.exists(): + paths.extend(path for path in output_path.parent.glob("*.jsonl") if path.name.startswith(prefix)) + + deduped: list[Path] = [] + seen: set[Path] = set() + for path in paths: + try: + key = path.resolve() + except OSError: + key = path + if key in seen: + continue + seen.add(key) + deduped.append(path) + return deduped + + +def _expected_route_keys( + *, + sample_id: str, + layers: list[int], + write_phases: str, + decode_steps: int, +) -> set[tuple[str, str, int, int | None]]: + keys: set[tuple[str, str, int, int | None]] = set() + if write_phases in ("prefill", "both"): + keys.update((sample_id, "prefill", layer, None) for layer in layers) + if write_phases in ("decode", "both"): + for step in range(1, decode_steps + 1): + keys.update((sample_id, "decode", layer, step) for layer in layers) + return keys + + +def _batch_log_row( + *, + args: argparse.Namespace, + batch: list[dict[str, Any]], + status: str, + seconds: float, + processed_samples: int = 0, + processed_tokens: int = 0, + error: BaseException | None = None, +) -> dict[str, Any]: + row = { + "benchmark": args.benchmark, + "model_key": args.model_key, + "sample_ids": [item["sample_id"] for item in batch], + "sample_indices": [item["sample_index"] for item in batch], + "input_tokens": [len(item["input_ids"]) for item in batch], + "status": status, + "seconds": seconds, + "processed_samples": processed_samples, + "processed_prefill_tokens": processed_tokens, + "layers": args.layers, + "decode_steps": args.decode_steps, + "write_phases": args.write_phases, + } + if error is not None: + row["error_type"] = type(error).__name__ + row["error"] = str(error) + return row + + +def _gini(values: list[int]) -> float: + total = sum(values) + if total == 0: + return 0.0 + sorted_values = sorted(values) + n = len(sorted_values) + weighted = sum((idx + 1) * value for idx, value in enumerate(sorted_values)) + return (2 * weighted) / (n * total) - (n + 1) / n + + +def _safe_layers(requested: list[int], layer_count: int) -> list[int]: + out = [] + for layer in requested: + if layer < 0: + layer = layer_count + layer + if 0 <= layer < layer_count: + out.append(layer) + return sorted(set(out)) + + +def _make_padded_batch(batch: list[dict[str, Any]], pad_token_id: int) -> tuple[torch.Tensor, torch.Tensor, list[int]]: + lengths = [len(item["input_ids"]) for item in batch] + padded = max(lengths) + input_ids = torch.full((len(batch), padded), pad_token_id, dtype=torch.long) + attention_mask = torch.zeros((len(batch), padded), dtype=torch.long) + for row, item in enumerate(batch): + values = torch.tensor(item["input_ids"], dtype=torch.long) + input_ids[row, : values.numel()] = values + attention_mask[row, : values.numel()] = 1 + return input_ids, attention_mask, lengths + + +def _base_forward(model, **kwargs): + base = getattr(model, "model", None) + if base is None: + return model(**kwargs) + return base(**kwargs) + + +def _next_token_from_lengths(model, hidden_states: torch.Tensor, lengths: list[int]) -> torch.Tensor: + lm_head = getattr(model, "lm_head", None) + if lm_head is None: + raise ValueError("model has no lm_head for greedy decode token selection") + batch_idx = torch.arange(hidden_states.shape[0], device=hidden_states.device) + last_idx = torch.tensor([length - 1 for length in lengths], device=hidden_states.device) + with torch.inference_mode(): + logits = lm_head(hidden_states[batch_idx, last_idx, :].unsqueeze(1)) + return logits.argmax(dim=-1) + + +def _next_token_from_hidden(model, hidden_states: torch.Tensor) -> torch.Tensor: + lm_head = getattr(model, "lm_head", None) + if lm_head is None: + raise ValueError("model has no lm_head for greedy decode token selection") + with torch.inference_mode(): + logits = lm_head(hidden_states[:, -1:, :]) + return logits.argmax(dim=-1) + + +def _split_router_logits(router_logits: torch.Tensor, lengths: list[int], padded_length: int) -> list[torch.Tensor]: + logits = router_logits.detach() + batch_size = len(lengths) + if logits.ndim == 3: + if logits.shape[0] != batch_size: + raise ValueError(f"unexpected 3D router logits shape: {tuple(logits.shape)}") + return [logits[idx, : lengths[idx], :] for idx in range(batch_size)] + if logits.ndim != 2: + raise ValueError(f"unexpected router logits shape: {tuple(logits.shape)}") + if logits.shape[0] == batch_size * padded_length: + reshaped = logits.reshape(batch_size, padded_length, logits.shape[-1]) + return [reshaped[idx, : lengths[idx], :] for idx in range(batch_size)] + if logits.shape[0] == sum(lengths): + chunks = [] + start = 0 + for length in lengths: + chunks.append(logits[start : start + length, :]) + start += length + return chunks + raise ValueError( + f"cannot split router logits shape={tuple(logits.shape)} " + f"batch={batch_size} padded={padded_length} sum_lengths={sum(lengths)}" + ) + + +def _routing_payload(logits: torch.Tensor, top_k: int, num_experts: int) -> dict[str, Any]: + logits_f = logits.float() + top_values, top_indices = torch.topk(logits_f, k=top_k, dim=-1) + top_weights = torch.softmax(top_values, dim=-1) + counts = torch.bincount(top_indices.reshape(-1).cpu(), minlength=num_experts).tolist() + tokens = int(logits_f.shape[0]) + pair_count = tokens * top_k + return { + "tokens": tokens, + "pair_count": pair_count, + "counts": counts, + "unique_experts": sum(1 for value in counts if value > 0), + "zero_experts": sum(1 for value in counts if value == 0), + "tail_lt4_experts": sum(1 for value in counts if 0 < value < 4), + "hot_ge16_experts": sum(1 for value in counts if value >= 16), + "gini": _gini(counts), + "routes": top_indices.cpu().tolist(), + "route_weights": [[float(x) for x in row] for row in top_weights.cpu().tolist()], + } + + +def _write_phase_rows( + *, + out_path: Path, + batch: list[dict[str, Any]], + router_logits: tuple[torch.Tensor, ...], + layers: list[int], + lengths: list[int], + padded_length: int, + top_k: int, + num_experts: int, + phase: str, + routing_source: str, + decode_step: int | None = None, +) -> None: + for layer in layers: + split = _split_router_logits(router_logits[layer], lengths, padded_length) + for item, logits in zip(batch, split, strict=True): + row = { + "benchmark": item["benchmark"], + "model_key": item["model_key"], + "model": MODEL_CONFIGS[item["model_key"]]["name"], + "sample_id": item["sample_id"], + "sample_index": item["sample_index"], + "category": item.get("category", ""), + "phase": phase, + "layer": layer, + "input_tokens": len(item["input_ids"]), + "routing_source": routing_source, + **_routing_payload(logits, top_k, num_experts), + } + if decode_step is not None: + row["decode_step"] = decode_step + append_jsonl(out_path, row) + + +def _process_batch( + *, + model, + batch: list[dict[str, Any]], + layers: list[int], + top_k: int, + num_experts: int, + pad_token_id: int, + decode_steps: int, + write_phases: str, + out_path: Path, +) -> tuple[int, int]: + if decode_steps > 0 and write_phases in ("decode", "both") and len(batch) > 1: + raise ValueError("decode routing is only supported with --batch-size 1") + input_ids, attention_mask, lengths = _make_padded_batch(batch, pad_token_id) + need_prefill_router_logits = write_phases in ("prefill", "both") + with torch.inference_mode(): + prefill = _base_forward( + model, + input_ids=input_ids, + attention_mask=attention_mask, + output_router_logits=need_prefill_router_logits, + use_cache=decode_steps > 0, + return_dict=True, + ) + safe = layers + if write_phases in ("prefill", "both"): + _write_phase_rows( + out_path=out_path, + batch=batch, + router_logits=tuple(prefill.router_logits), + layers=safe, + lengths=lengths, + padded_length=input_ids.shape[1], + top_k=top_k, + num_experts=num_experts, + phase="prefill", + routing_source="true_hidden_state_router_logits_topk_softmax", + ) + + past_key_values = getattr(prefill, "past_key_values", None) + next_token = _next_token_from_lengths(model, prefill.last_hidden_state, lengths) if decode_steps > 0 else None + decode_attention_mask = attention_mask + for step in range(decode_steps): + decode_attention_mask = torch.cat( + [decode_attention_mask, torch.ones((decode_attention_mask.shape[0], 1), dtype=decode_attention_mask.dtype)], + dim=1, + ) + with torch.inference_mode(): + decoded = _base_forward( + model, + input_ids=next_token, + attention_mask=decode_attention_mask, + past_key_values=past_key_values, + output_router_logits=True, + use_cache=True, + return_dict=True, + ) + past_key_values = decoded.past_key_values + if write_phases in ("decode", "both"): + _write_phase_rows( + out_path=out_path, + batch=batch, + router_logits=tuple(decoded.router_logits), + layers=safe, + lengths=[1] * len(batch), + padded_length=1, + top_k=top_k, + num_experts=num_experts, + phase="decode", + routing_source="true_hidden_state_router_logits_topk_softmax_greedy_decode", + decode_step=step + 1, + ) + next_token = _next_token_from_hidden(model, decoded.last_hidden_state) + return len(batch), sum(lengths) + + +def run(args: argparse.Namespace) -> dict[str, Any]: + ensure_paths() + torch.set_num_threads(args.threads) + model_cfg = MODEL_CONFIGS[args.model_key] + output_path = args.output or (OUT_ROOT / "true_routing" / f"{args.benchmark}_{args.model_key}_pilot.jsonl") + if output_path.exists() and not args.resume: + output_path.unlink() + + config = AutoConfig.from_pretrained(model_cfg["path"], local_files_only=True) + tokenizer = AutoTokenizer.from_pretrained(model_cfg["path"], local_files_only=True) + top_k = int(getattr(config, "num_experts_per_tok", None) or getattr(config, "experts_per_token")) + num_experts = int(getattr(config, "num_experts", None) or getattr(config, "num_local_experts")) + sample_filter = _parse_csv_strings(args.sample_ids) + safe_layers = _safe_layers(args.layers, int(model_cfg["layers"])) + resume_from = [Path(path) for path in (getattr(args, "resume_from", None) or [])] + resume_paths = _resume_scan_paths( + output_path=output_path, + resume_from=resume_from, + benchmark=args.benchmark, + model_key=args.model_key, + ) + existing_keys = _existing_route_keys_from_paths(resume_paths) if args.resume else set() + + records = [] + selected_total = 0 + skipped_existing = 0 + for sample_index, record in enumerate(_input_records(INPUT_FILES[args.benchmark], args.model_key)): + sample_id = str(record["sample_id"]) + if sample_filter is not None and sample_id not in sample_filter: + continue + input_ids = list(record["input_ids"]) + if args.max_input_tokens is not None: + input_ids = input_ids[: args.max_input_tokens] + if not input_ids: + continue + selected_total += 1 + if args.limit is not None and selected_total > args.limit: + break + if args.resume: + expected = _expected_route_keys( + sample_id=sample_id, + layers=safe_layers, + write_phases=args.write_phases, + decode_steps=args.decode_steps, + ) + if expected and expected.issubset(existing_keys): + skipped_existing += 1 + continue + records.append( + { + "benchmark": args.benchmark, + "model_key": args.model_key, + "sample_id": sample_id, + "sample_index": sample_index, + "category": record.get("category", ""), + "input_ids": input_ids, + } + ) + + if args.sort_by_length: + records.sort(key=lambda item: (len(item["input_ids"]), item["sample_index"])) + + print(f"model_path={model_cfg['path']}") + print(f"benchmark={args.benchmark} records={len(records)} output={output_path}") + if not records: + return { + "schema_version": 1, + "output": str(output_path), + "benchmark": args.benchmark, + "model_key": args.model_key, + "selected_samples_total": selected_total, + "selected_samples_after_resume": 0, + "skipped_existing_samples": skipped_existing, + "processed_samples": 0, + "processed_prefill_tokens": 0, + "run_seconds": 0.0, + "layers": args.layers, + "decode_steps": args.decode_steps, + "write_phases": args.write_phases, + "contains_route_weights": True, + "failed_batches": 0, + "failure_log": str(args.failure_log) if args.failure_log else None, + "sample_log": str(args.sample_log) if args.sample_log else None, + "resume_from": [str(path) for path in resume_from], + } + print("loading model", flush=True) + t_load = time.time() + model = AutoModelForCausalLM.from_pretrained( + model_cfg["path"], + local_files_only=True, + dtype="auto", + low_cpu_mem_usage=True, + attn_implementation="eager", + ) + model.eval() + print(f"model_loaded_sec={time.time() - t_load:.2f}", flush=True) + + pad_token_id = tokenizer.pad_token_id + if pad_token_id is None: + pad_token_id = tokenizer.eos_token_id + if pad_token_id is None: + pad_token_id = 0 + + processed = 0 + tokens = 0 + failures: list[dict[str, Any]] = [] + batch: list[dict[str, Any]] = [] + t_run = time.time() + for item in records: + if args.resume: + expected = _expected_route_keys( + sample_id=str(item["sample_id"]), + layers=safe_layers, + write_phases=args.write_phases, + decode_steps=args.decode_steps, + ) + live_resume_paths = _resume_scan_paths( + output_path=output_path, + resume_from=resume_from, + benchmark=args.benchmark, + model_key=args.model_key, + ) + live_existing_keys = _existing_route_keys_from_paths(live_resume_paths) + if expected and expected.issubset(live_existing_keys): + skipped_existing += 1 + print(f"skip_existing_dynamic id={item['sample_id']}", flush=True) + continue + print(f"sample={processed + len(batch) + 1} id={item['sample_id']} tokens={len(item['input_ids'])}", flush=True) + batch.append(item) + if len(batch) >= args.batch_size: + batch_start = time.time() + try: + done, seen = _process_batch( + model=model, + batch=batch, + layers=safe_layers, + top_k=top_k, + num_experts=num_experts, + pad_token_id=pad_token_id, + decode_steps=args.decode_steps, + write_phases=args.write_phases, + out_path=output_path, + ) + processed += done + tokens += seen + if args.sample_log: + append_jsonl( + args.sample_log, + _batch_log_row( + args=args, + batch=batch, + status="passed", + seconds=time.time() - batch_start, + processed_samples=done, + processed_tokens=seen, + ), + ) + except Exception as exc: + batch_seconds = time.time() - batch_start + failure = { + "benchmark": args.benchmark, + "model_key": args.model_key, + "sample_ids": [item["sample_id"] for item in batch], + "sample_indices": [item["sample_index"] for item in batch], + "seconds": batch_seconds, + "error_type": type(exc).__name__, + "error": str(exc), + "traceback": traceback.format_exc(), + } + failures.append(failure) + if args.failure_log: + append_jsonl(args.failure_log, failure) + if args.sample_log: + append_jsonl( + args.sample_log, + _batch_log_row( + args=args, + batch=batch, + status="failed", + seconds=batch_seconds, + error=exc, + ), + ) + print(f"failed batch sample_ids={failure['sample_ids']}: {exc}", flush=True) + if not args.keep_going: + raise + batch = [] + if batch: + batch_start = time.time() + try: + done, seen = _process_batch( + model=model, + batch=batch, + layers=safe_layers, + top_k=top_k, + num_experts=num_experts, + pad_token_id=pad_token_id, + decode_steps=args.decode_steps, + write_phases=args.write_phases, + out_path=output_path, + ) + processed += done + tokens += seen + if args.sample_log: + append_jsonl( + args.sample_log, + _batch_log_row( + args=args, + batch=batch, + status="passed", + seconds=time.time() - batch_start, + processed_samples=done, + processed_tokens=seen, + ), + ) + except Exception as exc: + batch_seconds = time.time() - batch_start + failure = { + "benchmark": args.benchmark, + "model_key": args.model_key, + "sample_ids": [item["sample_id"] for item in batch], + "sample_indices": [item["sample_index"] for item in batch], + "seconds": batch_seconds, + "error_type": type(exc).__name__, + "error": str(exc), + "traceback": traceback.format_exc(), + } + failures.append(failure) + if args.failure_log: + append_jsonl(args.failure_log, failure) + if args.sample_log: + append_jsonl( + args.sample_log, + _batch_log_row( + args=args, + batch=batch, + status="failed", + seconds=batch_seconds, + error=exc, + ), + ) + print(f"failed batch sample_ids={failure['sample_ids']}: {exc}", flush=True) + if not args.keep_going: + raise + summary = { + "schema_version": 1, + "output": str(output_path), + "benchmark": args.benchmark, + "model_key": args.model_key, + "selected_samples_total": selected_total, + "selected_samples_after_resume": len(records), + "skipped_existing_samples": skipped_existing, + "processed_samples": processed, + "processed_prefill_tokens": tokens, + "run_seconds": time.time() - t_run, + "layers": args.layers, + "decode_steps": args.decode_steps, + "write_phases": args.write_phases, + "contains_route_weights": True, + "failed_batches": len(failures), + "failure_log": str(args.failure_log) if args.failure_log else None, + "sample_log": str(args.sample_log) if args.sample_log else None, + "resume_from": [str(path) for path in resume_from], + } + return summary + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--benchmark", choices=sorted(INPUT_FILES), required=True) + parser.add_argument("--model-key", choices=sorted(MODEL_CONFIGS), default="qwen3") + parser.add_argument("--limit", type=int) + parser.add_argument("--sample-ids") + parser.add_argument("--max-input-tokens", type=int) + parser.add_argument("--layers", type=_parse_csv_ints, default=_parse_csv_ints("0,12,23")) + parser.add_argument("--decode-steps", type=int, default=1) + parser.add_argument("--threads", type=int, default=8) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--output", type=Path) + parser.add_argument("--resume", action="store_true") + parser.add_argument("--sort-by-length", action="store_true") + parser.add_argument("--write-phases", choices=("prefill", "decode", "both"), default="both") + parser.add_argument("--keep-going", action="store_true") + parser.add_argument("--failure-log", type=Path) + parser.add_argument("--sample-log", type=Path) + parser.add_argument("--resume-from", type=Path, action="append", default=[]) + args = parser.parse_args() + summary = run(args) + print(summary) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/transactional_emulator/testbench/window1_p2/p2_utils.py b/transactional_emulator/testbench/window1_p2/p2_utils.py new file mode 100644 index 00000000..df7f12cc --- /dev/null +++ b/transactional_emulator/testbench/window1_p2/p2_utils.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +import gzip +import json +import os +import sys +from pathlib import Path +from typing import Any, Iterable + +from transactional_emulator.testbench.window1_p1.p1_utils import ( + PLENA_ROOT, + REPO_ROOT, + TESTBENCH_ROOT, + gini_from_counts, + write_csv, + write_json, +) + + +OUT_ROOT = PLENA_ROOT / "outputs" / "window1_p2" +PAPER_ARTIFACTS = PLENA_ROOT / "paper_artifacts" +WEIGHTS_ROOT = PLENA_ROOT / "weights" +COMPILER_ROOT = PLENA_ROOT / "repos" / "PLENA_Compiler" +TOOLS_ROOT = PLENA_ROOT / "repos" / "PLENA_Tools" + +MODEL_CONFIGS = { + "qwen3": { + "name": "Qwen3-30B-A3B", + "path": WEIGHTS_ROOT / "qwen3-30b-a3b", + "hidden_size": 2048, + "intermediate_size": 768, + "num_experts": 128, + "top_k": 8, + "layers": 48, + "policy_name": "qwen3_moe", + "activation_policy": "standard_swiglu", + }, +} + +INPUT_FILES = { + "gpqa_diamond": PAPER_ARTIFACTS / "gpqa_diamond_model_inputs.jsonl.gz", + "bfcl_v3": PAPER_ARTIFACTS / "bfcl_v3_model_inputs.jsonl.gz", +} + + +def ensure_paths() -> None: + for path in (REPO_ROOT, COMPILER_ROOT, TOOLS_ROOT, TESTBENCH_ROOT): + text = str(path) + if text not in sys.path: + sys.path.insert(0, text) + os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") + os.environ["PYTHONPATH"] = ":".join( + [ + str(REPO_ROOT), + str(COMPILER_ROOT), + str(TOOLS_ROOT), + str(TESTBENCH_ROOT), + os.environ.get("PYTHONPATH", ""), + ] + ) + + +def open_text(path: Path, mode: str = "rt"): + if path.suffix == ".gz": + return gzip.open(path, mode, encoding="utf-8") + return path.open(mode, encoding="utf-8") + + +def iter_jsonl(path: Path) -> Iterable[dict[str, Any]]: + with open_text(path, "rt") as handle: + for line in handle: + line = line.strip() + if line: + yield json.loads(line) + + +def append_jsonl(path: Path, row: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with open_text(path, "at") as handle: + handle.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n") + + +def load_json(path: Path) -> dict[str, Any]: + with path.open(encoding="utf-8") as handle: + return json.load(handle) + + +def routing_stats(indices: list[list[int]], num_experts: int) -> dict[str, Any]: + counts = [0 for _ in range(num_experts)] + for row in indices: + for expert_id in row: + counts[int(expert_id)] += 1 + nonzero = [idx for idx, value in enumerate(counts) if value > 0] + total = sum(counts) + return { + "expert_counts": counts, + "active_experts": nonzero, + "hot_experts": sorted(nonzero, key=lambda idx: (-counts[idx], idx)), + "cold_experts": [idx for idx, value in enumerate(counts) if value == 0], + "duplicate_factor": float(total) / float(len(nonzero)) if nonzero else 0.0, + "gini": gini_from_counts(counts), + } + + +def stable_id(text: str) -> str: + return "".join(ch if ch.isalnum() or ch in ("_", "-") else "_" for ch in text) + + +__all__ = [ + "INPUT_FILES", + "MODEL_CONFIGS", + "OUT_ROOT", + "PAPER_ARTIFACTS", + "append_jsonl", + "ensure_paths", + "iter_jsonl", + "load_json", + "routing_stats", + "stable_id", + "write_csv", + "write_json", +] + diff --git a/transactional_emulator/testbench/window1_p2/qwen3_trace_replay_test.py b/transactional_emulator/testbench/window1_p2/qwen3_trace_replay_test.py new file mode 100644 index 00000000..c515fbff --- /dev/null +++ b/transactional_emulator/testbench/window1_p2/qwen3_trace_replay_test.py @@ -0,0 +1,404 @@ +#!/usr/bin/env python3 +"""Replay a Qwen3 true route trace through a fixed-route MoE emulator program.""" + +from __future__ import annotations + +import argparse +import json +import math +import sys +from pathlib import Path +from typing import Any + +import torch + +_REPOS_ROOT = Path(__file__).resolve().parents[4] +_TOP_LEVEL_COMPILER = _REPOS_ROOT / "PLENA_Compiler" +if _TOP_LEVEL_COMPILER.exists(): + sys.path.insert(0, str(_TOP_LEVEL_COMPILER)) + +from compiler.aten.plena import PlenaCompiler +from transactional_emulator.testbench.aten.configurable import add_hw_args, setup_hw +from transactional_emulator.testbench.emulator_runner import ( + compare_emulator_output, + run_emulator, + run_emulator_repeat_gate, +) +from transactional_emulator.testbench.layout_utils import infer_hbm_tensor_layouts, prestage_bf16_vram_matrix +from transactional_emulator.testbench.models.gpt_oss.attention_semantics_test import _comparison_params +from transactional_emulator.testbench.routed_moe.gpt_oss_moe_gather_scatter_test import _align_to +from transactional_emulator.testbench.sim_env_utils import create_mem_for_sim +from transactional_emulator.testbench.window1_p1.validate_route_trace import validate_trace +from transactional_emulator.testbench.window1_p2.p2_utils import OUT_ROOT, ensure_paths, load_json, write_json +from transactional_emulator.tools.create_sim_env import create_sim_env + + +def _flatten_int(rows: list[list[int]]) -> list[int]: + return [int(value) for row in rows for value in row] + + +def _flatten_float(rows: list[list[float]]) -> list[float]: + return [float(value) for row in rows for value in row] + + +def _expert_stride(prog: PlenaCompiler, shape: tuple[int, int]) -> int: + raw_size = int(shape[0] * shape[1] * prog.real_data_ratio) + return _align_to(raw_size, prog.mlen) + + +def _build_selected_dummy_weight_table( + prog: PlenaCompiler, + *, + prefix: str, + selected_experts: list[int], + num_experts: int, + shape: tuple[int, int], + input_tensors: dict[str, torch.Tensor], +) -> tuple[list[Any], int, int]: + stride = _expert_stride(prog, shape) + base = prog._allocate_hbm(stride * num_experts) + inputs = [] + zero = torch.zeros(shape, dtype=torch.bfloat16) + for expert_id in selected_experts: + name = f"{prefix}_e{expert_id}" + inputs.append(prog.input(name, shape=shape, hbm_addr=base + expert_id * stride)) + input_tensors[name] = zero + if not inputs: + raise ValueError("selected_experts cannot be empty") + return inputs, base, stride + + +def _synthetic_x(rows: int, hidden: int, *, mode: str, seed: int) -> torch.Tensor: + if mode == "zeros": + return torch.zeros(rows, hidden, dtype=torch.bfloat16) + if mode == "random": + torch.manual_seed(seed) + return (torch.randn(rows, hidden) * 0.02).to(torch.bfloat16) + raise ValueError(f"unknown input mode {mode!r}") + + +def build_artifacts(args: argparse.Namespace) -> dict[str, Any]: + ensure_paths() + trace = load_json(args.trace) + errors = validate_trace(trace, allow_missing_artifacts=True) + if errors: + raise ValueError("Invalid route trace:\n" + "\n".join(errors)) + + model = trace["model"] + workload = trace["workload"] + routing = trace["routing"] + if model["name"] != "Qwen3-30B-A3B": + raise ValueError(f"qwen3_trace_replay_test supports Qwen3-30B-A3B, got {model['name']!r}") + + build_dir = args.build_dir or (OUT_ROOT / "trace_replay" / trace["trace_id"]) + build_dir = build_dir.expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + hw = setup_hw(args, build_dir) + + rows = int(workload["token_count"]) + hidden = int(model["hidden_size"]) + intermediate = int(model["intermediate_size"]) + num_experts = int(model["num_experts"]) + top_k = int(model["top_k"]) + topk_indices = [[int(value) for value in row] for row in routing["topk_indices"]] + topk_weights = [[float(value) for value in row] for row in routing["topk_weights"]] + pair_count = rows * top_k + if len(topk_indices) != rows or len(topk_weights) != rows: + raise ValueError("trace topk row count does not match token_count") + selected_experts = sorted({expert_id for row in topk_indices for expert_id in row}) + + prog = PlenaCompiler(mlen=args.mlen, blen=args.blen, real_data_ratio=hw.real_data_ratio) + input_tensors: dict[str, torch.Tensor] = {} + + physical_rows = max(args.blen, math.ceil(rows / args.blen) * args.blen) + vram_preload = torch.zeros(physical_rows * hidden, dtype=torch.bfloat16) + x = _synthetic_x(rows, hidden, mode=args.input_mode, seed=args.seed) + x_vram = prestage_bf16_vram_matrix( + prog=prog, + name="TraceReplayX", + tensor=x, + vram_addr=0, + physical_shape=(physical_rows, hidden), + vram_preload=vram_preload, + ) + + gate_inputs, gate_base, gate_stride = _build_selected_dummy_weight_table( + prog, + prefix="QwenGate", + selected_experts=selected_experts, + num_experts=num_experts, + shape=(hidden, intermediate), + input_tensors=input_tensors, + ) + up_inputs, up_base, up_stride = _build_selected_dummy_weight_table( + prog, + prefix="QwenUp", + selected_experts=selected_experts, + num_experts=num_experts, + shape=(hidden, intermediate), + input_tensors=input_tensors, + ) + down_inputs, down_base, down_stride = _build_selected_dummy_weight_table( + prog, + prefix="QwenDown", + selected_experts=selected_experts, + num_experts=num_experts, + shape=(intermediate, hidden), + input_tensors=input_tensors, + ) + weight_templates = (gate_inputs[0], up_inputs[0], down_inputs[0]) + weight_table_bases = (gate_base, up_base, down_base) + weight_table_strides = (gate_stride, up_stride, down_stride) + + zero = prog.fp_var("decoder_zero", size=1) + one = prog.fp_var("decoder_one", size=args.blen) + neg_alpha = prog.fp_var("decoder_neg_alpha", size=args.blen) + limit_pos = prog.fp_var("decoder_unused_limit_pos", size=args.blen) + limit_neg = prog.fp_var("decoder_unused_limit_neg", size=args.blen) + shared_zero_row = prog.fp_var("decoder_shared_zero_row", size=args.mlen) + topk_weight_var = prog.fp_var("trace_topk_weights", size=pair_count) + route_fp_scratch = prog.fp_var("trace_route_fp_scratch", size=args.mlen) + topk_weights_fp_base = topk_weight_var.address + topk_indices_int_base = 0 + + accumulator = prog.alloc( + "TraceReplayAccumulator", + rows=rows, + cols=hidden, + strict=False, + physical_shape=(physical_rows, hidden), + ) + prog.moe_true_zero_vram_rows_v0( + accumulator, + rows=list(range(rows)), + hidden=hidden, + zero_row=shared_zero_row, + policy_name="qwen3_moe", + name="trace_acc_zero", + ) + + for pair_idx in range(pair_count): + token_idx = pair_idx // top_k + gathered = prog.moe_gather_token_rows_from_vram_v0( + x_vram, + token_indices=[token_idx], + hidden=hidden, + zero_row=shared_zero_row, + policy_name="qwen3_moe", + name=f"trace_pair{pair_idx}_vram_gather_t{token_idx}", + ) + expert_out = prog.moe_dynamic_expert_pair_v0( + gathered, + weight_templates, + weight_table_bases=weight_table_bases, + weight_table_strides=weight_table_strides, + expert_indices_int_base=topk_indices_int_base, + weights_fp_base=topk_weights_fp_base, + pair_idx=pair_idx, + bias_tables=None, + rows=args.blen, + intermediate=intermediate, + constants=(zero, limit_pos, limit_neg, one, neg_alpha), + zero_row=shared_zero_row, + route_fp_scratch=route_fp_scratch, + policy_name="qwen3_moe", + activation_policy="standard_swiglu", + name=f"trace_pair{pair_idx}", + ) + prog.moe_scatter_add_active_rows_v0( + accumulator, + expert_out, + token_indices=[token_idx], + active_rows=[0], + hidden=hidden, + policy_name="qwen3_moe", + name=f"trace_pair{pair_idx}_scatter", + ) + + isa = prog.compile() + fp_preload_len = max( + neg_alpha.address + neg_alpha.size, + topk_weight_var.address + topk_weight_var.size, + route_fp_scratch.address + route_fp_scratch.size, + shared_zero_row.address + shared_zero_row.size, + ) + fp_preload = [0.0] * fp_preload_len + for idx in range(one.size): + fp_preload[one.address + idx] = 1.0 + for idx in range(neg_alpha.size): + fp_preload[neg_alpha.address + idx] = -1.0 + flat_weights = _flatten_float(topk_weights) + for idx, value in enumerate(flat_weights): + fp_preload[topk_weights_fp_base + idx] = value + + int_preload = torch.tensor(_flatten_int(topk_indices), dtype=torch.int32) + golden = torch.zeros(rows, hidden, dtype=torch.bfloat16) + comparison_params = _comparison_params( + prog.get_vram_addr(accumulator.name), + rows, + hidden, + args.mlen, + physical_rows=accumulator.physical_shape[0], + ) + tensor_layouts = infer_hbm_tensor_layouts(input_tensors) + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + data_order = sorted(input_tensors, key=lambda name: hbm_addrs[name]) + + create_sim_env( + input_tensors, + isa, + { + "original_output": golden, + "compile_info": { + "trace_id": trace["trace_id"], + "measurement_note": trace.get("measurement_note"), + "input_mode": args.input_mode, + "selected_experts": selected_experts, + }, + }, + fp_preload=fp_preload, + int_preload=int_preload, + build_dir=str(build_dir), + vram_preload=vram_preload, + tensor_layouts=tensor_layouts, + ) + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="qwen3_trace_replay", + specified_data_order=data_order, + build_path=build_dir, + input_tensors=input_tensors, + tensor_layouts=tensor_layouts, + hbm_addrs=hbm_addrs, + ) + (build_dir / "comparison_params.json").write_text(json.dumps(comparison_params, indent=2) + "\n") + (build_dir / "generated_asm_code.asm").write_text(isa) + write_json(build_dir / "trace.json", trace) + manifest = { + "schema_version": 1, + "trace_id": trace["trace_id"], + "trace_path": str(args.trace), + "benchmark": workload["benchmark"], + "sample_id": workload["sample_id"], + "phase": workload["phase"], + "layer": model["layer_index"], + "rows": rows, + "hidden": hidden, + "intermediate": intermediate, + "num_experts": num_experts, + "top_k": top_k, + "pair_count": pair_count, + "selected_experts": selected_experts, + "selected_expert_count": len(selected_experts), + "mlen": args.mlen, + "blen": args.blen, + "input_mode": args.input_mode, + "topk_indices_int_base": topk_indices_int_base, + "topk_weights_fp_base": topk_weights_fp_base, + "weight_table_bases": {"gate": gate_base, "up": up_base, "down": down_base}, + "weight_table_strides": {"gate": gate_stride, "up": up_stride, "down": down_stride}, + "hbm_input_tensor_count": len(input_tensors), + "asm_lines": len(isa.splitlines()), + "measurement_note": "self-consistent upper bound, absolute accuracy pending RTL (Window 2)", + "comparison_params": comparison_params, + } + write_json(build_dir / "qwen3_trace_replay_manifest.json", manifest) + return {"trace": trace, "build_dir": build_dir, "manifest": manifest} + + +def run_trace(args: argparse.Namespace) -> dict[str, Any]: + built = build_artifacts(args) + build_dir: Path = built["build_dir"] + manifest = built["manifest"] + if args.no_run: + return {"schema_version": 1, "build_dir": str(build_dir), "manifest": manifest, "ran": False} + + metrics = run_emulator( + build_dir, + threads=args.emu_threads, + stage_profile=args.stage_profile, + dump_cwd=build_dir, + overlap_prefetch_compute=args.experimental_overlap_prefetch_compute, + ) + results, params = compare_emulator_output(build_dir) + gate = { + "passed": bool(results.get("test_pass", results.get("allclose_pass", False))), + "allclose_pass": bool(results.get("allclose_pass", False)), + "relative_match_rate": results.get("relative_match_rate"), + "max_error": results.get("max_error"), + "relative_error": results.get("relative_error"), + "zero_input_gate": manifest["input_mode"] == "zeros", + } + repeat_summary = None + if args.repeat_gate: + repeat_summary = run_emulator_repeat_gate( + build_dir, + repeats=args.repeat_gate, + threads=args.emu_threads, + stage_profile=False, + overlap_prefetch_compute=args.experimental_overlap_prefetch_compute, + ) + summary = { + **manifest, + "run_metrics": metrics, + "comparison_params_runtime": params, + "emulator_compare_raw": { + key: results[key] + for key in ( + "mse", + "mae", + "max_error", + "relative_error", + "relative_match_rate", + "allclose_match_rate", + "match_rate", + "allclose_pass", + "test_pass", + "atol", + "rtol", + ) + if key in results + }, + "full_vram_gate": gate, + "repeat_gate": repeat_summary, + } + write_json(build_dir / "qwen3_trace_replay_results.json", summary) + write_json(build_dir / "gather_scatter_results.json", summary) + if args.cleanup_dumps: + removed = [] + for name in ("mram_dump.bin", "vram_dump.bin", "hbm_dump.bin", "fpsram_dump.bin", "intsram_dump.bin"): + path = build_dir / name + if path.exists(): + path.unlink() + removed.append(name) + summary["cleanup_removed_dumps"] = removed + write_json(build_dir / "qwen3_trace_replay_results.json", summary) + write_json(build_dir / "gather_scatter_results.json", summary) + print(json.dumps(summary, indent=2, sort_keys=True)) + if not gate["passed"]: + raise AssertionError(f"trace replay functional gate failed: {gate}") + return summary + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + add_hw_args(parser) + parser.add_argument("trace", type=Path) + parser.add_argument("--build-dir", type=Path) + parser.add_argument("--emu-threads", type=int, default=1) + parser.add_argument("--input-mode", choices=("zeros", "random"), default="zeros") + parser.add_argument("--stage-profile", action="store_true") + parser.add_argument("--repeat-gate", type=int, default=0) + parser.add_argument("--experimental-overlap-prefetch-compute", action="store_true") + parser.add_argument("--keep-dumps", dest="cleanup_dumps", action="store_false") + parser.add_argument("--no-run", action="store_true") + parser.set_defaults(cleanup_dumps=True) + parser.set_defaults(mlen=128, blen=4) + args = parser.parse_args() + run_trace(args) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/transactional_emulator/testbench/window1_p2/run_trace_batch.py b/transactional_emulator/testbench/window1_p2/run_trace_batch.py new file mode 100644 index 00000000..e4922c31 --- /dev/null +++ b/transactional_emulator/testbench/window1_p2/run_trace_batch.py @@ -0,0 +1,145 @@ +#!/usr/bin/env python3 +"""Run a batch of P2 route traces through the Qwen3 replay harness.""" + +from __future__ import annotations + +import argparse +import subprocess +import sys +import time +from pathlib import Path + +from transactional_emulator.testbench.window1_p2.p2_utils import OUT_ROOT, REPO_ROOT, load_json, write_json + + +def _prune_success_artifacts(build_dir: Path) -> dict[str, int]: + prune_suffixes = {".pt"} + prune_names = { + "hbm_for_behave_sim.bin", + "fp_sram.bin", + "int_sram.bin", + "vram_preload.bin", + "vector_result.mem", + } + pruned_files = 0 + pruned_bytes = 0 + for path in sorted(child for child in build_dir.iterdir() if child.is_file()): + if path.suffix not in prune_suffixes and path.name not in prune_names: + continue + size = path.stat().st_size + path.unlink() + pruned_files += 1 + pruned_bytes += size + return {"pruned_file_count": pruned_files, "pruned_bytes": pruned_bytes} + + +def _trace_paths(args: argparse.Namespace) -> list[Path]: + paths: list[Path] = [] + for pattern in args.trace_glob: + paths.extend(sorted(Path().glob(pattern) if not pattern.startswith("/") else Path("/").glob(pattern[1:]))) + if args.limit is not None: + paths = paths[: args.limit] + return paths + + +def run(args: argparse.Namespace) -> dict: + traces = _trace_paths(args) + manifest = { + "schema_version": 1, + "started_at": time.time(), + "trace_count": len(traces), + "runs": [], + } + args.manifest_out.parent.mkdir(parents=True, exist_ok=True) + for idx, trace_path in enumerate(traces, start=1): + trace = load_json(trace_path) + build_dir = args.root / trace["trace_id"] + result_path = build_dir / "qwen3_trace_replay_results.json" + if args.skip_existing and result_path.exists(): + row = { + "trace": str(trace_path), + "trace_id": trace["trace_id"], + "build_dir": str(build_dir), + "status": "skipped_existing", + } + manifest["runs"].append(row) + write_json(args.manifest_out, manifest) + print(f"[{idx}/{len(traces)}] skip {trace['trace_id']}", flush=True) + continue + + build_dir.mkdir(parents=True, exist_ok=True) + cmd = [ + sys.executable, + "-m", + "transactional_emulator.testbench.window1_p2.qwen3_trace_replay_test", + str(trace_path), + "--build-dir", + str(build_dir), + "--mlen", + str(args.mlen), + "--blen", + str(args.blen), + "--emu-threads", + str(args.emu_threads), + ] + if args.stage_profile: + cmd.append("--stage-profile") + if args.experimental_overlap_prefetch_compute: + cmd.append("--experimental-overlap-prefetch-compute") + if args.keep_dumps: + cmd.append("--keep-dumps") + print(f"[{idx}/{len(traces)}] run {trace['trace_id']}", flush=True) + start = time.time() + log_path = build_dir / "p2_batch_stdout.log" + with log_path.open("w", encoding="utf-8") as log: + proc = subprocess.run(cmd, cwd=args.cwd, text=True, stdout=log, stderr=subprocess.STDOUT) + row = { + "trace": str(trace_path), + "trace_id": trace["trace_id"], + "build_dir": str(build_dir), + "command": cmd, + "return_code": proc.returncode, + "status": "passed" if proc.returncode == 0 else "failed", + "host_batch_seconds": time.time() - start, + "log_path": str(log_path), + } + if row["status"] == "passed" and args.prune_success_artifacts: + row.update(_prune_success_artifacts(build_dir)) + manifest["runs"].append(row) + write_json(args.manifest_out, manifest) + print( + f"[{idx}/{len(traces)}] {row['status']} {trace['trace_id']} " + f"({row['host_batch_seconds']:.1f}s)", + flush=True, + ) + if proc.returncode != 0 and not args.keep_going: + return manifest + manifest["ended_at"] = time.time() + write_json(args.manifest_out, manifest) + return manifest + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--trace-glob", action="append", required=True) + parser.add_argument("--root", type=Path, default=OUT_ROOT / "trace_replay") + parser.add_argument("--manifest-out", type=Path, default=OUT_ROOT / "trace_replay_batch_manifest.json") + parser.add_argument("--cwd", type=Path, default=REPO_ROOT) + parser.add_argument("--mlen", type=int, default=128) + parser.add_argument("--blen", type=int, default=4) + parser.add_argument("--emu-threads", type=int, default=1) + parser.add_argument("--limit", type=int) + parser.add_argument("--stage-profile", action="store_true") + parser.add_argument("--skip-existing", action="store_true") + parser.add_argument("--keep-going", action="store_true") + parser.add_argument("--keep-dumps", action="store_true") + parser.add_argument("--experimental-overlap-prefetch-compute", action="store_true") + parser.add_argument("--prune-success-artifacts", action="store_true") + args = parser.parse_args() + manifest = run(args) + failed = [row for row in manifest["runs"] if row.get("status") == "failed"] + return 1 if failed else 0 + + +if __name__ == "__main__": + raise SystemExit(main())