MoE: shared experts, arbitrary V_TOPK expert shapes, and an enforced stage-marker contract

[qichao-arlo-wang]

What this PR does

Adds the shared-expert branch of MoE, generalizes V_TOPK to arbitrary expert
shapes, and replaces the undeclared ASM-comment contract that stage profiling
depended on — then makes that contract impossible to violate silently.

Companion to AICrossSim/PLENA_Simulator#97,
which carries the emulator half and the submodule pin. Neither is useful without
the other.

Not ready to merge. The pin in #97 climbs through this branch in six steps;
merging either side alone leaves the other pointing at a commit that is no
longer on a branch. Land them together, this one first.

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. Reading top to bottom gives
the argument rather than the chronology: what the ISA needed, what the attribution
contract is, what shared experts are, how the contract is enforced, what enforcing
it turned up, and what runs it.

Group	Commits	What it establishes
A — ISA: routing at arbitrary expert shapes	1	C_SET_TOPK_REG, so a new expert shape is not an ISA revision
B — Stage attribution becomes a declared contract	1	@stage= markers and MOE_STAGES replace substring-matched prose
C — Shared experts	3	The emitters, the fusion identity, the caveat consumers need
D — The contract is enforced, not just documented	5	stage becomes required; a lint that covers every route around it
E — Attribution corrections found by review	1	residual_setup, and the Qwen router GEMMs that had no marker
F — CI: the guard actually runs	1	The job, and the trigger a stacked branch needs

Groups C and E depend on B; D depends on C; the simulator's pin ladder steps
through A→C→D→E→F in that order. That ordering is the reason the two PRs can be
read side by side at any commit rather than only at the tip.

---

Group A — C_SET_TOPK_REG: routing at any expert shape

V_TOPK's rmask was a two-entry table: 0 = 32 experts/top-4, 1 = 128/top-8.
That covers GPT-OSS and Qwen3-30B-A3B and nothing else — every other production
MoE falls outside it:

model	experts	top_k
Llama-4 Scout	16	1
Qwen2-MoE	60	4
DeepSeek-V2-Lite	64	6
DeepSeek-V3 / Kimi K2	256	8

Each would have cost an ISA revision plus a matching emulator change. New opcode
C_SET_TOPK_REG (6'h38) is a sticky control register in the same family as
C_SET_SCALE_REG / C_SET_STRIDE_REG / C_SET_V_MASK_REG, holding
(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.

The 8-bit shift is chosen so every shape up to 16383 experts packs into a single
22-bit S_ADDI_INT immediate — no S_LUI_INT pair. It bounds top_k at 255,
which no published MoE approaches.

C_SET_TOPK_REG is an rd-only form, so parse_asm_file already handles it and
only the encoder set needed an entry.

Group B — explicit @stage= markers

Stage attribution was an undeclared cross-repo contract: the emulator
substring-matched prose like "GPT-OSS gather token rows" out of emitted
comments. Rewording silently reclassified instructions, and there was no way at
all to introduce a stage the substring table did not already know — which is
exactly what made shared experts unmeasurable.

Emitters now emit ; @stage=<name>, validated against MOE_STAGES at ASM-gen
time so a typo fails immediately instead of collapsing a region into other.

Markers are authoritative and sticky, so stage became a parameter wherever
one emitter serves several phases:

• moe_true_zero_vram_rows_v0 zeroes combine accumulators, gather padding and
  route-weight tiles — the emulator previously guessed between them from the
  preceding comment's stage, a stateful rule that could not express "this clear
  belongs to the shared branch" at all.
• moe_materialize_route_weights_for_active_rows_v0 serves both route weights and
  the shared gate. Without the parameter every gate instruction bills to
  expert_route_weight, making a program with no routing appear to spend time
  computing route weights.

MOE_STAGES names the three shared_expert_* stages one commit before the
emitters that use them arrive. The vocabulary is the contract, and the emulator
holds its StageKind enum equal to it in both directions, so declaring a
stage and emitting it have to be one step or the guard fails from both sides.

This group also completes the moe_* API migration: the simulator testbench had
already been written against a generalized moe_* API — three of its files call
eight methods that did not exist in this repo at all, plus a policy_name
parameter no method accepted. Every old name survives as an alias
(_DEPRECATED_METHOD_ALIASES), so in-flight callers and the GPT-OSS bring-up
tests that predate the rename keep working unchanged. moe_router_select_v0 is
new: it generalizes the top-k emitter to arbitrary (num_experts, top_k) via
Group A's escape, preferring the fixed rmask table where it applies.

Group C — 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 Scout/Maverick	1	== routed	none
GLM-4.5, Qwen3-Next	1	independent	none
GPT-OSS, Qwen3-MoE, Mixtral	0	—	—

New emitters:

• moe_shared_expert_v0 — dense FFN over all rows with static weights
• moe_shared_gate_v0 — Qwen2-MoE's sigmoid(x @ w_gate)
• moe_combine_shared_and_routed_v0 — the unweighted add
• fused_shared_intermediate() — DeepSeek n_shared → fused MLP width

n_shared_experts is deliberately not an emitter parameter. DeepSeek stores
its shared experts pre-concatenated and instantiates a single DeepseekV2MLP
of width moe_intermediate_size × n_shared_experts. That is a checkpoint-loading
concern, not an emission one.

The fusion identity, stated precisely. Because SwiGLU is elementwise along the
intermediate axis while the down projection sums over it, the fused MLP equals the
sum of the individual shared experts in exact real arithmetic — an algebraic
identity, not an approximation. As emitted it is not bit-identical, for two
reasons unrelated to the algebra:

• the down projection reduces over the concatenated axis in one pass rather than
  summing n_shared separate reductions, and FP addition is not associative;
• expert weights are MXFP8 (e4m3, one e8m0 scale per block of 8 along that axis),
  so when moe_intermediate_size is not a multiple of the block, the fused
  tensor's blocks straddle expert boundaries and pick up scales derived from two
  different experts — the quantised weights themselves differ, not merely their sum.

Neither makes the fusion wrong. They are stated because #97's shared-expert test
is bit-exact at atol=rtol=0, so a test written against "exactly equal" would
fail for entirely correct reasons and send someone hunting a bug that is not there.

SHARED_VS_ROUTED_NOTE is declared here as canonical text. The warning that a
shared-vs-routed cycle ratio measures this compiler's lowering rather than MoE
hardware previously existed as prose in this module and, in different words, as
a hardcoded string in the emulator that ships it in the profile JSON — with
nothing checking they agreed. The wording of a caveat is the whole of its content.
The emulator keeps its own copy (it writes that JSON with no Python available) and
a test on that side asserts byte equality against this constant.

The gate reuses moe_materialize_route_weights_for_active_rows_v0: one FP scalar
per token broadcast across hidden is structurally identical to a V_TOPK route
weight, only the scalar's origin differs. The sigmoid itself runs in the scalar FP
unit (S_SUB_FP / S_EXP_FP / S_ADD_FP / S_RECI_FP) — there is one value per
token, and the vector form would compute MLEN copies of it.

Group D — stage becomes required, and stays required

Group B gave the reused emitters a stage parameter. This group makes it
impossible to not supply one.

Markers are sticky, so an emitter called from inside a marked region inherits the
enclosing marker: a default for that parameter is a silent wrong answer rather
than a missing one. This is not hypothetical — the shared-expert sigmoid gate
inherited expert_route_weight and misattributed 999 instructions while every
total still added up and every test stayed green.

The default is dropped from moe_materialize_route_weights_for_active_rows_v0,
moe_true_zero_vram_rows_v0 and moe_expert_activation_v0. This is a breaking
signature change; the six simulator testbench call sites it turns into hard
errors are fixed in #97, in the commit that moves the pin onto it.

test_moe_stage_attribution.py enforces the rule, and the four commits after it
are each about a way the lint was not yet doing its job:

commit	the gap it closed
self-check through the helpers	the self-check reimplemented the walk inline and had already drifted — it selected ast.FunctionDef, and ast.AsyncFunctionDef is a sibling class, not a subclass, so async def emitters were outside it entirely
scope the lint to MoE callees	stage is not a reserved word; qkt_multiply(stage="decode") is correct attention code that the unscoped lint reported as an unknown MoE stage. A lint that fires on correct code gets suppressed, and then protects nothing
match positional stage names	_moe_stage_arguments inspected node.keywords only, and all 19 moe_stage_marker(...) call sites pass the stage positionally — the lint was scanning three keyword arguments and none of the markers. Coverage went from 3 stage names to 25
close indirect routes	lambdas, functools.partial/partialmethod bindings and decorators handed stage= each supply the argument without any call site naming it

The third row is the one worth pausing on: for several commits this branch shipped
a lint that provably could not catch the defect it was written for.

A decorator that injects a stage without naming it in its own call cannot be
detected without resolving what the decorator does. That limit is stated at the
helper, with a TODO saying the answer is an explicit allowlist rather than deeper
analysis.

Group E — attribution corrections found by review

Two holes on the same stretch of the decoder MoE sublayer. Neither was findable
before Groups B and D: the lint checks that a stage name is declared, not that
the region it labels is the region it describes.

• residual_setup joins MOE_STAGES. MoE input preparation — the residual
  buffer zero, the residual copy, the input RMSNorm and its norm-weight multiply —
  was being labelled accumulator_init, which is the combine accumulator that
  expert outputs are scattered into. Different thing, and the cost scales with
  rows × hidden rather than with the accumulator.
• The two qwen3_router_logits_* emitters gain the router_topk marker. They
  had none, and both lower through general-purpose projection helpers that have
  none either, so the entire router GEMM was billed to whatever stage preceded the
  call. moe_router_logits_bf16_v0 already marked itself; these now match. This is
  part of the same fix rather than a separate one: relabelling the residual zero
  without it would have moved the Qwen router GEMM into residual_setup — a new
  wrong attribution instead of the old one.

Markers are comments (IsaBuilder().comment(...)), so no instruction count, cycle
total or numerical result changes — only which stage the profile bills.

Group F — CI

Before this, ci.yml ran the generator tests and a codegen smoke and nothing
under aten/tests at all — so the guard could not have failed a pull request,
which is indistinguishable from not having written it. The moe-stage-guard job
needs neither torch nor a checkpoint (the lint parses source with ast), so it is
pytest + pyyaml and a couple of seconds.

Naming one file in a workflow step recreates the same hole one file over, so
test_every_test_file_here_is_wired_into_ci asserts every aten/tests/test_*.py
is either named in ci.yml or listed in _UNWIRED_TESTS. The five files on that
list all import torch; pinning them is what stops a new unwired file hiding among
them, and the list is checked in both directions.

workflow_dispatch is kept for a different reason: pull_request is filtered to
branches: [main], so a branch stacked on a parent branch fires no jobs at all.
This branch targets main directly and no longer needs it — the next stacked branch
will.

---

Validation

pytest aten/tests/test_moe_stage_attribution.py — 7/7, on pytest + pyyaml with no
torch installed. Verified per commit across the branch: 0 syntax failures at all
12, and the guard passes at every commit where it exists (3 → 4 → 6 → 7 tests as
each gap closes).

Each guard was also verified to fail correctly: an injected defaulted-stage
lambda, a partial binding and a partialmethod binding are each reported with
file and line; a stage="acumulator_init" typo on a MoE callee is reported while
qkt_multiply(stage="decode") is not; an injected moe_stage_marker("expert_bais", ...)
is reported at program_routed_moe.py:921 (and was not reported before the
positional-matching commit); an unwired test file and a stale _UNWIRED_TESTS
entry each fail by name.

From the simulator side (see #97): three shared-expert configurations bit-exact at
atol=rtol=0; six routing shapes end to end, including DeepSeek-V3's 256 experts
across four MLEN-wide logit blocks selecting expert 255; all six pre-existing
routed-MoE CI tests still pass, including gpt_oss_topk_test which now exercises
the alias path; 87 emulator unit tests.

Known limitation, not introduced here

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, its matrix-op count
is strictly linear in tokens × top_k and matrix-engine row utilisation is fixed
at 1/64 regardless of token count. The shared branch added here is properly
batched (3 M_MM for 1 token or 64).

So a shared-vs-routed cycle ratio measured today reflects the routed lowering's
per-pair overhead at least as much as the architecture. That is what
SHARED_VS_ROUTED_NOTE says, and #97 now emits it as
classification.attribution_notes.shared_vs_routed so the caveat reaches the
people reading the numbers rather than only the people reading this source.
Expert-major batching is follow-up work.

Boundary with Simulator #97

• Pair-major routed utilisation is this PR's follow-up, not #97's. Fixing it
  means re-lowering the routed path to token-major / expert-batched, which is
  entirely a compiler concern. #97 explicitly does not touch it.
• The shared gate's emit-time unrolled token loop (moe_shared_gate_v0,
  for token_idx in range(rows)) is a shared-path codegen follow-up tracked on
  the #97 side, not part of the pair-major re-lowering. Static instruction count
  there grows with batch size; it is unrelated to routed placement.
• MOE_STAGES is a cross-repo contract, and it is now a checked one. #97's
  emulator-side guard parses this declaration out of the pinned submodule with
  Python's own ast and compares it against StageKind in both directions. It
  used to compare against a hand-copied 13-entry array — a copy only fails when
  somebody remembers to update it, which is the same failure mode as having no
  guard, and the compiler side is the side that moves. Adding residual_setup
  here would not have failed the old test.
• Scope split, so the two are reviewable independently. This PR owns the
  emitters, the MOE_STAGES vocabulary and the lint over aten/plena; #97 owns
  the emulator's stage profile, the testbenches, the timing harnesses, and the
  same lint pointed at testbench/ across the repo boundary. The submodule pin is
  the only coupling, and it moves in that repository.

Note on branch naming

The branch is still called feat/shared-expert-moe, which now under-describes it —
it carries the ISA change, the attribution contract and its enforcement too.
Renaming would close this PR, so the name stays and this paragraph exists instead.

🤖 Generated with Claude Code