MoE: timing-replay harnesses, shared experts, parameterized V_TOPK routing

[qichao-arlo-wang]

What this PR does

Two related pieces of MoE work:

1. Timing-replay harnesses (moe_timing/{replay,qwen,campaign}) — capture real
   MoE routing, replay it through the emulator, measure cycles and HBM bytes.
2. Shared-expert support — DeepSeek / Qwen2-MoE / Llama-4 / GLM, plus the two
   things it turned out to require: routing at arbitrary expert shapes, and an
   explicit stage-attribution contract to replace substring-matching the
   compiler's ASM comments.

Requires AICrossSim/PLENA_Compiler#69;
the submodule pin bump is included here. Neither is useful without the other.

---

Part 1 — Timing-replay harnesses

• replay/ — dependency-free route-trace schema and validator, sample-trace
  generator, replay runner, timing validation gates, result export.
• qwen/ — Qwen3 router-trace generation from tokenized inputs and local
  weights, conversion to the validated schema, replay, resumable batch runner,
  pilot export.
• campaign/ — stratified largest-remainder subset selection,
  serial-vs-parallel determinism gate, checkpointed parallel replay.

Rebased onto the merged emulator work

This branch originally carried a fourth commit, "emulator: deterministic timing,
stage profiling, and overlap model". That work landed on main independently as
2539533 (#98), was hardened further in fc57756 (#99), and the copy here was a
superseded duplicate — it has been dropped.

fea3dc5 migrates the harnesses to the resulting stage-profile schema v3.
Three fields these files read no longer exist, and one of their sums is now
explicitly forbidden:

was	now
total_stage_wall_cycles	removed as tautological (always identical to total_profiled_cycles). summarize_run reports the profiled figure, in picoseconds and cycles.
cycle_fraction	renamed time_fraction — once both derived from picoseconds they were the same ratio, so only the exact one survives.
resource_proxy_cycles.ramulator_proxy	removed in #98 — incremented with the identical value as dma, so it carried no information.
Σ per-stage wall_cycles in _sum_stages	sums wall_picos and rounds once. Under v3 each stage rounds up independently, so n stages over-report by up to n−1 cycles.

All three removed fields were read through .get(), so before this commit they
had been silently producing None columns rather than raising.

Also passes dump_cwd to run_emulator_repeat_gate in the Qwen3 replay. #99
added that parameter for this exact caller; without it the repeat runs fall back
to the shared emulator directory, so concurrent campaign workers race on
vram_dump.bin / fpsram_dump.bin.

---

Part 2 — Shared experts

A shared expert is an FFN every token passes through, summed into the routed
output unweighted:

y = shared(x) + Σ_k route_weight_k · routed_expert_k(x)

model	n_shared	shared intermediate	gate
DeepSeek-V2/V3, Kimi K2	1–2	moe_inter × n_shared	none
Qwen2-MoE	1	independent	sigmoid
Llama-4, GLM-4.5, Qwen3-Next	1	independent	none
GPT-OSS, Qwen3-MoE, Mixtral	0	—	—

n_shared_experts is not an emitter parameter: DeepSeek stores its shared experts
pre-concatenated and instantiates a single MLP of width
moe_intermediate_size × n_shared_experts. Because SwiGLU is elementwise along
the intermediate axis while the down projection sums over it, the fused MLP is
exactly equal to the sum of the individual shared experts — not an
approximation.

C_SET_TOPK_REG — routing at any expert shape

V_TOPK's rmask was a two-entry table (32/top-4, 128/top-8) covering GPT-OSS and
Qwen3-30B-A3B and nothing else. Llama-4 Scout is 16/top-1, Qwen2-MoE 60/top-4,
DeepSeek-V2-Lite 64/top-6, DeepSeek-V3 and Kimi K2 256/top-8 — each would have
cost an ISA revision.

topk_softmax was already generic in both counts; only the decode table was
not. New sticky control register holds (num_experts << 8) | top_k; rmask=15
escapes to it, 0 and 1 keep their exact previous meaning so existing programs
emit byte-identical ASM.

topk_policy is Option<u32> rather than a u32 defaulting to 0 so V_TOPK can
trap naming the missing C_SET_TOPK_REG. A 0 default would unpack to top_k=0
and abort inside topk_softmax with "topk must be positive", pointing at the
wrong thing.

classify_timing_access's V_TOPK arm derives the logit read extent from the
expert count; its pre-existing _ => 128 fallback would have under-reported every
model wider than 128 experts, crediting the unmodelled rows to the overlay as free
hiding capacity.

Stage attribution: @stage= markers replace substring sniffing

classify_comment guessed stages by substring-matching the compiler's generated
ASM prose. Unfixable for shared experts: there was no way to express a stage the
substring table did not already contain, and the compiler's policy_name
parameter would have broken the GPT-OSS literals the table keys on.

; @stage=<name> markers now win. A program carrying any marker is classified
only by markers — mixing is not a conservative middle ground but wrong, because
a marked region's body still contains ordinary comments from general-purpose
helpers ("sub projection", "subblock [") that a live substring rule would
immediately steal back. Legacy classification stays for unmarked programs,
including the ASM already checked in under moe_timing.

Three new stages: shared_expert_projection / _activation / _gate, so the
profile can answer what fraction of MoE time is architecturally dense versus
routing-dependent. Schema version 4, with two new reported fields:

• classification.stage_attribution — which mode ran. Every vocabulary_* field
  only means anything in legacy mode; under markers they describe rules that
  never ran.
• classification.unresolved_stage_markers — names the compiler emitted that no
  StageKind matches. The compiler validates against its own MOE_STAGES before
  emitting, so anything here means the vocabularies drifted, and those regions
  silently inherit the previous stage.

Route-trace schema v2

Adds shared_experts / shared_intermediate_size / shared_gate to model. A
v1 trace is a valid v2 trace with no shared branch, so both are accepted.
shared_intermediate_size is the fused width. The validator rejects
half-specified combinations: a shared_experts with no width would replay as
routed-only while still labelled shared, and the timing split would attribute
nothing while looking healthy.

summarize_run gains a shared-vs-routed picosecond split. gather /
scatter_combine / accumulator_init are excluded from both terms — combine
plumbing shared by the branches, and billing it to "routed" would overstate its
share.

Tests

recipe	covers
just test-shared-moe	DeepSeek: ungated SwiGLU shared expert
just test-shared-moe-gated	Qwen2-MoE: adds the sigmoid shared-expert gate
just test-shared-moe-deepseek-fused	n_shared=2 fused + routed-accumulator combine
just test-router-policy-all	six routing shapes, both encodings

Shared-expert comparisons are bit-exact at atol=rtol=0. MXFP8-representable
inputs make weight quantization the identity, leaving only BF16 rounding, which
_shared_moe_reference.py mirrors step by step — silu(g)*u in one rounding does
not equal the hardware's six.

The gate weight is scaled by a power of two so the sigmoid stays strictly
interior; at full scale the logits reach ±16 and the gate saturates to exactly 1.0
for most rows, which would let a completely unapplied gate still pass. A guard
fails the test if any scalar saturates.

Router policies assert which encoding they took, not just that routing was
correct — otherwise a shape could quietly fall back to a fixed rmask, produce
correct indices anyway, and leave the escape path untested while appearing
covered. DeepSeek-V3 spans four MLEN-wide logit blocks and selects expert 255.

Negative-tested

injected defect	caught by
packing shift 8 → 7	router policy test, wrong indices
@stage= markers not emitted	attribution fell back to legacy substrings
shared gate multiply dropped	numeric, 0.00% match rate, MSE 1.4e6
shared branch marked as routed	{'expert_projection': 568} in an unrouted program

The last is not hypothetical — the guard caught it live during bring-up at 999
misattributed instructions, because the broadcast emitter the gate reuses hardcoded
expert_route_weight, and re-marking afterwards is too late for a sticky marker.

Side effect: three previously-unrunnable files now resolve

gpt_oss_moe_gather_scatter_test.py, models/gpt_oss/attention_semantics_test.py
and moe_timing/qwen/qwen3_trace_replay.py called eight compiler methods that did
not exist in the pinned submodule. AST-verified: all now resolve.

They still cannot run — they need the gpt-oss-20b checkpoint plus
huggingface_hub/safetensors, which is the pre-existing reason the justfile
excludes them from CI. That half is unchanged.

---

Scope

These harnesses measure the timing of MoE layers in isolation. End-to-end
model timing belongs in analytic_models/ as a composition layer. Numerical
correctness of the MoE ops is validated by the routed-MoE op tests and the new
shared-expert tests, not by the timing harnesses: the replay drives the emulator
with dummy zero expert weights, so its gate is a shape/no-crash smoke,
deliberately named zero_input_smoke_gate. Overlap-adjusted cycle counts are
opt-in and are an estimate, not a bound in either direction.

Known limitations

The routed path is pair-major. One (token, expert) pair at a time in a
BLEN-row slot, re-fetching expert weights per pair. Measured by emitting ASM and
counting matrix ops (MLEN=64, hidden=intermediate=64, top_k=2):

lowering	tokens	M_MM	H_PREFETCH_M	useful rows/tile	utilisation
routed (pair-major)	1	6	6	1 / 64	1.56%
routed (pair-major)	4	24	24	1 / 64	1.56%
routed (pair-major)	16	96	96	1 / 64	1.56%
shared (batched)	1	3	3	1 / 64	1.56%
shared (batched)	8	3	3	8 / 64	12.5%
shared (batched)	64	3	3	64 / 64	100%

Routed cost is strictly linear in tokens × top_k with zero batching benefit; the
shared branch costs the same for 1 token or 64. The compiler calls the pair-major
lowering "intentionally wasteful but keeps the first L2 correctness path exact" —
the right call for bring-up, but now the dominant term in anything measuring
prefill.

So a shared-vs-routed cycle ratio measured today reflects the routed lowering's
per-pair overhead at least as much as it reflects the architecture. Documented
on _branch_split and in program_moe_shared.py's module docstring rather than
left for whoever reads the numbers to discover. Expert-major batching is follow-up
work; note it gains nothing at batch=1 decode and is a prefill / large-batch
concern.

Carried over from the original review, none introduced by this branch:

• moe_timing has no CI wiring and no tests of its own. ~4.4k lines of Python
  with no justfile target and no assertions. ruff covers formatting and lint;
  nothing covers execution. (The shared-expert and routing tests added here are
  in CI.)
• campaign/run_subset.py's _result_ok reads a stale gate key
  (result.get("functional_gate", True), defaulting to True). Since the results
  file is written before the gate assertion fires, --skip-existing treats a
  previously failed trace as complete and never retries it.
• qwen/utils.py hard-codes a machine layout: PLENA_ROOT = REPO_ROOT.parents[1],
  with no environment override (unlike PLENA_OUT_ROOT).
• replay/timing_validation_gates.py cannot fail. main() computes the gates,
  writes the summary, prints them, and returns 0 unconditionally.
• Phase-based naming survives in docstrings and output filenames
  (Window 1 P1/P2/P3, p3_rev_*.json).

Validation

• 3 shared-expert variants bit-exact, 6 routing policies pass end to end
• All 6 pre-existing routed-MoE CI tests still pass, including gpt_oss_topk_test
  which now exercises the compiler's deprecated-alias path
• 87 emulator unit tests
• python -m compileall, ruff format --check, ruff check — clean
• Call sites checked against main's emulator_runner.py signatures

When to reconsider adding an ISA for shared expert

Shared-expert support here is a software-only decomposition: no new opcode, no
new register. The only ISA addition in this PR is C_SET_TOPK_REG (0x38), which
belongs to the routed branch. The shared expert lowers to existing
linear_projection + moe_expert_activation_v0, and fused_shared_intermediate
folds n_shared > 1 into one wide MLP that is exactly equal to the sum of the
individual shared experts — SwiGLU is elementwise along the intermediate axis and
the down projection sums over it, so this is not an approximation.

This is deliberate, and it matches practice: no mainstream accelerator gives the
shared expert its own datapath, because it is the dense path the hardware already
covers. Hardware assist goes to routing/gather-scatter — which is exactly where
V_TOPK / C_SET_TOPK_REG sit.

It is also worth being explicit that an ISA would not have solved this PR's actual
pain points. Of the four (marker hijack, the gate's cost shape, pair-major
utilization, profile-contract fragility), a shared_expert_v0 macro-op addresses
none: the first and last are compiler↔emulator attribution contracts, and the
third belongs to the routed lowering. A macro-op would arguably make the first
worse, by hiding the misattribution under one opcode.

So "no ISA" is a decision with an expiry condition, not inertia. Revisit if any of:

• T1 — performance. shared_expert_gate exceeds ~10% of MoE cycles at
  realistic prefill (rows >= 512), or the gate's emit-time-unrolled token loop
  (program_moe_shared.py, for token_idx in range(rows)) overruns the
  instruction-memory budget. Note the fix then is a general V_GEMV /
  vector-reduce-with-activation, not a shared-expert instruction: attention
  scores, router logits and norm statistics all want the same shape. The sigmoid
  itself is not a candidate — it is already 4 scalar-FP ops and correctly kept off
  the vector path.
• T2 — attribution. A marker-hijack-class misattribution recurs a third time.
  The response is still not an ISA: it is moving attribution from ASM comments
  to hardware events (tag the stage at dispatch), which needs no new opcode.
• T3 — expressiveness. A shared-expert variant appears that the current
  instruction set cannot express exactly — e.g. per-token dynamic shared width.
  None exists today; the variation across DeepSeek / Qwen2 / GLM / Llama-4 is
  entirely in width, gating and count, all of which lower exactly as-is.

The bar used above: an ISA addition should (1) express something SW composition
cannot, (2) pay off across models rather than one checkpoint shape, (3) reduce
verification surface, and (4) encode something semantically settled.
C_SET_TOPK_REG meets all four — one control register retires a whole class of
per-shape ISA revisions. A shared_expert_v0 macro-op meets none, and would have
to encode n_shared, activation policy, gate presence and bias presence in
hardware — the same trap C_SET_TOPK_REG exists to avoid.

Staying SW-only carries no forward-compat cost: nothing shared-specific enters the
ISA, so adding hardware later cannot break existing programs.

Boundary with Compiler PR #69

• Pair-major routed utilization (~1.56%) is PR refactor: split accelerator op dispatch #69's scope, not this PR's. The
  routed path processes one (token, expert) pair per BLEN-row slot and re-fetches
  expert weights per pair; fixing it means re-lowering routed to token-major /
  expert-batched. This PR must not touch it.
• Consequently, do not read the shared-vs-routed cycle ratio as an architectural
  result until that lands. Shared is batched and routed is per-pair, so the ratio
  reflects the routed lowering's overhead at least as much as the architecture.
  This caveat is now emitted in the profile itself as
  classification.attribution_notes.shared_vs_routed.
• The gate's emit-time unrolled token loop is follow-up to this line of work,
  not refactor: split accelerator op dispatch #69. It is a shared-path codegen issue (static instruction count grows with
  batch), unrelated to routed placement.
• PR refactor: split accelerator op dispatch #69 carries the reciprocal note.

Review hardening (stacked)

Defensive follow-ups from architectural review live on feat/moe-timing-review-hardening
(stacked on this branch, not yet part of it):

• the compiler-vocabulary guard now parses MOE_STAGES out of the pinned submodule
  instead of comparing against a hand-copied list that could never fail in the
  direction it existed to protect;
• classification.attribution_notes.shared_vs_routed carries the caveat above into
  the profile JSON, where consumers actually read it;
• stage becomes a required argument on the three stage-polymorphic emitters
  (PLENA_Compiler feat/moe-stage-required-arg), so the marker-hijack bug class
  fails at emit time instead of silently misattributing.

🤖 Generated with Claude Code