feat: add binary activation wire format

This commit is contained in:
Dobromir Popov
2026-06-29 14:58:55 +03:00
parent b17deb4d0e
commit 2a4383e353
6 changed files with 413 additions and 28 deletions

View File

@@ -1,4 +1,4 @@
"""US-002 integration test: two-node shard pipeline with activation tensors flowing A B."""
"""Integration tests for two-node shard pipelines with activation tensors flowing A to B."""
import json
import urllib.request
@@ -62,3 +62,54 @@ def test_two_node_activation_pipeline():
gateway.stop()
node_a.stop()
node_b.stop()
def test_binary_activation_pipeline_chunks_512_token_prefill():
node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False)
port_a = node_a.start()
node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True)
port_b = node_b.start()
gateway = GatewayServer(
inference_route=[
f"http://127.0.0.1:{port_a}",
f"http://127.0.0.1:{port_b}",
]
)
gateway_port = gateway.start()
try:
prompt = " ".join(f"tok{i}" for i in range(512))
payload = json.dumps({
"model": "stub-model",
"messages": [{"role": "user", "content": prompt}],
}).encode()
req = urllib.request.Request(
f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
data=payload,
headers={"Content-Type": "application/json"},
method="POST",
)
with urllib.request.urlopen(req) as resp:
assert resp.status == 200
body = json.loads(resp.read())
assert body["choices"][0]["message"]["content"] == f"stub response to: {prompt}"
assert node_a.forward_chunk_count == 4
assert node_b.forward_chunk_count == 4
assert not node_a.received_activations
assert node_b.received_activations
chunk_responses = gateway.last_binary_chunk_responses
assert len(chunk_responses) == 4
for index, chunk in enumerate(chunk_responses):
assert chunk.shape == [1, 128, 64]
assert chunk.dtype == "bfloat16"
assert chunk.chunk_index == index
assert chunk.chunk_total == 4
assert chunk.headers["x-meshnet-wire"] == "2"
assert len(chunk.body) == 1 * 128 * 64 * 2
finally:
gateway.stop()
node_a.stop()
node_b.stop()