139 lines
5.0 KiB
Python
139 lines
5.0 KiB
Python
"""Certified architecture adapters for the public TensorBundle boundary.
|
|
|
|
The adapter is intentionally small: it owns boundary names and endpoint rules,
|
|
not transformer execution. llama.cpp owns local graphs; callers select a
|
|
certified adapter before accepting an activation from another Shard.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
import struct
|
|
from typing import Callable, Mapping, Sequence
|
|
|
|
from .native_protocol import HIDDEN_STATES, ProtocolError, encode_bundle, encode_tensor, pb
|
|
|
|
|
|
class Architecture(str, Enum):
|
|
DENSE = "dense"
|
|
MOE = "moe"
|
|
MLA = "mla"
|
|
|
|
|
|
class BoundaryStage(str, Enum):
|
|
HEAD = "head"
|
|
MIDDLE = "middle"
|
|
TAIL = "tail"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ProtocolIdentity:
|
|
request_id: str
|
|
runtime_recipe_digest: str
|
|
chat_template_id: str
|
|
chat_template_version: str
|
|
reasoning_mode: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SamplingParameters:
|
|
temperature: float
|
|
top_p: float
|
|
top_k: int
|
|
seed: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TailOutput:
|
|
kind: str
|
|
value: int | object
|
|
|
|
@classmethod
|
|
def sampled_token(cls, token_id: int) -> "TailOutput":
|
|
if token_id < 0:
|
|
raise ProtocolError("sampled token id must be non-negative")
|
|
return cls("sampled_token", token_id)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TypedTailResult:
|
|
identity: ProtocolIdentity
|
|
sampling: SamplingParameters
|
|
output_kind: str
|
|
sampled_token_id: int | None = None
|
|
logits: object | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ArchitectureBoundaryAdapter:
|
|
architecture: Architecture
|
|
required_names: frozenset[str]
|
|
|
|
def bundle_from_token_ids(
|
|
self,
|
|
token_ids: Sequence[int],
|
|
token_embedding: Callable[[int], Sequence[float]],
|
|
):
|
|
"""Head-only embedding entry point; middle/tail never receive IDs."""
|
|
if self.architecture is not Architecture.DENSE:
|
|
raise ProtocolError("head token embedding is not certified for this architecture")
|
|
if not token_ids:
|
|
raise ProtocolError("head requires at least one token id")
|
|
rows = [tuple(token_embedding(token)) for token in token_ids]
|
|
if not rows or not rows[0] or any(len(row) != len(rows[0]) for row in rows):
|
|
raise ProtocolError("token embedding returned inconsistent hidden widths")
|
|
payload = struct.pack("<" + "f" * (len(rows) * len(rows[0])), *(x for row in rows for x in row))
|
|
return self.bundle_from_named_payloads({HIDDEN_STATES: payload}, shape=[1, len(rows), len(rows[0])])
|
|
|
|
def bundle_from_named_payloads(
|
|
self, payloads: Mapping[str, bytes], *, shape: Sequence[int] | None = None
|
|
):
|
|
names = set(payloads)
|
|
if not self.required_names <= names:
|
|
missing = sorted(self.required_names - names)
|
|
raise ProtocolError(f"{self.architecture.value} boundary requires {missing}")
|
|
tensors = []
|
|
for name, payload in payloads.items():
|
|
tensor_shape = list(shape) if name == HIDDEN_STATES and shape else [len(payload) // 4]
|
|
if len(payload) % 4:
|
|
raise ProtocolError(f"{name!r} F32 fixture payload is not word aligned")
|
|
tensors.append(encode_tensor(name, payload, tensor_shape, pb.DTYPE_FLOAT32))
|
|
return encode_bundle(tensors)
|
|
|
|
def input_for(self, stage: BoundaryStage, bundle):
|
|
"""Accept architecture state only after the head embedding boundary."""
|
|
if stage is BoundaryStage.HEAD:
|
|
raise ProtocolError("head accepts token ids and owns token embedding")
|
|
if bundle is None:
|
|
raise ProtocolError(f"{stage.value} requires a TensorBundle")
|
|
from .native_protocol import decode_bundle
|
|
|
|
payloads = decode_bundle(bundle)
|
|
if not self.required_names <= set(payloads):
|
|
raise ProtocolError(f"{self.architecture.value} boundary requires {sorted(self.required_names)}")
|
|
return bundle
|
|
|
|
def tail_result(
|
|
self, *, identity: ProtocolIdentity, sampling: SamplingParameters, output: TailOutput
|
|
) -> TypedTailResult:
|
|
if not identity.request_id or not identity.runtime_recipe_digest:
|
|
raise ProtocolError("tail result requires exact request and recipe identity")
|
|
if output.kind != "sampled_token":
|
|
raise ProtocolError("uncertified tail output kind")
|
|
return TypedTailResult(identity, sampling, output.kind, sampled_token_id=int(output.value))
|
|
|
|
|
|
_ADAPTERS = {
|
|
Architecture.DENSE: ArchitectureBoundaryAdapter(Architecture.DENSE, frozenset({HIDDEN_STATES})),
|
|
Architecture.MOE: ArchitectureBoundaryAdapter(Architecture.MOE, frozenset({HIDDEN_STATES, "router_logits"})),
|
|
Architecture.MLA: ArchitectureBoundaryAdapter(Architecture.MLA, frozenset({HIDDEN_STATES, "mla_position_state"})),
|
|
}
|
|
|
|
|
|
def adapter_for(architecture: Architecture | str) -> ArchitectureBoundaryAdapter:
|
|
try:
|
|
return _ADAPTERS[Architecture(architecture)]
|
|
except (KeyError, ValueError):
|
|
raise ProtocolError(f"unsupported architecture {architecture!r}") from None
|