recipes/docker-images/mia-tp4/tests/test_moe_b12x_next_cpu.py

"""CPU checks of adapter/moe_b12x_next.py's host-side logic (no GPU, no engine, no b12x_next build).

- hook order: install_method -> install_runner and install_runner -> install_method both import
  b12x_next through ``_b12x()`` (pinned-commit and drift checks), and a wrong pin or a drifted
  b12x source refuses in both orders (fake ``b12x_next`` / ``sglang`` packages on disk).
- a layer that reaches the fused function without a conversion refuses instead of running the
  stock FlashInfer kernel on checkpoint-layout weights.
- chunk selection: the chunk is the largest prepared capacity one launch can run; with
  CHUNKED_PREFILL_SIZE=65536 (ladder top 65536 above the E=384 launch limit of 61,717 tokens) the
  65536 plan is dropped, the chunk is 4096, and ``forward`` never hands a plan more rows than its
  capacity for any M up to 140k; ``_run`` refuses an oversized batch.
- CUDA_GRAPH_MAX_BS_DECODE above the adapter's graph list names the uncovered row counts.
- DSV41_MOE_B12X_NEXT_DETERMINISTIC=1 moves the <= 8 row plan to the internal route planner and
  races nothing.

usage: python3 tests/test_moe_b12x_next_cpu.py   (needs torch; runs each scenario in a subprocess)
"""
import json
import os
import subprocess
import sys
import tempfile
import textwrap

HERE = os.path.dirname(os.path.abspath(__file__))
ADAPTER = os.path.join(HERE, "..", "adapter")

FAKE_FILES = {
    "b12x_next/__init__.py": "",
    "b12x_next/preparation/__init__.py": "class PreparationSession:\n    pass\n",
    "b12x_next/moe/__init__.py": "",
    "b12x_next/moe/fused_moe/__init__.py": "def run(binding):\n    pass\n",
    "b12x_next/moe/fused_moe/_preparation.py": textwrap.dedent("""\
        class _FusedMoeState:
            def bind(self, **kwargs):
                kwargs["experts"] = self.experts._impl
                b = self.scratch.bind(**kwargs)
                return b, dict(compact_launches=self.compact_launches), "unit_scale_contract", "_w4a16_launches"
        """),
    "b12x_next/moe/fused_moe/api.py": textwrap.dedent("""\
        from dataclasses import replace
        def bind(state, plan, **kwargs):
            return replace(state.bind(**kwargs), plan=plan)
        """),
    "sglang/__init__.py": "",
    "sglang/srt/__init__.py": "",
    "sglang/srt/distributed/__init__.py": "def get_tp_group():\n    return None\n",
    "sglang/srt/distributed/device_communicators/__init__.py": "",
    "sglang/srt/distributed/device_communicators/pynccl_allocator.py": "def use_symmetric_memory(g):\n    pass\n",
    "sglang/srt/layers/__init__.py": "",
    "sglang/srt/layers/dp_attention.py": "def is_allocation_symmetric():\n    return False\n",
    "sglang/srt/layers/quantization/__init__.py": "",
    "sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py": textwrap.dedent("""\
        import torch
        # the stock method dispatches to fused_experts_none_to_flashinfer_mxfp4
        class Mxfp4FlashinferCutlassMoEMethod:
            load_up_proj_weight_first = True
            def create_weights(self, layer, e, h, n, dt, **kw):
                assert n % 128 == 0
                self._fp8.create_weights(layer, e, h, n, dt, fp4_scale_dtype=torch.float8_e8m0fnu, **kw)
            def process_weights_after_loading(self, layer):
                self._fp8.process_weights_after_loading(layer)
                block_scale_interleave(layer)
            def create_moe_runner(self, layer, cfg):
                pass
            def apply(self, layer, d):
                pass
        """),
    "sglang/srt/layers/moe/__init__.py": "",
    "sglang/srt/layers/moe/topk.py": "class TopKOutputChecker:\n    pass\n",
    "sglang/srt/layers/moe/token_dispatcher/__init__.py": "",
    "sglang/srt/layers/moe/token_dispatcher/standard.py": "class StandardCombineInput:\n    pass\n",
    "sglang/srt/layers/moe/moe_runner/__init__.py": "",
    "sglang/srt/layers/moe/moe_runner/base.py": textwrap.dedent("""\
        def fused_experts_none_to_flashinfer_mxfp4(dispatch_output, quant_info, runner_config):
            return "stock"
        class FusedOpPool:
            _fused_funcs = {("none", "flashinfer_mxfp4"): fused_experts_none_to_flashinfer_mxfp4}
        """),
    "sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py": "",
    "sglang/srt/layers/moe/fused_moe_triton/__init__.py": "",
    "sglang/srt/layers/moe/fused_moe_triton/layer.py": textwrap.dedent("""\
        class FusedMoE:
            def _load_w13(self, loaded_weight, shard_dim):
                if getattr(self.quant_method, "load_up_proj_weight_first", False):
                    pass
                shard_size = loaded_weight.shape[shard_dim] // self.moe_tp_size
                start = shard_size
            def _load_w2(self, loaded_weight, shard_dim):
                shard_size = loaded_weight.shape[shard_dim] // self.moe_tp_size
        """),
}

def make_tree(root, commit, drift=False):
    for rel, text in FAKE_FILES.items():
        path = os.path.join(root, rel)
        os.makedirs(os.path.dirname(path), exist_ok=True)
        if drift and rel.endswith("_preparation.py"):
            text = text.replace("unit_scale_contract", "unit_scale")
        with open(path, "w") as f:
            f.write(text)
    with open(os.path.join(root, "b12x_next", "SOURCE_COMMIT"), "w") as f:
        f.write(commit + "\n")


def run(code, env=None, path=None):
    """Runs ``code`` in a fresh interpreter with the adapter (and a fake tree) importable; returns stdout."""
    e = dict(os.environ)
    for k in list(e):
        # PYTHONPATH: the image's adapter dir carries sitecustomize.py, whose import hooks would
        # install the adapter on the fake modules' import and take the hook order out of our hands
        if k == "PYTHONPATH" or k.startswith(("DSV41_", "CUDA_GRAPH_MAX_BS_DECODE", "CHUNKED_PREFILL_SIZE",
                                               "DSPARK_", "B12X_")):
            del e[k]
    e.update({"DSV41_MOE_B12X_NEXT": "1", "STATE_PATH": tempfile.gettempdir()}, **(env or {}))
    pre = f"import sys; sys.path[:0] = {[p for p in (path, ADAPTER) if p]!r}\n"
    r = subprocess.run([sys.executable, "-c", pre + textwrap.dedent(code)], env=e, capture_output=True, text=True,
                       timeout=300)
    if r.returncode != 0:
        raise AssertionError(f"scenario failed:\n{code}\n--- stdout\n{r.stdout}\n--- stderr\n{r.stderr[-3000:]}")
    return r.stdout


HOOKS = """
import moe_b12x_next as ad
import sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe as mq
import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass as fic
steps = {"method": lambda: ad.install_method(mq), "runner": lambda: ad.install_runner(fic)}
out = []
for name in ORDER:
    try:
        steps[name]()
        out.append([name, "ok", sorted(ad._B), sorted(ad._ENG)])
    except RuntimeError as exc:
        out.append([name, "refused", str(exc)[:200]])
        break
print("RESULT " + __import__("json").dumps(out))
"""


def hooks(tree, order):
    out = run(f"ORDER = {order!r}\n" + HOOKS, path=tree)
    return json.loads(out.split("RESULT ", 1)[1])


def check_hook_order():
    pinned = run("import moe_b12x_next as ad; print(ad.PINNED_COMMIT)").strip()
    with tempfile.TemporaryDirectory() as good, tempfile.TemporaryDirectory() as bad, \
            tempfile.TemporaryDirectory() as drift:
        make_tree(good, pinned)
        make_tree(bad, "0" * 40)
        make_tree(drift, pinned, drift=True)
        for order in (["method", "runner"], ["runner", "method"]):
            res = hooks(good, order)
            assert [r[1] for r in res] == ["ok", "ok"], res
            for r in res:                       # b12x modules present from the first hook on
                assert {"fm", "prep", "pkg"} <= set(r[2]), (order, r)
            assert {"get_tp_group", "TopKOutputChecker"} <= set(res[-1][3]), res
            assert not {"fm", "prep"} & set(res[-1][3]) and "get_tp_group" not in res[-1][2], res
            for tree, why in ((bad, "is built from 0000"), (drift, "drifted")):
                res = hooks(tree, order)
                assert res[0][1] == "refused" and why in res[0][2], (order, why, res)
        print("hook order: method->runner and runner->method both run the pin + drift checks; "
              "a wrong pin / drifted source refuses in both orders")


def check_unconverted_layer():
    out = run("""
        import torch, types, moe_b12x_next as ad
        qi = types.SimpleNamespace(w13_weight=torch.zeros(2, 4, dtype=torch.uint8))
        ad._state["orig_fused"] = lambda *a: "stock kernel ran"
        try:
            ad._fused_b12x_next(None, qi, None)
            print("RAN")
        except RuntimeError as exc:
            print("REFUSED", exc)
    """)
    assert out.startswith("REFUSED") and "without a b12x_next conversion" in out, out
    print("unconverted layer: refuses instead of running the stock kernel")


CHUNK = """
import torch, types, moe_b12x_next as ad
LIMIT = 61717                      # _safe_dynamic_token_chunk at E=384, K=5120, N=576, top-6
g = ad._Geometry.__new__(ad._Geometry)
caps = g.capacities()
# dynamic plans launch up to LIMIT tokens; b12x's micro plan (tiny caps) launches its capacity
limits = {c: (c if c <= 8 else LIMIT) for c in caps}
chunk, kept, dropped = ad.select_chunk(limits)
g.caps, g.chunk, g.direct = kept, chunk, True
g.table = [None] + [g._lookup(m) for m in range(1, ad.EXACT_MAX + 1)]
seen, worst = set(), 0
def fake_run(cap, x, ids, w, out, impl):
    global worst
    assert x.shape[0] <= cap and cap in kept, (x.shape[0], cap)
    worst = max(worst, x.shape[0] - cap)
    seen.add(cap)
g._run = fake_run
for m in list(range(1, 5000)) + [8191, 8192, 61717, 61718, 65535, 65536, 65537, 140000]:
    x = torch.empty(m, 0)
    ids = torch.empty(m, 6, dtype=torch.int32)
    w = torch.empty(m, 6, dtype=torch.float32)
    g.forward(None, x, ids, w, x)
# the real _run refuses a batch larger than the plan's capacity before touching the plan
try:
    ad._Geometry._run(g, 4096, torch.empty(4097, 0), None, None, None, None)
    raised = False
except RuntimeError:
    raised = True
# a launch limit between two ladder steps (Fable M2): chunk = the largest capacity under it
c2 = ad.select_chunk({c: (c if c <= 8 else 3000) for c in caps})
try:
    ad.select_chunk({4096: 100})
    none_ok = True
except RuntimeError:
    none_ok = False
print("RESULT " + __import__("json").dumps(dict(ladder=ad.LADDER, caps=caps, chunk=chunk, kept_top=kept[-3:],
      dropped=dropped, raised=raised, c2=[c2[0], c2[2]], none_ok=none_ok, seen_top=max(seen))))
"""


def check_chunk_selection():
    res = json.loads(run(CHUNK, env={"CHUNKED_PREFILL_SIZE": "65536"}).split("RESULT ", 1)[1])
    assert res["ladder"][-1] == 65536 and 65536 in res["caps"], res
    assert res["chunk"] == 4096 and res["dropped"] == [65536] and res["seen_top"] == 4096, res
    assert res["raised"] and res["c2"] == [2048, [4096, 65536]] and not res["none_ok"], res
    base = json.loads(run(CHUNK, env={"CHUNKED_PREFILL_SIZE": "4096"}).split("RESULT ", 1)[1])
    assert base["chunk"] == 4096 and base["dropped"] == [] and base["caps"] == res["caps"][:-1], base
    print(f"chunk selection: CHUNKED_PREFILL_SIZE=65536 -> capacity 65536 dropped (limit 61717), chunk "
          f"{res['chunk']}; every chunk of M in 1..140000 fits its plan; oversized _run refuses; limit 3000 -> "
          f"chunk {res['c2'][0]}; production ladder (4096) unchanged")


GRAPH = """
import moe_b12x_next as ad
tuned = sorted(ad._Geometry.tuned_caps(type("G", (), {"topk": 6})()))
print("RESULT " + __import__("json").dumps(dict(rows=ad.uncovered_rows(), planner=ad.SMALL_PLANNER, bs=ad.GRAPH_BS,
                                                tuned=tuned)))
"""


def graph(env):
    return json.loads(run(GRAPH, env=env).split("RESULT ", 1)[1])


def check_graph_bs_and_planner():
    assert graph({})["rows"] == [] and graph({"CUDA_GRAPH_MAX_BS_DECODE": "16"})["rows"] == []
    r = graph({"CUDA_GRAPH_MAX_BS_DECODE": "32"})
    want = sorted({k * b for k in (5, 6) for b in range(17, 33)})
    assert r["rows"] == want and r["bs"][-1] == 16, r
    r = graph({"CUDA_GRAPH_MAX_BS_DECODE": "32", "DSV41_MOE_B12X_NEXT_GRAPH_BS": ",".join(map(str, range(1, 33)))})
    assert r["rows"] == [], r
    r = graph({"CUDA_GRAPH_MAX_BS_DECODE": "8"})
    assert r["rows"] == [] and r["bs"] == [1, 2, 3, 4, 5, 6, 7, 8], r
    d = graph({})
    assert d["planner"] == "triton" and d["tuned"] == [6 * b for b in d["bs"]], d
    det = graph({"DSV41_MOE_B12X_NEXT_DETERMINISTIC": "1"})
    assert det["planner"] == "internal" and det["tuned"] == [], det
    assert graph({"DSV41_MOE_B12X_NEXT_SMALL_PLAN": "internal:none:16"})["planner"] == "internal"
    print(f"graph bs: CUDA_GRAPH_MAX_BS_DECODE=32 names {len(want)} uncovered row counts ({want[0]}..{want[-1]}), "
          f"none at 16 or with a covering list; DETERMINISTIC=1 -> internal route planner, no race")


TABLE = """
import moe_b12x_next as ad
print("RESULT " + __import__("json").dumps(dict(
    table=ad.PLAN_TABLE, p6=[ad.table_plan(c, 6) for c in (1, 6, 8, 9, 96, 4096)],
    p3=[ad.table_plan(c, 3) for c in (5, 40, 80)])))
"""


def table(env):
    return json.loads(run(TABLE, env=env).split("RESULT ", 1)[1])


def check_plan_table():
    d = table({})
    assert d["p6"] == [["triton", 48, 16]] * 3 + [None] * 3 and d["p3"] == [["triton", 48, 16], None, None], d
    det = table({"DSV41_MOE_B12X_NEXT_DETERMINISTIC": "1"})
    assert det["p6"][:3] == [["internal", 48, 16]] * 3 and det["p6"][3:] == [None] * 3, det
    assert table({"DSV41_MOE_B12X_NEXT_SMALL_PLAN": "internal:none:16"})["p6"][0] == ["internal", None, 16]
    t = table({"DSV41_MOE_B12X_NEXT_DETERMINISTIC": "1",
               "DSV41_MOE_B12X_NEXT_PLAN_TABLE": "1-8=triton:48:16, 40-80@k3=internal:32:16, 96=heur, 9-96=internal:none:16"})
    assert t["p6"] == [["internal", 48, 16]] * 3 + [["internal", None, 16], None, None], t
    assert t["p3"] == [["internal", 48, 16], ["internal", 32, 16], ["internal", 32, 16]], t
    for bad in ("1-8=cutlass:48:16", "8-1=internal:48:16", "x=heur", "1-8=internal:48"):
        try:
            run(TABLE, env={"DSV41_MOE_B12X_NEXT_PLAN_TABLE": bad})
            raise RuntimeError(f"{bad!r} was accepted")
        except AssertionError as exc:
            assert "DSV41_MOE_B12X_NEXT_PLAN_TABLE: bad" in str(exc), (bad, str(exc)[-300:])
    print("plan table: default = <= 8 row pin only (unchanged plans), SMALL_PLAN still honoured, top-k scoping, "
          "first match wins, heur entries, deterministic triton -> internal, malformed entries refuse to load")


def main():
    check_hook_order()
    check_unconverted_layer()
    check_chunk_selection()
    check_graph_bs_and_planner()
    check_plan_table()
    print("moe_b12x_next CPU OK")


if __name__ == "__main__":
    main()