From 8a6b94b9badd87c8b5dad8614b4a5c0a3bd253f0 Mon Sep 17 00:00:00 2001 From: "Qichao (Arlo) Wang" Date: Sun, 9 Aug 2026 13:07:17 +0100 Subject: [PATCH] feat(moe): a terminator for the sticky stage marker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Markers are sticky, and a program does not stop where its MoE region does. The last MoE marker therefore ran to the end of the file, and every instruction after it -- an lm_head, the next sublayer -- was billed to a MoE stage, feeding the shared-vs-routed ratio the profile exists to report. This is a hazard the marker mechanism introduced: under the legacy substring rules the epilogue fell back to `other`, a conservative miss. Turning markers on converted it into a confident wrong answer. `non_moe` joins MOE_STAGES and `moe_end_marker()` emits it -- a separate entry point from `moe_stage_marker(MOE_END_STAGE)` because closing a region and setting one are different acts, and the caller who has to remember it is assembling a decoder program, not writing an emitter. The name is spelled literally inside MOE_STAGES rather than as `MOE_END_STAGE`. Both repositories recover that set with a parser that refuses anything it cannot evaluate, so a name reference fails to parse rather than resolving. A test holds the constant and the entry equal instead. Also documents two limits on `moe_shared_gate_v0` that were only discoverable by hitting them: `name` must be unique within a program, and `rows` is bounded near 297 by FPRAM -- `one` and `neg_one` must each be at least `rows` long alongside the per-token gate scalars, so the ceiling is near `rows * 3`. Cannot land alone: the emulator holds `StageKind` equal to MOE_STAGES in both directions, and fails with "compiler emits @stage=non_moe but no StageKind matches it" until the companion adds the variant. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 5 --- .github/workflows/ci.yml | 2 + aten/plena/program_moe_shared.py | 19 ++++++++ aten/plena/program_routed_moe.py | 41 ++++++++++++++++ aten/tests/test_moe_stage_terminator.py | 63 +++++++++++++++++++++++++ 4 files changed, 125 insertions(+) create mode 100644 aten/tests/test_moe_stage_terminator.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 809afc1..3732abc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -55,6 +55,8 @@ jobs: - name: Run the top-k policy encoding guard run: PYTHONPATH=. python3 -m pytest aten/tests/test_moe_topk_policy_encoding.py -v + - name: Run the stage-terminator guard + run: PYTHONPATH=. python3 -m pytest aten/tests/test_moe_stage_terminator.py -v unit-tests: name: Generator unit tests diff --git a/aten/plena/program_moe_shared.py b/aten/plena/program_moe_shared.py index 2c8421f..f441b9c 100644 --- a/aten/plena/program_moe_shared.py +++ b/aten/plena/program_moe_shared.py @@ -127,6 +127,25 @@ def moe_shared_gate_v0( broadcast across the hidden axis, ready to multiply the shared expert output. That is the same shape and mechanism the routed branch uses for route weights, so the two are combined identically downstream. + + Two limits worth knowing before you reach them: + + - ``name`` must be unique within a program. The scratch buffers are + derived from it, so calling this twice with the default is a collision, + not a second gate. + - ``rows`` is bounded by FPRAM. One gate scalar per token, plus the + ``one`` and ``neg_one`` constants which must themselves be at least + ``rows`` long, out of 1024 f16 slots shared with everything else -- so + the ceiling is near ``rows * 3``, not ``rows``. Measured at 297 with a + caller that sizes every other constant to 1; a realistically-shaped + caller runs out sooner. + + Static instruction count grows with ``rows`` -- the token loop is unrolled + at emit time -- and the per-token cost grows with ``hidden`` as well, but + that second part is the broadcast helper's, not this loop's: measured 27 + instructions per token at ``hidden=64`` against 182 at ``hidden=2048``, + while the dot-product loop itself is one ``C_LOOP`` regardless. A codegen + follow-up, not a correctness limit. """ if hidden % self.mlen != 0: raise ValueError(f"{name}: hidden={hidden} must be divisible by MLEN={self.mlen}") diff --git a/aten/plena/program_routed_moe.py b/aten/plena/program_routed_moe.py index f5c226b..9720c6c 100644 --- a/aten/plena/program_routed_moe.py +++ b/aten/plena/program_routed_moe.py @@ -22,6 +22,12 @@ #: Every stage name the emulator's ``StageKind`` understands. #: +#: Stage name that closes the MoE region. The emulator resolves it to a distinct +#: ``StageKind`` of its own -- not the unclassified fallback -- so an epilogue is +#: told apart from a region the classifier had no opinion about. Declared as a +#: constant because both repositories key on the exact string. +MOE_END_STAGE = "non_moe" + #: Marker emission is validated against this set, so a typo fails at ASM-gen time #: instead of quietly collapsing a region into ``other``. MOE_STAGES: frozenset[str] = frozenset( @@ -45,6 +51,18 @@ "shared_expert_projection", "shared_expert_activation", "shared_expert_gate", + # Terminator. Markers are sticky and there was no way to say "the MoE + # region is over", so every instruction after the last MoE marker -- + # the lm_head, the next sublayer, anything at all -- kept that marker + # to the end of the program, and that cost lands in the shared-vs-routed + # ratio. A marker like any other, so it is emitted where the region ends + # rather than inferred. + # + # Spelled literally, not as `MOE_END_STAGE`: both repositories recover + # this set with a parser that refuses anything it cannot evaluate, so a + # name reference here fails to parse rather than resolving. The constant + # and this entry are held equal by a test instead. + "non_moe", } ) @@ -134,12 +152,33 @@ def moe_stage_marker(stage: str, detail: str = "") -> str: ``expert_weight_prefetch``. That is why the routed path marks its own dynamic prefetch explicitly; the shared path deliberately does not, folding weight traffic into ``shared_expert_projection`` where it belongs. + + Because markers are sticky and a program does not end where its MoE region + does, the last MoE marker otherwise runs to the end of the file. Emit + :func:`moe_end_marker` at the point the region closes; see + :data:`MOE_END_STAGE`. """ if stage not in MOE_STAGES: raise ValueError(f"unknown MoE stage {stage!r}; expected one of {sorted(MOE_STAGES)}") return f"{MOE_STAGE_MARKER_PREFIX}{stage}" + (f" {detail}" if detail else "") +def moe_end_marker(detail: str = "") -> str: + """Close the MoE region, so what follows is billed to no MoE stage. + + A separate entry point rather than ``moe_stage_marker(MOE_END_STAGE)`` + because closing a region and setting one are different acts, and the caller + that has to remember this is the one assembling a decoder program, not the + one writing an emitter. Emit it once, after the last MoE work in the + program -- the combine, or the shared/routed add. + + Omitting it is not silent: the profile reports how many instructions carry + the final marker, so an epilogue billed to a MoE stage is visible in the + JSON without reading the source that produced it. + """ + return moe_stage_marker(MOE_END_STAGE, detail) + + class ProgramRoutedMoeMixin: """Routed-MoE v0 emit helpers used by GPT-OSS and Qwen bring-up. @@ -1738,8 +1777,10 @@ def moe_expert_activation_v0( __all__ = [ + "MOE_END_STAGE", "MOE_STAGES", "MOE_STAGE_MARKER_PREFIX", "ProgramRoutedMoeMixin", + "moe_end_marker", "moe_stage_marker", ] diff --git a/aten/tests/test_moe_stage_terminator.py b/aten/tests/test_moe_stage_terminator.py new file mode 100644 index 0000000..660781d --- /dev/null +++ b/aten/tests/test_moe_stage_terminator.py @@ -0,0 +1,63 @@ +"""Guards on the MoE region terminator. + +These cover a failure that produces a plausible-looking wrong answer rather than +an error, which is why it is pinned here rather than left to a numerical test. + +The terminator exists because markers are sticky and a program does not end +where its MoE region does. Without one, the last MoE marker runs to the end of +the file and every instruction after it -- an lm_head, the next sublayer -- is +billed to a MoE stage, and that cost lands in the shared-vs-routed ratio. The +size of the effect is whatever the epilogue happens to be, so it is measured in +the pull request rather than pinned here, where it would rot. +""" + +from __future__ import annotations + +import pytest + +from compiler.aten.plena.program_routed_moe import ( + MOE_END_STAGE, + MOE_STAGES, + moe_end_marker, + moe_stage_marker, +) + + +def test_the_terminator_is_part_of_the_declared_vocabulary() -> None: + """It has to be, or the emulator's both-directions guard rejects it.""" + assert MOE_END_STAGE in MOE_STAGES + + +def test_the_terminator_is_not_a_stage_a_caller_could_mistake_for_work() -> None: + """No other declared stage may be confusable with the terminator. + + Asserting `MOE_END_STAGE == "non_moe"` and then that it does not start with + a work prefix is two ways of restating the literal on the line above. What + is worth pinning is the relationship to the rest of the vocabulary: the + terminator must not be a prefix of, or prefixed by, any stage that names + real work, or a substring match somewhere would pick the wrong one. + """ + others = MOE_STAGES - {MOE_END_STAGE} + assert others, "MOE_STAGES holds nothing but the terminator" + for stage in others: + assert not stage.startswith(MOE_END_STAGE), f"{stage!r} is prefixed by the terminator" + assert not MOE_END_STAGE.startswith(stage), f"the terminator is prefixed by {stage!r}" + + +def test_moe_end_marker_emits_the_terminator() -> None: + assert moe_end_marker() == f"@stage={MOE_END_STAGE}" + assert moe_end_marker("after the combine") == f"@stage={MOE_END_STAGE} after the combine" + + +def test_a_typo_of_the_terminator_fails_like_any_other() -> None: + """It is a marker, not a special case. + + `moe_stage_marker(MOE_END_STAGE) == moe_end_marker()` was the previous + assertion here; it is a tautology, since `moe_end_marker` is defined as that + call. What is checkable is that the terminator is not exempt from the + vocabulary check -- a near miss must be rejected, not silently accepted as + "close enough to the end marker". + """ + for typo in ("non-moe", "nonmoe", "non_moe_", "NON_MOE"): + with pytest.raises(ValueError, match="unknown MoE stage"): + moe_stage_marker(typo)