| """Wire codec of offload/dspark/CONTRACT.md v1.3 (lane O3): little-endian, offsets as in the contract's tables. |
| |
| Pure Python + numpy views; no torch. Every decode checks lengths and raises ProtocolError (-> frame MALFORMED).""" |
|
|
| from __future__ import annotations |
|
|
| import struct |
| from dataclasses import dataclass |
|
|
| import numpy as np |
|
|
| MAGIC = 0x31505344 |
| VERSION = 1 |
| OP_HELLO, OP_STEP, OP_RESET, OP_STATS, OP_PING = 1, 2, 3, 4, 5 |
|
|
| |
| OK, BAD_VERSION, IDENTITY_MISMATCH, CONFIG_MISMATCH, NO_SESSION, STALE_SEQ, MALFORMED, CAPACITY, INTERNAL, \ |
| UNSUPPORTED_OP = range(10) |
| FRAME_STATUS = {0: "OK", 1: "BAD_VERSION", 2: "IDENTITY_MISMATCH", 3: "CONFIG_MISMATCH", 4: "NO_SESSION", |
| 5: "STALE_SEQ", 6: "MALFORMED", 7: "CAPACITY", 8: "INTERNAL", 9: "UNSUPPORTED_OP"} |
| |
| R_OK, R_PARTIAL_WINDOW, R_STALE_GENERATION, R_BAD_ENTRY, R_INTERNAL = range(5) |
|
|
| |
| PAYLOAD_TAPS_BF16, PAYLOAD_MAINX_BF16 = 1, 2 |
| SAMPLE_GREEDY = 1 |
| SAMPLE_TOP32, CAP_TOP32, STEP_SAMPLED, REPLY_SAMPLED = 2, 8, 2, 2 |
| SAMPLE_META = struct.Struct("<qI4x") |
| CAP_TAPS_BF16, CAP_MAINX_BF16, CAP_HAS_MARGIN = 1, 2, 4 |
| |
| STEP_WANT_MARGIN = 1 |
| REQ_NEW_GEN, REQ_NEED_DRAFT, REQ_RESYNC = 1, 2, 4 |
| REPLY_HAS_MARGIN = 1 |
|
|
| APP_HDR = struct.Struct("<IHHQIIiI") |
| HELLO_REQ = struct.Struct("<6I4I10I32s32s48s") |
| HELLO_REP = struct.Struct("<6I32s32s40s") |
| STEP_HDR = struct.Struct("<QHHIII8x") |
| REQ_ENTRY = struct.Struct("<BBBBIqIiq") |
| REPLY_HDR = struct.Struct("<QHHIII8x") |
| REQ_STATUS = struct.Struct("<HHI") |
|
|
| assert APP_HDR.size == 32 and HELLO_REQ.size == 192 and HELLO_REP.size == 128 |
| assert STEP_HDR.size == 32 and REQ_ENTRY.size == 32 and REPLY_HDR.size == 32 and REQ_STATUS.size == 8 |
|
|
|
|
| class ProtocolError(ValueError): |
| pass |
|
|
|
|
| @dataclass |
| class AppHeader: |
| magic: int |
| version: int |
| op: int |
| seq: int |
| session: int |
| body_len: int |
| status: int |
| flags: int |
|
|
| @classmethod |
| def decode(cls, buf) -> "AppHeader": |
| if len(buf) < APP_HDR.size: |
| raise ProtocolError(f"frame of {len(buf)} B is shorter than the 32 B header") |
| return cls(*APP_HDR.unpack_from(buf, 0)) |
|
|
| def encode(self) -> bytes: |
| return APP_HDR.pack(self.magic, self.version, self.op, self.seq, self.session, self.body_len, self.status, |
| self.flags) |
|
|
|
|
| def reply(req: AppHeader, status: int, body: bytes = b"") -> bytes: |
| return AppHeader(MAGIC, VERSION, req.op, req.seq, req.session, len(body), status, 0).encode() + body |
|
|
|
|
| @dataclass |
| class Hello: |
| proto_min: int |
| proto_max: int |
| num_spec: int |
| window: int |
| hidden: int |
| n_taps: int |
| tap_layer_ids: tuple[int, int, int, int] |
| payload_kind: int |
| row_bytes: int |
| noise_token_id: int |
| vocab_size: int |
| max_model_len: int |
| max_slots: int |
| max_rows_per_step: int |
| sample_mode: int |
| tp_size: int |
| drafter_identity: bytes |
| markov_identity: bytes |
| client_name: str |
|
|
| @classmethod |
| def decode(cls, body) -> "Hello": |
| if len(body) != HELLO_REQ.size: |
| raise ProtocolError(f"HELLO body {len(body)} B, want {HELLO_REQ.size}") |
| v = HELLO_REQ.unpack_from(body, 0) |
| return cls(v[0], v[1], v[2], v[3], v[4], v[5], tuple(v[6:10]), v[10], v[11], v[12], v[13], v[14], v[15], |
| v[16], v[17], v[18], v[20], v[21], v[22].rstrip(b"\0").decode("ascii", "replace")) |
|
|
| def encode(self) -> bytes: |
| return HELLO_REQ.pack(self.proto_min, self.proto_max, self.num_spec, self.window, self.hidden, self.n_taps, |
| *self.tap_layer_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, self.drafter_identity, self.markov_identity, |
| self.client_name.encode()[:48]) |
|
|
|
|
| def hello_reply(proto: int, max_slots: int, max_rows: int, caps: int, window: int, num_spec: int, |
| drafter_identity: bytes, markov_identity: bytes, name: str) -> bytes: |
| return HELLO_REP.pack(proto, max_slots, max_rows, caps, window, num_spec, drafter_identity, markov_identity, |
| name.encode()[:40]) |
|
|
|
|
| @dataclass |
| class ReqEntry: |
| slot: int |
| req_flags: int |
| k_window: int |
| generation: int |
| first_pos: int |
| n_rows: int |
| anchor_token: int |
| anchor_pos: int |
| n_out: int = 0 |
| row_start: int = 0 |
|
|
|
|
| @dataclass |
| class Step: |
| step_id: int |
| n_req: int |
| step_flags: int |
| n_rows: int |
| row_bytes: int |
| rows_offset: int |
| entries: list[ReqEntry] |
| sample_meta: list[tuple[int, int]] |
| rows: np.ndarray |
|
|
|
|
| def decode_step(body, expected_row_bytes: int, max_slots: int, max_rows: int) -> Step: |
| """STEP request body -> Step. Raises ProtocolError for anything structurally wrong (frame MALFORMED); |
| CAPACITY problems raise CapacityError.""" |
| if len(body) < STEP_HDR.size: |
| raise ProtocolError("STEP body shorter than its 32 B header") |
| step_id, n_req, step_flags, n_rows, row_bytes, rows_offset = STEP_HDR.unpack_from(body, 0) |
| if n_req < 1: |
| raise ProtocolError("STEP with no request") |
| if n_req > max_slots or n_rows > max_rows: |
| raise CapacityError(f"STEP with {n_req} requests / {n_rows} rows (limits {max_slots} / {max_rows})") |
| if row_bytes != expected_row_bytes: |
| raise ProtocolError(f"row_bytes {row_bytes}, HELLO said {expected_row_bytes}") |
| entries_end = STEP_HDR.size + n_req * REQ_ENTRY.size |
| meta_end = entries_end + (n_req * SAMPLE_META.size if step_flags & STEP_SAMPLED else 0) |
| if rows_offset % 64 or rows_offset < meta_end: |
| raise ProtocolError(f"rows_offset {rows_offset} (entries end at {entries_end}, must be a multiple of 64)") |
| if len(body) != rows_offset + n_rows * row_bytes: |
| raise ProtocolError(f"STEP body {len(body)} B, header implies {rows_offset + n_rows * row_bytes}") |
| entries = [] |
| start = 0 |
| slots = set() |
| for i in range(n_req): |
| v = REQ_ENTRY.unpack_from(body, STEP_HDR.size + i * REQ_ENTRY.size) |
| e = ReqEntry(v[0], v[1], v[2], v[4], v[5], v[6], v[7], v[8], n_out=v[3]) |
| if e.slot in slots: |
| raise ProtocolError(f"slot {e.slot} twice in one STEP") |
| slots.add(e.slot) |
| e.row_start = start |
| start += e.n_rows |
| entries.append(e) |
| if start != n_rows: |
| raise ProtocolError(f"entries carry {start} rows, header says {n_rows}") |
| mv = memoryview(body)[rows_offset: rows_offset + n_rows * row_bytes] |
| rows = np.frombuffer(mv, dtype=np.uint8).reshape(n_rows, row_bytes) if n_rows else np.empty((0, row_bytes), |
| np.uint8) |
| meta = [] |
| if step_flags & STEP_SAMPLED: |
| for i in range(n_req): |
| seed, mask = SAMPLE_META.unpack_from(body, entries_end + i * SAMPLE_META.size) |
| if mask not in (0, 1): |
| raise ProtocolError("sample mask must be 0/1") |
| meta.append((seed, mask)) |
| return Step(step_id, n_req, step_flags, n_rows, row_bytes, rows_offset, entries, meta, rows) |
|
|
|
|
| class CapacityError(ValueError): |
| pass |
|
|
|
|
| def encode_step(step_id: int, entries: list[ReqEntry], rows: bytes, row_bytes: int, step_flags: int = 0) -> bytes: |
| """Client side (tests, tools): build a STEP body; rows = the row block, entries in order.""" |
| n_rows = sum(e.n_rows for e in entries) |
| entries_end = STEP_HDR.size + len(entries) * REQ_ENTRY.size |
| rows_offset = (entries_end + 63) // 64 * 64 |
| out = bytearray(STEP_HDR.pack(step_id, len(entries), step_flags, n_rows, row_bytes, rows_offset)) |
| for e in entries: |
| out += REQ_ENTRY.pack(e.slot, e.req_flags, e.k_window, e.n_out, e.generation, e.first_pos, e.n_rows, |
| e.anchor_token, e.anchor_pos) |
| out += b"\0" * (rows_offset - len(out)) |
| assert len(rows) == n_rows * row_bytes |
| return bytes(out) + rows |
|
|
|
|
| def encode_step_reply(step_id: int, num_spec: int, statuses: list[tuple[int, int, int]], draft_ids: np.ndarray, |
| conf: np.ndarray, margin: np.ndarray | None, worker_us: int, recv_to_send_us: int) -> bytes: |
| n = len(statuses) |
| flags = REPLY_HAS_MARGIN if margin is not None else 0 |
| out = bytearray(REPLY_HDR.pack(step_id, n, num_spec, flags, worker_us & 0xFFFFFFFF, |
| recv_to_send_us & 0xFFFFFFFF)) |
| for st, wr, gen in statuses: |
| out += REQ_STATUS.pack(st, wr, gen & 0xFFFFFFFF) |
| out += np.ascontiguousarray(draft_ids, dtype="<i4").tobytes() |
| out += np.ascontiguousarray(conf, dtype="<f4").tobytes() |
| if margin is not None: |
| out += np.ascontiguousarray(margin, dtype="<f4").tobytes() |
| return bytes(out) |
|
|
|
|
| @dataclass |
| class StepReply: |
| step_id: int |
| n_req: int |
| num_spec: int |
| reply_flags: int |
| worker_us: int |
| recv_to_send_us: int |
| statuses: list[tuple[int, int, int]] |
| draft_ids: np.ndarray |
| conf: np.ndarray |
| margin: np.ndarray | None |
|
|
|
|
| def decode_step_reply(body) -> StepReply: |
| step_id, n, N, flags, wus, rts = REPLY_HDR.unpack_from(body, 0) |
| off = REPLY_HDR.size |
| sts = [REQ_STATUS.unpack_from(body, off + i * REQ_STATUS.size) for i in range(n)] |
| off += n * REQ_STATUS.size |
| ids = np.frombuffer(body, dtype="<i4", count=n * N, offset=off).reshape(n, N) |
| off += 4 * n * N |
| conf = np.frombuffer(body, dtype="<f4", count=n * N, offset=off).reshape(n, N) |
| off += 4 * n * N |
| margin = None |
| if flags & REPLY_HAS_MARGIN: |
| margin = np.frombuffer(body, dtype="<f4", count=n * N, offset=off).reshape(n, N) |
| off += 4 * n * N |
| if off != len(body): |
| raise ProtocolError(f"STEP reply {len(body)} B, parsed {off}") |
| return StepReply(step_id, n, N, flags, wus, rts, sts, ids, conf, margin) |
|
|
|
|
| def serial_newer(a: int, b: int) -> bool: |
| """a is newer than b in u32 serial-number arithmetic (RFC 1982).""" |
| d = (a - b) & 0xFFFFFFFF |
| return d != 0 and d < 0x80000000 |
|
|
