MoE: required stage argument, a lint that runs, and two attribution fixes - #70
Closed
qichao-arlo-wang wants to merge 11 commits into
Closed
MoE: required stage argument, a lint that runs, and two attribution fixes#70qichao-arlo-wang wants to merge 11 commits into
stage argument, a lint that runs, and two attribution fixes#70qichao-arlo-wang wants to merge 11 commits into
Conversation
V_TOPK's rmask was a two-entry table: 0 = 32 experts/top-4 (GPT-OSS), 1 = 128/top-8 (Qwen3-30B-A3B). Every other production MoE shape falls outside it -- Qwen2-MoE 60/top-4, DeepSeek-V2-Lite 64/top-6, DeepSeek-V3 and Kimi K2 256/top-8, Llama-4 Scout 16/top-1 -- so each new architecture would have cost an ISA revision plus a matching emulator change. Add C_SET_TOPK_REG (6'h38), 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. The 8-bit shift keeps the packed value inside a single 22-bit S_ADDI_INT immediate for every shape up to 16383 experts, so no S_LUI_INT pair is needed. It bounds top_k at 255, which no published MoE approaches. The register is a rd-only form, so parse_asm_file already handles it and only the encoder set needed the entry. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Adds the shared-expert branch of MoE -- an FFN every token passes through,
summed into the routed output unweighted:
y = shared(x) + sum_k route_weight_k * routed_expert_k(x)
DeepSeek-V2/V3, Kimi K2, Llama-4, GLM-4.5 and Qwen3-Next all have one;
GPT-OSS, Qwen3-MoE and Mixtral do not, which is why the substrate had no
notion of it.
moe_shared_expert_v0 dense FFN over all rows, 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 MLP of width
moe_intermediate_size * n_shared_experts. That is a checkpoint-loading
concern, and SwiGLU being elementwise along the intermediate axis makes the
fused form exactly equal to the sum -- 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, since there is one value per token and the
vector form would compute MLEN copies of it.
Explicit @stage= markers
------------------------
Stage attribution was an undeclared cross-repo contract: the emulator
substring-matched prose like "GPT-OSS gather token rows" out of the 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 rather than 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
accumulators, gather padding and route tiles; the row-weight broadcast serves
both route weights and the shared gate).
moe_* API migration
-------------------
Renames the gpt_oss_*_v0 methods the simulator testbench had already been
written against -- three of its files called eight methods that did not
exist in this repo at all -- and adds the policy_name parameter they pass.
moe_router_select_v0 is new: it generalizes the top-k emitter to arbitrary
(num_experts, top_k) via C_SET_TOPK_REG, keeping the two hardwired rmask
policies where they apply so existing programs emit byte-identical ASM.
Every old name survives as an alias, so nothing breaks mid-flight.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…emitters Marker attribution is sticky: a helper called from inside a marked region inherits the enclosing marker, so an emitter reused across stages must be told which one it serves. When that parameter has a default, forgetting it is not a missing answer but a silently wrong one -- instruction counts, cycle totals and numerics are all unchanged, so nothing downstream fails. That is not hypothetical. The shared-expert sigmoid gate reused `moe_materialize_route_weights_for_active_rows_v0`, inherited its `expert_route_weight` default, and misattributed 999 instructions with the whole suite green. Adding the `stage` parameter fixed that call site; it left the mechanism that produced it in place. Drop the default from all three stage-polymorphic emitters: moe_materialize_route_weights_for_active_rows_v0 (was expert_route_weight) moe_true_zero_vram_rows_v0 (was accumulator_init) moe_expert_activation_v0 (was expert_activation) `moe_true_zero_vram_rows_v0` is the one that mattered most: it serves the most stages and had the most callers, so its default silently absorbed every caller that never thought about attribution. The scope rule is "takes a `stage` parameter", not "is reused across >= 2 stages", because the two coincide by construction -- an emitter serving exactly one stage hardcodes its marker and never takes the parameter. One internal caller (moe_dynamic_expert_pair_v0 -> moe_expert_activation_v0) was relying on the default and now names `expert_activation` explicitly. All 7 compiler call sites pass `stage`; behaviour is unchanged. test_moe_stage_attribution.py locks the rule in: no defaulted `stage` parameters, every literal `stage=` argument is a declared MOE_STAGES member, and a self-check proving the AST walk can still fail -- without which a broken scanner would pass over zero findings. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The lint landed in 209e299 with no workflow invoking it. `ci.yml` ran the generator parser tests and a codegen smoke, and nothing under `aten/tests` at all -- so the guard against a defaulted `stage` parameter could not have failed a pull request. A test nobody runs is indistinguishable from a test nobody wrote. Add a `moe-stage-guard` job. It needs neither torch nor a checkpoint, because the guard reads sources with `ast` rather than importing them; pyyaml is the one dependency, and only because collecting anything under `aten/tests` imports `aten/__init__.py`, which pulls the op registry in. That keeps the job at pytest + pyyaml and a couple of seconds, so it does not sit behind `needs: syntax-check`. Also add `test_every_test_file_here_is_wired_into_ci`, because naming a single file in a workflow step reintroduces the same hole one directory over: the next test file added here would again run nowhere. It 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 and some want a real checkpoint, so wiring them is separate work -- pinning them is what stops a new unwired file from hiding among them. The list is checked in both directions, so an entry that gets wired up or deleted has to leave it. Verified: the guard passes at 4/4 on pytest + pyyaml with no torch present; dropping an unwired file into the directory fails it by name; and a stale `_UNWIRED_TESTS` entry fails it too. Addresses Agent 2 CRIT#1. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`test_guard_would_catch_a_defaulted_stage_parameter` exists to prove the lint can fail. It did that by reimplementing the argument walk inline, which proves only that the copy agrees with itself -- and the copy had already drifted: it selected functions with `isinstance(node, ast.FunctionDef)`, while `ast.AsyncFunctionDef` is a sibling class rather than a subclass. So `async def` emitters were outside the self-check entirely. The real lint iterates `_functions`, which handles both, so the two had silently diverged in exactly the direction this test is supposed to rule out. Extract `_defaulted_stage_params(func)` -- the kwonly/positional/posonly default alignment, in one place -- and call it from both the lint and the self-check. The self-check now writes its fixture to `tmp_path` and iterates `_functions`, so it exercises the same read-parse-select path as the lint instead of a parallel one. The fixture grows an `async def` offender and an `async def` clean case, so the drift above would now be caught. Assert on sorted names rather than walk order, which is an implementation detail of `ast.walk`. Verified: 4/4 pass; the shared helper flags defaulted `stage` for async kwonly, async positional and sync posonly parameters, and leaves a non-defaulted `stage` and a stage-less function clean. Addresses Agent 2 CRIT#2. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`test_stage_arguments_are_declared_moe_stages` matched every `stage=` keyword argument in the tree and checked the value against MOE_STAGES. `stage` is not a reserved word: the attention path calls `qkt_multiply(stage="decode")`, where it selects prefill vs decode. Within `aten/plena` today every `stage=` callee happens to be a `moe_*` emitter, so the lint is green by luck -- widen the scan by one directory and it reports correct attention code as passing an unknown MoE stage. A lint that fires on correct code gets suppressed, and then it protects nothing. Match on the callee instead: `_is_moe_callee` accepts the `moe_` and `gpt_oss_` prefixes (`moe_stage_marker` falls under `moe_`), and `_moe_stage_arguments` yields `(callee, keyword)` pairs so the scan is one helper rather than a nested walk inlined into the test. `_callee_name` resolves `f(...)` and `obj.f(...)` alike. Add `test_non_moe_stage_arguments_are_not_flagged` over a fixture holding `qkt_multiply(stage="decode")`, `attention_softmax(stage="attn_input")` and `builder.flash_attention(stage="prefill")` next to one real MoE call, asserting only the MoE call is matched. Offender lines now name the callee and point at the keyword rather than the call, which for these multi-line emitter calls is a different line. Verified: 5/5 pass; a deliberate `stage="acumulator_init"` typo on `moe_true_zero_vram_rows_v0` is still reported, and `qkt_multiply` is not. Addresses Agent 2 SF S2. Landed before the simulator-side scan because that scan reaches `flashattn_qkt_test.py`, where the false positive is real. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Two attribution holes on the same stretch of the decoder MoE sublayer. **No stage describes MoE input preparation.** Between the last attention marker and the router, a program zeroes the residual buffer, copies the post-attention hidden state into it, and runs the MoE input RMSNorm with its norm-weight multiply. Testbenches were labelling the residual zero `accumulator_init`, which is a different thing: `accumulator_init` is the combine accumulator that expert outputs are scattered into, while these rows are a residual copy added back *after* the experts run. The cost scales with rows x hidden, and none of it is accumulator initialisation. None of the other twelve stages fits either, so add `residual_setup` rather than keep folding it into a stage it contradicts. **The two Qwen router emitters emit no marker at all.** `moe_router_logits_bf16_v0` marks itself `router_topk`; `qwen3_router_logits_matrix_bf16_rowpacked_v0` and `qwen3_router_logits_packed_skinny_bf16_rowpacked_v0` do not, and both lower through `linear_projection_bf16*` / the packed-skinny helper, which are general-purpose and carry no marker either. Markers being sticky, the entire router GEMM was billed to whatever stage preceded the call. That makes this part of the same fix rather than a separate one: relabelling the residual zero without marking these would move the Qwen router GEMM from `accumulator_init` into `residual_setup` -- a new wrong attribution instead of the old one. Both paths now mark themselves like the vector-dot path. Markers are comments (`IsaBuilder().comment(...)`), so no instruction count, cycle total or numerical result changes; only which stage the profile bills the region to. The emulator side of `residual_setup` -- the `StageKind` variant and the testbench call sites that use it -- lands with the pin bump in PLENA_Simulator, since `stage_marker_names_match_the_compiler_vocabulary` requires the two vocabularies to move together. Verified: the module parses; the stage extractor recovers 14 stages including `residual_setup`; all four router emitters now emit `router_topk`; the attribution guard passes 5/5. Addresses Agent 3 SF S3 (the misattribution) and a marker gap found while fixing it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…efault `test_stage_parameters_have_no_default` walks `def` and `async def` signatures. Making `stage` required is only enforced to the extent the lint covers every way of supplying one, and three constructs are outside it: - a lambda, which carries the same `arguments` node with defaults and none of the coverage; - `functools.partial(emit, stage=...)` and `partialmethod`, which pre-bind the argument at the binding site so every call through the result omits it; - a decorator handed `stage=`, which can inject it into the wrapped call. Each leaves call sites that never name a stage, which is precisely the condition the required parameter exists to prevent -- reached by a route the lint was not looking at. None is hypothetical enough to ignore: `partial` is the ordinary way to specialise an emitter, and `_DEPRECATED_METHOD_ALIASES` already rebinds emitters onto the mixin. `_stage_defaulting_lambdas` reuses `_defaulted_stage_params` unchanged, since `ast.Lambda` and `ast.FunctionDef` share the `arguments` node. The decorator check catches the declared shape, `@with_stage(stage="gather")`. A decorator that injects a stage without naming it in its own call -- from a closure, a registry, an attribute -- cannot be detected without resolving what the decorator does, which is past what a source lint can do. That limit is stated at the helper with a TODO saying the answer would be an explicit allowlist of stage-injecting decorators, not deeper analysis. `test_guard_would_catch_an_indirectly_supplied_stage` covers all three against a fixture, with clean counterparts for each. Verified by injecting each construct into `program_routed_moe.py` and running the lint: a defaulted-stage lambda, a `partial` binding and a `partialmethod` binding are each reported with file and line; a lambda whose `stage` has no default, and a `partial` binding something other than `stage`, are not. 7/7 pass with the file restored. Addresses Agent 2's NIT on AST bypasses. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…two places The warning that a shared-vs-routed cycle ratio measures this compiler's lowering rather than MoE hardware exists twice: as a paragraph in this module's docstring, and as `SHARED_VS_ROUTED_NOTE` in the emulator's `stage_profile.rs`, which puts it in the profile JSON. The two say the same thing in different words, with nothing checking that they keep doing so. That is two failure modes at once. A caveat living only in compiler source is a caveat nobody applies, because the JSON is what consumers read -- which is why the emulator restated it. A caveat restated by hand across two repositories is one that quietly stops matching, and the wording of a caveat is the whole of its content. Declare `SHARED_VS_ROUTED_NOTE` here as the canonical text. The docstring keeps the module-specific detail -- the per-pair gather and the "intentionally wasteful" quote -- and defers the consequence to the constant instead of restating it. The emulator keeps its own `const` rather than reading this at runtime: the profile is written by the emulator with no Python available, and shelling out to produce a JSON field would be worse than a checked copy. A test on that side parses this constant and asserts the two are byte-identical, so the copy cannot drift silently. It lands with the pin bump. Verified: the module parses; the two strings are byte-identical at 432 characters; the attribution guard passes 7/7. Addresses Agent 1 NIT#6. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`fused_shared_intermediate` said the fused MLP "is exactly equal to the sum of the individual shared experts -- no approximation". The algebra is right: SwiGLU is elementwise along the intermediate axis and the down projection sums over it, so concatenating the shared experts and running one MLP is an identity. But "exactly equal" reads as a claim about the emitted program, and the emitted program is not bit-identical, for two reasons that have nothing to do with the algebra: - The down projection reduces over the whole concatenated axis in one pass, where the unfused form does `n_shared_experts` reductions and adds them. Floating-point addition is not associative, so the roundings differ before quantisation enters into it. - Expert weights are stored MXFP8 -- e4m3 elements with one e8m0 scale per block of 8 along that axis (`load_config.rs:717`). When `moe_intermediate_size` is not a multiple of that block, the fused tensor's blocks straddle expert boundaries, so a block's scale is derived from values belonging to two different experts and the quantised weights themselves differ, not merely their sum. Both are small and neither makes the fusion wrong. The reason to state them is that a bit-exactness test written against the old wording would fail for entirely correct reasons, and whoever wrote it would go looking for a bug that is not there. Addresses Agent 3 N1. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This was referenced Jul 30, 2026
`test_stage_arguments_are_declared_moe_stages` justifies itself with "a typo on a rarely-exercised branch would reach the emulator, land in `unresolved_stage_markers`, and leave that region inheriting the previous marker". It could not have caught one. `_moe_stage_arguments` inspected `node.keywords` only, and all 19 `moe_stage_marker(...)` call sites under `aten/plena` pass the stage **positionally** — so the one construct that actually emits a marker was entirely outside the lint. It was scanning three `stage=` keyword arguments and none of the markers. `moe_stage_marker` validates its argument at runtime, but only on paths a test executes, which is exactly the gap this lint is supposed to cover. Match the first positional argument for callees in `_POSITIONAL_STAGE_CALLEES` as well as `stage=` keywords. Coverage across `aten/plena` goes from 3 stage names to 25. `_moe_stage_arguments` now yields the `ast.Constant` rather than the keyword node, since the two spellings have no common wrapper. `test_non_moe_stage_arguments_are_not_flagged` grows a positional marker in its fixture and a standalone assertion that a positional typo is visible, so the blind spot cannot come back silently. It compares sorted results: `ast.walk` is breadth-first and its order is not a property worth pinning. Also adds `workflow_dispatch` to `ci.yml`. `pull_request` is filtered to `branches: [main]`, so a stacked PR — which targets its parent branch — fires no jobs at all, and the guard this file added in 74ec261 could not run on the PR that adds it. Manual dispatch is the smallest fix that makes a stacked PR verifiable before it retargets main. Verified: 7/7 pass. Injecting `moe_stage_marker("expert_bais", ...)` into `program_routed_moe.py` is now reported as `program_routed_moe.py:921`; before this commit it was not reported at all. Addresses review findings on lint coverage and CI triggering. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
qichao-arlo-wang
force-pushed
the
feat/shared-expert-moe
branch
from
July 31, 2026 00:56
311561e to
60239b0
Compare
Collaborator
Author
|
Closing: everything on this branch is now in #69, which targets The two branches have been merged into one history and reorganized by intent rather than by review round, so this PR's diff no longer means anything — its base ( What changed in the process, beyond the merge:
The emulator half, PLENA_Simulator#101, is closed into PLENA_Simulator#97 the same way. Nothing is lost: the pre-reorganization history is preserved at |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this PR does
Stacked review hardening for the MoE stage-attribution contract. Makes
stagearequired argument on the stage-polymorphic emitters, adds a lint that enforces it
(and the routes around it), and closes two attribution gaps found while doing so.
1.
stagebecomes required — and stays required209e299drops the default frommoe_materialize_route_weights_for_active_rows_v0,moe_true_zero_vram_rows_v0andmoe_expert_activation_v0. Markers are sticky, soan 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_weightand misattributed 999 instructions while every total stilladded up and every test stayed green.
A defaulted
stageis only one of several ways to supply one without the callernaming it.
7b9d46ealso rejects:argumentsnode, none of the coverage;functools.partial/partialmethodbindings that pre-bindstage;stage=, e.g.@with_stage(stage="gather").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.
2. The lint now runs
74ec261adds amoe-stage-guardjob. Before it,ci.ymlran the generator testsand a codegen smoke and nothing under
aten/testsat all — so the guard couldnot have failed a pull request, which is indistinguishable from not having written
it. The 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_ciasserts everyaten/tests/test_*.pyis either named in
ci.ymlor listed in_UNWIRED_TESTS. The five files on thatlist all import torch; pinning them is what stops a new unwired file hiding among
them, and the list is checked in both directions.
3. Attribution fixes
residual_setupjoinsMOE_STAGES(7565629). MoE input preparation — theresidual buffer zero, the residual copy, the input RMSNorm and its norm-weight
multiply — was being labelled
accumulator_init, which is the combineaccumulator that expert outputs are scattered into. Different thing, and the cost
scales with
rows × hiddenrather than with the accumulator.qwen3_router_logits_*emitters gain therouter_topkmarker. Theyhad 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_v0already marked itself; these now match. This ispart 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 newwrong attribution instead of the old one.
Markers are comments (
IsaBuilder().comment(...)), so no instruction count, cycletotal or numerical result changes — only which stage the profile bills.
4. Caveat de-duplication
SHARED_VS_ROUTED_NOTEbecomes a module constant inprogram_moe_shared.py(
916f990). The warning that a shared-vs-routed ratio measures this compiler'slowering rather than MoE hardware previously existed as prose here 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.
a387f98also qualifiesfused_shared_intermediate's "exactly equal" as holdingin exact real arithmetic: as emitted, FP associativity and MXFP8 block alignment
(e4m3, one e8m0 scale per block of 8) both make it non-bit-identical. Small, and it
does not make the fusion wrong — but a bit-exactness test written against the old
wording would have failed for entirely correct reasons.
Validation
pytest aten/tests/test_moe_stage_attribution.py— 7/7, on pytest + pyyaml with notorch installed. Each guard was also verified to fail correctly: an injected
defaulted-stage lambda, a
partialbinding and apartialmethodbinding are eachreported with file and line; a
stage="acumulator_init"typo on a MoE callee isreported while
qkt_multiply(stage="decode")is not; an unwired test file and astale
_UNWIRED_TESTSentry each fail by name.🤖 Generated with Claude Code