MoE: a terminator for the sticky stage marker

.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

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}")

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",
 ]

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)