popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: implement DGR-006 tensor bundle boundary

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-14 13:44:37 +03:00
parent 91c450840d
commit 7925e5253d
16 changed files with 802 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
packages/node/native/tests/test_shard_protocol_conformance.cpp
View File

@@ -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();