MoE: shared experts, parameterized V_TOPK routing, explicit stage attribution

[qichao-arlo-wang]

What this PR does

Emulator and testbench side of shared-expert MoE support. Companion to
AICrossSim/PLENA_Compiler#69,
which this PR bumps the submodule pin to. Neither is useful without the other.

Stacked on #97. Base is feat/moe-timing, so the diff shows only this
work (13 files, +1419/−27). Against main it would swallow #97's 19 files.
Merge #97 first, or retarget this to main after.

1. 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. 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", which points at the
wrong thing.

classify_timing_access now takes the policy as a second accessor closure. Its
V_TOPK arm derives the logit read extent from the expert count, and the
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.

2. Stage attribution: @stage= markers replace substring sniffing

classify_comment guessed stages by substring-matching the compiler's generated
ASM prose. That is 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. The split
exists so the profile can answer what fraction of MoE time is architecturally
dense (shared, every token) versus routing-dependent.

Two new reported fields, schema version 4:

• classification.stage_attribution — "explicit_stage_markers" or
  "legacy_comment_substrings". Every vocabulary_* field only means anything in
  the latter 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.

3. Tests

moe_shared_expert_test.py — bit-exact, three architectures

variant	recipe
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

All at atol=rtol=0. Exactness comes from MXFP8-representable inputs making
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.

moe_router_policy_test.py — every production routing shape

Six policies through compiler → assembler → emulator, including DeepSeek-V3's 256
experts spanning four MLEN-wide logit blocks and selecting expert 255 (unencodable
before C_SET_TOPK_REG).

Each case asserts which encoding it took, not just that routing was correct.
Without that a shape could quietly fall back to a fixed rmask, produce correct
indices anyway, and leave the escape path untested while appearing covered.

Negative-tested

Each guard against the defect it claims to catch:

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

4. Route-trace schema v2

Adds the optional shared-expert block to model: shared_experts,
shared_intermediate_size, shared_gate. A v1 trace is a valid v2 trace with no
shared branch, so both are accepted rather than forcing every existing trace to be
regenerated.

shared_intermediate_size is the fused width, already multiplied by
n_shared — DeepSeek replays a single MLP, not n of them. The validator rejects
the half-specified combinations: a shared_experts with no width would replay as
routed-only while still being labelled shared, and the timing split would
attribute nothing while looking healthy.

summarize_run gains a shared-vs-routed picosecond split, summed exactly and
rounded once. gather / scatter_combine / accumulator_init are excluded from
both terms: they are combine plumbing shared by the branches, and billing them to
"routed" would overstate its share.

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

CI

New step just test-moe-shared-all — three shared-expert variants plus six
routing policies. Synthetic throughout: no checkpoint, no HF libraries.

Validation

• 3 shared-expert variants bit-exact, 6 routing policies pass
• 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
• ruff format / ruff check clean

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

The routed cost is strictly linear in tokens × top_k with zero batching benefit;
the shared branch added here 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", which was the right call for bring-up but is 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. This is
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
(one token, top_k groups of one row each) — it is a prefill / large-batch-decode
concern.

🤖 Generated with Claude Code