MoE: shared experts, arbitrary V_TOPK routing, and a checked stage-attribution contract

[qichao-arlo-wang]

What this PR does

Three related pieces of MoE work, plus the cross-repo contract that turned out to
be load-bearing for all of them:

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.
3. That contract, made checkable — the emulator reads the compiler's stage
   vocabulary and its caveat text rather than mirroring them, and the compiler's
   stage lint is pointed at the testbenches on this side of the repo boundary.

Requires AICrossSim/PLENA_Compiler#69;
the submodule pin climbs through that branch here, in six steps.

Not ready to merge. Merging either side alone leaves the other pointing at a
commit that is no longer on a branch. Land #69 first, then this.

How to read the commits

The history is grouped by intent, not by the order the work happened in. Each
commit's body opens with its group and position, and every commit that moves the
submodule pin says which step of the ladder it is and why it has to move there.

Group	Commits	What it establishes
A — ISA execution: arbitrary expert shapes	2	Decode C_SET_TOPK_REG, execute the rmask=15 escape
B — MoE timing-replay harnesses	4	replay/ qwen/ campaign/, and the v3 schema migration
C — Stage attribution becomes a declared contract	2	@stage= markers win over substring sniffing; schema v4
D — The profile's cross-repo claims are checked, not mirrored	3	MOE_STAGES read with ast; the caveat held byte-identical
E — The contract is enforced across the repo boundary	2	Testbench call sites name their stage; the compiler's lint reaches them
F — Attribution corrections found by review	2	residual_setup; the router joins the routed branch
G — Pin at the companion PR's head	1	The one pin bump with no code change on this side

The submodule pin ladder

Six steps, each in the same commit as the code that needs it, each strictly
forward along #69's history. This is the whole of the coupling between the two PRs:

#	moves the pin onto (#69)	because, here
1	shared-expert emitters	StageKind gains the three shared_expert_* stages, and the guards hold the two vocabularies equal in both directions — so they cannot land separately
2	SHARED_VS_ROUTED_NOTE	the byte-equality test has nothing to compare against until the constant exists
3	stage becomes required	a breaking signature change; six testbench call sites become hard errors and are fixed in the same commit
4	lint matches positional args	that commit changes _moe_stage_arguments to yield the ast.Constant, and this repo imports that helper rather than copying it
5	residual_setup stage	same both-directions vocabulary lock as step 1
6	CI wiring	nothing here depends on it; the pin points at the head of the branch under review rather than five commits into it

---

Group A — executing C_SET_TOPK_REG

Companion to #69's ISA change. topk_softmax was already generic in
(expert_count, top_k) — only the rmask decode table was not — so this is a
decode and plumbing change, not a semantics change.

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 — DeepSeek-V3 and Kimi K2 are 256 — crediting the
unmodelled rows to the overlay as free hiding capacity.

The second commit moves the 0x38 => C_SET_TOPK_REG decode arm out from between
0x2E and 0x2F, making the control-op block strictly ascending. Match arms are
disjoint constants, so this is behaviour-neutral.

Group B — 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.

The emulator work these were written against landed on main independently as
#98 and was hardened in #99, which changed what the profile emits. The fourth
commit migrates 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 that commit they
had been silently producing None columns rather than raising.

Group C — @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. (router_topk is
excluded from both terms here, and that turns out to be wrong — Group F fixes it
and adds the test that would have caught it.)

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.
#69's Group D exists to retire that whole bug class rather than the one instance.

Group D — the profile's cross-repo claims are checked, not mirrored

The stage vocabulary is read, not copied. stage_marker_names_match_the_compiler_vocabulary
used to compare StageKind against COMPILER_MOE_STAGES, a 13-entry array in
stage_profile.rs commented "Mirrors MOE_STAGES in …". A copy only fails when
someone remembers to update it — the same failure mode as no guard at all, since
the compiler is the side that moves.

It now parses MOE_STAGES out of the pinned submodule with python3 -c "import ast; ...", emitting sorted JSON that serde_json decodes. Not with a
hand-rolled scanner: a scanner that takes the first textual hit, matches a
brace-balanced region and pulls double-quoted runs out of it fits today's
declaration and is silently wrong on four shapes — returning a truncated set
rather than failing:

shape	scanner behaviour
single-quoted names	panics
a comment containing }	truncates at the comment
frozenset({..}) | frozenset({..})	keeps only the first operand
{..} | {..}	keeps only the first operand
a docstring mentioning MOE_STAGES	parses the docstring's decoy set

The silent ones are what matter: the guard compares StageKind against whatever
comes back, so a truncated vocabulary lets it pass while real drift goes
unreported — the exact vacuous green it exists to prevent. ast binds the
declaration at module scope rather than finding it by text, so docstring mentions
and same-named locals are excluded by construction, and quote style, comments,
unions, frozenset(...), list arguments and implicit concatenation all fall out
without special cases. Ten adversarial shapes and six unprovable ones are pinned.

The caveat consumers actually read is now emitted, and checked.
classification.attribution_notes.shared_vs_routed is a new consumer-visible
field carrying the routed-lowering caveat to where numbers are actually read,
rather than leaving it in compiler docstrings. The emulator keeps its own const
— it writes that JSON with no Python available — but a test parses #69's canonical
SHARED_VS_ROUTED_NOTE and asserts byte equality. STRING_CONSTANT_EXTRACTOR
generalises the ast approach to any module-level string and is pinned to fail
rather than invent a value for an absent constant, a value-less annotation, a
non-string, a computed value, or a function-local of the same name.

schema_version stays at 4: v4 is introduced by Group C on this same unmerged
branch, so this extends a schema that has never shipped. The field is purely
additive, but a consumer cannot detect its presence from the version alone.

A missing submodule checkout or absent python3 skips with a reason rather than
panicking — but only outside CI, keyed on $CI. CI checks out submodules: recursive and runs inside nix develop, so a missing prerequisite there is a
broken pipeline, and skipping would restore the vacuous green these guards exist
to prevent.

Group E — the contract crosses the repo boundary

The compiler lints its own emitters and stops at PLENA_Compiler. The testbenches
here call the same emitters from the other side of it, so a wrong stage name in a
testbench was exactly as invisible as one in the compiler used to be.

The first commit is the evidence: making stage required turned six testbench
call sites into hard errors, every one of which had been inheriting
moe_true_zero_vram_rows_v0's accumulator_init default. None had ever stated
an attribution; they simply inherited one. That they surfaced at all is luck — a
signature change stopped them running. A stage name that was wrong but spelled
like a real one would have run fine and changed no instruction count, cycle total
or numerical result.

test_testbench_stage_attribution.py closes that. The matcher is imported from
the compiler's guard by file path, not copied — a second copy of "which callees
take a MoE stage argument" is the same mirroring failure this whole branch is
about. Loading by path also sidesteps aten/__init__.py, so the job needs pytest
and nothing else: no nix, no yaml, no torch.

test_the_known_call_sites_are_actually_scanned pins the sites by path and count,
because both lints pass by finding nothing and a glob that stopped matching would
look identical to being correct. Keyed by path relative to the testbench root
rather than by basename: two files with the same name in different directories
would otherwise collapse into one key and silently drop a site from the count this
exists to pin.

Group F — attribution corrections found by review

• residual_setup. DecoderMoeResidual holds the pre-MoE hidden state added
  back after the experts run — not the combine accumulator, which is zeroed
  separately further down and was already correct. Markers being sticky,
  accumulator_init was also absorbing the residual copy, the input RMSNorm and
  its norm-weight multiply. ResidualSetup also joins the pair-label reset, for
  the same reason RouterTopk and AccumulatorInit are in it: it runs before any
  routing, so it must not inherit the previous (token, expert) pair's label.
  Latent in a single-layer program; real the moment a second MoE layer is emitted.
• The HBM store slot is zeroed, immediately scatter-added into, then stored, so
  it belongs to scatter_combine rather than accumulator_init.
• router_topk joins the routed branch. It does not scale with top_k —
  which is why it was excluded — but it scales with num_experts and has no
  counterpart at all in a dense or shared-only layer. Excluding it understated
  routing by the whole router GEMM plus top-k selection. Input and combine plumbing
  stays in neither term, unchanged.

test_branch_split.py asserts every StageKind the emulator can bill to is
classified exactly once across shared, routed and plumbing — read out of
stage_profile.rs rather than mirrored, so a stage added on the Rust side cannot
quietly fall out of the split and read as free.

CI

stage-attribution-guard runs the testbench guard and the branch-split unit tests.
PLENA_Compiler joins the workflow's path filters: a compiler pin bump is precisely
what moves MOE_STAGES, and it was not triggering this workflow at all.
test_every_guard_file_is_wired_into_ci requires every torch-free test_*.py under
testbench/ to be named in the job, so the next guard added here cannot run nowhere.

Unlike #69, the CI wiring here is not collected into one commit at the end. Each
job is added alongside the guard it runs, because test_every_guard_file_is_wired_into_ci
asserts the workflow's file list is complete — deferring the wiring would make the
intervening commits red.

---

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
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. That is what
classification.attribution_notes.shared_vs_routed now says, in the JSON, held
byte-identical to the compiler's canonical wording. Expert-major batching is
follow-up work on #69; 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 ~4.5k lines of Python with thin CI coverage. ruff covers
  formatting and lint. Partly closed here: replay/utils.py's branch split now has
  unit tests in CI, and 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.
• replay/utils.py:13 hard-codes a machine layout: PLENA_ROOT = REPO_ROOT.parents[1],
  with no environment override (unlike PLENA_OUT_ROOT). qwen/utils.py:13
  imports the name and is a consumer, not the definition — both being line 13 is a
  coincidence worth not repeating.
• 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

cargo test cannot link ramulator outside nix develop, which is not available
on the machine this was written on. So, explicitly:

• Per commit across the branch: cargo fmt --all -- --check and cargo check --tests pass at all 7 Rust-touching commits; the pytest guards pass at every
  commit where they exist (4 → 10 tests), run with CI=1 so none of them can take
  a prerequisite-missing skip.
• The StageKind ↔ MOE_STAGES lockstep was checked directly at every commit,
  since cargo test cannot run locally: the two sets match in both directions from
  the moment the marker mechanism lands, 13 names through Group E and 14 from the
  residual_setup commit onward, stepping exactly where the pin does.
• The Rust↔Python plumbing was executed for real in a standalone crate using the
  exact source text of the extractors: the adversarial and negative canaries pass,
  the real compiler source yields 14 stages, the two caveat strings match verbatim
  at 432 characters, and a reworded caveat compares unequal.
• index() is dense and unique over 0..=14 (15 variants, Other last).
• Each Python guard was verified to fail correctly — a stage-name typo, a removed
  call site, an unclassified stage, an unwired guard file and a missing submodule
  under $CI each fail by name.
• 3 shared-expert variants bit-exact, 6 routing policies pass end to end, all 6
  pre-existing routed-MoE CI tests still pass, 87 emulator unit tests, and
  python -m compileall / ruff format --check / ruff check clean.

The full cargo test --workspace --release and the just integration suites need
nix develop and are covered by CI on this PR, which — unlike the stacked branch
this work previously lived on — now triggers automatically, since the base is main.

When to reconsider adding an ISA for shared expert

Shared-expert support is a software-only decomposition: no new opcode, no new
register. The only ISA addition across both PRs 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 equal to the sum of the individual shared
experts in exact real arithmetic (see #69 on FP associativity and MXFP8 block
alignment, neither of which affects the decomposition's validity).

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.

An ISA would also not have solved this work'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
  overruns the instruction-memory budget. 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.

Note on branch naming

The branch is still called feat/moe-timing, which now under-describes it — the
timing harnesses are one of four groups. Renaming would close this PR, so the name
stays and this paragraph exists instead.

🤖 Generated with Claude Code