MoE: shared experts, arbitrary V_TOPK expert shapes, explicit stage markers

[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.

Companion to AICrossSim/PLENA_Simulator#100, which bumps the submodule pin.
Neither is useful without the other; land this first.

1. 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.

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	—	—

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 — and 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.

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.

3. 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. This was caught by the new test's stage-distribution
  guard during bring-up, at 999 misattributed instructions.

4. 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. This completes that
migration.

moe_router_select_v0 is new: it generalizes the top-k emitter to arbitrary
(num_experts, top_k) via C_SET_TOPK_REG, preferring the fixed rmask table
where it applies.

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.

Validation

Verified from the simulator side (see the companion PR):

• Three shared-expert configurations pass bit-exactly at atol=rtol=0
• Six routing shapes pass 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 it
reflects the architecture. This is documented in program_moe_shared.py's module
docstring. Expert-major batching is follow-up work.

Boundary with Simulator PR #97

PLENA_Simulator #97 pins
this branch and carries the emulator half (stage attribution, timing-replay
harnesses). The split, so neither side waits on the other:

• Pair-major routed utilisation (the fixed 1/64 ≈ 1.56% row utilisation noted
  above) 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.
• Because of that, the shared-vs-routed cycle ratio is not an architectural
  result yet. #97 now emits this caveat in the profile JSON itself as
  classification.attribution_notes.shared_vs_routed, so a consumer reading only
  the profile still sees it — previously it lived solely in
  program_moe_shared.py's module docstring, i.e. nowhere a numbers-consumer looks.
• 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 in program_routed_moe.py is a cross-repo contract. #97's
  emulator-side guard now parses this set directly out of the pinned submodule, so
  adding or renaming a stage here without a matching StageKind fails that test.
  Previously it compared against a hand-copied list, which could not fail in the
  direction that actually moves.

Stacked follow-up

feat/moe-stage-required-arg (stacked on this branch) makes stage a required
argument on the three stage-polymorphic emitters —
moe_materialize_route_weights_for_active_rows_v0, moe_true_zero_vram_rows_v0
and moe_expert_activation_v0. The shared-gate misattribution fixed in this PR
came from inheriting a stage default; removing the defaults retires the bug class
rather than the single instance. Note this is a breaking signature change — the
matching simulator call-site updates ride along with the pin bump in #97.

🤖 Generated with Claude Code