Skip to content

MoE timing-replay: deterministic emulator timing + route-trace harnesses#97

Open
qichao-arlo-wang wants to merge 4 commits into
mainfrom
feat/moe-timing
Open

MoE timing-replay: deterministic emulator timing + route-trace harnesses#97
qichao-arlo-wang wants to merge 4 commits into
mainfrom
feat/moe-timing

Conversation

@qichao-arlo-wang

Copy link
Copy Markdown
Collaborator

What this PR does

Consolidates the MoE timing-replay stack — previously split across #93, #94, #95, #96 — into a single branch against main, with the review findings from those PRs fixed. It supersedes and replaces all four (which are being closed).

Organized as four thematic commits:

  1. emulator: deterministic timing, stage profiling, and overlap model — deterministic event ordering (monotonic timer sequence ids; address-ordered HBM gather / ramulator transfers), cycle-attributed stage profile (schema v2, serde), an opt-in prefetch/compute overlap timing model, and a repeat gate + per-run dump cwd in the Python runner.
  2. testbench: MoE route-trace replay harness and validation gates (P1) — dependency-free route-trace schema + validator, sample-trace generator, replay runner, timing validation gates, result export.
  3. testbench: Qwen3 MoE route-trace generation and replay (P2) — true Qwen3 router-trace generation from tokenized inputs + local weights, conversion to the validated schema, replay, resumable batch runner, pilot export.
  4. testbench: representative MoE timing runner (P3) — stratified largest-remainder subset selection, serial-vs-parallel determinism gate, checkpointed parallel replay.

Fixes applied during consolidation

  • Dropped the timing goldens entirely (check_timing_goldens.py, overlap_golden_workloads.json, run_timing_smoke.sh, the timing-smoke recipe): they depended on out-of-repo build artifacts and were not CI-runnable. The overlap-timing engine (timing_overlay.rs, unit-tested) and the standalone timing_validation_gates.py are kept.
  • Removed the atomic-global sentinel for overlap report cycles in favor of a documented Mutex<Option<u64>> (a module cell is still needed because Executor::spawn requires Future<Output = ()>), and clarified the muddled overlap log line.
  • Output roots no longer escape the repo: window1_p{1,2,3} outputs default to a git-ignored in-repo outputs/ (override the base with PLENA_OUT_ROOT); added /outputs/ to .gitignore.
  • Trimmed P3: dropped the superseded run_p3_scale_batch.py and the two personal tmux babysitter scripts (wait_p3_rev_stage_b.py, watch_p3_rev_raw_build.py); removed dead code in run_p3_rev_subset.py.
  • Renamed qwen3_trace_replay_test.pyqwen3_trace_replay.py (it is a CLI, not a pytest module) and updated all references.
  • Parenthesized a precedence-fragile validator condition; resolved all ruff findings (E402, RUF005).

Validation (local, in nix develop)

  • cargo fmt --all --check, cargo build --release, cargo test --workspace --release — all green.
  • just test-aten-linear — PASSED, 100% match / 0.0 error.
  • Stage-profile schema-v2 emitted end-to-end; total_simulation_cycles == total_profiled_cycles invariant holds.
  • ruff check . + ruff format --check . — clean.

Compiler submodule stays pinned at a765612 (row-granular ABI), identical to main.

qichao-arlo-wang and others added 4 commits July 23, 2026 23:15
Deterministic event ordering (monotonic timer sequence ids; address-ordered HBM
gather / ramulator transfers), a cycle-attributed serde stage profile (schema
v2) with a routed-MoE classification-coverage guard that warns / fails loudly
when the compiler's ASM comment vocabulary drifts, and an opt-in prefetch/
compute overlap model that accumulates in the picosecond domain (rounding to
cycles once) and treats write-out ops as dependency-free. A repeat gate and
per-run dump cwd round out the Python runner.

Co-authored-by: Michael C Li <mcl123@ee-beholder0.ee.ic.ac.uk>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…iming/replay)

A dependency-free route-trace schema and validator (hardened to append errors
rather than crash on malformed input), a sample-trace generator, a replay
runner, timing validation gates (structural, not magic-number), and result
export. Outputs default to a git-ignored in-repo location; PLENA_OUT_ROOT
overrides the base.

Co-authored-by: Michael C Li <mcl123@ee-beholder0.ee.ic.ac.uk>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Generate true Qwen3 router traces from tokenized inputs and local weights,
convert them to the replay schema, and replay through the emulator. The replay
gate is an honest zero-input shape smoke (numerical correctness is validated by
the routed-MoE op tests, not here); the resumable batch runner skips only
genuinely-passed traces; the compiler submodule resolves on a standard checkout.

Co-authored-by: Michael C Li <mcl123@ee-beholder0.ee.ic.ac.uk>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Stratified largest-remainder subset selection, a serial-vs-parallel determinism
gate that no longer false-negatives when stage profiling is off, and
checkpointed parallel replay that is resilient to individual unreadable traces.

Co-authored-by: Michael C Li <mcl123@ee-beholder0.ee.ic.ac.uk>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant