feat: implement DGR-006 tensor bundle boundary
This commit is contained in:
138
packages/node/meshnet_node/architecture_boundary.py
Normal file
138
packages/node/meshnet_node/architecture_boundary.py
Normal file
@@ -0,0 +1,138 @@
|
||||
"""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
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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,
|
||||
*,
|
||||
|
||||
@@ -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)
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -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",)
|
||||
|
||||
@@ -32,6 +32,15 @@ Both `--check` modes run in CI via `tests/test_native_shard_protocol.py`, so a
|
||||
schema edit that is not accompanied by regenerated output fails the suite rather
|
||||
than shipping stubs that disagree with the schema they claim to implement.
|
||||
|
||||
## DGR-006 decode and tail compatibility
|
||||
|
||||
`DecodeStep.bundle` is the versioned `TensorBundle` fast-path boundary. It is
|
||||
authoritative whenever present and supports architecture sidebands. The original
|
||||
`DecodeStep.tensor` remains readable as the compact one-tensor encoding for
|
||||
certified boundaries that need only one tensor; new readers wrap it into a
|
||||
one-member bundle. Tail completions use `TailResult`, which binds logits or a
|
||||
sampled token to request/recipe identity and sampling/template/reasoning inputs.
|
||||
|
||||
## Building and running the C++ conformance test
|
||||
|
||||
If the machine has no protobuf C++ toolchain:
|
||||
|
||||
@@ -33,6 +33,16 @@ enum SchemaVersion {
|
||||
SCHEMA_VERSION_1 = 1;
|
||||
}
|
||||
|
||||
// Certified transformer boundary family. Names alone are not an adapter: a
|
||||
// Node must select an explicit architecture contract before interpreting a
|
||||
// TensorBundle.
|
||||
enum ArchitectureType {
|
||||
ARCHITECTURE_TYPE_UNSPECIFIED = 0;
|
||||
ARCHITECTURE_TYPE_DENSE = 1;
|
||||
ARCHITECTURE_TYPE_MOE = 2;
|
||||
ARCHITECTURE_TYPE_MLA = 3;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tensor bundle — the public activation boundary
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -142,6 +152,11 @@ message TensorBundle {
|
||||
// boundary payload can evolve without a whole-protocol version bump.
|
||||
uint32 bundle_version = 1;
|
||||
repeated NamedTensor tensors = 2;
|
||||
// Explicit adapter selection; names are never interpreted through unchecked
|
||||
// substitutions. UNSPECIFIED is accepted only for legacy DGR-002 bundles.
|
||||
ArchitectureType architecture = 3;
|
||||
// Adapter-owned semantic boundary, such as "pre_tail_residual".
|
||||
string boundary_point = 4;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -393,8 +408,9 @@ message ActivationChunk {
|
||||
// A decode step is one token: the envelope's identity fields are already fixed
|
||||
// for the life of the stream, so repeating them per token is pure overhead on
|
||||
// the hottest path. A DecodeStep carries only what changes — the step, the
|
||||
// position, and one tensor — and inherits the rest from the SessionOpen
|
||||
// handshake. A peer may always fall back to ActivationChunk with PHASE_DECODE.
|
||||
// position, and one compact boundary encoding — and inherits the rest from the
|
||||
// SessionOpen handshake. A Node may always fall back to ActivationChunk with
|
||||
// PHASE_DECODE.
|
||||
message DecodeStep {
|
||||
// Idempotency step within the session; also orders the stream.
|
||||
uint64 idempotency_step = 1;
|
||||
@@ -402,11 +418,50 @@ message DecodeStep {
|
||||
uint64 position = 2;
|
||||
// Tokens the receiver's cache must already hold. A mismatch is a CACHE_MISS.
|
||||
uint64 expected_past_len = 3;
|
||||
// The single boundary tensor for this token, typically [1, 1, hidden].
|
||||
// Legacy compact one-tensor boundary, typically [1, 1, hidden]. New readers
|
||||
// accept it as a TensorBundle with one member; writers retain it where the
|
||||
// selected architecture genuinely requires only one tensor.
|
||||
NamedTensor tensor = 4;
|
||||
// Work id for attribution; the session/epoch/fingerprint/range are inherited.
|
||||
string work_id = 5;
|
||||
int64 deadline_unix_nanos = 6;
|
||||
// Versioned architecture boundary. This carries sidebands such as MoE router
|
||||
// state or MLA/IndexShare state; when both fields are set, bundle is
|
||||
// authoritative and tensor is retained only for older-Node forwarding.
|
||||
TensorBundle bundle = 7;
|
||||
}
|
||||
|
||||
// Exact request-level identity for a tail result. Runtime recipe identity is
|
||||
// deliberately repeated here: sampling the right logits with the wrong chat
|
||||
// template or reasoning mode is a different request, not a compatible result.
|
||||
message RequestRecipeIdentity {
|
||||
string request_id = 1;
|
||||
string runtime_recipe_digest = 2;
|
||||
string chat_template_id = 3;
|
||||
string chat_template_version = 4;
|
||||
string reasoning_mode = 5;
|
||||
ArchitectureType architecture = 6;
|
||||
}
|
||||
|
||||
// Sampling is an explicit tail-only operation. A non-tail Shard never applies
|
||||
// final norm, LM head, row pruning, or sampling to its boundary output.
|
||||
message SamplingParameters {
|
||||
float temperature = 1;
|
||||
float top_p = 2;
|
||||
uint32 top_k = 3;
|
||||
uint64 seed = 4;
|
||||
bool greedy = 5;
|
||||
}
|
||||
|
||||
// Typed tail completion. The oneof keeps token id 0 unambiguous and prevents a
|
||||
// caller from treating an activation tensor as an inferred completion.
|
||||
message TailResult {
|
||||
RequestRecipeIdentity identity = 1;
|
||||
SamplingParameters sampling = 2;
|
||||
oneof output {
|
||||
TensorBundle logits = 3;
|
||||
uint32 sampled_token_id = 4;
|
||||
}
|
||||
}
|
||||
|
||||
// Drop session state for a Route Session. Bounded memory does not depend on
|
||||
@@ -471,7 +526,8 @@ message SessionResponse {
|
||||
ActivationChunk chunk = 2;
|
||||
Ack ack = 3;
|
||||
FlowControl flow_control = 4;
|
||||
ShardStatus status = 5;
|
||||
ShardStatus status = 5;
|
||||
TailResult tail_result = 6;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
BIN
packages/node/native/testdata/decode_step_golden.binpb
vendored
Normal file
BIN
packages/node/native/testdata/decode_step_golden.binpb
vendored
Normal file
Binary file not shown.
@@ -265,7 +265,9 @@ void TestDecodeFastPathIsSmall() {
|
||||
step->set_position(1024);
|
||||
step->set_expected_past_len(1024);
|
||||
step->set_work_id("work-7f3a");
|
||||
sp::NamedTensor *tensor = step->mutable_tensor();
|
||||
sp::TensorBundle *bundle = step->mutable_bundle();
|
||||
bundle->set_bundle_version(1);
|
||||
sp::NamedTensor *tensor = bundle->add_tensors();
|
||||
tensor->set_name("hidden_states");
|
||||
tensor->add_shape(1);
|
||||
tensor->add_shape(1);
|
||||
@@ -290,7 +292,27 @@ void TestDecodeFastPathIsSmall() {
|
||||
CHECK(parsed.ParseFromString(bytes));
|
||||
CHECK(parsed.kind_case() == sp::SessionRequest::kDecode);
|
||||
CHECK_EQ(parsed.decode().position(), 1024u);
|
||||
CHECK_EQ(parsed.decode().tensor().total_bytes(), 16u);
|
||||
CHECK_EQ(parsed.decode().bundle().tensors_size(), 1);
|
||||
CHECK_EQ(parsed.decode().bundle().tensors(0).total_bytes(), 16u);
|
||||
}
|
||||
|
||||
void TestDecodeBundleVector(const std::string &bytes) {
|
||||
sp::SessionRequest request;
|
||||
CHECK(request.ParseFromString(bytes));
|
||||
CHECK(request.kind_case() == sp::SessionRequest::kDecode);
|
||||
const sp::DecodeStep &step = request.decode();
|
||||
CHECK_EQ(step.idempotency_step(), 43u);
|
||||
CHECK_EQ(step.position(), 384u);
|
||||
CHECK_EQ(step.expected_past_len(), 384u);
|
||||
CHECK_EQ(step.work_id(), std::string("decode-7f3a"));
|
||||
CHECK(step.has_bundle());
|
||||
CHECK_EQ(step.bundle().bundle_version(), 1u);
|
||||
CHECK_EQ(step.bundle().architecture(), sp::ARCHITECTURE_TYPE_MLA);
|
||||
CHECK_EQ(step.bundle().boundary_point(), std::string("pre_tail_residual"));
|
||||
CHECK_EQ(step.bundle().tensors_size(), 2);
|
||||
CHECK_EQ(step.bundle().tensors(0).name(), std::string("hidden_states"));
|
||||
CHECK_EQ(step.bundle().tensors(1).name(), std::string("index_topk"));
|
||||
CHECK_EQ(step.bundle().tensors(1).dtype(), sp::DTYPE_INT32);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
@@ -308,9 +330,11 @@ int main(int argc, char **argv) {
|
||||
ReadFile(testdata / "session_request_golden.binpb");
|
||||
const std::string capability_bytes =
|
||||
ReadFile(testdata / "capability_report_golden.binpb");
|
||||
const std::string decode_bytes = ReadFile(testdata / "decode_step_golden.binpb");
|
||||
|
||||
TestSessionRequestVector(session_bytes);
|
||||
TestCapabilityReportVector(capability_bytes);
|
||||
TestDecodeBundleVector(decode_bytes);
|
||||
TestUnknownFieldsArePreserved(session_bytes);
|
||||
TestSparseMessageParses();
|
||||
TestDecodeFastPathIsSmall();
|
||||
|
||||
Reference in New Issue
Block a user