feat: implement DGR-006 tensor bundle boundary

packages/node/meshnet_node/architecture_boundary.py

@@ -0,0 +1,186 @@
+"""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,
+)
+
+
+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)
+
+
+@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":
+            raise ProtocolError("uncertified tail output kind")
+        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=int(output.value),
+        )
+        validate_tail_result(message)
+        return TypedTailResult(identity, sampling, "sampled_token_id", message)
+
+
+_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