launcher/docker-images/mia-tp4/tests/test_prefill_sp.py

"""CPU check for adapter/prefill_sp.py: the row partition, the tail-row gathers, and the whole
sharded layer loop against the engine's stock loop, with four TP ranks emulated by threads over
a fake process group (all-gather = concatenation, all-reduce = rank-order bf16 sum, reduce-scatter
= a different summation order, like NCCL's).

The model is a toy with the engine's structure: residual [M, 4, D], a row-mixing "attention"
whose wo_b all-reduces unless told to skip, a hash-routed "MoE" that honours mlp_reduce_scatter,
an Engram (owned-row lookup + all-reduce, wkv, gate), bounded replay (late layers on the tail
rows), DSpark aux capture and the vision `where`. Its hc combine picks one of two numerically
different reductions by row count (as the engine's >= 4096 fused norm does), keyed on
sp_logical_rows(): a gate keyed on the shard size would change bits, and the test shows it.

Checks, for several chunk sizes (4096, 2052 = odd shards, 2048 with a full-row tail, 2050 = not
divisible, below MIN_ROWS) and tail layouts (one request, requests straddling shard boundaries,
many short requests, no bounded replay):
  * stage 1 EXACT == stock, bit for bit, on every rank;
  * stage 1 fast == P0 (comm), bit for bit (same reduce-scatters on the same tensors);
  * P0 differs from stock only through the summation order (and matches it with an exact RS);
  * the wkv self-check falls back to the full-M GEMM when the shard GEMM is not bit-exact.

  PYTHONPATH=adapter python tests/test_prefill_sp.py
"""
import os
import threading
from contextlib import contextmanager, nullcontext
from types import SimpleNamespace

os.environ.setdefault("DSV41_PREFILL_SP", "1")
import torch  # noqa: E402

import prefill_sp as sp  # noqa: E402

W = 4
D = 64
BF = torch.bfloat16


# ------------------------------------------------------------------------------------------
# fake process group over threads
# ------------------------------------------------------------------------------------------
class ThreadGroup:
    def __init__(self, world, rs_order="ring"):
        self.world_size = world
        self.unique_name = "tp"
        self.device_group = None
        self._bar = threading.Barrier(world)
        self._slots = [None] * world
        self._tl = threading.local()
        self.rs_order = rs_order
        self.calls = {"ag": 0, "ar": 0, "rs": 0}

    @property
    def rank_in_group(self):
        return self._tl.rank

    def bind(self, rank):
        self._tl.rank = rank

    def _exchange(self, t):
        r = self.rank_in_group
        self._slots[r] = t.clone()
        self._bar.wait()
        vals = list(self._slots)
        self._bar.wait()
        return vals

    def all_gather_into_tensor(self, out, x):
        vals = self._exchange(x.contiguous())
        if self.rank_in_group == 0:
            self.calls["ag"] += 1
        out.view(-1).copy_(torch.cat([v.reshape(-1) for v in vals]))

    def all_reduce(self, x):
        vals = self._exchange(x)
        if self.rank_in_group == 0:
            self.calls["ar"] += 1
        acc = vals[0].clone()
        for v in vals[1:]:
            acc = acc + v                               # rounds to the dtype at every hop
        return acc

    def reduce_scatter_tensor(self, out, x):
        vals = self._exchange(x.contiguous())
        r, s = self.rank_in_group, out.shape[0]
        if r == 0:
            self.calls["rs"] += 1
        parts = [v[r * s:(r + 1) * s] for v in vals]
        order = list(range(self.world_size)) if self.rs_order == "same" else \
            [(r + 1 + k) % self.world_size for k in range(self.world_size)]
        acc = parts[order[0]].clone()
        for o in order[1:]:
            acc = acc + parts[o]
        out.copy_(acc)


def run_ranks(group, fn):
    out, errs = [None] * group.world_size, []

    def body(r):
        group.bind(r)
        try:
            out[r] = fn(r)
        except BaseException as exc:  # noqa: BLE001
            errs.append((r, exc))
            group._bar.abort()

    ts = [threading.Thread(target=body, args=(r,)) for r in range(group.world_size)]
    for t in ts:
        t.start()
    for t in ts:
        t.join()
    if errs:
        raise errs[0][1]
    return out


# ------------------------------------------------------------------------------------------
# fake engine
# ------------------------------------------------------------------------------------------
_fwd = threading.local()


class _Forward:
    @property
    def mlp_reduce_scatter(self):
        return getattr(_fwd, "mrs", False)

    @contextmanager
    def scoped(self, mlp_reduce_scatter=False):
        prev = getattr(_fwd, "mrs", False)
        _fwd.mrs = mlp_reduce_scatter
        try:
            yield
        finally:
            _fwd.mrs = prev


FORWARD = _Forward()
GATE = {"logical": True}


def hc_rows(x):
    return sp.sp_logical_rows(x) if GATE["logical"] else x.shape[0]


class RowParallel(torch.nn.Module):
    def __init__(self, group, weights):
        super().__init__()
        self.group, self.w = group, weights
        self.reduce_results, self.tp_size = True, group.world_size

    def forward(self, input_, skip_all_reduce=False):
        out = (input_.float() @ self.w[self.group.rank_in_group]).to(BF)
        if not skip_all_reduce:
            out = self.group.all_reduce(out)
        return out, None


class FakeAttn(torch.nn.Module):
    def __init__(self, group, g):
        super().__init__()
        self.group = group
        self.scale = torch.randn(group.world_size, D, generator=g)
        self.wo_b = RowParallel(group, torch.randn(group.world_size, D, D, generator=g) / D ** 0.5)
        self.attn_tp_size = group.world_size
        self.seen_rows = []

    def forward(self, x, positions, forward_batch, x_quant=None):
        assert positions.shape[0] == x.shape[0], "attention needs the positions of every row"
        self.seen_rows.append(x.shape[0])
        # causal row mixing: attention needs every earlier row
        h = x.float().cumsum(0) / torch.arange(1, x.shape[0] + 1).unsqueeze(1)
        h = (h * self.scale[self.group.rank_in_group]).to(BF)
        o, _ = self.wo_b(h)
        return o

    def maybe_use_decode_attn_tp(self, fb):
        return nullcontext()


class FakeMoE(torch.nn.Module):
    def __init__(self, group, g):
        super().__init__()
        self.group = group
        self.table = torch.randn(16, group.world_size, D, generator=g)
        self.tp_size = group.world_size
        self._shared_expert_tp1 = False
        self._enable_a2a_moe = False

    def forward(self, hidden_states, forward_batch=None, gemm_output_zero_allocator=None,
                input_ids=None, input_ids_global=None, skip_shared_experts=False):
        assert input_ids.shape[0] == hidden_states.shape[0], "hash routing needs every row's id"
        w = self.table[input_ids % 16, self.group.rank_in_group]
        out = (hidden_states.float() * torch.tanh(w)).to(BF)
        if not FORWARD.mlp_reduce_scatter:
            out = self.group.all_reduce(out)
        return out


class FakeEmbed:
    def __init__(self, group, g, vocab, cols, dim):
        self.group, self.tp_size, self._shared = group, group.world_size, False
        self.table = torch.randn(vocab, cols, dim, generator=g).to(BF)
        self.vocab = vocab

    def _owned_rows(self, ids):
        r = self.group.rank_in_group
        lo, hi = self.vocab * r // self.tp_size, self.vocab * (r + 1) // self.tp_size
        owned = (ids >= lo) & (ids < hi)                                 # [T, cols]
        vals = self.table[ids, torch.arange(ids.shape[1])]              # [T, cols, dim]
        return vals.masked_fill(~owned.unsqueeze(-1), 0)


class FakeEngram:
    def __init__(self, group, g, hash_index, m_dependent=False):
        self.group, self.layer_hash_index = group, hash_index
        self.embed = FakeEmbed(group, g, 97, 3, 8)
        self.wkv_w = torch.randn(24, D, generator=g) / 5
        self.m_dependent = m_dependent
        self.wkv_calls = []

    def wkv(self, t):
        self.wkv_calls.append(t.shape[0])
        out = (t.float() @ self.wkv_w)
        if self.m_dependent and t.shape[0] >= 2048:   # a GEMM whose tactic changes with M
            out = out * (1 + 2 ** -7)
        return out.to(BF), None

    def apply_gate(self, x, kv):
        return (x.float() + torch.sigmoid(kv.float()).unsqueeze(1)).to(x.dtype)

    def __call__(self, x, ids, forward_batch=None, cp_all_tokens=False):
        emb = self.group.all_reduce(self.embed._owned_rows(ids))
        kv, _ = self.wkv(emb.flatten(-2))
        return self.apply_gate(x, kv)


class FakeLayer:
    def __init__(self, group, g, layer_id, engram=None):
        self.layer_id = layer_id
        self.self_attn = FakeAttn(group, g)
        self.mlp = FakeMoE(group, g)
        self.engram = engram
        self.fn = torch.randn(2, 4, D, generator=g)
        self.norm_w = torch.randn(D, generator=g)
        self.dsa_enable_prefill_cp = False
        self.use_fused_mhc_post_pre = False

    def _hc_combine(self, R, pre):
        if pre is None:
            return R[:, 0, :].contiguous()
        p = pre.float()
        if hc_rows(R) >= 4096:          # "fused" reduction order
            y = ((R[:, 0].float() * p[:, 0:1] + R[:, 1].float() * p[:, 1:2])
                 + (R[:, 2].float() * p[:, 2:3] + R[:, 3].float() * p[:, 3:4]))
        else:                           # "unfused": sequential, rounded before the norm
            y = R[:, 0].float() * p[:, 0:1]
            for c in (1, 2, 3):
                y = y + R[:, c].float() * p[:, c:c + 1]
            y = y.to(BF).float()
        y = y * torch.rsqrt(y.square().mean(-1, keepdim=True) + 1e-6) * self.norm_w
        return y.to(BF)

    def _stats(self, R, which):
        mix = (R.float() * self.fn[which]).sum(-1)                         # [N, 4]
        pre = torch.sigmoid(mix)
        post = 2 * torch.sigmoid(mix.flip(-1))
        comb = torch.softmax(mix.unsqueeze(-1) * mix.unsqueeze(-2) / 8, dim=-1)
        return pre, post, comb

    def hc_post(self, x, R, post, comb):
        out = post.unsqueeze(-1) * x.float().unsqueeze(1) + (comb.unsqueeze(-1) * R.float().unsqueeze(2)).sum(1)
        return out.to(BF)

    def forward_hc_pre_from_prev(self, positions, hidden_states, input_ids, forward_batch,
                                 input_ids_global, prev_pre, precomputed_attn=None,
                                 next_norm=None, next_input=None):
        R = hidden_states
        x = self._hc_combine(R, prev_pre)
        x = self.self_attn(x=x, positions=positions, forward_batch=forward_batch, x_quant=None)
        attn_pre, post, comb = self._stats(R, 0)
        R = self.hc_post(x, R, post, comb)
        x = self._hc_combine(R, attn_pre)
        with FORWARD.scoped(mlp_reduce_scatter=False):
            x = self.mlp(x, forward_batch, input_ids=input_ids, input_ids_global=input_ids_global,
                         skip_shared_experts=False)
        ffn_pre, post, comb = self._stats(R, 1)
        R = self.hc_post(x, R, post, comb)
        return R, ffn_pre


class Tail:
    def __init__(self, lens, window=128, pad_rows=0):
        starts, idx = [], []
        pos = 0
        for n in lens:
            k = min(window, n)
            idx.extend(range(pos + n - k, pos + n))
            pos += n
        self.token_indices = torch.tensor(idx, dtype=torch.int64)
        self.contiguous_start = (pos - min(window, lens[0])) if len(lens) == 1 else None
        self.positions = self.token_indices.clone()
        self.pad_rows = pad_rows
        self.cp_metadata = None

    def real_rows(self, t):
        return t[self.contiguous_start:] if self.contiguous_start is not None else t[self.token_indices]

    def rows(self, t):
        r = self.real_rows(t)
        if self.pad_rows:
            r = torch.cat([r, r.new_zeros((self.pad_rows, *r.shape[1:]))])
        return r


class Backend:
    def __init__(self, tail):
        self.tail_forward_metadata = SimpleNamespace(late_layer_tail=tail)
        self.entered = 0

    def enter_late_layer_tail(self, fb):
        self.entered += 1
        return "saved"

    def exit_late_layer_tail(self, saved, fb):
        assert saved == "saved"


class FakeModel:
    def __init__(self, group, seed, n_layers=6, late=4, vision=False, m_dependent_wkv=False):
        g = torch.Generator().manual_seed(seed)
        self.pp_group = SimpleNamespace(world_size=1)
        self.hidden_size, self.hc_mult, self.hc_pre_from_prev_sublayer = D, 4, True
        self.start_layer, self.end_layer, self.late_layer_start = 0, n_layers, late
        self.config = SimpleNamespace(model_type="deepseek_v41", vision_n_layers=1 if vision else 0,
                                      image_token_id=7)
        eng = {1: FakeEngram(group, g, 0, m_dependent_wkv), 3: FakeEngram(group, g, 1, m_dependent_wkv)}
        self.layers = [FakeLayer(group, g, i, eng.get(i)) for i in range(n_layers)]
        self.hash_a = torch.randint(1, 97, (2, 3), generator=g)
        self.dspark_layers_to_capture = [2, 5]

    def engram_hasher(self, input_ids, fb):
        return (input_ids[:, None, None] * self.hash_a + torch.arange(3)) % 97

    def _check_late_layer_tail_readers(self, fb):
        pass


def fake_engine(group, backend):
    return SimpleNamespace(
        get_attn_backend=lambda: backend,
        check_cuda_graph_backend=lambda *a: False,
        Phase=SimpleNamespace(PREFILL="prefill"),
        Backend=SimpleNamespace(TC_PIECEWISE="pcg"),
        get_global_expert_distribution_recorder=lambda: SimpleNamespace(
            with_current_layer=lambda i: nullcontext()),
        get_forward=lambda: FORWARD,
        is_in_breakable_cuda_graph=lambda: False,
        is_cp_active=lambda fb: False,
        get_attn_tp_context=lambda: SimpleNamespace(input_scattered=False),
        get_tp_group=lambda: group,
        get_parallel=lambda: SimpleNamespace(attn_dp_size=1, tp_size=group.world_size,
                                             attn_tp_size=group.world_size),
        get_moe_a2a_backend=lambda: SimpleNamespace(is_none=lambda: True),
        get_platform=lambda: SimpleNamespace(is_sm100=False),
        envs=SimpleNamespace(SGLANG_ENABLE_DETERMINISTIC_INFERENCE=SimpleNamespace(get=lambda: False)),
        MQALayer=FakeAttn,
    )


FakeAttn.forward = sp._make_attn_forward(FakeAttn.forward)
FakeMoE.forward = sp._make_moe_forward(FakeMoE.forward)
FakeLayer.forward_hc_pre_from_prev = sp._make_layer_forward(FakeLayer.forward_hc_pre_from_prev)
sp._REQUIRE_CUDA = False


def _agree_min(p, local):
    vals = p.group._exchange(torch.tensor([1 if local else 0]))
    return bool(min(int(v) for v in vals))


sp._agree_min = _agree_min


# ------------------------------------------------------------------------------------------
# reference: the engine's stock loop, same fake model
# ------------------------------------------------------------------------------------------
def stock_forward_layers(self, positions, hidden_states, forward_batch, input_ids, input_ids_global,
                         capture_dspark, aux_out):
    hash_ids = self.engram_hasher(input_ids, forward_batch)
    tail = None
    if self.late_layer_start is not None:
        backend = sp._M["v4"].get_attn_backend()
        tail = backend.tail_forward_metadata.late_layer_tail
    prev_pre = None
    saved = None
    for i in range(self.start_layer, self.end_layer):
        if tail is not None and i == self.late_layer_start:
            saved = backend.enter_late_layer_tail(forward_batch)
            hidden_states, prev_pre, input_ids, input_ids_global = (
                tail.rows(hidden_states), tail.rows(prev_pre), tail.rows(input_ids),
                tail.rows(input_ids_global))
            positions = tail.positions
            hash_ids = tail.rows(hash_ids)
        engram = self.layers[i].engram
        if engram is not None:
            before = hidden_states
            hidden_states = engram(hidden_states, hash_ids[:, engram.layer_hash_index], forward_batch)
            if self.config.vision_n_layers > 0:
                hidden_states = torch.where((input_ids == self.config.image_token_id)[:, None, None],
                                            before, hidden_states)
        if capture_dspark and i in self.dspark_layers_to_capture:
            aux = hidden_states
            if tail is not None and i < self.late_layer_start:
                aux = tail.rows(aux)
            aux_out.append(aux.mean(dim=1))
        hidden_states, prev_pre = self.layers[i].forward_hc_pre_from_prev(
            positions=positions, hidden_states=hidden_states, input_ids=input_ids,
            forward_batch=forward_batch, input_ids_global=input_ids_global, prev_pre=prev_pre)
    if saved is not None:
        backend.exit_late_layer_tail(saved, forward_batch)
        return hidden_states, prev_pre, tail
    return hidden_states, prev_pre, None


FB = SimpleNamespace(forward_mode=SimpleNamespace(is_extend_without_speculative=lambda: True))


def run(mode, rows, lens, *, exact=False, late=4, vision=False, rs_order="ring", m_dep=False, seed=1,
        debug=False):
    """Every rank's (hidden, pre, aux list) for one forward in `mode` (stock | comm | shard)."""
    group = ThreadGroup(W, rs_order)
    model = FakeModel(group, seed, late=late, vision=vision, m_dependent_wkv=m_dep)
    tail = Tail(lens) if late is not None else None
    backend = Backend(tail)
    sp._M.clear()
    sp._M.update(v4=fake_engine(group, backend), v2=SimpleNamespace(DeepseekV2MoE=FakeMoE))
    sp._STATIC["checked"] = False
    sp._STATIC["tail_checked"] = True         # the engine's tail helpers are hash-checked on GPU
    sp._WKV_OK.clear()
    g = torch.Generator().manual_seed(seed + 100)
    R0 = torch.randn(rows, 4, D, generator=g).to(BF)
    ids = torch.randint(0, 40, (rows,), generator=g)
    pos = torch.arange(rows)
    old = sp.MODE, sp.EXACT

    def body(r):
        aux = []
        if mode == "stock":
            out = stock_forward_layers(model, pos, R0.clone(), FB, ids, ids, True, aux)
            return out[0], out[1], aux, None
        p = sp._plan(model, R0, FB)
        if p is None:
            out = stock_forward_layers(model, pos, R0.clone(), FB, ids, ids, True, aux)
            return out[0], out[1], aux, None
        d = sp._Dbg(1, rows, p, r) if debug else None
        sp._ctx.dbg = d
        try:
            out = sp._sp_forward_layers(model, p, pos, R0.clone(), FB, ids, ids, True, aux)
        finally:
            sp._ctx.dbg = None
        if d is not None:
            sp._debug_end(d, out)
            return out[0], out[1], aux, (p, d)
        return out[0], out[1], aux, p

    sp.MODE, sp.EXACT = ("off" if mode == "stock" else mode), exact
    try:
        res = run_ranks(group, body)
    finally:
        sp.MODE, sp.EXACT = old
    return res, model, group


def same(a, b):
    ha, pa, xa, _ = a
    hb, pb, xb, _ = b
    return (torch.equal(ha, hb) and torch.equal(pa, pb) and len(xa) == len(xb)
            and all(torch.equal(u, v) for u, v in zip(xa, xb)))


def main():
    torch.set_num_threads(1)
    # --- partition -------------------------------------------------------------------------
    for rows in (4096, 2048, 2052, 1500, 128, 65536):
        if not sp.eligible_rows(rows, W, 128):
            continue
        cover = []
        for r in range(W):
            lo, hi = sp.shard_range(rows, W, r)
            cover.extend(range(lo, hi))
        assert cover == list(range(rows)), rows
    assert not sp.eligible_rows(2050, W, 1024) and not sp.eligible_rows(2044, W, 2048)
    assert sp.eligible_rows(2048, W, 2048) and not sp.eligible_rows(4096, 1, 2048)
    idx = torch.tensor([0, 511, 512, 1023, 1024, 4095])
    o, l = sp.owner_local(idx, 1024)
    assert o.tolist() == [0, 0, 0, 0, 1, 3] and l.tolist() == [0, 511, 512, 1023, 0, 1023]
    assert sp.tail_plan(4096, 4, 128) == "owners" and sp.tail_plan(2048, 4, 512) == "full"

    # --- tail-row gathers against tail.rows(full) ---------------------------------------------
    for rows, lens, pad in ((4096, [4096], 0), (4096, [1000, 50, 2046, 1000], 0),
                            (2052, [513, 1, 1025, 513], 0), (2048, [64] * 32, 0),
                            (4096, [4000, 96], 3), (8192, [2047, 2049, 4096], 0)):
        group = ThreadGroup(W)
        tail = Tail(lens, pad_rows=pad)
        full = torch.randn(rows, 4, D).to(BF)
        pre = torch.rand(rows, 4)

        def body(r, rows=rows, tail=tail, full=full, pre=pre, group=group):
            p = sp._Plan(rows, W, r, group, True, False)
            return sp._tail_rows(p, tail, full[p.lo:p.hi]), sp._tail_rows(p, tail, pre[p.lo:p.hi])

        for got_h, got_p in run_ranks(group, body):
            assert torch.equal(got_h, tail.rows(full)) and torch.equal(got_p, tail.rows(pre)), (rows, lens)

    # --- whole loop ----------------------------------------------------------------------------
    MIN = sp.MIN_ROWS
    sp.MIN_ROWS = 1024
    try:
        cases = [
            ("one request", 4096, [4096], dict()),
            ("straddling requests", 4096, [1000, 50, 2046, 1000], dict()),
            ("odd shards + vision", 2052, [513, 1, 1025, 513], dict(vision=True)),
            ("tail == all rows", 2048, [64] * 32, dict()),
            ("no bounded replay", 4096, [4096], dict(late=None)),
            ("not divisible (stock)", 2050, [2050], dict()),
            ("below MIN_ROWS (stock)", 1020, [1020], dict()),
        ]
        for name, rows, lens, kw in cases:
            stock, _, _ = run("stock", rows, lens, **kw)
            exact, m_ex, g_ex = run("shard", rows, lens, exact=True, **kw)
            fast, m_fast, g_fast = run("shard", rows, lens, **kw)
            comm, _, g_comm = run("comm", rows, lens, **kw)
            comm_same, _, _ = run("comm", rows, lens, rs_order="same", **kw)
            planned = exact[0][3] is not None
            for r in range(W):
                assert same(exact[r], stock[r]), (name, "exact != stock", r)
                assert same(fast[r], comm[r]), (name, "fast != P0", r)
                assert same(comm_same[r], stock[r]), (name, "P0 with a same-order RS != stock", r)
            if planned:
                assert not same(fast[0], stock[0]), (name, "the RS order should show")
                attn = m_fast.layers[0].self_attn
                assert set(attn.seen_rows) == {rows}, (name, attn.seen_rows)
                late = kw.get("late", 4)
                tail_ars = 0 if late is None else 2 * (len(m_fast.layers) - late)   # stock late layers
                assert g_fast.calls["ar"] == tail_ars == g_comm.calls["ar"], (name, g_fast.calls)
                eng_shard = m_fast.layers[1].engram.wkv_calls
                assert rows // W in eng_shard, (name, eng_shard)
            else:
                assert same(fast[0], stock[0]) and same(comm[0], stock[0]), name
            print(f"  {name:24s} rows={rows}: exact==stock, fast==P0"
                  f"{'' if planned else ' (stock path)'}; collectives fast {g_fast.calls}")

        # gates keyed on the shard size instead of the chunk: bits change (the test is sensitive)
        GATE["logical"] = False
        try:
            wrong, _, _ = run("shard", 4096, [4096], exact=True)
        finally:
            GATE["logical"] = True
        stock, _, _ = run("stock", 4096, [4096])
        assert not same(wrong[0], stock[0]), "a shard-keyed gate should change the result"

        # wkv whose result depends on M: the self-check must fall back and stay exact
        stock, _, _ = run("stock", 4096, [1000, 3096], m_dep=True)
        exact, m_ex, _ = run("shard", 4096, [1000, 3096], exact=True, m_dep=True)
        assert all(same(exact[r], stock[r]) for r in range(W)), "wkv fallback not exact"
        assert sp._WKV_OK and not any(sp._WKV_OK.values()), sp._WKV_OK
        stock, _, _ = run("stock", 4096, [1000, 3096])
        exact, _, _ = run("shard", 4096, [1000, 3096], exact=True)
        assert all(same(exact[r], stock[r]) for r in range(W))
        assert sp._WKV_OK and all(sp._WKV_OK.values()), sp._WKV_OK
    finally:
        sp.MIN_ROWS = MIN
    # debug compare + fingerprints: on (exact and fast), results unchanged, nothing flagged
    old_dbg = set(sp.DEBUG)
    sp.DEBUG.clear()
    sp.DEBUG.update({"compare", "fp"})
    try:
        stock, _, _ = run("stock", 4096, [1000, 50, 2046, 1000])
        for exact_ in (True, False):
            dbg, _, _ = run("shard", 4096, [1000, 50, 2046, 1000], exact=exact_, debug=True)
            ref, _, _ = run("shard", 4096, [1000, 50, 2046, 1000], exact=exact_)
            for r in range(W):
                d = dbg[r][3][1]
                assert same(dbg[r][:3] + (None,), ref[r][:3] + (None,)), "debug changed the result"
                assert d.checked >= 10 and not d.bad, (d.checked, d.bad)
                assert {k for _, k, _ in d.fps} == {"Rin", "ain", "aout", "min", "mout", "Rout", "pre"}
                assert {layer for layer, _, _ in d.fps} == set(range(6)), d.fps[:3]
            if exact_:
                assert all(same(dbg[r][:3] + (None,), stock[r]) for r in range(W))
    finally:
        sp.DEBUG.clear()
        sp.DEBUG.update(old_dbg)
    print("test_prefill_sp: ok")


if __name__ == "__main__":
    main()