122 lines
4.3 KiB
Python
122 lines
4.3 KiB
Python
"""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"
|