"""Encode and decode the native Shard protocol's named-tensor bundles. The generated stubs give us message *structure*; they cannot enforce the invariants that keep a distributed forward correct. A bundle whose declared shape disagrees with its byte count, whose fragments leave a hole, or whose checksum does not match is not a slightly-wrong activation — it is silently wrong tokens for the rest of the generation. So decoding is validating: every path into a tensor's bytes goes through :func:`decode_tensor`, which refuses a payload it cannot fully account for. Compression is a transport optimisation and is decided by the same policy layer the existing HTTP seam already uses (``activation_compression``), so a node's tuned thresholds apply to both transports. """ from __future__ import annotations from dataclasses import dataclass import struct from typing import Iterable, Sequence from ..activation_compression import ( CompressionPolicy, compress_activation, decompress_activation, ) from .generated import shard_runtime_pb2 as pb # The schema generation this build speaks. A peer offering something else is # rejected at the handshake rather than being half-understood. SCHEMA_VERSION = pb.SCHEMA_VERSION_1 # Generation of the tensor-bundle layout, versioned independently of the # protocol so a boundary payload can evolve without a protocol bump. BUNDLE_VERSION = 1 # Token-aligned prefill chunk bound. 128 tokens is the size ADR-0008 already # uses on the HTTP seam; keeping it identical means seam bytes stay comparable # across transports. DEFAULT_MAX_PREFILL_CHUNK_TOKENS = 128 # gRPC's default maximum receive size. Fragmenting below it keeps us inside the # default limits of any conformant peer instead of requiring every client to # raise its window. DEFAULT_MAX_CHUNK_BYTES = 4 * 1024 * 1024 # Leave room for envelope and framing overhead inside one chunk message. DEFAULT_MAX_FRAGMENT_BYTES = 1024 * 1024 DEFAULT_MAX_INFLIGHT_CHUNKS = 8 # Canonical boundary tensor name for a dense transformer hidden state. HIDDEN_STATES = "hidden_states" _DTYPE_ITEMSIZE: dict[int, int] = { pb.DTYPE_BFLOAT16: 2, pb.DTYPE_FLOAT16: 2, pb.DTYPE_FLOAT32: 4, pb.DTYPE_INT32: 4, pb.DTYPE_INT64: 8, pb.DTYPE_UINT8: 1, pb.DTYPE_INT8: 1, pb.DTYPE_BOOL: 1, } class ProtocolError(Exception): """A peer sent something this build cannot safely interpret.""" class PayloadCorrupt(ProtocolError): """A tensor payload failed validation: size, coverage, or checksum.""" def itemsize(dtype: int) -> int: try: return _DTYPE_ITEMSIZE[dtype] except KeyError: raise ProtocolError(f"unsupported dtype {dtype}") from None def expected_bytes(shape: Sequence[int], dtype: int) -> int: """Byte count a tensor of `shape` and `dtype` must occupy.""" if any(dim < 0 for dim in shape): raise ProtocolError(f"negative dimension in shape {list(shape)}") count = 1 for dim in shape: count *= dim return count * itemsize(dtype) # --- CRC32C ---------------------------------------------------------------- # # CRC32C (Castagnoli), not zlib's CRC32: it is the checksum gRPC, and the # storage systems these payloads pass through, already use, and hardware # implements it. `google_crc32c` is used when present; the table fallback keeps # the default test suite dependency-free. _CRC32C_POLY = 0x82F63B78 _CRC32C_TABLE: list[int] = [] for _i in range(256): _c = _i for _ in range(8): _c = (_c >> 1) ^ (_CRC32C_POLY if _c & 1 else 0) _CRC32C_TABLE.append(_c) try: # pragma: no cover - depends on an optional native package from google_crc32c import value as _fast_crc32c except ImportError: # pragma: no cover _fast_crc32c = None def crc32c(data: bytes) -> int: if _fast_crc32c is not None: # pragma: no cover - optional fast path return _fast_crc32c(data) crc = 0xFFFFFFFF for byte in data: crc = (crc >> 8) ^ _CRC32C_TABLE[(crc ^ byte) & 0xFF] return crc ^ 0xFFFFFFFF def checksum_of(data: bytes) -> pb.Checksum: return pb.Checksum( algorithm=pb.CHECKSUM_ALGORITHM_CRC32C, value=struct.pack(">I", crc32c(data)), ) # --- Tensors --------------------------------------------------------------- def encode_tensor( name: str, data: bytes, shape: Sequence[int], dtype: int = pb.DTYPE_BFLOAT16, *, policy: CompressionPolicy | None = None, max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES, ) -> pb.NamedTensor: """Build a NamedTensor, compressing and fragmenting as needed. `data` is the uncompressed little-endian payload. The checksum is taken over it *before* compression so it stays valid whichever framing a hop chooses. """ if max_fragment_bytes <= 0: raise ProtocolError("max_fragment_bytes must be positive") declared = expected_bytes(shape, dtype) if len(data) != declared: raise ProtocolError( f"tensor {name!r} declares shape {list(shape)} ({declared} bytes) " f"but carries {len(data)} bytes" ) body = data compression = pb.COMPRESSION_NONE if policy is not None: result = compress_activation(data, policy) if result.compressed: body = result.body compression = pb.COMPRESSION_ZSTD tensor = pb.NamedTensor( name=name, shape=list(shape), dtype=dtype, byte_order=pb.BYTE_ORDER_LITTLE_ENDIAN, total_bytes=len(data), compression=compression, checksum=checksum_of(data), ) # Fragment the wire body (compressed if we compressed). Offsets walk the # wire body so a receiver can verify coverage without assuming arrival # order; a zstd frame is not decodable per fragment, so reassembly comes # first and decompression happens once, in decode_tensor. slices = [body[i : i + max_fragment_bytes] for i in range(0, len(body), max_fragment_bytes)] if not slices: # A zero-element tensor is legal (e.g. an empty mask) and still needs a # fragment, so coverage checks have something to verify. slices = [b""] offset = 0 for index, piece in enumerate(slices): tensor.fragments.append( pb.TensorFragment( fragment_index=index, fragment_count=len(slices), byte_offset=offset, payload=piece, ) ) offset += len(piece) return tensor def decode_tensor(tensor: pb.NamedTensor) -> bytes: """Reassemble, decompress and validate a NamedTensor's payload. Raises PayloadCorrupt rather than returning a payload it cannot fully account for: a hole in the fragments or a bad checksum means the activation is not what the sender computed, and continuing would corrupt the route. """ if tensor.byte_order == pb.BYTE_ORDER_BIG_ENDIAN: raise ProtocolError(f"tensor {tensor.name!r} is big-endian; wire order is little-endian") if tensor.byte_order != pb.BYTE_ORDER_LITTLE_ENDIAN: raise ProtocolError(f"tensor {tensor.name!r} declares no byte order") if not tensor.fragments: raise PayloadCorrupt(f"tensor {tensor.name!r} carries no fragments") fragments = sorted(tensor.fragments, key=lambda f: f.byte_offset) count = fragments[0].fragment_count if any(f.fragment_count != count for f in fragments): raise PayloadCorrupt(f"tensor {tensor.name!r} has inconsistent fragment_count") if len(fragments) != count: raise PayloadCorrupt( f"tensor {tensor.name!r} expects {count} fragments but carries {len(fragments)}" ) if {f.fragment_index for f in fragments} != set(range(count)): raise PayloadCorrupt(f"tensor {tensor.name!r} has duplicate or missing fragment indices") # Contiguity: offsets must tile the body exactly, with no hole and no overlap. body = bytearray() for fragment in fragments: if fragment.byte_offset != len(body): raise PayloadCorrupt( f"tensor {tensor.name!r} fragment {fragment.fragment_index} starts at " f"{fragment.byte_offset}, expected {len(body)}" ) body.extend(fragment.payload) if tensor.compression == pb.COMPRESSION_ZSTD: data = decompress_activation(bytes(body), "zstd").body elif tensor.compression in (pb.COMPRESSION_NONE, pb.COMPRESSION_UNSPECIFIED): data = bytes(body) else: raise ProtocolError(f"tensor {tensor.name!r} uses unsupported compression") if len(data) != tensor.total_bytes: raise PayloadCorrupt( f"tensor {tensor.name!r} declares {tensor.total_bytes} bytes but " f"reassembled {len(data)}" ) declared = expected_bytes(tensor.shape, tensor.dtype) if declared != tensor.total_bytes: raise PayloadCorrupt( f"tensor {tensor.name!r} shape {list(tensor.shape)} implies {declared} bytes " f"but declares {tensor.total_bytes}" ) algorithm = tensor.checksum.algorithm if algorithm == pb.CHECKSUM_ALGORITHM_CRC32C: if tensor.checksum.value != struct.pack(">I", crc32c(data)): raise PayloadCorrupt(f"tensor {tensor.name!r} failed its CRC32C check") elif algorithm not in (pb.CHECKSUM_ALGORITHM_NONE, pb.CHECKSUM_ALGORITHM_UNSPECIFIED): raise ProtocolError(f"tensor {tensor.name!r} uses unsupported checksum algorithm") return data def encode_bundle(tensors: Iterable[pb.NamedTensor]) -> pb.TensorBundle: return pb.TensorBundle(bundle_version=BUNDLE_VERSION, tensors=list(tensors)) def decode_bundle(bundle: pb.TensorBundle) -> dict[str, bytes]: """Validate every tensor in a bundle and return name -> payload.""" if bundle.bundle_version > BUNDLE_VERSION: raise ProtocolError( f"bundle version {bundle.bundle_version} is newer than this build " f"understands ({BUNDLE_VERSION})" ) payloads: dict[str, bytes] = {} for tensor in bundle.tensors: if not tensor.name: raise ProtocolError("bundle carries an unnamed tensor") if tensor.name in payloads: raise ProtocolError(f"bundle carries duplicate tensor {tensor.name!r}") payloads[tensor.name] = decode_tensor(tensor) return payloads # --- Bounded prefill chunking ---------------------------------------------- @dataclass(frozen=True) class PrefillChunk: """One token-aligned slice of a prefill.""" chunk_index: int chunk_count: int first_position: int token_count: int @property def final_chunk(self) -> bool: return self.chunk_index == self.chunk_count - 1 def chunk_info(self) -> pb.ChunkInfo: return pb.ChunkInfo( chunk_index=self.chunk_index, chunk_count=self.chunk_count, final_chunk=self.final_chunk, ) def position(self) -> pb.PositionSpan: return pb.PositionSpan( first_position=self.first_position, token_count=self.token_count ) def plan_prefill_chunks( total_tokens: int, *, first_position: int = 0, max_tokens: int = DEFAULT_MAX_PREFILL_CHUNK_TOKENS, ) -> list[PrefillChunk]: """Split a prefill into bounded, token-aligned chunks. Splits fall on token boundaries only (ADR-0008): a fragment of a token's hidden state is not a thing a receiver can execute. """ if total_tokens <= 0: raise ProtocolError("a prefill must carry at least one token") if max_tokens <= 0: raise ProtocolError("max_tokens must be positive") count = (total_tokens + max_tokens - 1) // max_tokens chunks = [] for index in range(count): offset = index * max_tokens chunks.append( PrefillChunk( chunk_index=index, chunk_count=count, first_position=first_position + offset, token_count=min(max_tokens, total_tokens - offset), ) ) return chunks def default_flow_control() -> pb.FlowControl: return pb.FlowControl( credits_granted=DEFAULT_MAX_INFLIGHT_CHUNKS, max_inflight_chunks=DEFAULT_MAX_INFLIGHT_CHUNKS, max_chunk_bytes=DEFAULT_MAX_CHUNK_BYTES, max_prefill_chunk_tokens=DEFAULT_MAX_PREFILL_CHUNK_TOKENS, ) def negotiate_flow_control( proposed: pb.FlowControl, limits: pb.FlowControl ) -> pb.FlowControl: """Settle a stream's limits: the strictest bound of either peer wins. Taking the minimum means neither peer can raise the other's ceiling, so a misconfigured — or hostile — sender cannot talk a worker into unbounded queues by proposing a large window. """ def _min(a: int, b: int, fallback: int) -> int: candidates = [v for v in (a, b) if v > 0] return min(candidates) if candidates else fallback return pb.FlowControl( credits_granted=_min( proposed.credits_granted, limits.credits_granted, DEFAULT_MAX_INFLIGHT_CHUNKS ), max_inflight_chunks=_min( proposed.max_inflight_chunks, limits.max_inflight_chunks, DEFAULT_MAX_INFLIGHT_CHUNKS, ), max_chunk_bytes=_min( proposed.max_chunk_bytes, limits.max_chunk_bytes, DEFAULT_MAX_CHUNK_BYTES ), max_prefill_chunk_tokens=_min( proposed.max_prefill_chunk_tokens, limits.max_prefill_chunk_tokens, DEFAULT_MAX_PREFILL_CHUNK_TOKENS, ), )