popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: implement DGR-006 tensor bundle boundary

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-14 13:44:37 +03:00
parent 91c450840d
commit 7925e5253d
16 changed files with 802 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

186
packages/node/meshnet_node/architecture_boundary.py Normal file
View File

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

6
packages/node/meshnet_node/native_protocol/__init__.py
View File

@@ -31,6 +31,8 @@ from .codec import (
checksum_of,
crc32c,
decode_bundle,
decode_step_bundle,
encode_decode_step,
decode_tensor,
default_flow_control,
encode_bundle,
@@ -40,6 +42,7 @@ from .codec import (
negotiate_flow_control,
plan_prefill_chunks,
validate_session_message_size,
validate_tail_result,
)
__all__ = [
@@ -60,6 +63,8 @@ __all__ = [
"checksum_of",
"crc32c",
"decode_bundle",
"decode_step_bundle",
"encode_decode_step",
"decode_tensor",
"default_flow_control",
"encode_bundle",
@@ -70,4 +75,5 @@ __all__ = [
"pb",
"plan_prefill_chunks",
"validate_session_message_size",
"validate_tail_result",
]

83
packages/node/meshnet_node/native_protocol/codec.py
View File

@@ -328,6 +328,8 @@ def decode_tensor(
def encode_bundle(
tensors: Iterable[pb.NamedTensor],
*,
architecture: int = pb.ARCHITECTURE_TYPE_UNSPECIFIED,
boundary_point: str = "",
max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
max_tensors: int = DEFAULT_MAX_TENSORS_PER_BUNDLE,
) -> pb.TensorBundle:
@@ -338,7 +340,12 @@ def encode_bundle(
raise ProtocolError(
f"bundle carries {len(tensor_list)} tensors, exceeding limit {max_tensors}"
)
bundle = pb.TensorBundle(bundle_version=BUNDLE_VERSION, tensors=tensor_list)
bundle = pb.TensorBundle(
bundle_version=BUNDLE_VERSION,
tensors=tensor_list,
architecture=architecture,
boundary_point=boundary_point,
)
if bundle.ByteSize() > max_chunk_bytes:
raise ProtocolError(
f"serialized tensor bundle exceeds the {max_chunk_bytes}-byte "
@@ -392,6 +399,80 @@ def decode_bundle(
return payloads
def encode_decode_step(
bundle: pb.TensorBundle,
*,
idempotency_step: int,
position: int,
expected_past_len: int,
work_id: str,
deadline_unix_nanos: int = 0,
prefer_compact_one_tensor: bool = True,
) -> pb.DecodeStep:
"""Encode a decode boundary, retaining the deliberate compact fallback."""
step = pb.DecodeStep(
idempotency_step=idempotency_step,
position=position,
expected_past_len=expected_past_len,
work_id=work_id,
deadline_unix_nanos=deadline_unix_nanos,
)
if prefer_compact_one_tensor and len(bundle.tensors) == 1:
step.tensor.CopyFrom(bundle.tensors[0])
else:
step.bundle.CopyFrom(bundle)
return step
def validate_tail_result(result: pb.TailResult) -> None:
"""Fail closed unless a tail completion is bound to its exact recipe."""
identity = result.identity
required = (
identity.request_id,
identity.runtime_recipe_digest,
identity.chat_template_id,
identity.chat_template_version,
identity.reasoning_mode,
)
if not all(required) or identity.architecture == pb.ARCHITECTURE_TYPE_UNSPECIFIED:
raise ProtocolError("tail result lacks exact request/recipe/template identity")
if result.WhichOneof("output") not in {"logits", "sampled_token_id"}:
raise ProtocolError("tail result lacks logits or sampled token output")
def decode_step_bundle(
step: pb.DecodeStep,
*,
max_chunk_bytes: int = DEFAULT_MAX_CHUNK_BYTES,
max_fragment_bytes: int = DEFAULT_MAX_FRAGMENT_BYTES,
max_fragments: int = DEFAULT_MAX_FRAGMENTS_PER_TENSOR,
max_tensors: int = DEFAULT_MAX_TENSORS_PER_BUNDLE,
) -> dict[str, bytes]:
"""Decode a fast-path boundary with the DGR-006 compatibility rule.
`bundle` is authoritative because it can carry architecture sidebands. The
old `tensor` field remains the compact representation for a certified
one-tensor boundary and is accepted by new readers during rollout.
"""
if step.HasField("bundle"):
return decode_bundle(
step.bundle,
max_chunk_bytes=max_chunk_bytes,
max_fragment_bytes=max_fragment_bytes,
max_fragments=max_fragments,
max_tensors=max_tensors,
)
if step.HasField("tensor"):
return decode_bundle(
encode_bundle([step.tensor], max_chunk_bytes=max_chunk_bytes),
max_chunk_bytes=max_chunk_bytes,
max_fragment_bytes=max_fragment_bytes,
max_fragments=max_fragments,
max_tensors=max_tensors,
)
raise ProtocolError("decode step carries neither TensorBundle nor legacy tensor")
def validate_session_message_size(
message: pb.SessionRequest | pb.SessionResponse,
*,

25
packages/node/meshnet_node/native_protocol/conformance.py
View File

@@ -27,6 +27,7 @@ TESTDATA_DIR = pathlib.Path(__file__).resolve().parents[2] / "native/testdata"
GOLDEN_SESSION_REQUEST = "session_request_golden.binpb"
GOLDEN_CAPABILITY_REPORT = "capability_report_golden.binpb"
GOLDEN_DECODE_STEP = "decode_step_golden.binpb"
# Written by the C++ conformance test into its build tree; the Python test picks
# it up when present to prove the two languages agree byte-for-byte.
@@ -136,6 +137,30 @@ def canonical_capability_report() -> pb.CapabilityReport:
)
def canonical_decode_step() -> pb.SessionRequest:
"""The DGR-006 multi-tensor decode boundary vector."""
hidden = codec.encode_tensor(
codec.HIDDEN_STATES, bytes(range(16)), [1, 1, 4], pb.DTYPE_FLOAT32
)
index_topk = codec.encode_tensor(
"index_topk", (3).to_bytes(4, "little"), [1], pb.DTYPE_INT32
)
return pb.SessionRequest(
decode=pb.DecodeStep(
idempotency_step=43,
position=384,
expected_past_len=384,
work_id="decode-7f3a",
deadline_unix_nanos=DEADLINE_UNIX_NANOS,
bundle=codec.encode_bundle(
[hidden, index_topk],
architecture=pb.ARCHITECTURE_TYPE_MLA,
boundary_point="pre_tail_residual",
),
)
)
def serialize(message) -> bytes:
"""Serialize deterministically, so committed golden bytes are stable."""
return message.SerializeToString(deterministic=True)

162
packages/node/meshnet_node/native_protocol/generated/shard_runtime_pb2.py
View File

File diff suppressed because one or more lines are too long

73
packages/node/meshnet_node/native_protocol/generated/shard_runtime_pb2.pyi
View File

@@ -12,6 +12,13 @@ class SchemaVersion(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
SCHEMA_VERSION_UNSPECIFIED: _ClassVar[SchemaVersion]
SCHEMA_VERSION_1: _ClassVar[SchemaVersion]
class ArchitectureType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
__slots__ = ()
ARCHITECTURE_TYPE_UNSPECIFIED: _ClassVar[ArchitectureType]
ARCHITECTURE_TYPE_DENSE: _ClassVar[ArchitectureType]
ARCHITECTURE_TYPE_MOE: _ClassVar[ArchitectureType]
ARCHITECTURE_TYPE_MLA: _ClassVar[ArchitectureType]
class DType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
__slots__ = ()
DTYPE_UNSPECIFIED: _ClassVar[DType]
@@ -80,6 +87,10 @@ class ServingState(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
SERVING_STATE_NOT_SERVING: _ClassVar[ServingState]
SCHEMA_VERSION_UNSPECIFIED: SchemaVersion
SCHEMA_VERSION_1: SchemaVersion
ARCHITECTURE_TYPE_UNSPECIFIED: ArchitectureType
ARCHITECTURE_TYPE_DENSE: ArchitectureType
ARCHITECTURE_TYPE_MOE: ArchitectureType
ARCHITECTURE_TYPE_MLA: ArchitectureType
DTYPE_UNSPECIFIED: DType
DTYPE_BFLOAT16: DType
DTYPE_FLOAT16: DType
@@ -165,12 +176,16 @@ class NamedTensor(_message.Message):
def __init__(self, name: _Optional[str] = ..., shape: _Optional[_Iterable[int]] = ..., dtype: _Optional[_Union[DType, str]] = ..., byte_order: _Optional[_Union[ByteOrder, str]] = ..., total_bytes: _Optional[int] = ..., compression: _Optional[_Union[Compression, str]] = ..., checksum: _Optional[_Union[Checksum, _Mapping]] = ..., fragments: _Optional[_Iterable[_Union[TensorFragment, _Mapping]]] = ...) -> None: ...
class TensorBundle(_message.Message):
__slots__ = ("bundle_version", "tensors")
__slots__ = ("bundle_version", "tensors", "architecture", "boundary_point")
BUNDLE_VERSION_FIELD_NUMBER: _ClassVar[int]
TENSORS_FIELD_NUMBER: _ClassVar[int]
ARCHITECTURE_FIELD_NUMBER: _ClassVar[int]
BOUNDARY_POINT_FIELD_NUMBER: _ClassVar[int]
bundle_version: int
tensors: _containers.RepeatedCompositeFieldContainer[NamedTensor]
def __init__(self, bundle_version: _Optional[int] = ..., tensors: _Optional[_Iterable[_Union[NamedTensor, _Mapping]]] = ...) -> None: ...
architecture: ArchitectureType
boundary_point: str
def __init__(self, bundle_version: _Optional[int] = ..., tensors: _Optional[_Iterable[_Union[NamedTensor, _Mapping]]] = ..., architecture: _Optional[_Union[ArchitectureType, str]] = ..., boundary_point: _Optional[str] = ...) -> None: ...
class Fingerprint(_message.Message):
__slots__ = ("model_artifact_digest", "runtime_recipe_digest", "recipe_id", "recipe_version", "catalogue_version")
@@ -327,20 +342,64 @@ class ActivationChunk(_message.Message):
def __init__(self, envelope: _Optional[_Union[Envelope, _Mapping]] = ..., bundle: _Optional[_Union[TensorBundle, _Mapping]] = ...) -> None: ...
class DecodeStep(_message.Message):
__slots__ = ("idempotency_step", "position", "expected_past_len", "tensor", "work_id", "deadline_unix_nanos")
__slots__ = ("idempotency_step", "position", "expected_past_len", "tensor", "work_id", "deadline_unix_nanos", "bundle")
IDEMPOTENCY_STEP_FIELD_NUMBER: _ClassVar[int]
POSITION_FIELD_NUMBER: _ClassVar[int]
EXPECTED_PAST_LEN_FIELD_NUMBER: _ClassVar[int]
TENSOR_FIELD_NUMBER: _ClassVar[int]
WORK_ID_FIELD_NUMBER: _ClassVar[int]
DEADLINE_UNIX_NANOS_FIELD_NUMBER: _ClassVar[int]
BUNDLE_FIELD_NUMBER: _ClassVar[int]
idempotency_step: int
position: int
expected_past_len: int
tensor: NamedTensor
work_id: str
deadline_unix_nanos: int
def __init__(self, idempotency_step: _Optional[int] = ..., position: _Optional[int] = ..., expected_past_len: _Optional[int] = ..., tensor: _Optional[_Union[NamedTensor, _Mapping]] = ..., work_id: _Optional[str] = ..., deadline_unix_nanos: _Optional[int] = ...) -> None: ...
bundle: TensorBundle
def __init__(self, idempotency_step: _Optional[int] = ..., position: _Optional[int] = ..., expected_past_len: _Optional[int] = ..., tensor: _Optional[_Union[NamedTensor, _Mapping]] = ..., work_id: _Optional[str] = ..., deadline_unix_nanos: _Optional[int] = ..., bundle: _Optional[_Union[TensorBundle, _Mapping]] = ...) -> None: ...
class RequestRecipeIdentity(_message.Message):
__slots__ = ("request_id", "runtime_recipe_digest", "chat_template_id", "chat_template_version", "reasoning_mode", "architecture")
REQUEST_ID_FIELD_NUMBER: _ClassVar[int]
RUNTIME_RECIPE_DIGEST_FIELD_NUMBER: _ClassVar[int]
CHAT_TEMPLATE_ID_FIELD_NUMBER: _ClassVar[int]
CHAT_TEMPLATE_VERSION_FIELD_NUMBER: _ClassVar[int]
REASONING_MODE_FIELD_NUMBER: _ClassVar[int]
ARCHITECTURE_FIELD_NUMBER: _ClassVar[int]
request_id: str
runtime_recipe_digest: str
chat_template_id: str
chat_template_version: str
reasoning_mode: str
architecture: ArchitectureType
def __init__(self, request_id: _Optional[str] = ..., runtime_recipe_digest: _Optional[str] = ..., chat_template_id: _Optional[str] = ..., chat_template_version: _Optional[str] = ..., reasoning_mode: _Optional[str] = ..., architecture: _Optional[_Union[ArchitectureType, str]] = ...) -> None: ...
class SamplingParameters(_message.Message):
__slots__ = ("temperature", "top_p", "top_k", "seed", "greedy")
TEMPERATURE_FIELD_NUMBER: _ClassVar[int]
TOP_P_FIELD_NUMBER: _ClassVar[int]
TOP_K_FIELD_NUMBER: _ClassVar[int]
SEED_FIELD_NUMBER: _ClassVar[int]
GREEDY_FIELD_NUMBER: _ClassVar[int]
temperature: float
top_p: float
top_k: int
seed: int
greedy: bool
def __init__(self, temperature: _Optional[float] = ..., top_p: _Optional[float] = ..., top_k: _Optional[int] = ..., seed: _Optional[int] = ..., greedy: _Optional[bool] = ...) -> None: ...
class TailResult(_message.Message):
__slots__ = ("identity", "sampling", "logits", "sampled_token_id")
IDENTITY_FIELD_NUMBER: _ClassVar[int]
SAMPLING_FIELD_NUMBER: _ClassVar[int]
LOGITS_FIELD_NUMBER: _ClassVar[int]
SAMPLED_TOKEN_ID_FIELD_NUMBER: _ClassVar[int]
identity: RequestRecipeIdentity
sampling: SamplingParameters
logits: TensorBundle
sampled_token_id: int
def __init__(self, identity: _Optional[_Union[RequestRecipeIdentity, _Mapping]] = ..., sampling: _Optional[_Union[SamplingParameters, _Mapping]] = ..., logits: _Optional[_Union[TensorBundle, _Mapping]] = ..., sampled_token_id: _Optional[int] = ...) -> None: ...
class ReleaseSignal(_message.Message):
__slots__ = ("route_session_id", "route_epoch", "work_id")
@@ -409,18 +468,20 @@ class SessionRequest(_message.Message):
def __init__(self, open: _Optional[_Union[SessionOpen, _Mapping]] = ..., chunk: _Optional[_Union[ActivationChunk, _Mapping]] = ..., decode: _Optional[_Union[DecodeStep, _Mapping]] = ..., flow_control: _Optional[_Union[FlowControl, _Mapping]] = ..., release: _Optional[_Union[ReleaseSignal, _Mapping]] = ..., cancel: _Optional[_Union[CancelSignal, _Mapping]] = ...) -> None: ...
class SessionResponse(_message.Message):
__slots__ = ("accepted", "chunk", "ack", "flow_control", "status")
__slots__ = ("accepted", "chunk", "ack", "flow_control", "status", "tail_result")
ACCEPTED_FIELD_NUMBER: _ClassVar[int]
CHUNK_FIELD_NUMBER: _ClassVar[int]
ACK_FIELD_NUMBER: _ClassVar[int]
FLOW_CONTROL_FIELD_NUMBER: _ClassVar[int]
STATUS_FIELD_NUMBER: _ClassVar[int]
TAIL_RESULT_FIELD_NUMBER: _ClassVar[int]
accepted: SessionAccepted
chunk: ActivationChunk
ack: Ack
flow_control: FlowControl
status: ShardStatus
def __init__(self, accepted: _Optional[_Union[SessionAccepted, _Mapping]] = ..., chunk: _Optional[_Union[ActivationChunk, _Mapping]] = ..., ack: _Optional[_Union[Ack, _Mapping]] = ..., flow_control: _Optional[_Union[FlowControl, _Mapping]] = ..., status: _Optional[_Union[ShardStatus, _Mapping]] = ...) -> None: ...
tail_result: TailResult
def __init__(self, accepted: _Optional[_Union[SessionAccepted, _Mapping]] = ..., chunk: _Optional[_Union[ActivationChunk, _Mapping]] = ..., ack: _Optional[_Union[Ack, _Mapping]] = ..., flow_control: _Optional[_Union[FlowControl, _Mapping]] = ..., status: _Optional[_Union[ShardStatus, _Mapping]] = ..., tail_result: _Optional[_Union[TailResult, _Mapping]] = ...) -> None: ...
class CapabilityRequest(_message.Message):
__slots__ = ("schema_version",)