neuron-tai/packages/node/meshnet_node/architecture_boundary.py

"""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