325 lines
13 KiB
Python
325 lines
13 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,
|
|
validate_tail_result,
|
|
)
|
|
from .shard_engine import BoundaryBundle, EngineTensor
|
|
|
|
|
|
# This is deliberately an execution-boundary name, not a transport name. It
|
|
# identifies the value *before* final norm/output projection. A future wire
|
|
# codec may rename its field, but cannot reinterpret this value as logits.
|
|
DENSE_LLAMA_ARCHITECTURE = "dense-llama"
|
|
DENSE_RESIDUAL_BOUNDARY_V1 = "dense.residual.v1"
|
|
|
|
|
|
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
|
|
architecture: Architecture
|
|
|
|
|
|
@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)
|
|
|
|
@classmethod
|
|
def logits(cls, logits: object) -> "TailOutput":
|
|
"""Return raw logits under the explicit tail-only output contract."""
|
|
return cls("logits", logits)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TypedTailResult:
|
|
identity: ProtocolIdentity
|
|
sampling: SamplingParameters
|
|
output_kind: str
|
|
message: pb.TailResult
|
|
|
|
@property
|
|
def sampled_token_id(self) -> int | None:
|
|
return self.message.sampled_token_id if self.output_kind == "sampled_token_id" else None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ArchitectureBoundaryAdapter:
|
|
architecture: Architecture
|
|
required_names: frozenset[str]
|
|
|
|
@property
|
|
def protocol_architecture(self) -> int:
|
|
return {
|
|
Architecture.DENSE: pb.ARCHITECTURE_TYPE_DENSE,
|
|
Architecture.MOE: pb.ARCHITECTURE_TYPE_MOE,
|
|
Architecture.MLA: pb.ARCHITECTURE_TYPE_MLA,
|
|
}[self.architecture]
|
|
|
|
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,
|
|
architecture=self.protocol_architecture,
|
|
boundary_point="pre_tail_residual",
|
|
)
|
|
|
|
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 bundle.architecture != self.protocol_architecture:
|
|
raise ProtocolError("boundary architecture does not match certified adapter")
|
|
if bundle.boundary_point != "pre_tail_residual":
|
|
raise ProtocolError("unsupported architecture boundary point")
|
|
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 identity.architecture is not self.architecture:
|
|
raise ProtocolError("tail result architecture does not match certified adapter")
|
|
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":
|
|
if not isinstance(output.value, int):
|
|
raise ProtocolError("sampled tail output must carry an integer token id")
|
|
message = pb.TailResult(
|
|
identity=pb.RequestRecipeIdentity(
|
|
request_id=identity.request_id,
|
|
runtime_recipe_digest=identity.runtime_recipe_digest,
|
|
chat_template_id=identity.chat_template_id,
|
|
chat_template_version=identity.chat_template_version,
|
|
reasoning_mode=identity.reasoning_mode,
|
|
architecture=self.protocol_architecture,
|
|
),
|
|
sampling=pb.SamplingParameters(
|
|
temperature=sampling.temperature,
|
|
top_p=sampling.top_p,
|
|
top_k=sampling.top_k,
|
|
seed=sampling.seed,
|
|
greedy=sampling.temperature == 0.0,
|
|
),
|
|
sampled_token_id=output.value,
|
|
)
|
|
elif output.kind == "logits":
|
|
if not isinstance(output.value, pb.TensorBundle):
|
|
raise ProtocolError("logits tail output must carry a TensorBundle")
|
|
# Validate the logits bundle before putting it in the result; this
|
|
# rejects an incompatible boundary schema rather than passing an
|
|
# opaque tensor on to sampling.
|
|
from .native_protocol import decode_bundle
|
|
|
|
decode_bundle(output.value)
|
|
message = pb.TailResult(
|
|
identity=pb.RequestRecipeIdentity(
|
|
request_id=identity.request_id,
|
|
runtime_recipe_digest=identity.runtime_recipe_digest,
|
|
chat_template_id=identity.chat_template_id,
|
|
chat_template_version=identity.chat_template_version,
|
|
reasoning_mode=identity.reasoning_mode,
|
|
architecture=self.protocol_architecture,
|
|
),
|
|
sampling=pb.SamplingParameters(
|
|
temperature=sampling.temperature,
|
|
top_p=sampling.top_p,
|
|
top_k=sampling.top_k,
|
|
seed=sampling.seed,
|
|
greedy=sampling.temperature == 0.0,
|
|
),
|
|
logits=output.value,
|
|
)
|
|
else:
|
|
raise ProtocolError("uncertified tail output kind")
|
|
validate_tail_result(message)
|
|
return TypedTailResult(identity, sampling, message.WhichOneof("output"), message)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DenseLayerRange:
|
|
"""A certified, inclusive dense-Llama range within one loaded model."""
|
|
|
|
start_layer: int
|
|
end_layer: int
|
|
total_layers: int
|
|
architecture: str = DENSE_LLAMA_ARCHITECTURE
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.architecture != DENSE_LLAMA_ARCHITECTURE:
|
|
raise ProtocolError("dense boundary executor only certifies dense-llama")
|
|
if self.start_layer < 0 or self.end_layer < self.start_layer:
|
|
raise ProtocolError("dense range is empty or inverted")
|
|
if self.total_layers <= self.end_layer:
|
|
raise ProtocolError("dense range lies outside the model")
|
|
|
|
@property
|
|
def is_head(self) -> bool:
|
|
return self.start_layer == 0
|
|
|
|
@property
|
|
def is_tail(self) -> bool:
|
|
return self.end_layer == self.total_layers - 1
|
|
|
|
|
|
class DenseRangeBoundaryExecutor:
|
|
"""Execute one dense range without leaking endpoint ownership.
|
|
|
|
``run_layers`` owns only the local transformer blocks and receives/returns
|
|
the raw residual. It never receives a final norm/head callback. Only a
|
|
tail range receives ``tail_output``; consequently row pruning and logits
|
|
projection cannot accidentally happen before the final stage.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
layer_range: DenseLayerRange,
|
|
*,
|
|
embed_tokens: Callable[[tuple[int, ...]], EngineTensor],
|
|
run_layers: Callable[[EngineTensor], EngineTensor],
|
|
tail_output: Callable[[EngineTensor], TailOutput] | None = None,
|
|
) -> None:
|
|
if layer_range.is_tail != (tail_output is not None):
|
|
raise ProtocolError("only a dense tail range may own final norm/output")
|
|
self._range = layer_range
|
|
self._embed_tokens = embed_tokens
|
|
self._run_layers = run_layers
|
|
self._tail_output = tail_output
|
|
|
|
def execute(
|
|
self,
|
|
*,
|
|
token_ids: tuple[int, ...] | None = None,
|
|
boundary: BoundaryBundle | None = None,
|
|
) -> BoundaryBundle | TailOutput:
|
|
if self._range.is_head:
|
|
if token_ids is None or boundary is not None or not token_ids:
|
|
raise ProtocolError("dense head accepts non-empty token ids and no boundary bundle")
|
|
residual = self._embed_tokens(token_ids)
|
|
else:
|
|
if token_ids is not None or boundary is None:
|
|
raise ProtocolError("dense middle/tail requires a named residual boundary bundle")
|
|
residual = self._residual_from_boundary(boundary)
|
|
|
|
residual = self._run_layers(residual)
|
|
if residual.name != HIDDEN_STATES:
|
|
raise ProtocolError("dense range must return hidden_states residual")
|
|
|
|
if self._range.is_tail:
|
|
assert self._tail_output is not None
|
|
output = self._tail_output(residual)
|
|
if output.kind not in {"logits", "sampled_token"}:
|
|
raise ProtocolError("dense tail returned an uncertified output kind")
|
|
return output
|
|
|
|
# Do not normalize, project, sample, or prune rows here: this exact
|
|
# raw output becomes the next range's input.
|
|
return BoundaryBundle(
|
|
tensors=(residual,),
|
|
architecture=DENSE_LLAMA_ARCHITECTURE,
|
|
boundary_point=DENSE_RESIDUAL_BOUNDARY_V1,
|
|
)
|
|
|
|
@staticmethod
|
|
def _residual_from_boundary(boundary: BoundaryBundle) -> EngineTensor:
|
|
if boundary.architecture != DENSE_LLAMA_ARCHITECTURE:
|
|
raise ProtocolError("boundary architecture is not certified dense-llama")
|
|
if boundary.boundary_point != DENSE_RESIDUAL_BOUNDARY_V1:
|
|
raise ProtocolError("incompatible dense residual boundary schema")
|
|
if len(boundary.tensors) != 1 or boundary.tensors[0].name != HIDDEN_STATES:
|
|
raise ProtocolError("dense residual boundary requires exactly one hidden_states tensor")
|
|
return boundary.tensors[0]
|
|
|
|
|
|
_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
|