PLI Labs public snapshot 2026-10-28 (history squashed)

e1acc15 7 days ago

20 kB

# SPDX-License-Identifier: Apache-2.0
"""Wire codec for the remote DSpark contract (offload/dspark/CONTRACT.md, v1.0).
Pure Python + numpy, no torch and no vLLM: shared by the Spark proxy, the mock worker
and the offline tests. Every layout here mirrors a table in CONTRACT.md section 6; keep
them in step (and bump CONTRACT_VERSION / the change log) on any change.
"""

from __future__ import annotations

import hashlib
import struct
from dataclasses import dataclass, field

try:  # numpy is needed for STEP bodies; HELLO/PING/STATS framing works without it
    import numpy as np  # (the probe tool runs on Spark hosts whose python3 has no numpy)
except ImportError:  # pragma: no cover
    np = None  # type: ignore[assignment]

CONTRACT_VERSION = "1.8"

MAGIC = 0x31505344  # bytes "DSP1" little-endian
VERSION = 1

OP_HELLO = 1
OP_STEP = 2
OP_RESET = 3
OP_STATS = 4
OP_PING = 5
OP_NAMES = {OP_HELLO: "HELLO", OP_STEP: "STEP", OP_RESET: "RESET",
            OP_STATS: "STATS", OP_PING: "PING"}

# Frame status (CONTRACT 9.1)
ST_OK = 0
ST_BAD_VERSION = 1
ST_IDENTITY_MISMATCH = 2
ST_CONFIG_MISMATCH = 3
ST_NO_SESSION = 4
ST_STALE_SEQ = 5
ST_MALFORMED = 6
ST_CAPACITY = 7
ST_INTERNAL = 8
ST_UNSUPPORTED_OP = 9
FRAME_STATUS_NAMES = {
    ST_OK: "OK", ST_BAD_VERSION: "BAD_VERSION", ST_IDENTITY_MISMATCH: "IDENTITY_MISMATCH",
    ST_CONFIG_MISMATCH: "CONFIG_MISMATCH", ST_NO_SESSION: "NO_SESSION",
    ST_STALE_SEQ: "STALE_SEQ", ST_MALFORMED: "MALFORMED", ST_CAPACITY: "CAPACITY",
    ST_INTERNAL: "INTERNAL", ST_UNSUPPORTED_OP: "UNSUPPORTED_OP",
}
# Frame errors after which the remote must not be retried (configuration errors).
FATAL_FRAME_STATUS = frozenset(
    {ST_BAD_VERSION, ST_IDENTITY_MISMATCH, ST_CONFIG_MISMATCH, ST_UNSUPPORTED_OP}
)

# Per-request status (CONTRACT 9.2)
RS_OK = 0
RS_PARTIAL_WINDOW = 1
RS_STALE_GENERATION = 2
RS_BAD_ENTRY = 3
RS_INTERNAL = 4
REQ_STATUS_NAMES = {RS_OK: "OK", RS_PARTIAL_WINDOW: "PARTIAL_WINDOW",
                    RS_STALE_GENERATION: "STALE_GENERATION", RS_BAD_ENTRY: "BAD_ENTRY",
                    RS_INTERNAL: "INTERNAL"}
# Per-request statuses whose drafts may be used.
USABLE_REQ_STATUS = frozenset({RS_OK, RS_PARTIAL_WINDOW})

# ReqEntry.req_flags
RF_NEW_GEN = 0x1
RF_NEED_DRAFT = 0x2
RF_RESYNC = 0x4

# StepHeader.step_flags / ReplyHeader.reply_flags
SF_WANT_MARGIN = 0x1
RPF_HAS_MARGIN = 0x1

PAYLOAD_TAPS_BF16 = 1
PAYLOAD_MAINX_BF16 = 2
SAMPLE_GREEDY = 1
SAMPLE_TOP32 = 2
CAP_TOP32 = 8
SF_SAMPLED = 2
RPF_SAMPLED = 2

CAP_TAPS_BF16 = 0x1
CAP_MAINX_BF16 = 0x2
CAP_HAS_MARGIN = 0x4

APP_HEADER = struct.Struct("<IHHQIIiI")          # 32 B
HELLO_REQ = struct.Struct("<20I32s32s48s")       # 192 B
HELLO_REP = struct.Struct("<6I32s32s40s")        # 128 B
STEP_HEADER = struct.Struct("<QHHIII8x")         # 32 B
REQ_ENTRY = struct.Struct("<BBBBIqIiq")          # 32 B
REPLY_HEADER = struct.Struct("<QHHIII8x")        # 32 B
REQ_STATUS = struct.Struct("<HHI")               # 8 B
RESET_BODY = struct.Struct("<I")                 # 4 B

assert APP_HEADER.size == 32
assert HELLO_REQ.size == 192
assert HELLO_REP.size == 128
assert STEP_HEADER.size == 32
assert REQ_ENTRY.size == 32
assert REPLY_HEADER.size == 32
assert REQ_STATUS.size == 8

APP_HEADER_SIZE = APP_HEADER.size
U32_MASK = 0xFFFFFFFF


class WireError(ValueError):
    """A malformed or inconsistent message."""


def align64(n: int) -> int:
    return (n + 63) & ~63


def rows_offset_for(n_req: int, sampled: bool = False) -> int:
    """Offset of the row block from the start of a STEP body (multiple of 64)."""
    return align64(STEP_HEADER.size + n_req * REQ_ENTRY.size + (16 * n_req if sampled else 0))


def gen_newer(a: int, b: int) -> bool:
    """Serial-number arithmetic on u32 generations: is a newer than b?"""
    d = (a - b) & U32_MASK
    return d != 0 and d < 0x80000000


# --------------------------------------------------------------------------- headers


@dataclass
class AppHeader:
    op: int
    seq: int
    session: int
    body_len: int
    status: int = 0
    flags: int = 0
    magic: int = MAGIC
    version: int = VERSION

    def pack_into(self, buf, off: int = 0) -> None:
        APP_HEADER.pack_into(buf, off, self.magic, self.version, self.op, self.seq,
                             self.session, self.body_len, self.status, self.flags)

    def pack(self) -> bytes:
        b = bytearray(APP_HEADER_SIZE)
        self.pack_into(b)
        return bytes(b)

    @classmethod
    def unpack_from(cls, buf, off: int = 0) -> "AppHeader":
        if len(buf) - off < APP_HEADER_SIZE:
            raise WireError(f"short frame: {len(buf) - off} < {APP_HEADER_SIZE}")
        magic, version, op, seq, session, body_len, status, flags = APP_HEADER.unpack_from(
            buf, off)
        h = cls(op=op, seq=seq, session=session, body_len=body_len, status=status,
                flags=flags, magic=magic, version=version)
        return h

    def check(self, frame_len: int) -> None:
        if self.magic != MAGIC:
            raise WireError(f"bad magic 0x{self.magic:08x}")
        if self.version != VERSION:
            raise WireError(f"unsupported version {self.version}")
        if APP_HEADER_SIZE + self.body_len != frame_len:
            raise WireError(
                f"body_len {self.body_len} does not match frame length {frame_len}")


# --------------------------------------------------------------------------- HELLO


@dataclass
class Hello:
    num_spec: int = 5
    window: int = 128
    hidden: int = 5120
    n_taps: int = 3
    tap_layer_ids: tuple[int, ...] = (37, 38, 39)
    payload_kind: int = PAYLOAD_TAPS_BF16
    row_bytes: int = 3 * 5120 * 2
    noise_token_id: int = 128799
    vocab_size: int = 129280
    max_model_len: int = 327680
    max_slots: int = 4
    max_rows_per_step: int = 4 * 128
    sample_mode: int = SAMPLE_GREEDY
    tp_size: int = 2
    drafter_identity: bytes = bytes(32)
    markov_identity: bytes = bytes(32)
    client_name: str = ""
    proto_min: int = 1
    proto_max: int = 1

    def pack(self) -> bytes:
        ids = list(self.tap_layer_ids) + [0] * (4 - len(self.tap_layer_ids))
        if len(ids) != 4:
            raise WireError("at most 4 tap layers")
        return HELLO_REQ.pack(
            self.proto_min, self.proto_max, self.num_spec, self.window, self.hidden,
            self.n_taps, *ids, self.payload_kind, self.row_bytes, self.noise_token_id,
            self.vocab_size, self.max_model_len, self.max_slots, self.max_rows_per_step,
            self.sample_mode, self.tp_size, 0,
            _b32(self.drafter_identity), _b32(self.markov_identity),
            self.client_name.encode("ascii", "replace")[:48],
        )

    @classmethod
    def unpack(cls, body) -> "Hello":
        if len(body) != HELLO_REQ.size:
            raise WireError(f"HELLO body {len(body)} != {HELLO_REQ.size}")
        v = HELLO_REQ.unpack(bytes(body))
        n_taps = v[5]
        return cls(
            proto_min=v[0], proto_max=v[1], num_spec=v[2], window=v[3], hidden=v[4],
            n_taps=n_taps, tap_layer_ids=tuple(v[6:6 + n_taps]), payload_kind=v[10],
            row_bytes=v[11], noise_token_id=v[12], vocab_size=v[13], max_model_len=v[14],
            max_slots=v[15], max_rows_per_step=v[16], sample_mode=v[17], tp_size=v[18],
            drafter_identity=v[20], markov_identity=v[21],
            client_name=v[22].rstrip(b"\0").decode("ascii", "replace"),
        )

    def config_tuple(self) -> tuple:
        """Fields that client and worker must agree on (CONFIG_MISMATCH otherwise)."""
        return (self.num_spec, self.window, self.hidden, self.n_taps,
                tuple(self.tap_layer_ids), self.payload_kind, self.row_bytes,
                self.noise_token_id, self.vocab_size, self.sample_mode)


@dataclass
class HelloReply:
    proto: int = 1
    max_slots: int = 4
    max_rows_per_step: int = 512
    capabilities: int = CAP_TAPS_BF16
    window: int = 128
    num_spec: int = 5
    drafter_identity: bytes = bytes(32)
    markov_identity: bytes = bytes(32)
    worker_name: str = ""

    def pack(self) -> bytes:
        return HELLO_REP.pack(self.proto, self.max_slots, self.max_rows_per_step,
                              self.capabilities, self.window, self.num_spec,
                              _b32(self.drafter_identity), _b32(self.markov_identity),
                              self.worker_name.encode("ascii", "replace")[:40])

    @classmethod
    def unpack(cls, body) -> "HelloReply":
        if len(body) != HELLO_REP.size:
            raise WireError(f"HELLO reply body {len(body)} != {HELLO_REP.size}")
        v = HELLO_REP.unpack(bytes(body))
        return cls(proto=v[0], max_slots=v[1], max_rows_per_step=v[2], capabilities=v[3],
                   window=v[4], num_spec=v[5], drafter_identity=v[6], markov_identity=v[7],
                   worker_name=v[8].rstrip(b"\0").decode("ascii", "replace"))


def _b32(b: bytes) -> bytes:
    if len(b) != 32:
        raise WireError(f"identity must be 32 bytes, got {len(b)}")
    return b


# --------------------------------------------------------------------------- STEP


@dataclass
class ReqEntry:
    slot: int
    generation: int
    first_pos: int
    n_rows: int
    anchor_token: int
    anchor_pos: int
    flags: int = 0
    k_window: int = 0
    n_out: int = 0   # CONTRACT 1.3: draft positions the Spark consumes (0 = all N)

    def pack_into(self, buf, off: int) -> None:
        REQ_ENTRY.pack_into(buf, off, self.slot, self.flags, self.k_window, self.n_out,
                            self.generation & U32_MASK, self.first_pos, self.n_rows,
                            self.anchor_token, self.anchor_pos)

    @classmethod
    def unpack_from(cls, buf, off: int) -> "ReqEntry":
        slot, flags, k_window, n_out, gen, first_pos, n_rows, anchor_token, anchor_pos = (
            REQ_ENTRY.unpack_from(buf, off))
        return cls(slot=slot, generation=gen, first_pos=first_pos, n_rows=n_rows,
                   anchor_token=anchor_token, anchor_pos=anchor_pos, flags=flags,
                   k_window=k_window, n_out=n_out)


@dataclass
class StepRequest:
    """A decoded STEP body. ``rows`` is a [n_rows, row_bytes] uint8 view (no copy)."""
    step_id: int
    step_flags: int
    row_bytes: int
    entries: list[ReqEntry]
    rows: np.ndarray

    @property
    def n_req(self) -> int:
        return len(self.entries)


def step_body_len(n_req: int, n_rows: int, row_bytes: int) -> int:
    return rows_offset_for(n_req) + n_rows * row_bytes


def pack_step_prefix(buf, off: int, step_id: int, entries: list[ReqEntry],
                     n_rows: int, row_bytes: int, step_flags: int = 0, sample_meta=None) -> int:
    """Write StepHeader + entries + zero padding at ``buf[off:]`` (``off`` = start of the
    STEP body). Returns the offset of the row block relative to the body start. The
    caller places the rows at ``buf[off + rows_offset:]`` (e.g. a pinned D2H target)."""
    n_req = len(entries)
    ro = rows_offset_for(n_req, sample_meta is not None)
    STEP_HEADER.pack_into(buf, off, step_id, n_req, step_flags, n_rows, row_bytes, ro)
    p = off + STEP_HEADER.size
    total = 0
    for e in entries:
        e.pack_into(buf, p)
        p += REQ_ENTRY.size
        total += e.n_rows
    if total != n_rows:
        raise WireError(f"entries carry {total} rows, header says {n_rows}")
    pad_end = off + ro
    if pad_end > p:
        memoryview(buf).cast("B")[p:pad_end] = bytes(pad_end - p)
    if sample_meta is not None:
        from .sampled_wire import META
        if len(sample_meta) != n_req or not step_flags & SF_SAMPLED:
            raise WireError("missing sample flag/context")
        for i, (seed, mask) in enumerate(sample_meta):
            META.pack_into(buf, off + STEP_HEADER.size + n_req * REQ_ENTRY.size + i * META.size,
                           int(seed), int(mask))
    return ro


def encode_step(step_id: int, entries: list[ReqEntry], rows: np.ndarray | None,
                row_bytes: int, step_flags: int = 0) -> bytes:
    """Convenience (tests, mock): full STEP body as bytes. ``rows`` is [n, row_bytes]
    uint8 (or any array whose bytes are the rows in order)."""
    n_rows = sum(e.n_rows for e in entries)
    body = bytearray(step_body_len(len(entries), n_rows, row_bytes))
    ro = pack_step_prefix(body, 0, step_id, entries, n_rows, row_bytes, step_flags)
    if n_rows:
        raw = np.ascontiguousarray(rows).view(np.uint8).reshape(-1)
        if raw.size != n_rows * row_bytes:
            raise WireError(f"rows hold {raw.size} bytes, expected {n_rows * row_bytes}")
        body[ro:ro + raw.size] = raw.tobytes()
    return bytes(body)


def decode_step(body, *, max_slots: int, max_rows: int, expected_row_bytes: int,
                window: int) -> StepRequest:
    mv = memoryview(body)
    if len(mv) < STEP_HEADER.size:
        raise WireError("STEP body shorter than its header")
    step_id, n_req, step_flags, n_rows, row_bytes, ro = STEP_HEADER.unpack_from(mv, 0)
    if row_bytes != expected_row_bytes:
        raise WireError(f"row_bytes {row_bytes} != {expected_row_bytes}")
    if not 1 <= n_req <= max_slots:
        raise WireError(f"n_req {n_req} outside 1..{max_slots}")
    if n_rows > max_rows:
        raise WireError(f"n_rows {n_rows} > max_rows_per_step {max_rows}")
    if ro != rows_offset_for(n_req, bool(step_flags & SF_SAMPLED)):
        raise WireError(f"rows_offset {ro} != {rows_offset_for(n_req)}")
    if len(mv) != ro + n_rows * row_bytes:
        raise WireError(f"STEP body {len(mv)} != {ro + n_rows * row_bytes}")
    entries = [ReqEntry.unpack_from(mv, STEP_HEADER.size + i * REQ_ENTRY.size)
               for i in range(n_req)]
    seen = set()
    total = 0
    for e in entries:
        if e.slot >= max_slots or e.slot in seen:
            raise WireError(f"slot {e.slot} invalid or repeated")
        seen.add(e.slot)
        total += e.n_rows
    if total != n_rows:
        raise WireError(f"entries carry {total} rows, header says {n_rows}")
    rows = np.frombuffer(mv, dtype=np.uint8, count=n_rows * row_bytes, offset=ro)
    rows = rows.reshape(n_rows, row_bytes) if n_rows else rows.reshape(0, row_bytes)
    return StepRequest(step_id=step_id, step_flags=step_flags, row_bytes=row_bytes,
                       entries=entries, rows=rows)


@dataclass
class StepReply:
    step_id: int
    num_spec: int
    status: np.ndarray        # [n_req] uint16
    window_rows: np.ndarray   # [n_req] uint16
    generation: np.ndarray    # [n_req] uint32
    draft_ids: np.ndarray     # [n_req, N] int32
    confidence: np.ndarray    # [n_req, N] float32
    margin: np.ndarray | None = None  # [n_req, N] float32
    worker_us: int = 0
    recv_to_send_us: int = 0
    reply_flags: int = 0

    support_ids: np.ndarray | None = None
    proposal_q: np.ndarray | None = None
    sampled_mask: np.ndarray | None = None

    @property
    def n_req(self) -> int:
        return int(self.status.shape[0])

    def pack(self) -> bytes:
        n, N = self.n_req, self.num_spec
        flags = self.reply_flags | (RPF_HAS_MARGIN if self.margin is not None else 0)
        if self.support_ids is not None:
            flags |= RPF_SAMPLED
        parts = [REPLY_HEADER.pack(self.step_id, n, N, flags, self.worker_us,
                                   self.recv_to_send_us)]
        st = np.zeros(n, dtype=np.dtype([("s", "<u2"), ("w", "<u2"), ("g", "<u4")]))
        st["s"], st["w"], st["g"] = self.status, self.window_rows, self.generation
        parts.append(st.tobytes())
        parts.append(np.ascontiguousarray(self.draft_ids, dtype="<i4").reshape(n, N).tobytes())
        parts.append(np.ascontiguousarray(self.confidence, dtype="<f4").reshape(n, N).tobytes())
        if self.margin is not None:
            parts.append(np.ascontiguousarray(self.margin, dtype="<f4").reshape(n, N).tobytes())
        body = b"".join(parts)
        if flags & RPF_SAMPLED:
            from .sampled_wire import append_q
            if self.support_ids is None or self.proposal_q is None or self.sampled_mask is None:
                raise WireError("sampled reply requires support/q/mask")
            body = append_q(body, self.sampled_mask, self.support_ids, self.proposal_q)
        return body

    @classmethod
    def unpack(cls, body, *, expect_n_req: int | None = None,
               expect_num_spec: int | None = None) -> "StepReply":
        mv = memoryview(body)
        if len(mv) < REPLY_HEADER.size:
            raise WireError("STEP reply shorter than its header")
        step_id, n, N, flags, worker_us, r2s = REPLY_HEADER.unpack_from(mv, 0)
        if expect_n_req is not None and n != expect_n_req:
            raise WireError(f"reply n_req {n} != {expect_n_req}")
        if expect_num_spec is not None and N != expect_num_spec:
            raise WireError(f"reply num_spec {N} != {expect_num_spec}")
        has_margin = bool(flags & RPF_HAS_MARGIN)
        need = REPLY_HEADER.size + 8 * n + 8 * n * N + (4 * n * N if has_margin else 0)
        support_ids = proposal_q = sampled_mask = None
        if flags & RPF_SAMPLED:
            from .sampled_wire import unpack_q
            try:
                sampled_mask, support_ids, proposal_q = unpack_q(mv, need, n, N)
            except ValueError as exc:
                raise WireError(str(exc)) from exc
        elif len(mv) != need:
            raise WireError(f"STEP reply {len(mv)} bytes, expected {need}")
        o = REPLY_HEADER.size
        st = np.frombuffer(mv, dtype=np.dtype([("s", "<u2"), ("w", "<u2"), ("g", "<u4")]),
                           count=n, offset=o)
        o += 8 * n
        ids = np.frombuffer(mv, dtype="<i4", count=n * N, offset=o).reshape(n, N)
        o += 4 * n * N
        conf = np.frombuffer(mv, dtype="<f4", count=n * N, offset=o).reshape(n, N)
        o += 4 * n * N
        margin = None
        if has_margin:
            margin = np.frombuffer(mv, dtype="<f4", count=n * N, offset=o).reshape(n, N)
        return cls(step_id=step_id, num_spec=N, status=st["s"].copy(),
                   window_rows=st["w"].copy(), generation=st["g"].copy(),
                   draft_ids=ids.copy(), confidence=conf.copy(),
                   margin=None if margin is None else margin.copy(),
                   worker_us=worker_us, recv_to_send_us=r2s, reply_flags=flags,
                   support_ids=support_ids, proposal_q=proposal_q, sampled_mask=sampled_mask)


# --------------------------------------------------------------------------- frames


def frame(op: int, seq: int, session: int, body: bytes = b"", status: int = 0) -> bytes:
    return AppHeader(op=op, seq=seq, session=session, body_len=len(body),
                     status=status).pack() + body


def split_frame(buf) -> tuple[AppHeader, memoryview]:
    mv = memoryview(buf)
    h = AppHeader.unpack_from(mv)
    h.check(len(mv))
    return h, mv[APP_HEADER_SIZE:]


# --------------------------------------------------------------------------- identity


def drafter_identity(file_hashes: dict[str, str]) -> bytes:
    """CONTRACT 6.7: sha256 over "<basename> <sha256-hex>\\n" sorted by basename."""
    lines = []
    for name in sorted(file_hashes):
        h = file_hashes[name].strip().lower()
        if len(h) != 64 or any(c not in "0123456789abcdef" for c in h):
            raise ValueError(f"{name}: not a sha256 hex digest: {h!r}")
        if "/" in name:
            raise ValueError(f"{name}: basenames only")
        lines.append(f"{name} {h}\n")
    return hashlib.sha256("".join(lines).encode()).digest()


def parse_hex32(text: str, what: str) -> bytes:
    t = (text or "").strip().lower()
    if len(t) != 64:
        raise ValueError(f"{what}: expected 64 hex characters, got {len(t)}")
    return bytes.fromhex(t)


@dataclass
class Counters:
    """Small helper for per-reason counters (shared by client and mock worker)."""
    values: dict[str, int] = field(default_factory=dict)

    def inc(self, key: str, n: int = 1) -> None:
        self.values[key] = self.values.get(key, 0) + n

	# SPDX-License-Identifier: Apache-2.0
	"""Wire codec for the remote DSpark contract (offload/dspark/CONTRACT.md, v1.0).

	Pure Python + numpy, no torch and no vLLM: shared by the Spark proxy, the mock worker
	and the offline tests. Every layout here mirrors a table in CONTRACT.md section 6; keep
	them in step (and bump CONTRACT_VERSION / the change log) on any change.
	"""

	from __future__ import annotations

	import hashlib
	import struct
	from dataclasses import dataclass, field

	try: # numpy is needed for STEP bodies; HELLO/PING/STATS framing works without it
	import numpy as np # (the probe tool runs on Spark hosts whose python3 has no numpy)
	except ImportError: # pragma: no cover
	np = None # type: ignore[assignment]

	CONTRACT_VERSION = "1.8"

	MAGIC = 0x31505344 # bytes "DSP1" little-endian
	VERSION = 1

	OP_HELLO = 1
	OP_STEP = 2
	OP_RESET = 3
	OP_STATS = 4
	OP_PING = 5
	OP_NAMES = {OP_HELLO: "HELLO", OP_STEP: "STEP", OP_RESET: "RESET",
	OP_STATS: "STATS", OP_PING: "PING"}

	# Frame status (CONTRACT 9.1)
	ST_OK = 0
	ST_BAD_VERSION = 1
	ST_IDENTITY_MISMATCH = 2
	ST_CONFIG_MISMATCH = 3
	ST_NO_SESSION = 4
	ST_STALE_SEQ = 5
	ST_MALFORMED = 6
	ST_CAPACITY = 7
	ST_INTERNAL = 8
	ST_UNSUPPORTED_OP = 9
	FRAME_STATUS_NAMES = {
	ST_OK: "OK", ST_BAD_VERSION: "BAD_VERSION", ST_IDENTITY_MISMATCH: "IDENTITY_MISMATCH",
	ST_CONFIG_MISMATCH: "CONFIG_MISMATCH", ST_NO_SESSION: "NO_SESSION",
	ST_STALE_SEQ: "STALE_SEQ", ST_MALFORMED: "MALFORMED", ST_CAPACITY: "CAPACITY",
	ST_INTERNAL: "INTERNAL", ST_UNSUPPORTED_OP: "UNSUPPORTED_OP",
	}
	# Frame errors after which the remote must not be retried (configuration errors).
	FATAL_FRAME_STATUS = frozenset(
	{ST_BAD_VERSION, ST_IDENTITY_MISMATCH, ST_CONFIG_MISMATCH, ST_UNSUPPORTED_OP}
	)

	# Per-request status (CONTRACT 9.2)
	RS_OK = 0
	RS_PARTIAL_WINDOW = 1
	RS_STALE_GENERATION = 2
	RS_BAD_ENTRY = 3
	RS_INTERNAL = 4
	REQ_STATUS_NAMES = {RS_OK: "OK", RS_PARTIAL_WINDOW: "PARTIAL_WINDOW",
	RS_STALE_GENERATION: "STALE_GENERATION", RS_BAD_ENTRY: "BAD_ENTRY",
	RS_INTERNAL: "INTERNAL"}
	# Per-request statuses whose drafts may be used.
	USABLE_REQ_STATUS = frozenset({RS_OK, RS_PARTIAL_WINDOW})

	# ReqEntry.req_flags
	RF_NEW_GEN = 0x1
	RF_NEED_DRAFT = 0x2
	RF_RESYNC = 0x4

	# StepHeader.step_flags / ReplyHeader.reply_flags
	SF_WANT_MARGIN = 0x1
	RPF_HAS_MARGIN = 0x1

	PAYLOAD_TAPS_BF16 = 1
	PAYLOAD_MAINX_BF16 = 2
	SAMPLE_GREEDY = 1
	SAMPLE_TOP32 = 2
	CAP_TOP32 = 8
	SF_SAMPLED = 2
	RPF_SAMPLED = 2

	CAP_TAPS_BF16 = 0x1
	CAP_MAINX_BF16 = 0x2
	CAP_HAS_MARGIN = 0x4

	APP_HEADER = struct.Struct("<IHHQIIiI") # 32 B
	HELLO_REQ = struct.Struct("<20I32s32s48s") # 192 B
	HELLO_REP = struct.Struct("<6I32s32s40s") # 128 B
	STEP_HEADER = struct.Struct("<QHHIII8x") # 32 B
	REQ_ENTRY = struct.Struct("<BBBBIqIiq") # 32 B
	REPLY_HEADER = struct.Struct("<QHHIII8x") # 32 B
	REQ_STATUS = struct.Struct("<HHI") # 8 B
	RESET_BODY = struct.Struct("<I") # 4 B

	assert APP_HEADER.size == 32
	assert HELLO_REQ.size == 192
	assert HELLO_REP.size == 128
	assert STEP_HEADER.size == 32
	assert REQ_ENTRY.size == 32
	assert REPLY_HEADER.size == 32
	assert REQ_STATUS.size == 8

	APP_HEADER_SIZE = APP_HEADER.size
	U32_MASK = 0xFFFFFFFF


	class WireError(ValueError):
	"""A malformed or inconsistent message."""


	def align64(n: int) -> int:
	return (n + 63) & ~63


	def rows_offset_for(n_req: int, sampled: bool = False) -> int:
	"""Offset of the row block from the start of a STEP body (multiple of 64)."""
	return align64(STEP_HEADER.size + n_req * REQ_ENTRY.size + (16 * n_req if sampled else 0))


	def gen_newer(a: int, b: int) -> bool:
	"""Serial-number arithmetic on u32 generations: is a newer than b?"""
	d = (a - b) & U32_MASK
	return d != 0 and d < 0x80000000


	# --------------------------------------------------------------------------- headers


	@dataclass
	class AppHeader:
	op: int
	seq: int
	session: int
	body_len: int
	status: int = 0
	flags: int = 0
	magic: int = MAGIC
	version: int = VERSION

	def pack_into(self, buf, off: int = 0) -> None:
	APP_HEADER.pack_into(buf, off, self.magic, self.version, self.op, self.seq,
	self.session, self.body_len, self.status, self.flags)

	def pack(self) -> bytes:
	b = bytearray(APP_HEADER_SIZE)
	self.pack_into(b)
	return bytes(b)

	@classmethod
	def unpack_from(cls, buf, off: int = 0) -> "AppHeader":
	if len(buf) - off < APP_HEADER_SIZE:
	raise WireError(f"short frame: {len(buf) - off} < {APP_HEADER_SIZE}")
	magic, version, op, seq, session, body_len, status, flags = APP_HEADER.unpack_from(
	buf, off)
	h = cls(op=op, seq=seq, session=session, body_len=body_len, status=status,
	flags=flags, magic=magic, version=version)
	return h

	def check(self, frame_len: int) -> None:
	if self.magic != MAGIC:
	raise WireError(f"bad magic 0x{self.magic:08x}")
	if self.version != VERSION:
	raise WireError(f"unsupported version {self.version}")
	if APP_HEADER_SIZE + self.body_len != frame_len:
	raise WireError(
	f"body_len {self.body_len} does not match frame length {frame_len}")


	# --------------------------------------------------------------------------- HELLO


	@dataclass
	class Hello:
	num_spec: int = 5
	window: int = 128
	hidden: int = 5120
	n_taps: int = 3
	tap_layer_ids: tuple[int, ...] = (37, 38, 39)
	payload_kind: int = PAYLOAD_TAPS_BF16
	row_bytes: int = 3 * 5120 * 2
	noise_token_id: int = 128799
	vocab_size: int = 129280
	max_model_len: int = 327680
	max_slots: int = 4
	max_rows_per_step: int = 4 * 128
	sample_mode: int = SAMPLE_GREEDY
	tp_size: int = 2
	drafter_identity: bytes = bytes(32)
	markov_identity: bytes = bytes(32)
	client_name: str = ""
	proto_min: int = 1
	proto_max: int = 1

	def pack(self) -> bytes:
	ids = list(self.tap_layer_ids) + [0] * (4 - len(self.tap_layer_ids))
	if len(ids) != 4:
	raise WireError("at most 4 tap layers")
	return HELLO_REQ.pack(
	self.proto_min, self.proto_max, self.num_spec, self.window, self.hidden,
	self.n_taps, *ids, self.payload_kind, self.row_bytes, self.noise_token_id,
	self.vocab_size, self.max_model_len, self.max_slots, self.max_rows_per_step,
	self.sample_mode, self.tp_size, 0,
	_b32(self.drafter_identity), _b32(self.markov_identity),
	self.client_name.encode("ascii", "replace")[:48],
	)

	@classmethod
	def unpack(cls, body) -> "Hello":
	if len(body) != HELLO_REQ.size:
	raise WireError(f"HELLO body {len(body)} != {HELLO_REQ.size}")
	v = HELLO_REQ.unpack(bytes(body))
	n_taps = v[5]
	return cls(
	proto_min=v[0], proto_max=v[1], num_spec=v[2], window=v[3], hidden=v[4],
	n_taps=n_taps, tap_layer_ids=tuple(v[6:6 + n_taps]), payload_kind=v[10],
	row_bytes=v[11], noise_token_id=v[12], vocab_size=v[13], max_model_len=v[14],
	max_slots=v[15], max_rows_per_step=v[16], sample_mode=v[17], tp_size=v[18],
	drafter_identity=v[20], markov_identity=v[21],
	client_name=v[22].rstrip(b"\0").decode("ascii", "replace"),
	)

	def config_tuple(self) -> tuple:
	"""Fields that client and worker must agree on (CONFIG_MISMATCH otherwise)."""
	return (self.num_spec, self.window, self.hidden, self.n_taps,
	tuple(self.tap_layer_ids), self.payload_kind, self.row_bytes,
	self.noise_token_id, self.vocab_size, self.sample_mode)


	@dataclass
	class HelloReply:
	proto: int = 1
	max_slots: int = 4
	max_rows_per_step: int = 512
	capabilities: int = CAP_TAPS_BF16
	window: int = 128
	num_spec: int = 5
	drafter_identity: bytes = bytes(32)
	markov_identity: bytes = bytes(32)
	worker_name: str = ""

	def pack(self) -> bytes:
	return HELLO_REP.pack(self.proto, self.max_slots, self.max_rows_per_step,
	self.capabilities, self.window, self.num_spec,
	_b32(self.drafter_identity), _b32(self.markov_identity),
	self.worker_name.encode("ascii", "replace")[:40])

	@classmethod
	def unpack(cls, body) -> "HelloReply":
	if len(body) != HELLO_REP.size:
	raise WireError(f"HELLO reply body {len(body)} != {HELLO_REP.size}")
	v = HELLO_REP.unpack(bytes(body))
	return cls(proto=v[0], max_slots=v[1], max_rows_per_step=v[2], capabilities=v[3],
	window=v[4], num_spec=v[5], drafter_identity=v[6], markov_identity=v[7],
	worker_name=v[8].rstrip(b"\0").decode("ascii", "replace"))


	def _b32(b: bytes) -> bytes:
	if len(b) != 32:
	raise WireError(f"identity must be 32 bytes, got {len(b)}")
	return b


	# --------------------------------------------------------------------------- STEP


	@dataclass
	class ReqEntry:
	slot: int
	generation: int
	first_pos: int
	n_rows: int
	anchor_token: int
	anchor_pos: int
	flags: int = 0
	k_window: int = 0
	n_out: int = 0 # CONTRACT 1.3: draft positions the Spark consumes (0 = all N)

	def pack_into(self, buf, off: int) -> None:
	REQ_ENTRY.pack_into(buf, off, self.slot, self.flags, self.k_window, self.n_out,
	self.generation & U32_MASK, self.first_pos, self.n_rows,
	self.anchor_token, self.anchor_pos)

	@classmethod
	def unpack_from(cls, buf, off: int) -> "ReqEntry":
	slot, flags, k_window, n_out, gen, first_pos, n_rows, anchor_token, anchor_pos = (
	REQ_ENTRY.unpack_from(buf, off))
	return cls(slot=slot, generation=gen, first_pos=first_pos, n_rows=n_rows,
	anchor_token=anchor_token, anchor_pos=anchor_pos, flags=flags,
	k_window=k_window, n_out=n_out)


	@dataclass
	class StepRequest:
	"""A decoded STEP body. ``rows`` is a [n_rows, row_bytes] uint8 view (no copy)."""
	step_id: int
	step_flags: int
	row_bytes: int
	entries: list[ReqEntry]
	rows: np.ndarray

	@property
	def n_req(self) -> int:
	return len(self.entries)


	def step_body_len(n_req: int, n_rows: int, row_bytes: int) -> int:
	return rows_offset_for(n_req) + n_rows * row_bytes


	def pack_step_prefix(buf, off: int, step_id: int, entries: list[ReqEntry],
	n_rows: int, row_bytes: int, step_flags: int = 0, sample_meta=None) -> int:
	"""Write StepHeader + entries + zero padding at ``buf[off:]`` (``off`` = start of the
	STEP body). Returns the offset of the row block relative to the body start. The
	caller places the rows at ``buf[off + rows_offset:]`` (e.g. a pinned D2H target)."""
	n_req = len(entries)
	ro = rows_offset_for(n_req, sample_meta is not None)
	STEP_HEADER.pack_into(buf, off, step_id, n_req, step_flags, n_rows, row_bytes, ro)
	p = off + STEP_HEADER.size
	total = 0
	for e in entries:
	e.pack_into(buf, p)
	p += REQ_ENTRY.size
	total += e.n_rows
	if total != n_rows:
	raise WireError(f"entries carry {total} rows, header says {n_rows}")
	pad_end = off + ro
	if pad_end > p:
	memoryview(buf).cast("B")[p:pad_end] = bytes(pad_end - p)
	if sample_meta is not None:
	from .sampled_wire import META
	if len(sample_meta) != n_req or not step_flags & SF_SAMPLED:
	raise WireError("missing sample flag/context")
	for i, (seed, mask) in enumerate(sample_meta):
	META.pack_into(buf, off + STEP_HEADER.size + n_req * REQ_ENTRY.size + i * META.size,
	int(seed), int(mask))
	return ro


	def encode_step(step_id: int, entries: list[ReqEntry], rows: np.ndarray \| None,
	row_bytes: int, step_flags: int = 0) -> bytes:
	"""Convenience (tests, mock): full STEP body as bytes. ``rows`` is [n, row_bytes]
	uint8 (or any array whose bytes are the rows in order)."""
	n_rows = sum(e.n_rows for e in entries)
	body = bytearray(step_body_len(len(entries), n_rows, row_bytes))
	ro = pack_step_prefix(body, 0, step_id, entries, n_rows, row_bytes, step_flags)
	if n_rows:
	raw = np.ascontiguousarray(rows).view(np.uint8).reshape(-1)
	if raw.size != n_rows * row_bytes:
	raise WireError(f"rows hold {raw.size} bytes, expected {n_rows * row_bytes}")
	body[ro:ro + raw.size] = raw.tobytes()
	return bytes(body)


	def decode_step(body, *, max_slots: int, max_rows: int, expected_row_bytes: int,
	window: int) -> StepRequest:
	mv = memoryview(body)
	if len(mv) < STEP_HEADER.size:
	raise WireError("STEP body shorter than its header")
	step_id, n_req, step_flags, n_rows, row_bytes, ro = STEP_HEADER.unpack_from(mv, 0)
	if row_bytes != expected_row_bytes:
	raise WireError(f"row_bytes {row_bytes} != {expected_row_bytes}")
	if not 1 <= n_req <= max_slots:
	raise WireError(f"n_req {n_req} outside 1..{max_slots}")
	if n_rows > max_rows:
	raise WireError(f"n_rows {n_rows} > max_rows_per_step {max_rows}")
	if ro != rows_offset_for(n_req, bool(step_flags & SF_SAMPLED)):
	raise WireError(f"rows_offset {ro} != {rows_offset_for(n_req)}")
	if len(mv) != ro + n_rows * row_bytes:
	raise WireError(f"STEP body {len(mv)} != {ro + n_rows * row_bytes}")
	entries = [ReqEntry.unpack_from(mv, STEP_HEADER.size + i * REQ_ENTRY.size)
	for i in range(n_req)]
	seen = set()
	total = 0
	for e in entries:
	if e.slot >= max_slots or e.slot in seen:
	raise WireError(f"slot {e.slot} invalid or repeated")
	seen.add(e.slot)
	total += e.n_rows
	if total != n_rows:
	raise WireError(f"entries carry {total} rows, header says {n_rows}")
	rows = np.frombuffer(mv, dtype=np.uint8, count=n_rows * row_bytes, offset=ro)
	rows = rows.reshape(n_rows, row_bytes) if n_rows else rows.reshape(0, row_bytes)
	return StepRequest(step_id=step_id, step_flags=step_flags, row_bytes=row_bytes,
	entries=entries, rows=rows)


	@dataclass
	class StepReply:
	step_id: int
	num_spec: int
	status: np.ndarray # [n_req] uint16
	window_rows: np.ndarray # [n_req] uint16
	generation: np.ndarray # [n_req] uint32
	draft_ids: np.ndarray # [n_req, N] int32
	confidence: np.ndarray # [n_req, N] float32
	margin: np.ndarray \| None = None # [n_req, N] float32
	worker_us: int = 0
	recv_to_send_us: int = 0
	reply_flags: int = 0

	support_ids: np.ndarray \| None = None
	proposal_q: np.ndarray \| None = None
	sampled_mask: np.ndarray \| None = None

	@property
	def n_req(self) -> int:
	return int(self.status.shape[0])

	def pack(self) -> bytes:
	n, N = self.n_req, self.num_spec
	flags = self.reply_flags \| (RPF_HAS_MARGIN if self.margin is not None else 0)
	if self.support_ids is not None:
	flags \|= RPF_SAMPLED
	parts = [REPLY_HEADER.pack(self.step_id, n, N, flags, self.worker_us,
	self.recv_to_send_us)]
	st = np.zeros(n, dtype=np.dtype([("s", "<u2"), ("w", "<u2"), ("g", "<u4")]))
	st["s"], st["w"], st["g"] = self.status, self.window_rows, self.generation
	parts.append(st.tobytes())
	parts.append(np.ascontiguousarray(self.draft_ids, dtype="<i4").reshape(n, N).tobytes())
	parts.append(np.ascontiguousarray(self.confidence, dtype="<f4").reshape(n, N).tobytes())
	if self.margin is not None:
	parts.append(np.ascontiguousarray(self.margin, dtype="<f4").reshape(n, N).tobytes())
	body = b"".join(parts)
	if flags & RPF_SAMPLED:
	from .sampled_wire import append_q
	if self.support_ids is None or self.proposal_q is None or self.sampled_mask is None:
	raise WireError("sampled reply requires support/q/mask")
	body = append_q(body, self.sampled_mask, self.support_ids, self.proposal_q)
	return body

	@classmethod
	def unpack(cls, body, *, expect_n_req: int \| None = None,
	expect_num_spec: int \| None = None) -> "StepReply":
	mv = memoryview(body)
	if len(mv) < REPLY_HEADER.size:
	raise WireError("STEP reply shorter than its header")
	step_id, n, N, flags, worker_us, r2s = REPLY_HEADER.unpack_from(mv, 0)
	if expect_n_req is not None and n != expect_n_req:
	raise WireError(f"reply n_req {n} != {expect_n_req}")
	if expect_num_spec is not None and N != expect_num_spec:
	raise WireError(f"reply num_spec {N} != {expect_num_spec}")
	has_margin = bool(flags & RPF_HAS_MARGIN)
	need = REPLY_HEADER.size + 8 * n + 8 * n * N + (4 * n * N if has_margin else 0)
	support_ids = proposal_q = sampled_mask = None
	if flags & RPF_SAMPLED:
	from .sampled_wire import unpack_q
	try:
	sampled_mask, support_ids, proposal_q = unpack_q(mv, need, n, N)
	except ValueError as exc:
	raise WireError(str(exc)) from exc
	elif len(mv) != need:
	raise WireError(f"STEP reply {len(mv)} bytes, expected {need}")
	o = REPLY_HEADER.size
	st = np.frombuffer(mv, dtype=np.dtype([("s", "<u2"), ("w", "<u2"), ("g", "<u4")]),
	count=n, offset=o)
	o += 8 * n
	ids = np.frombuffer(mv, dtype="<i4", count=n * N, offset=o).reshape(n, N)
	o += 4 * n * N
	conf = np.frombuffer(mv, dtype="<f4", count=n * N, offset=o).reshape(n, N)
	o += 4 * n * N
	margin = None
	if has_margin:
	margin = np.frombuffer(mv, dtype="<f4", count=n * N, offset=o).reshape(n, N)
	return cls(step_id=step_id, num_spec=N, status=st["s"].copy(),
	window_rows=st["w"].copy(), generation=st["g"].copy(),
	draft_ids=ids.copy(), confidence=conf.copy(),
	margin=None if margin is None else margin.copy(),
	worker_us=worker_us, recv_to_send_us=r2s, reply_flags=flags,
	support_ids=support_ids, proposal_q=proposal_q, sampled_mask=sampled_mask)


	# --------------------------------------------------------------------------- frames


	def frame(op: int, seq: int, session: int, body: bytes = b"", status: int = 0) -> bytes:
	return AppHeader(op=op, seq=seq, session=session, body_len=len(body),
	status=status).pack() + body


	def split_frame(buf) -> tuple[AppHeader, memoryview]:
	mv = memoryview(buf)
	h = AppHeader.unpack_from(mv)
	h.check(len(mv))
	return h, mv[APP_HEADER_SIZE:]


	# --------------------------------------------------------------------------- identity


	def drafter_identity(file_hashes: dict[str, str]) -> bytes:
	"""CONTRACT 6.7: sha256 over "<basename> <sha256-hex>\\n" sorted by basename."""
	lines = []
	for name in sorted(file_hashes):
	h = file_hashes[name].strip().lower()
	if len(h) != 64 or any(c not in "0123456789abcdef" for c in h):
	raise ValueError(f"{name}: not a sha256 hex digest: {h!r}")
	if "/" in name:
	raise ValueError(f"{name}: basenames only")
	lines.append(f"{name} {h}\n")
	return hashlib.sha256("".join(lines).encode()).digest()


	def parse_hex32(text: str, what: str) -> bytes:
	t = (text or "").strip().lower()
	if len(t) != 64:
	raise ValueError(f"{what}: expected 64 hex characters, got {len(t)}")
	return bytes.fromhex(t)


	@dataclass
	class Counters:
	"""Small helper for per-reason counters (shared by client and mock worker)."""
	values: dict[str, int] = field(default_factory=dict)

	def inc(self, key: str, n: int = 1) -> None:
	self.values[key] = self.values.get(key, 0) + n