story: DGR-035 Implement dense architecture boundary input/output

packages/node/meshnet_node/architecture_boundary.py

@@ -20,6 +20,14 @@ from .native_protocol import (
     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):
@@ -63,6 +71,11 @@ class TailOutput:
             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:
@@ -148,28 +161,153 @@ class ArchitectureBoundaryAdapter:
             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 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")
-        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)
+        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 = {