feat: implement DGR-006 tensor bundle boundary

tests/test_architecture_boundary.py

@@ -0,0 +1,121 @@
+"""DGR-006 architecture-defined activation-boundary contract."""
+
+from __future__ import annotations
+
+import struct
+
+import pytest
+
+from meshnet_node.architecture_boundary import (
+    Architecture,
+    BoundaryStage,
+    ProtocolIdentity,
+    SamplingParameters,
+    TailOutput,
+    adapter_for,
+)
+from meshnet_node.native_protocol import ProtocolError, decode_bundle
+
+
+def _f32(values: list[float]) -> bytes:
+    return struct.pack("<" + "f" * len(values), *values)
+
+
+def _values(payload: bytes) -> tuple[float, ...]:
+    return struct.unpack("<" + "f" * (len(payload) // 4), payload)
+
+
+def test_dense_whole_model_and_two_ranges_match_for_prefill_and_greedy_decode() -> None:
+    adapter = adapter_for(Architecture.DENSE)
+    embeddings = {3: [1.0, 2.0], 7: [3.0, 4.0]}
+    head = adapter.bundle_from_token_ids([3, 7], lambda token: embeddings[token])
+
+    def head_layers(bundle):
+        return [value + 10.0 for value in _values(decode_bundle(bundle)["hidden_states"])]
+
+    def tail_layers(residual: list[float]) -> tuple[list[float], int]:
+        # The tail alone applies this fixture's final norm/head and greedy argmax.
+        logits = [sum(residual) / len(residual) + 20.0, 0.0]
+        return logits, max(range(len(logits)), key=logits.__getitem__)
+
+    # Whole-model prefill retains the unnormalized residual locally. The
+    # two-range lane sends that same residual before the tail-only head.
+    whole_prefill = head_layers(head)
+    seam = adapter.bundle_from_named_payloads(
+        {"hidden_states": _f32(whole_prefill)}, shape=[1, 2, 2]
+    )
+    two_range_prefill = list(_values(decode_bundle(adapter.input_for(BoundaryStage.TAIL, seam))["hidden_states"]))
+    assert two_range_prefill == whole_prefill
+
+    whole_logits, whole_token = tail_layers(whole_prefill)
+    two_logits, two_token = tail_layers(two_range_prefill)
+    assert two_logits == whole_logits
+    assert two_token == whole_token == 0
+
+    # One greedy decode step is the same contract with [batch, token, hidden].
+    decode_seam = adapter.bundle_from_named_payloads(
+        {"hidden_states": _f32([5.0, 6.0])}, shape=[1, 1, 2]
+    )
+    assert tail_layers([5.0, 6.0]) == tail_layers(
+        list(_values(decode_bundle(adapter.input_for(BoundaryStage.TAIL, decode_seam))["hidden_states"]))
+    )
+
+
+def test_middle_and_tail_reject_token_ids_and_require_boundary_bundle() -> None:
+    adapter = adapter_for(Architecture.MOE)
+
+    with pytest.raises(ProtocolError, match="head"):
+        adapter.bundle_from_token_ids([1], lambda _: [1.0])
+
+    with pytest.raises(ProtocolError, match="TensorBundle"):
+        adapter.input_for(BoundaryStage.MIDDLE, None)
+
+
+@pytest.mark.parametrize(
+    ("architecture", "names"),
+    [
+        (Architecture.MOE, {"hidden_states", "router_logits"}),
+        (Architecture.MLA, {"hidden_states", "mla_position_state"}),
+    ],
+)
+def test_architecture_adapters_route_and_validate_their_named_sidebands(
+    architecture: Architecture, names: set[str]
+) -> None:
+    adapter = adapter_for(architecture)
+    bundle = adapter.bundle_from_named_payloads(
+        {
+            name: (_f32([1.0, 2.0]) if name == "hidden_states" else _f32([0.0]))
+            for name in names
+        }
+    )
+
+    assert set(decode_bundle(adapter.input_for(BoundaryStage.MIDDLE, bundle))) == names
+    with pytest.raises(ProtocolError, match="requires"):
+        adapter.bundle_from_named_payloads({"hidden_states": _f32([1.0, 2.0])})
+
+
+def test_unknown_architecture_fails_closed() -> None:
+    with pytest.raises(ProtocolError, match="unsupported architecture"):
+        adapter_for("unchecked-name-substitution")
+
+
+def test_typed_tail_result_binds_sampling_and_request_recipe_identity() -> None:
+    adapter = adapter_for(Architecture.DENSE)
+    identity = ProtocolIdentity(
+        request_id="request-1",
+        runtime_recipe_digest="sha256:recipe",
+        chat_template_id="llama3",
+        chat_template_version="2",
+        reasoning_mode="max",
+        architecture=Architecture.DENSE,
+    )
+    result = adapter.tail_result(
+        identity=identity,
+        sampling=SamplingParameters(temperature=0.0, top_p=1.0, top_k=0, seed=9),
+        output=TailOutput.sampled_token(42),
+    )
+
+    assert result.identity.request_id == "request-1"
+    assert result.sampled_token_id == 42
+    assert result.output_kind == "sampled_token_id"
+    assert result.message.WhichOneof("output") == "sampled_token_id"

tests/test_native_shard_protocol.py

@@ -34,6 +34,8 @@ from meshnet_node.native_protocol import (
     ProtocolError,
     checksum_of,
     decode_bundle,
+    decode_step_bundle,
+    encode_decode_step,
     decode_tensor,
     default_flow_control,
     encode_bundle,
@@ -429,7 +431,7 @@ def test_decode_fast_path_is_much_smaller_than_a_full_envelope_chunk():
             idempotency_step=9,
             position=1024,
             expected_past_len=1024,
-            tensor=tensor,
+            bundle=encode_bundle([tensor]),
             work_id="work-7f3a",
         )
     )
@@ -443,7 +445,48 @@ def test_decode_fast_path_is_much_smaller_than_a_full_envelope_chunk():
     )
 
     assert len(fast.SerializeToString()) * 2 < len(full.SerializeToString())
-    assert decode_tensor(fast.decode.tensor) == hidden
+    assert decode_step_bundle(fast.decode) == {HIDDEN_STATES: hidden}
+
+
+def test_decode_fast_path_preserves_legacy_one_tensor_compatibility():
+    tensor = encode_tensor(HIDDEN_STATES, b"\x01\x02" * 8, [1, 1, 8], pb.DTYPE_BFLOAT16)
+    legacy = pb.DecodeStep(tensor=tensor)
+
+    assert decode_step_bundle(legacy) == {HIDDEN_STATES: b"\x01\x02" * 8}
+
+
+def test_decode_bundle_wins_over_legacy_tensor_and_can_carry_sidebands():
+    hidden = encode_tensor(HIDDEN_STATES, b"\x01\x02" * 8, [1, 1, 8], pb.DTYPE_BFLOAT16)
+    sideband = encode_tensor("index_topk", b"\x00" * 4, [1], pb.DTYPE_INT32)
+    decode = pb.DecodeStep(tensor=hidden, bundle=encode_bundle([hidden, sideband]))
+
+    assert decode_step_bundle(decode) == {
+        HIDDEN_STATES: b"\x01\x02" * 8,
+        "index_topk": b"\x00" * 4,
+    }
+
+
+def test_decode_writer_uses_compact_tensor_only_for_a_certified_one_tensor_bundle():
+    hidden = encode_tensor(HIDDEN_STATES, b"\x01\x02" * 8, [1, 1, 8], pb.DTYPE_BFLOAT16)
+    compact = encode_decode_step(
+        encode_bundle([hidden]), idempotency_step=1, position=1, expected_past_len=1, work_id="w"
+    )
+    sideband = encode_tensor("index_topk", b"\x00" * 4, [1], pb.DTYPE_INT32)
+    expanded = encode_decode_step(
+        encode_bundle([hidden, sideband]), idempotency_step=1, position=1, expected_past_len=1, work_id="w"
+    )
+
+    assert compact.HasField("tensor") and not compact.HasField("bundle")
+    assert expanded.HasField("bundle") and not expanded.HasField("tensor")
+
+
+def test_schema_exposes_typed_tail_result_and_bound_sampling_identity():
+    assert {"identity", "sampling", "logits", "sampled_token_id"} <= set(
+        pb.TailResult.DESCRIPTOR.fields_by_name
+    )
+    assert {"request_id", "runtime_recipe_digest", "chat_template_id", "chat_template_version", "reasoning_mode"} <= set(
+        pb.RequestRecipeIdentity.DESCRIPTOR.fields_by_name
+    )
 
 
 def test_flow_control_defaults_bound_the_queue_and_the_message():
@@ -491,6 +534,20 @@ def test_committed_vectors_still_encode_as_promised():
     report = (conformance.TESTDATA_DIR / conformance.GOLDEN_CAPABILITY_REPORT).read_bytes()
     assert conformance.serialize(conformance.canonical_capability_report()) == report
 
+    decode = (conformance.TESTDATA_DIR / conformance.GOLDEN_DECODE_STEP).read_bytes()
+    assert conformance.serialize(conformance.canonical_decode_step()) == decode
+
+
+def test_decode_golden_preserves_the_multi_tensor_boundary():
+    golden = (conformance.TESTDATA_DIR / conformance.GOLDEN_DECODE_STEP).read_bytes()
+    request = pb.SessionRequest.FromString(golden)
+
+    assert request.decode.idempotency_step == 43
+    assert decode_step_bundle(request.decode) == {
+        HIDDEN_STATES: bytes(range(16)),
+        "index_topk": (3).to_bytes(4, "little"),
+    }
+
 
 def test_golden_session_request_round_trips_with_every_field_intact():
     golden = (conformance.TESTDATA_DIR / conformance.GOLDEN_SESSION_REQUEST).read_bytes()