From 11752f614bbd9019625aa803d765cce75f8b1a05 Mon Sep 17 00:00:00 2001 From: dlimpus Date: Thu, 25 Jun 2026 19:52:51 -0500 Subject: [PATCH 1/3] Fix(io): track RDMA notification completions in transfer status --- src/io/rdma/backend_impl.cpp | 43 +++++++++++++++++++++++++++--------- src/io/rdma/common.cpp | 22 +++++++++++++----- src/io/rdma/common.hpp | 4 +++- tests/cpp/io/test_engine.cpp | 30 +++++++++++++++++++++++++ 4 files changed, 83 insertions(+), 16 deletions(-) diff --git a/src/io/rdma/backend_impl.cpp b/src/io/rdma/backend_impl.cpp index 0e901438e..5a00cd156 100644 --- a/src/io/rdma/backend_impl.cpp +++ b/src/io/rdma/backend_impl.cpp @@ -738,13 +738,33 @@ NotifManager::FlushDrainStats NotifManager::ProcessOneCqe( struct ibv_recv_wr* bad = nullptr; SYSCALL_RETURN_ZERO(ibv_post_recv(ep.local.ibvHandle.qp, &wr, &bad)); } else if (wc[i].opcode == IBV_WC_SEND) { - if (!IsNotifSendWrId(wc[i].wr_id)) { - MORI_IO_WARN( - "ProcessOneCqe: unexpected SEND completion with non-notification wr_id {}; " - "releasing 1 sqDepth under current SEND invariant", - wc[i].wr_id); + if (IsNotifSendWrId(wc[i].wr_id)) { + // Backward-compatible handling for any notification SENDs posted by + // older code that used tagged transfer IDs instead of ledger records. + if (ep.sqDepth) ep.sqDepth->fetch_sub(1, std::memory_order_relaxed); + } else { + int mergedBatchSize = 0; + auto meta = ep.ledger + ? ep.ledger->ReleaseByCqe(wc[i].wr_id, ep.sqDepth.get(), &mergedBatchSize) + : nullptr; + if (meta) { + uint32_t finishedBefore = meta->finishedBatchSize.fetch_add(mergedBatchSize); + TransferStatus* statusPtr = meta->status; + if (statusPtr != nullptr && + (finishedBefore + mergedBatchSize) == meta->totalBatchSize) { + statusPtr->Update(StatusCode::SUCCESS, ibv_wc_status_str(wc[i].status)); + } + MORI_IO_TRACE( + "ProcessOneCqe: notification SEND CQE for task {} total={} finished={} cur={}", + meta->id, meta->totalBatchSize, finishedBefore, mergedBatchSize); + } else { + MORI_IO_WARN( + "ProcessOneCqe: notification SEND CQE has no ledger record for wr_id {}; " + "releasing 1 sqDepth under fallback path", + wc[i].wr_id); + if (ep.sqDepth) ep.sqDepth->fetch_sub(1, std::memory_order_relaxed); + } } - if (ep.sqDepth) ep.sqDepth->fetch_sub(1, std::memory_order_relaxed); } else { // Batch path: wr_id carries a recordId from the SubmissionLedger. uint64_t recordId = wc[i].wr_id; @@ -1198,7 +1218,8 @@ void RdmaBackendSession::ReadWrite(size_t localOffset, size_t remoteOffset, size TransferStatus* status, TransferUniqueId id, bool isRead) { MORI_IO_FUNCTION_TIMER; status->SetCode(StatusCode::IN_PROGRESS); - auto callbackMeta = std::make_shared(status, id, 1); + const int notifBatchSize = config.enableNotification ? static_cast(eps.size()) : 0; + auto callbackMeta = std::make_shared(status, id, 1 + notifBatchSize); internal::PublishCurrentIoCallDiagnostics(callbackMeta); RdmaOpRet ret = RdmaBatchReadWrite(eps, localMrPerEp, remoteMrPerEp, {localOffset}, @@ -1211,7 +1232,7 @@ void RdmaBackendSession::ReadWrite(size_t localOffset, size_t remoteOffset, size status->Update(ret.code, ret.message); } if (!ret.Failed() && config.enableNotification) { - RdmaOpRet notifRet = RdmaNotifyTransfer(eps, status, id); + RdmaOpRet notifRet = RdmaNotifyTransfer(eps, callbackMeta, id); if (notifRet.Failed()) { status->Update(notifRet.code, notifRet.message); } @@ -1223,7 +1244,9 @@ void RdmaBackendSession::BatchReadWrite(const SizeVec& localOffsets, const SizeV TransferUniqueId id, bool isRead) { MORI_IO_FUNCTION_TIMER; status->SetCode(StatusCode::IN_PROGRESS); - auto callbackMeta = std::make_shared(status, id, sizes.size()); + const int notifBatchSize = config.enableNotification ? static_cast(eps.size()) : 0; + auto callbackMeta = std::make_shared( + status, id, static_cast(sizes.size()) + notifBatchSize); internal::PublishCurrentIoCallDiagnostics(callbackMeta); RdmaOpRet ret; if (executor) { @@ -1241,7 +1264,7 @@ void RdmaBackendSession::BatchReadWrite(const SizeVec& localOffsets, const SizeV status->Update(ret.code, ret.message); } if (!ret.Failed() && config.enableNotification) { - RdmaOpRet notifRet = RdmaNotifyTransfer(eps, status, id); + RdmaOpRet notifRet = RdmaNotifyTransfer(eps, callbackMeta, id); if (notifRet.Failed()) { status->Update(notifRet.code, notifRet.message); } diff --git a/src/io/rdma/common.cpp b/src/io/rdma/common.cpp index 7e24dab91..64ccb79f8 100644 --- a/src/io/rdma/common.cpp +++ b/src/io/rdma/common.cpp @@ -370,9 +370,10 @@ static void ResetMergedWorkRequestPointers(MergedWorkRequest* wr) { /* Rdma Utilities */ /* ---------------------------------------------------------------------------------------------- */ -RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, TransferStatus* status, TransferUniqueId id) { +RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, + std::shared_ptr callbackMeta, + TransferUniqueId id) { MORI_IO_FUNCTION_TIMER; - (void)status; std::string reserveErr; int reserved = 0; @@ -388,6 +389,12 @@ RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, TransferStatus* status, Trans } for (size_t i = 0; i < eps.size(); i++) { + if (!eps[i].ledger) { + for (int j = static_cast(i); j < reserved; ++j) ReleaseSqDepth(eps[j], 1); + return {StatusCode::ERR_RDMA_OP, + "submission ledger is not initialized for notification SEND tracking"}; + } + const application::RdmaEndpoint& ep = eps[i].local; NotifMessage msg{id, static_cast(i), static_cast(eps.size())}; @@ -396,8 +403,10 @@ RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, TransferStatus* status, Trans sge.length = sizeof(NotifMessage); sge.lkey = 0; + const uint64_t recordId = eps[i].ledger->Insert(1, true, callbackMeta, 1); + struct ibv_send_wr wr{}; - wr.wr_id = MakeNotifSendWrId(id); + wr.wr_id = recordId; wr.opcode = IBV_WR_SEND; wr.send_flags = IBV_SEND_INLINE | IBV_SEND_SIGNALED; wr.sg_list = &sge; @@ -406,8 +415,11 @@ RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, TransferStatus* status, Trans struct ibv_send_wr* bad_wr = nullptr; int ret = ibv_post_send(ep.ibvHandle.qp, &wr, &bad_wr); if (ret != 0) { - // WR i was reserved but failed to post if bad_wr points at this WR. - if (bad_wr == &wr) ReleaseSqDepth(eps[i], 1); + // This call posts a single WR, so a non-zero return means the current + // notification was not accepted and no CQE will arrive for its ledger record. + ReleaseSqDepth(eps[i], 1); + int dummy = 0; + eps[i].ledger->ReleaseByCqe(recordId, nullptr, &dummy); // Any remaining endpoints are reserved but not posted yet. for (int j = i + 1; j < eps.size(); ++j) ReleaseSqDepth(eps[j], 1); std::string message = diff --git a/src/io/rdma/common.hpp b/src/io/rdma/common.hpp index ee30be208..449a421cb 100644 --- a/src/io/rdma/common.hpp +++ b/src/io/rdma/common.hpp @@ -242,7 +242,9 @@ struct RdmaOpRet { bool Failed() { return code > StatusCode::ERR_BEGIN; } }; -RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, TransferStatus* status, TransferUniqueId id); +RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, + std::shared_ptr callbackMeta, + TransferUniqueId id); RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, const std::vector& localMrPerEp, diff --git a/tests/cpp/io/test_engine.cpp b/tests/cpp/io/test_engine.cpp index 3ca816310..78b16a281 100644 --- a/tests/cpp/io/test_engine.cpp +++ b/tests/cpp/io/test_engine.cpp @@ -246,6 +246,35 @@ void CaseSubmissionLedgerBasic() { Require(sqDepth2.load(std::memory_order_relaxed) == 5, "sq depth after posted CQE release"); } +void CaseNotificationCompletionFanIn() { + TransferStatus status; + status.SetCode(StatusCode::IN_PROGRESS); + auto meta = std::make_shared(&status, 42, 3); + + uint32_t finishedBefore = meta->finishedBatchSize.fetch_add(2); + if (finishedBefore + 2 == meta->totalBatchSize) { + status.Update(StatusCode::SUCCESS, "data complete"); + } + Require(status.InProgress(), "data completion must not finish before notification SEND CQE"); + + finishedBefore = meta->finishedBatchSize.fetch_add(1); + if (finishedBefore + 1 == meta->totalBatchSize) { + status.Update(StatusCode::SUCCESS, "notification complete"); + } + Require(status.Succeeded(), "final notification completion should finish transfer"); + + TransferStatus failedStatus; + failedStatus.SetCode(StatusCode::IN_PROGRESS); + auto failedMeta = std::make_shared(&failedStatus, 43, 2); + (void)failedMeta->finishedBatchSize.fetch_add(1); + failedStatus.Update(StatusCode::ERR_RDMA_OP, "notification failed"); + finishedBefore = failedMeta->finishedBatchSize.fetch_add(1); + if (finishedBefore + 1 == failedMeta->totalBatchSize) { + failedStatus.Update(StatusCode::SUCCESS, "late success"); + } + Require(failedStatus.Failed(), "notification failure must not be overwritten by late success"); +} + void CaseWrIdNamespaceHelpers() { const uint64_t taggedZero = MakeNotifSendWrId(0); Require(taggedZero == kNotifSendWrIdTag, "tagged zero should only set the reserved high bit"); @@ -1451,6 +1480,7 @@ int main(int argc, char* argv[]) { SetLogLevel("info"); std::vector cases = { {"submission_ledger_basic", CaseSubmissionLedgerBasic}, + {"notification_completion_fan_in", CaseNotificationCompletionFanIn}, {"wr_id_namespace_helpers", CaseWrIdNamespaceHelpers}, {"rdma_backend_config_chunking_fields", CaseRdmaBackendConfigChunkingFields}, {"resolve_requested_nics", CaseResolveRequestedNics}, From df59641b344c6b0320d573024295deea23af3c23 Mon Sep 17 00:00:00 2001 From: dlimpus Date: Fri, 26 Jun 2026 17:03:33 -0500 Subject: [PATCH 2/3] Fix(io): preserve notify count in WR-counted batches When transfer chunking rewrites batch accounting from request count to final WR count, keep the notification SEND completions in the total so success cannot be skipped or reported early when CQEs arrive out of order. Signed-off-by: dlimpus --- src/io/rdma/backend_impl.cpp | 4 ++-- src/io/rdma/common.cpp | 7 ++++--- src/io/rdma/common.hpp | 3 +-- tests/cpp/io/test_engine.cpp | 11 +++++++++++ 4 files changed, 18 insertions(+), 7 deletions(-) diff --git a/src/io/rdma/backend_impl.cpp b/src/io/rdma/backend_impl.cpp index 5a00cd156..98b50e4dc 100644 --- a/src/io/rdma/backend_impl.cpp +++ b/src/io/rdma/backend_impl.cpp @@ -1245,8 +1245,8 @@ void RdmaBackendSession::BatchReadWrite(const SizeVec& localOffsets, const SizeV MORI_IO_FUNCTION_TIMER; status->SetCode(StatusCode::IN_PROGRESS); const int notifBatchSize = config.enableNotification ? static_cast(eps.size()) : 0; - auto callbackMeta = std::make_shared( - status, id, static_cast(sizes.size()) + notifBatchSize); + auto callbackMeta = + std::make_shared(status, id, static_cast(sizes.size()) + notifBatchSize); internal::PublishCurrentIoCallDiagnostics(callbackMeta); RdmaOpRet ret; if (executor) { diff --git a/src/io/rdma/common.cpp b/src/io/rdma/common.cpp index 64ccb79f8..b1d915106 100644 --- a/src/io/rdma/common.cpp +++ b/src/io/rdma/common.cpp @@ -370,8 +370,7 @@ static void ResetMergedWorkRequestPointers(MergedWorkRequest* wr) { /* Rdma Utilities */ /* ---------------------------------------------------------------------------------------------- */ -RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, - std::shared_ptr callbackMeta, +RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, std::shared_ptr callbackMeta, TransferUniqueId id) { MORI_IO_FUNCTION_TIMER; @@ -687,8 +686,10 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, if (mergedWrCount > static_cast(std::numeric_limits::max())) { return {StatusCode::ERR_INVALID_ARGS, "final WR count exceeds int range"}; } + const int notifBatchSize = + std::max(0, callbackMeta->totalBatchSize - static_cast(batchSize)); for (size_t k = 0; k < mergedWrCount; ++k) mergedWrs[k].mergedRequests = 1; - callbackMeta->totalBatchSize = static_cast(mergedWrCount); + callbackMeta->totalBatchSize = static_cast(mergedWrCount) + notifBatchSize; } size_t epNum = eps.size(); diff --git a/src/io/rdma/common.hpp b/src/io/rdma/common.hpp index 449a421cb..a8b853562 100644 --- a/src/io/rdma/common.hpp +++ b/src/io/rdma/common.hpp @@ -242,8 +242,7 @@ struct RdmaOpRet { bool Failed() { return code > StatusCode::ERR_BEGIN; } }; -RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, - std::shared_ptr callbackMeta, +RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, std::shared_ptr callbackMeta, TransferUniqueId id); RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, diff --git a/tests/cpp/io/test_engine.cpp b/tests/cpp/io/test_engine.cpp index 78b16a281..2f41f96ea 100644 --- a/tests/cpp/io/test_engine.cpp +++ b/tests/cpp/io/test_engine.cpp @@ -263,6 +263,17 @@ void CaseNotificationCompletionFanIn() { } Require(status.Succeeded(), "final notification completion should finish transfer"); + TransferStatus outOfOrderStatus; + outOfOrderStatus.SetCode(StatusCode::IN_PROGRESS); + auto outOfOrderMeta = std::make_shared(&outOfOrderStatus, 44, 5); + (void)outOfOrderMeta->finishedBatchSize.fetch_add(2); // notification SEND CQEs first + finishedBefore = outOfOrderMeta->finishedBatchSize.fetch_add(3); + if (finishedBefore + 3 == outOfOrderMeta->totalBatchSize) { + outOfOrderStatus.Update(StatusCode::SUCCESS, "data complete after notification"); + } + Require(outOfOrderStatus.Succeeded(), + "notification-first completion order must still finish at the exact total"); + TransferStatus failedStatus; failedStatus.SetCode(StatusCode::IN_PROGRESS); auto failedMeta = std::make_shared(&failedStatus, 43, 2); From 75888a46cb7ebe0ba178cd5dea221aead92d6c64 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 9 Jul 2026 00:17:05 +0000 Subject: [PATCH 3/3] feat(umbp): add TurboQuant KV encoding metadata Thread versioned KV encoding descriptors through UMBP metadata, routing, and client APIs so quantized KV payloads can move over Mori IO as opaque bytes. Co-authored-by: Cursor --- .../distributed/master/global_block_index.cpp | 6 + src/umbp/distributed/master/master_client.cpp | 8 +- src/umbp/distributed/master/master_server.cpp | 5 + .../distributed/peer/peer_dram_allocator.cpp | 25 +++- src/umbp/distributed/peer/peer_service.cpp | 7 +- .../distributed/peer/peer_ssd_manager.cpp | 17 ++- .../distributed/peer/ssd_copy_pipeline.cpp | 2 +- src/umbp/distributed/pool_client.cpp | 127 ++++++++++++---- src/umbp/distributed/proto/umbp.proto | 37 +++++ src/umbp/distributed/proto/umbp_peer.proto | 5 + src/umbp/include/umbp/codec/kv_encoding.h | 141 ++++++++++++++++++ .../umbp/distributed/master/master_client.h | 1 + .../distributed/peer/batch_resolve_codec.h | 107 ++++++++++++- .../distributed/peer/peer_dram_allocator.h | 11 +- .../umbp/distributed/peer/peer_ssd_manager.h | 3 + .../include/umbp/distributed/pool_client.h | 26 +++- src/umbp/include/umbp/distributed/types.h | 59 +++++++- src/umbp/tests/CMakeLists.txt | 11 ++ src/umbp/tests/test_batch_resolve_codec.cpp | 23 +++ src/umbp/tests/test_kv_encoding.cpp | 71 +++++++++ src/umbp/tests/test_peer_dram_allocator.cpp | 41 +++++ 21 files changed, 690 insertions(+), 43 deletions(-) create mode 100644 src/umbp/include/umbp/codec/kv_encoding.h create mode 100644 src/umbp/tests/test_kv_encoding.cpp diff --git a/src/umbp/distributed/master/global_block_index.cpp b/src/umbp/distributed/master/global_block_index.cpp index 247f8c963..49fd36f4f 100644 --- a/src/umbp/distributed/master/global_block_index.cpp +++ b/src/umbp/distributed/master/global_block_index.cpp @@ -130,6 +130,9 @@ size_t ApplyAddOrRemoveLocked(EntryMap& entries, NodeToKeys& node_to_keys, return 0; } loc->size = ev.size; + loc->encoding = ev.encoding; + if (loc->encoding.stored_bytes == 0) loc->encoding.stored_bytes = ev.size; + if (loc->encoding.logical_bytes == 0) loc->encoding.logical_bytes = loc->encoding.stored_bytes; return 1; } // REMOVE @@ -222,6 +225,9 @@ void GlobalBlockIndex::ReplaceNodeLocations(const std::string& node_id, auto [loc, inserted] = FindOrInsertLocation(entry, node_id, ev->tier); (void)inserted; loc->size = ev->size; + loc->encoding = ev->encoding; + if (loc->encoding.stored_bytes == 0) loc->encoding.stored_bytes = ev->size; + if (loc->encoding.logical_bytes == 0) loc->encoding.logical_bytes = loc->encoding.stored_bytes; sh.node_to_keys[node_id].insert(ev->key); } } diff --git a/src/umbp/distributed/master/master_client.cpp b/src/umbp/distributed/master/master_client.cpp index 59ddbe5cb..3a76cfa18 100644 --- a/src/umbp/distributed/master/master_client.cpp +++ b/src/umbp/distributed/master/master_client.cpp @@ -35,6 +35,7 @@ #include "umbp/common/env_time.h" #include "umbp/distributed/master/master_metrics.h" #include "umbp/distributed/master/rpc_latency_timer.h" +#include "umbp/distributed/peer/batch_resolve_codec.h" #include "umbp/distributed/peer/peer_dram_allocator.h" #include "umbp/distributed/peer/peer_ssd_manager.h" @@ -108,6 +109,7 @@ void FillBundle(::umbp::EventBundle* dst, const EventBundle& src) { pe->set_key(ev.key); pe->set_tier(ToProtoTier(ev.tier)); pe->set_size(ev.size); + FillProtoKvEncoding(ev.encoding, pe->mutable_encoding()); } } @@ -258,6 +260,7 @@ grpc::Status MasterClient::RouteGet(const std::string& key, r.node_id = resp.node_id(); r.tier = FromProtoTier(resp.tier()); r.size = resp.size(); + r.encoding = KvEncodingFromProto(resp.encoding(), r.size); r.peer_address = resp.peer_address(); *out_result = std::move(r); return grpc::Status::OK; @@ -333,7 +336,8 @@ grpc::Status MasterClient::BatchRouteGet(const std::vector& keys, // Columnar response: node_ref[i] is a 1-based index into resp.nodes() // (0 = not found); tier[i] / size[i] are parallel per-key arrays. const int n = resp.node_ref_size(); - if (resp.tier_size() != n || resp.size_size() != n) { + if (resp.tier_size() != n || resp.size_size() != n || + (resp.encoding_size() != 0 && resp.encoding_size() != n)) { return grpc::Status(grpc::StatusCode::INTERNAL, "BatchRouteGet: malformed columnar response (array length mismatch)"); } @@ -349,6 +353,8 @@ grpc::Status MasterClient::BatchRouteGet(const std::vector& keys, r.node_id = node.node_id(); r.tier = FromProtoTier(resp.tier(i)); r.size = resp.size(i); + r.encoding = + resp.encoding_size() == n ? KvEncodingFromProto(resp.encoding(i), r.size) : RawKvEncoding(r.size); r.peer_address = node.peer_address(); (*out)[i] = std::move(r); } diff --git a/src/umbp/distributed/master/master_server.cpp b/src/umbp/distributed/master/master_server.cpp index 3ebd9db5a..2a5f2920b 100644 --- a/src/umbp/distributed/master/master_server.cpp +++ b/src/umbp/distributed/master/master_server.cpp @@ -32,6 +32,7 @@ #include #include "mori/utils/mori_log.hpp" +#include "umbp/distributed/peer/batch_resolve_codec.h" #include "umbp.grpc.pb.h" #include "umbp/common/env_time.h" #include "umbp/distributed/master/external_kv_block_index.h" @@ -98,6 +99,7 @@ KvEvent FromProtoEvent(const ::umbp::KvEvent& pe) { ev.key = pe.key(); ev.tier = static_cast(pe.tier()); ev.size = pe.size(); + ev.encoding = KvEncodingFromProto(pe.encoding(), ev.size); return ev; } @@ -376,6 +378,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser response->set_tier(static_cast<::umbp::TierType>(result->location.tier)); response->set_size(result->location.size); response->set_peer_address(result->peer_address); + FillProtoKvEncoding(result->location.encoding, response->mutable_encoding()); if (metrics_) { metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_ROUTE_GET, @@ -439,6 +442,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser response->add_node_ref(0); response->add_tier(::umbp::TIER_UNKNOWN); response->add_size(0); + FillProtoKvEncoding(RawKvEncoding(0), response->add_encoding()); continue; } std::string node_key = opt->location.node_id; @@ -454,6 +458,7 @@ class MasterServer::UMBPMasterServiceImpl final : public ::umbp::UMBPMaster::Ser response->add_node_ref(it->second); response->add_tier(static_cast<::umbp::TierType>(opt->location.tier)); response->add_size(opt->location.size); + FillProtoKvEncoding(opt->location.encoding, response->add_encoding()); if (metrics_) { metrics_->addCounter(MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_GET, MORI_UMBP_METRIC_CLIENT_BATCH_ROUTE_GET_HELP, diff --git a/src/umbp/distributed/peer/peer_dram_allocator.cpp b/src/umbp/distributed/peer/peer_dram_allocator.cpp index 063786c0c..76267037c 100644 --- a/src/umbp/distributed/peer/peer_dram_allocator.cpp +++ b/src/umbp/distributed/peer/peer_dram_allocator.cpp @@ -28,6 +28,7 @@ #include #include "mori/utils/mori_log.hpp" +#include "umbp/codec/kv_encoding.h" namespace mori::umbp { @@ -101,12 +102,18 @@ const PageBitmapAllocator* PeerDramAllocator::AllocatorForLocked(TierType tier) PeerDramAllocator::AllocateResult PeerDramAllocator::Allocate(const std::string& key, uint64_t size, TierType tier) { + return Allocate(key, size, tier, RawKvEncoding(size)); +} + +PeerDramAllocator::AllocateResult PeerDramAllocator::Allocate( + const std::string& key, uint64_t size, TierType tier, const KvEncodingDescriptor& encoding) { std::lock_guard lock(mutex_); - return AllocateLocked(key, size, tier); + return AllocateLocked(key, size, tier, encoding); } PeerDramAllocator::AllocateResult PeerDramAllocator::AllocateLocked(const std::string& key, - uint64_t size, TierType tier) { + uint64_t size, TierType tier, + const KvEncodingDescriptor& encoding) { auto fail = [&](Outcome outcome, const char* reason) { MORI_UMBP_WARN("[PeerDramAllocator] Allocate reason={} key='{}' size={} tier={}", reason, key, size, static_cast(tier)); @@ -145,6 +152,9 @@ PeerDramAllocator::AllocateResult PeerDramAllocator::AllocateLocked(const std::s slot.tier = tier; slot.pages = std::move(*pages); slot.size = size; + slot.encoding = encoding; + if (slot.encoding.stored_bytes == 0) slot.encoding.stored_bytes = size; + if (slot.encoding.logical_bytes == 0) slot.encoding.logical_bytes = slot.encoding.stored_bytes; slot.deadline = std::chrono::steady_clock::now() + pending_ttl_; slot.generation = allocator_generation_; pending_[slot.slot_id] = slot; @@ -160,7 +170,7 @@ std::vector PeerDramAllocator::BatchAllo std::lock_guard lock(mutex_); for (size_t i = 0; i < entries.size(); ++i) { const auto& entry = entries[i]; - auto result = AllocateLocked(entry.key, entry.size, entry.tier); + auto result = AllocateLocked(entry.key, entry.size, entry.tier, entry.encoding); out[i].outcome = result.outcome; out[i].slot = std::move(result.slot); if (out[i].outcome == Outcome::kSuccessAllocated && out[i].slot.has_value()) { @@ -219,7 +229,8 @@ bool PeerDramAllocator::CommitLocked(uint64_t slot_id, const std::string& key, owned.tier = it->second.tier; owned.pages = std::move(it->second.pages); owned.size = it->second.size; - QueueEventLocked(KvEvent{KvEvent::Kind::ADD, key, owned.tier, owned.size}); + owned.encoding = it->second.encoding; + QueueEventLocked(KvEvent{KvEvent::Kind::ADD, key, owned.tier, owned.size, owned.encoding}); owned_[key] = std::move(owned); pending_.erase(it); bytes_committed = owned_[key].size; @@ -276,6 +287,7 @@ PeerDramAllocator::ResolveResult PeerDramAllocator::Resolve(const std::string& k r.tier = it->second.tier; r.pages = it->second.pages; r.size = it->second.size; + r.encoding = it->second.encoding; // Extend the read lease so concurrent Evict reports bytes_freed=0 for // this key. steady_clock is monotonic and read_lease_ttl_ is fixed, // so this assignment is always >= any previous deadline for the key. @@ -297,6 +309,7 @@ std::vector PeerDramAllocator::BatchResolve( entry.tier = it->second.tier; entry.pages = it->second.pages; entry.size = it->second.size; + entry.encoding = it->second.encoding; if (include_descs) { entry.descs = BuildBufferDescsLocked(it->second.tier, it->second.pages); } @@ -360,6 +373,7 @@ std::optional PeerDramAllocator::AcquireDramCopy DramCopyPin pin; pin.total_size = it->second.size; + pin.encoding = it->second.encoding; pin.segments = BuildCopySegmentsLocked(it->second.tier, it->second.pages, it->second.size); pin.pin_token = next_pin_token_++; pins_[key] = PinState{pin.pin_token, std::chrono::steady_clock::now()}; @@ -496,7 +510,8 @@ std::vector PeerDramAllocator::SnapshotOwnedKeysLocked() const { std::vector out; out.reserve(owned_.size()); for (const auto& kv : owned_) { - out.push_back(KvEvent{KvEvent::Kind::ADD, kv.first, kv.second.tier, kv.second.size}); + out.push_back( + KvEvent{KvEvent::Kind::ADD, kv.first, kv.second.tier, kv.second.size, kv.second.encoding}); } return out; } diff --git a/src/umbp/distributed/peer/peer_service.cpp b/src/umbp/distributed/peer/peer_service.cpp index 634dd5cd0..82323925a 100644 --- a/src/umbp/distributed/peer/peer_service.cpp +++ b/src/umbp/distributed/peer/peer_service.cpp @@ -319,8 +319,8 @@ class PeerServiceServer::UMBPPeerServiceImpl final : public ::umbp::UMBPPeer::Se response->set_outcome(::umbp::ALLOCATE_SLOT_OUTCOME_FAILED); return grpc::Status::OK; } - auto result = - dram_alloc_->Allocate(request->key(), request->size(), FromProtoTier(request->tier())); + auto result = dram_alloc_->Allocate(request->key(), request->size(), FromProtoTier(request->tier()), + KvEncodingFromProto(request->encoding(), request->size())); switch (result.outcome) { case PeerDramAllocator::Outcome::kSuccessAlreadyExists: response->set_outcome(::umbp::ALLOCATE_SLOT_OUTCOME_SUCCESS_ALREADY_EXISTS); @@ -381,6 +381,7 @@ class PeerServiceServer::UMBPPeerServiceImpl final : public ::umbp::UMBPPeer::Se auto descs = dram_alloc_->BufferDescsForPages(r.tier, r.pages); FillPagesAndDescs(response, r.pages, dram_alloc_->PageSize(), descs); response->set_size(r.size); + FillProtoKvEncoding(r.encoding, response->mutable_encoding()); RecordInboundGet(r.size, "remote"); return grpc::Status::OK; } @@ -421,6 +422,7 @@ class PeerServiceServer::UMBPPeerServiceImpl final : public ::umbp::UMBPPeer::Se alloc_entry.key = entry.key(); alloc_entry.size = entry.size(); alloc_entry.tier = FromProtoTier(entry.tier()); + alloc_entry.encoding = KvEncodingFromProto(entry.encoding(), entry.size()); entries.push_back(std::move(alloc_entry)); } @@ -525,6 +527,7 @@ class PeerServiceServer::UMBPPeerServiceImpl final : public ::umbp::UMBPPeer::Se if (r.found) { e.tier = r.tier; e.size = r.size; + e.encoding = r.encoding; e.pages = std::move(r.pages); total_bytes += r.size; if (!omit_descs) { diff --git a/src/umbp/distributed/peer/peer_ssd_manager.cpp b/src/umbp/distributed/peer/peer_ssd_manager.cpp index 8a8229b73..25eedb0d7 100644 --- a/src/umbp/distributed/peer/peer_ssd_manager.cpp +++ b/src/umbp/distributed/peer/peer_ssd_manager.cpp @@ -25,6 +25,7 @@ #include #include "mori/utils/mori_log.hpp" +#include "umbp/codec/kv_encoding.h" #include "umbp/local/tiers/spdk_proxy_tier.h" #include "umbp/local/tiers/ssd_tier.h" #include "umbp/local/tiers/tier_backend.h" @@ -132,6 +133,12 @@ void PeerSsdManager::TouchLocked(const std::string& key) { bool PeerSsdManager::Write(const std::string& key, const std::vector>& segments, size_t total_size) { + return Write(key, segments, total_size, RawKvEncoding(total_size)); +} + +bool PeerSsdManager::Write(const std::string& key, + const std::vector>& segments, + size_t total_size, const KvEncodingDescriptor& encoding) { if (!backend_) return false; // Optimization, not just defense: the DRAM pin only dedups *concurrent* @@ -197,8 +204,12 @@ bool PeerSsdManager::Write(const std::string& key, TouchLocked(key); } else { lru_.push_front(key); - owned_.emplace(key, OwnedEntry{total_size, lru_.begin()}); - pending_events_.push_back(KvEvent{KvEvent::Kind::ADD, key, TierType::SSD, total_size}); + KvEncodingDescriptor stored_encoding = encoding; + if (stored_encoding.stored_bytes == 0) stored_encoding.stored_bytes = total_size; + if (stored_encoding.logical_bytes == 0) stored_encoding.logical_bytes = stored_encoding.stored_bytes; + owned_.emplace(key, OwnedEntry{total_size, stored_encoding, lru_.begin()}); + pending_events_.push_back( + KvEvent{KvEvent::Kind::ADD, key, TierType::SSD, total_size, stored_encoding}); } } @@ -417,7 +428,7 @@ std::vector PeerSsdManager::SnapshotOwnedKeysLocked() const { std::vector out; out.reserve(owned_.size()); for (const auto& [key, entry] : owned_) { - out.push_back(KvEvent{KvEvent::Kind::ADD, key, TierType::SSD, entry.size}); + out.push_back(KvEvent{KvEvent::Kind::ADD, key, TierType::SSD, entry.size, entry.encoding}); } return out; } diff --git a/src/umbp/distributed/peer/ssd_copy_pipeline.cpp b/src/umbp/distributed/peer/ssd_copy_pipeline.cpp index ff223db6b..e5453cfc0 100644 --- a/src/umbp/distributed/peer/ssd_copy_pipeline.cpp +++ b/src/umbp/distributed/peer/ssd_copy_pipeline.cpp @@ -169,7 +169,7 @@ void SsdCopyPipeline::RunTask(const SsdCopyTask& task) { // Backend IO runs outside the allocator lock; the pin keeps the source // pages alive for the duration of this synchronous Write. - if (ssd_->Write(task.key, pin->segments, pin->total_size)) { + if (ssd_->Write(task.key, pin->segments, pin->total_size, pin->encoding)) { metrics_.copied_ok.fetch_add(1, std::memory_order_relaxed); } else { metrics_.failed.fetch_add(1, std::memory_order_relaxed); diff --git a/src/umbp/distributed/pool_client.cpp b/src/umbp/distributed/pool_client.cpp index 00e4c9aef..54204d155 100644 --- a/src/umbp/distributed/pool_client.cpp +++ b/src/umbp/distributed/pool_client.cpp @@ -42,6 +42,7 @@ #include "mori/io/backend.hpp" #include "mori/utils/mori_log.hpp" +#include "umbp/codec/kv_encoding.h" #include "umbp/common/env_time.h" #include "umbp/distributed/master/master_metrics.h" #include "umbp/distributed/peer/batch_resolve_codec.h" @@ -699,12 +700,13 @@ bool PoolClient::LocalGetPages(const std::vector& pages, uint64_t } PoolClient::PutAttemptOutcome PoolClient::ExecuteLocalPut(const std::string& key, const void* src, - size_t size, TierType tier) { + size_t size, TierType tier, + const KvEncodingDescriptor& encoding) { if (!peer_alloc_) { MORI_UMBP_ERROR("[PoolClient] Local Put requested but peer allocator unavailable"); return PutAttemptOutcome::kFatal; } - auto alloc_res = peer_alloc_->Allocate(key, size, tier); + auto alloc_res = peer_alloc_->Allocate(key, size, tier, encoding); switch (alloc_res.outcome) { case PeerDramAllocator::Outcome::kSuccessAlreadyExists: return PutAttemptOutcome::kSuccessAlreadyExists; @@ -739,7 +741,8 @@ PoolClient::PutAttemptOutcome PoolClient::ExecuteLocalPut(const std::string& key } PoolClient::GetAttemptOutcome PoolClient::ExecuteLocalGet(const std::string& key, void* dst, - size_t size) { + size_t size, + KvEncodingDescriptor* out_encoding) { if (!peer_alloc_) { MORI_UMBP_ERROR("[PoolClient] Local Get requested but peer allocator unavailable"); return GetAttemptOutcome::kFatal; @@ -751,6 +754,7 @@ PoolClient::GetAttemptOutcome PoolClient::ExecuteLocalGet(const std::string& key if (!LocalGetPages(resolved.pages, peer_alloc_->PageSize(), dst, size)) { return GetAttemptOutcome::kFatal; } + if (out_encoding != nullptr) *out_encoding = resolved.encoding; master_client_->AddCounter(MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL, MORI_UMBP_METRIC_CLIENT_OUTBOUND_GET_BYTES_TOTAL_HELP, {{"traffic", "local"}}, static_cast(size)); @@ -801,22 +805,38 @@ PoolClient::GetAttemptOutcome PoolClient::ExecuteLocalSsdGet(const std::string& // --------------------------------------------------------------------------- bool PoolClient::Put(const std::string& key, const void* src, size_t size) { + return PutEncoded(key, src, size, RawKvEncoding(size)); +} + +bool PoolClient::PutEncoded(const std::string& key, const void* src, size_t size, + const KvEncodingDescriptor& encoding) { std::vector keys{key}; std::vector srcs{src}; std::vector sizes{size}; - auto results = BatchPut(keys, srcs, sizes); + std::vector encodings{encoding}; + auto results = BatchPutEncoded(keys, srcs, sizes, encodings); return !results.empty() && results[0]; } std::vector PoolClient::BatchPut(const std::vector& keys, const std::vector& srcs, const std::vector& sizes) { + std::vector encodings; + encodings.reserve(sizes.size()); + for (size_t size : sizes) encodings.push_back(RawKvEncoding(size)); + return BatchPutEncoded(keys, srcs, sizes, encodings); +} + +std::vector PoolClient::BatchPutEncoded(const std::vector& keys, + const std::vector& srcs, + const std::vector& sizes, + const std::vector& encodings) { // NOTE: BatchPut only lands data in the DRAM/HBM tier. The SSD tier copy is // asynchronous and owner-driven: a successful slot Commit (local in // ExecuteLocalPut, remote on the peer's CommitSlot) enqueues the copy onto the // SsdCopyPipeline. There is no explicit SSD write on this path. const auto call_start = std::chrono::steady_clock::now(); - if (keys.size() != srcs.size() || keys.size() != sizes.size()) { + if (keys.size() != srcs.size() || keys.size() != sizes.size() || keys.size() != encodings.size()) { MORI_UMBP_ERROR("[PoolClient] BatchPut: vector length mismatch"); return std::vector(keys.size(), false); } @@ -830,6 +850,23 @@ std::vector PoolClient::BatchPut(const std::vector& keys, std::vector block_sizes(keys.size()); for (size_t i = 0; i < sizes.size(); ++i) block_sizes[i] = static_cast(sizes[i]); + for (size_t i = 0; i < encodings.size(); ++i) { + std::string err; + if (!ValidateKvEncodingDescriptor(encodings[i], &err)) { + MORI_UMBP_ERROR("[PoolClient] BatchPut: invalid encoding for key='{}': {}", keys[i], err); + return std::vector(keys.size(), false); + } + uint64_t expected_size = encodings[i].stored_bytes; + if (expected_size == 0 && encodings[i].kind == KvEncodingKind::TURBOQUANT) { + auto layout = ComputeTurboQuantLayout(encodings[i]); + if (layout.has_value()) expected_size = layout->stored_bytes; + } + if (expected_size != sizes[i]) { + MORI_UMBP_ERROR("[PoolClient] BatchPut: encoding size mismatch for key='{}' (expected {}, got {})", + keys[i], expected_size, sizes[i]); + return std::vector(keys.size(), false); + } + } std::vector> routes; std::unordered_set excludes; auto status = master_client_->BatchRoutePut(keys, block_sizes, excludes, &routes); @@ -839,7 +876,7 @@ std::vector PoolClient::BatchPut(const std::vector& keys, } if (routes.size() < keys.size()) routes.resize(keys.size()); - BatchPutPlan plan = PartitionBatchPutTargets(keys, srcs, sizes, routes, &outcomes); + BatchPutPlan plan = PartitionBatchPutTargets(keys, srcs, sizes, encodings, routes, &outcomes); ExecuteBatchPutPlan(plan, &outcomes); const auto call_end = std::chrono::steady_clock::now(); @@ -864,7 +901,8 @@ std::vector PoolClient::BatchPut(const std::vector& keys, PoolClient::BatchPutPlan PoolClient::PartitionBatchPutTargets( const std::vector& keys, const std::vector& srcs, - const std::vector& sizes, const std::vector>& routes, + const std::vector& sizes, const std::vector& encodings, + const std::vector>& routes, std::vector* results) { BatchPutPlan plan; const size_t count = keys.size(); @@ -885,12 +923,22 @@ PoolClient::BatchPutPlan PoolClient::PartitionBatchPutTargets( // Self-target: deferred (with its tier) so ExecuteBatchPutPlan can run the // local memcpy inside the remote-DRAM submit..wait window. plan.local_items.push_back(BatchPutItem{ - .index = i, .key = &keys[i], .src = srcs[i], .size = sizes[i], .route = route}); + .index = i, + .key = &keys[i], + .src = srcs[i], + .size = sizes[i], + .encoding = &encodings[i], + .route = route}); continue; } if (route.tier != TierType::DRAM && route.tier != TierType::HBM) continue; plan.remote_groups[route.node_id].push_back(BatchPutItem{ - .index = i, .key = &keys[i], .src = srcs[i], .size = sizes[i], .route = route}); + .index = i, + .key = &keys[i], + .src = srcs[i], + .size = sizes[i], + .encoding = &encodings[i], + .route = route}); } return plan; } @@ -907,7 +955,7 @@ void PoolClient::ExecuteBatchPutPlan(const BatchPutPlan& plan, const auto t0 = std::chrono::steady_clock::now(); LocalParallelFor(local.size(), nthr, [&](size_t k) { const auto& item = local[k]; - switch (ExecuteLocalPut(*item.key, item.src, item.size, item.route.tier)) { + switch (ExecuteLocalPut(*item.key, item.src, item.size, item.route.tier, *item.encoding)) { case PutAttemptOutcome::kSuccess: (*results)[item.index] = PutEntryOutcome::kSucceeded; break; @@ -1169,6 +1217,7 @@ bool PoolClient::AllocateRemotePutEntries(const std::vector& items entry->set_size(item.size); entry->set_tier(static_cast<::umbp::TierType>(item.route.tier)); entry->set_key(*item.key); + FillProtoKvEncoding(*item.encoding, entry->mutable_encoding()); } ::umbp::BatchAllocateSlotsResponse alloc_resp; @@ -1222,6 +1271,7 @@ bool PoolClient::AllocateRemotePutEntries(const std::vector& items entry.result_index = item.index; entry.item = &item; entry.slot_id = plan.slot_id; + entry.encoding = *item.encoding; entry.plan = std::move(plan); entries->push_back(std::move(entry)); } @@ -1408,16 +1458,30 @@ void PoolClient::FinalizeRemotePutEntries(std::vector& entries, // --------------------------------------------------------------------------- bool PoolClient::Get(const std::string& key, void* dst, size_t size) { + return GetEncoded(key, dst, size, nullptr); +} + +bool PoolClient::GetEncoded(const std::string& key, void* dst, size_t size, + KvEncodingDescriptor* out_encoding) { std::vector keys{key}; std::vector dsts{dst}; std::vector sizes{size}; - auto results = BatchGet(keys, dsts, sizes); + std::vector encodings; + auto results = BatchGetEncoded(keys, dsts, sizes, out_encoding == nullptr ? nullptr : &encodings); + if (out_encoding != nullptr && !encodings.empty()) *out_encoding = encodings.front(); return !results.empty() && results[0]; } std::vector PoolClient::BatchGet(const std::vector& keys, const std::vector& dsts, const std::vector& sizes) { + return BatchGetEncoded(keys, dsts, sizes, nullptr); +} + +std::vector PoolClient::BatchGetEncoded(const std::vector& keys, + const std::vector& dsts, + const std::vector& sizes, + std::vector* out_encodings) { const auto call_start = std::chrono::steady_clock::now(); std::vector results(keys.size(), false); if (keys.size() != dsts.size() || keys.size() != sizes.size()) { @@ -1428,6 +1492,9 @@ std::vector PoolClient::BatchGet(const std::vector& keys, MORI_UMBP_ERROR("[PoolClient] BatchGet: client not initialized"); return results; } + if (out_encodings != nullptr) { + out_encodings->assign(keys.size(), KvEncodingDescriptor{}); + } std::vector> routes; std::unordered_set excludes; @@ -1440,8 +1507,8 @@ std::vector PoolClient::BatchGet(const std::vector& keys, routes.resize(keys.size()); } - BatchGetPlan plan = PartitionBatchGetTargets(keys, dsts, sizes, routes); - ExecuteBatchGetPlan(plan, keys, dsts, sizes, &results); + BatchGetPlan plan = PartitionBatchGetTargets(keys, dsts, sizes, out_encodings, routes); + ExecuteBatchGetPlan(plan, keys, dsts, sizes, out_encodings, &results); const auto call_end = std::chrono::steady_clock::now(); const double seconds = @@ -1460,7 +1527,8 @@ std::vector PoolClient::BatchGet(const std::vector& keys, PoolClient::BatchGetPlan PoolClient::PartitionBatchGetTargets( const std::vector& keys, const std::vector& dsts, - const std::vector& sizes, const std::vector>& routes) { + const std::vector& sizes, std::vector* out_encodings, + const std::vector>& routes) { BatchGetPlan plan; for (size_t i = 0; i < keys.size(); ++i) { // Zero-size gets are rejected before local fallback or remote read: an @@ -1475,22 +1543,23 @@ PoolClient::BatchGetPlan PoolClient::PartitionBatchGetTargets( continue; } const auto& route = routes[i].value(); + BatchGetItem item{.index = i, + .key = &keys[i], + .dst = const_cast(dsts[i]), + .size = sizes[i], + .out_encoding = out_encodings == nullptr ? nullptr : &(*out_encodings)[i], + .route = route}; if (route.node_id == config_.master_config.node_id) { // Self-target: both DRAM/HBM and SSD are deferred (collected as indices) // so ExecuteBatchGetPlan can place them inside the remote-DRAM in-flight // window in the overlap path. if (route.tier == TierType::SSD) { - plan.local_ssd_indices.push_back(i); + plan.local_ssd_items.push_back(std::move(item)); } else { plan.local_dram_indices.push_back(i); } continue; } - BatchGetItem item{.index = i, - .key = &keys[i], - .dst = const_cast(dsts[i]), - .size = sizes[i], - .route = route}; if (route.tier == TierType::SSD) { plan.remote_ssd_groups[route.node_id].push_back(std::move(item)); } else { @@ -1502,7 +1571,9 @@ PoolClient::BatchGetPlan PoolClient::PartitionBatchGetTargets( void PoolClient::ExecuteBatchGetPlan(const BatchGetPlan& plan, const std::vector& keys, const std::vector& dsts, - const std::vector& sizes, std::vector* results) { + const std::vector& sizes, + std::vector* out_encodings, + std::vector* results) { // Parallel local DRAM reads: different threads handle different keys. Resolve // is mutex-serialized in the allocator; the per-key memcpy in // ExecuteLocalGet->LocalGetPages runs lock-free in parallel. results is @@ -1515,7 +1586,9 @@ void PoolClient::ExecuteBatchGetPlan(const BatchGetPlan& plan, const std::vector std::vector ok(idx.size(), 0); LocalParallelFor(idx.size(), nthr, [&](size_t k) { const size_t i = idx[k]; - if (ExecuteLocalGet(keys[i], const_cast(dsts[i]), sizes[i]) == + KvEncodingDescriptor* out_encoding = + out_encodings == nullptr ? nullptr : &(*out_encodings)[i]; + if (ExecuteLocalGet(keys[i], const_cast(dsts[i]), sizes[i], out_encoding) == GetAttemptOutcome::kSuccess) { ok[k] = 1; } @@ -1541,10 +1614,11 @@ void PoolClient::ExecuteBatchGetPlan(const BatchGetPlan& plan, const std::vector // thread, reading straight into the user buffer (no staging / RDMA / lease). // Independent per key, so its position in the schedule is correctness-neutral. auto run_local_ssd = [&]() { - for (size_t i : plan.local_ssd_indices) { - if (ExecuteLocalSsdGet(keys[i], const_cast(dsts[i]), sizes[i]) == + for (const auto& item : plan.local_ssd_items) { + if (ExecuteLocalSsdGet(*item.key, item.dst, item.size) == GetAttemptOutcome::kSuccess) { - (*results)[i] = true; + if (item.out_encoding != nullptr) *item.out_encoding = item.route.encoding; + (*results)[item.index] = true; } } }; @@ -1827,6 +1901,7 @@ bool PoolClient::PrepareRemoteGetEntries(const std::vector& items, entry.item = &item; entry.plan.page_size = decoded.page_size; entry.plan.pages = std::move(decoded.keys[i].pages); + entry.encoding = key.encoding; // Descriptors were hydrated batch-level above; the per-entry plan carries // none (BuildRemoteGetTransfers' EnsureBufferDescsCached call is a no-op on // an empty list and the read path resolves descriptors by buffer_index). @@ -1935,6 +2010,7 @@ void PoolClient::FinalizeRemoteGetEntries(std::vector& entries, master_client_->AddCounter(MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL, MORI_UMBP_METRIC_CLIENT_INBOUND_GET_BYTES_TOTAL_HELP, {{"traffic", "remote"}}, static_cast(entry.item->size)); + if (entry.item->out_encoding != nullptr) *entry.item->out_encoding = entry.encoding; (*results)[entry.result_index] = true; } } @@ -2478,6 +2554,7 @@ void PoolClient::ProcessRemoteSsdBatchGet(const std::vector& items switch (outcome) { case SsdGetOutcome::kSuccess: (*results)[item.index] = true; + if (item.out_encoding != nullptr) *item.out_encoding = item.route.encoding; done = true; break; case SsdGetOutcome::kMiss: // definitive miss diff --git a/src/umbp/distributed/proto/umbp.proto b/src/umbp/distributed/proto/umbp.proto index 8bc2805ab..4e50979ab 100644 --- a/src/umbp/distributed/proto/umbp.proto +++ b/src/umbp/distributed/proto/umbp.proto @@ -12,6 +12,40 @@ enum TierType { TIER_SSD = 3; // NVMe SSD } +enum KvEncodingKind { + KV_ENCODING_UNKNOWN = 0; + KV_ENCODING_RAW = 1; + KV_ENCODING_TURBOQUANT = 2; +} + +enum KvDType { + KV_DTYPE_UNKNOWN = 0; + KV_DTYPE_FP16 = 1; + KV_DTYPE_BF16 = 2; + KV_DTYPE_FP8 = 3; + KV_DTYPE_UINT8 = 4; +} + +message KvEncodingDescriptor { + uint32 schema_version = 1; + KvEncodingKind kind = 2; + KvDType original_dtype = 3; + string preset = 4; + uint32 key_bits = 5; + uint32 value_bits = 6; + uint32 head_dim = 7; + uint32 block_size = 8; + uint32 num_layers = 9; + uint32 num_tokens = 10; + uint32 num_heads = 11; + uint32 hidden_dim = 12; + uint32 skip_first_layers = 13; + uint32 skip_last_layers = 14; + uint32 packing_version = 15; + uint64 stored_bytes = 16; + uint64 logical_bytes = 17; +} + // Per-tier capacity reported by a peer node. Master treats every value // as ground truth and applies it unconditionally — the peer's allocator // is the single source of physical capacity. @@ -86,6 +120,7 @@ message KvEvent { string key = 2; TierType tier = 3; uint64 size = 4; // ADD only + KvEncodingDescriptor encoding = 5; // ADD only; absent/zero means raw bytes } message EventBundle { @@ -144,6 +179,7 @@ message RouteGetResponse { TierType tier = 3; uint64 size = 4; string peer_address = 5; + KvEncodingDescriptor encoding = 6; } // ============================================================ @@ -198,6 +234,7 @@ message BatchRouteGetResponse { repeated uint32 node_ref = 2; // per key; 1-based index into nodes, 0 = not found [packed] repeated TierType tier = 3; // per key [packed] repeated uint64 size = 4; // per key [packed] + repeated KvEncodingDescriptor encoding = 5; // per key; default raw when absent } // Read-only batched existence probe. Unlike BatchRouteGet, this does diff --git a/src/umbp/distributed/proto/umbp_peer.proto b/src/umbp/distributed/proto/umbp_peer.proto index ff37fda7a..4fe852707 100644 --- a/src/umbp/distributed/proto/umbp_peer.proto +++ b/src/umbp/distributed/proto/umbp_peer.proto @@ -96,6 +96,7 @@ message AllocateSlotRequest { uint64 size = 1; TierType tier = 2; string key = 3; // owned_ dedup (master-index-lag fallback) + KvEncodingDescriptor encoding = 4; } // AllocateSlot / BatchAllocateSlots per-entry outcome. // FAILED_NO_SPACE is split out so the writer can keep retrying with @@ -150,6 +151,7 @@ message ResolveKeyResponse { uint64 page_size = 3; repeated BufferMemoryDesc descs = 4; uint64 size = 5; + KvEncodingDescriptor encoding = 6; } // Master-driven eviction. Idempotent: already-evicted keys produce a @@ -228,4 +230,7 @@ message BatchResolveKeysResponse { // (buffer_index, page_index) pairs, walking keys in order. repeated uint32 buffer_index = 7; repeated uint32 page_index = 8; + // Per-key encoding descriptors, parallel to found/tier/size. Older peers may + // omit this; clients treat missing descriptors as raw bytes. + repeated KvEncodingDescriptor encoding = 9; } diff --git a/src/umbp/include/umbp/codec/kv_encoding.h b/src/umbp/include/umbp/codec/kv_encoding.h new file mode 100644 index 000000000..8ac90a77d --- /dev/null +++ b/src/umbp/include/umbp/codec/kv_encoding.h @@ -0,0 +1,141 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +#pragma once + +#include +#include +#include + +#include "umbp/distributed/types.h" + +namespace mori::umbp { + +struct TurboQuantLayout { + uint64_t key_packed_size = 0; + uint64_t value_packed_size = 0; + uint64_t slot_size = 0; + uint64_t slot_size_aligned = 0; + uint64_t num_blocks = 0; + uint64_t compressed_layers = 0; + uint64_t raw_layers = 0; + uint64_t compressed_bytes = 0; + uint64_t raw_bytes = 0; + uint64_t stored_bytes = 0; +}; + +inline uint64_t PackedBytes(uint64_t elements, uint32_t bits) { + return (elements * static_cast(bits) + 7) / 8; +} + +inline uint64_t DTypeBytes(KvDType dtype) { + switch (dtype) { + case KvDType::FP16: + case KvDType::BF16: + return 2; + case KvDType::FP8: + case KvDType::UINT8: + return 1; + default: + return 0; + } +} + +inline bool IsRawEncoding(const KvEncodingDescriptor& desc) { + return desc.kind == KvEncodingKind::RAW || desc.kind == KvEncodingKind::UNKNOWN; +} + +inline bool IsTurboQuantPresetSupported(const std::string& preset) { + return preset == "turboquant_k8v4" || preset == "turboquant_4bit_nc" || + preset == "turboquant_k3v4_nc" || preset == "turboquant_3bit_nc"; +} + +inline std::optional ComputeTurboQuantLayout( + const KvEncodingDescriptor& desc) { + if (desc.kind != KvEncodingKind::TURBOQUANT) return std::nullopt; + if (!IsTurboQuantPresetSupported(desc.preset)) return std::nullopt; + if (desc.head_dim == 0 || desc.block_size == 0 || desc.num_layers == 0 || + desc.num_tokens == 0 || desc.num_heads == 0) { + return std::nullopt; + } + + const uint32_t key_bits = + desc.key_bits != 0 ? desc.key_bits : (desc.preset == "turboquant_k8v4" ? 8 : 4); + const uint32_t value_bits = desc.value_bits != 0 ? desc.value_bits : 4; + + TurboQuantLayout layout; + const bool key_fp8 = key_bits == 8; + layout.key_packed_size = key_fp8 ? desc.head_dim : PackedBytes(desc.head_dim, key_bits) + 2; + layout.value_packed_size = PackedBytes(desc.head_dim, value_bits) + 4; + layout.slot_size = layout.key_packed_size + layout.value_packed_size; + layout.slot_size_aligned = layout.slot_size + (layout.slot_size % 2); + layout.num_blocks = (desc.num_tokens + desc.block_size - 1) / desc.block_size; + + const uint32_t quant_start = + desc.skip_first_layers < desc.num_layers ? desc.skip_first_layers : desc.num_layers; + const uint32_t raw_suffix = + desc.skip_last_layers < desc.num_layers - quant_start ? desc.skip_last_layers + : desc.num_layers - quant_start; + layout.compressed_layers = desc.num_layers - quant_start - raw_suffix; + layout.raw_layers = desc.num_layers - layout.compressed_layers; + + const uint64_t raw_dtype_bytes = DTypeBytes(desc.original_dtype); + if (layout.raw_layers > 0 && raw_dtype_bytes == 0) return std::nullopt; + + layout.raw_bytes = 2ull * layout.raw_layers * desc.num_tokens * + static_cast(desc.num_heads) * desc.head_dim * raw_dtype_bytes; + layout.compressed_bytes = layout.compressed_layers * layout.num_blocks * desc.block_size * + static_cast(desc.num_heads) * layout.slot_size_aligned; + layout.stored_bytes = layout.raw_bytes + layout.compressed_bytes; + return layout; +} + +inline bool ValidateKvEncodingDescriptor(const KvEncodingDescriptor& desc, + std::string* error = nullptr) { + auto fail = [&](const char* msg) { + if (error != nullptr) *error = msg; + return false; + }; + if (desc.schema_version == 0) return fail("schema_version must be non-zero"); + if (IsRawEncoding(desc)) { + if (desc.stored_bytes == 0) return fail("raw encoding stored_bytes must be non-zero"); + return true; + } + if (desc.kind != KvEncodingKind::TURBOQUANT) return fail("unsupported KV encoding kind"); + auto layout = ComputeTurboQuantLayout(desc); + if (!layout.has_value()) return fail("invalid TurboQuant descriptor"); + if (desc.stored_bytes != 0 && desc.stored_bytes != layout->stored_bytes) { + return fail("TurboQuant stored_bytes does not match descriptor layout"); + } + return true; +} + +inline KvEncodingDescriptor RawKvEncoding(uint64_t stored_bytes, + KvDType dtype = KvDType::UNKNOWN) { + KvEncodingDescriptor desc; + desc.kind = KvEncodingKind::RAW; + desc.original_dtype = dtype; + desc.stored_bytes = stored_bytes; + desc.logical_bytes = stored_bytes; + return desc; +} + +} // namespace mori::umbp diff --git a/src/umbp/include/umbp/distributed/master/master_client.h b/src/umbp/include/umbp/distributed/master/master_client.h index 4d329e05b..5a757c686 100644 --- a/src/umbp/include/umbp/distributed/master/master_client.h +++ b/src/umbp/include/umbp/distributed/master/master_client.h @@ -67,6 +67,7 @@ struct RouteGetResult { std::string node_id; TierType tier = TierType::UNKNOWN; uint64_t size = 0; + KvEncodingDescriptor encoding; std::string peer_address; }; diff --git a/src/umbp/include/umbp/distributed/peer/batch_resolve_codec.h b/src/umbp/include/umbp/distributed/peer/batch_resolve_codec.h index cbcc4ae48..e05c0a7a7 100644 --- a/src/umbp/include/umbp/distributed/peer/batch_resolve_codec.h +++ b/src/umbp/include/umbp/distributed/peer/batch_resolve_codec.h @@ -31,6 +31,7 @@ #include #include +#include "umbp/codec/kv_encoding.h" #include "umbp/distributed/types.h" #include "umbp_peer.pb.h" @@ -42,6 +43,7 @@ struct ResolvedKeyEntry { bool found = false; TierType tier = TierType::UNKNOWN; uint64_t size = 0; + KvEncodingDescriptor encoding; std::vector pages; }; @@ -51,6 +53,7 @@ struct DecodedResolveKey { bool found = false; TierType tier = TierType::UNKNOWN; uint64_t size = 0; + KvEncodingDescriptor encoding; std::vector pages; }; @@ -88,6 +91,103 @@ inline ::umbp::TierType BatchResolveTierToProto(TierType t) { } } +inline ::umbp::KvEncodingKind KvEncodingKindToProto(KvEncodingKind kind) { + switch (kind) { + case KvEncodingKind::RAW: + return ::umbp::KV_ENCODING_RAW; + case KvEncodingKind::TURBOQUANT: + return ::umbp::KV_ENCODING_TURBOQUANT; + default: + return ::umbp::KV_ENCODING_UNKNOWN; + } +} + +inline KvEncodingKind KvEncodingKindFromProto(::umbp::KvEncodingKind kind) { + switch (kind) { + case ::umbp::KV_ENCODING_RAW: + return KvEncodingKind::RAW; + case ::umbp::KV_ENCODING_TURBOQUANT: + return KvEncodingKind::TURBOQUANT; + default: + return KvEncodingKind::UNKNOWN; + } +} + +inline ::umbp::KvDType KvDTypeToProto(KvDType dtype) { + switch (dtype) { + case KvDType::FP16: + return ::umbp::KV_DTYPE_FP16; + case KvDType::BF16: + return ::umbp::KV_DTYPE_BF16; + case KvDType::FP8: + return ::umbp::KV_DTYPE_FP8; + case KvDType::UINT8: + return ::umbp::KV_DTYPE_UINT8; + default: + return ::umbp::KV_DTYPE_UNKNOWN; + } +} + +inline KvDType KvDTypeFromProto(::umbp::KvDType dtype) { + switch (dtype) { + case ::umbp::KV_DTYPE_FP16: + return KvDType::FP16; + case ::umbp::KV_DTYPE_BF16: + return KvDType::BF16; + case ::umbp::KV_DTYPE_FP8: + return KvDType::FP8; + case ::umbp::KV_DTYPE_UINT8: + return KvDType::UINT8; + default: + return KvDType::UNKNOWN; + } +} + +inline void FillProtoKvEncoding(const KvEncodingDescriptor& src, + ::umbp::KvEncodingDescriptor* dst) { + dst->set_schema_version(src.schema_version); + dst->set_kind(KvEncodingKindToProto(src.kind)); + dst->set_original_dtype(KvDTypeToProto(src.original_dtype)); + dst->set_preset(src.preset); + dst->set_key_bits(src.key_bits); + dst->set_value_bits(src.value_bits); + dst->set_head_dim(src.head_dim); + dst->set_block_size(src.block_size); + dst->set_num_layers(src.num_layers); + dst->set_num_tokens(src.num_tokens); + dst->set_num_heads(src.num_heads); + dst->set_hidden_dim(src.hidden_dim); + dst->set_skip_first_layers(src.skip_first_layers); + dst->set_skip_last_layers(src.skip_last_layers); + dst->set_packing_version(src.packing_version); + dst->set_stored_bytes(src.stored_bytes); + dst->set_logical_bytes(src.logical_bytes); +} + +inline KvEncodingDescriptor KvEncodingFromProto(const ::umbp::KvEncodingDescriptor& src, + uint64_t fallback_size = 0) { + KvEncodingDescriptor out; + out.schema_version = src.schema_version() == 0 ? 1 : src.schema_version(); + out.kind = KvEncodingKindFromProto(src.kind()); + if (out.kind == KvEncodingKind::UNKNOWN) out.kind = KvEncodingKind::RAW; + out.original_dtype = KvDTypeFromProto(src.original_dtype()); + out.preset = src.preset(); + out.key_bits = src.key_bits(); + out.value_bits = src.value_bits(); + out.head_dim = src.head_dim(); + out.block_size = src.block_size(); + out.num_layers = src.num_layers(); + out.num_tokens = src.num_tokens(); + out.num_heads = src.num_heads(); + out.hidden_dim = src.hidden_dim(); + out.skip_first_layers = src.skip_first_layers(); + out.skip_last_layers = src.skip_last_layers(); + out.packing_version = src.packing_version() == 0 ? 1 : src.packing_version(); + out.stored_bytes = src.stored_bytes() == 0 ? fallback_size : src.stored_bytes(); + out.logical_bytes = src.logical_bytes() == 0 ? out.stored_bytes : src.logical_bytes(); + return out; +} + // Encode resolved keys into the SoA response. `descs` are the batch-level // deduplicated buffer descriptors (by buffer_index); pass an empty vector to // omit them (honoring BatchResolveKeysRequest.omit_descs). Clears `resp` @@ -110,6 +210,7 @@ inline void EncodeBatchResolveResponse(const std::vector& keys resp->mutable_found()->Reserve(static_cast(keys.size())); resp->mutable_tier()->Reserve(static_cast(keys.size())); resp->mutable_size()->Reserve(static_cast(keys.size())); + resp->mutable_encoding()->Reserve(static_cast(keys.size())); resp->mutable_page_count()->Reserve(static_cast(keys.size())); resp->mutable_buffer_index()->Reserve(static_cast(total_pages)); resp->mutable_page_index()->Reserve(static_cast(total_pages)); @@ -118,6 +219,7 @@ inline void EncodeBatchResolveResponse(const std::vector& keys resp->add_found(k.found); resp->add_tier(BatchResolveTierToProto(k.tier)); resp->add_size(k.size); + FillProtoKvEncoding(k.encoding, resp->add_encoding()); resp->add_page_count(static_cast(k.pages.size())); for (const auto& p : k.pages) { resp->add_buffer_index(p.buffer_index); @@ -155,7 +257,8 @@ inline DecodedBatchResolve DecodeBatchResolveResponse( // hold at least the summed page_count. Anything else is a malformed peer // response; refuse to decode it. if (resp.tier_size() != n || resp.size_size() != n || resp.page_count_size() != n || - resp.buffer_index_size() != resp.page_index_size()) { + resp.buffer_index_size() != resp.page_index_size() || + (resp.encoding_size() != 0 && resp.encoding_size() != n)) { return out; } @@ -172,6 +275,8 @@ inline DecodedBatchResolve DecodeBatchResolveResponse( k.found = resp.found(i); k.tier = BatchResolveTierFromProto(resp.tier(i)); k.size = resp.size(i); + k.encoding = resp.encoding_size() == n ? KvEncodingFromProto(resp.encoding(i), k.size) + : RawKvEncoding(k.size); const uint32_t pc = resp.page_count(i); k.pages.reserve(pc); for (uint32_t p = 0; p < pc; ++p) { diff --git a/src/umbp/include/umbp/distributed/peer/peer_dram_allocator.h b/src/umbp/include/umbp/distributed/peer/peer_dram_allocator.h index 46383a71d..f6d861b98 100644 --- a/src/umbp/include/umbp/distributed/peer/peer_dram_allocator.h +++ b/src/umbp/include/umbp/distributed/peer/peer_dram_allocator.h @@ -98,6 +98,7 @@ class PeerDramAllocator : public OwnedLocationSource { TierType tier = TierType::UNKNOWN; std::vector pages; uint64_t size = 0; + KvEncodingDescriptor encoding; std::chrono::steady_clock::time_point deadline; // Snapshot of allocator_generation_ at Allocate(). Commit() // rejects slots whose generation no longer matches the current. @@ -108,6 +109,7 @@ class PeerDramAllocator : public OwnedLocationSource { TierType tier = TierType::UNKNOWN; std::vector pages; uint64_t size = 0; + KvEncodingDescriptor encoding; }; struct ResolveResult { @@ -115,6 +117,7 @@ class PeerDramAllocator : public OwnedLocationSource { TierType tier = TierType::UNKNOWN; std::vector pages; uint64_t size = 0; + KvEncodingDescriptor encoding; }; // ResolveResult + descs built under the same lock. @@ -123,6 +126,7 @@ class PeerDramAllocator : public OwnedLocationSource { TierType tier = TierType::UNKNOWN; std::vector pages; uint64_t size = 0; + KvEncodingDescriptor encoding; std::vector descs; }; @@ -151,6 +155,7 @@ class PeerDramAllocator : public OwnedLocationSource { std::string key; uint64_t size = 0; TierType tier = TierType::UNKNOWN; + KvEncodingDescriptor encoding; }; struct BatchAllocateResult { @@ -172,6 +177,8 @@ class PeerDramAllocator : public OwnedLocationSource { // Reserve `size` bytes on `tier` for `key`. `key` enables owned_ // dedup (master-index-lag fallback; primary dedup is at BatchRoutePut). AllocateResult Allocate(const std::string& key, uint64_t size, TierType tier); + AllocateResult Allocate(const std::string& key, uint64_t size, TierType tier, + const KvEncodingDescriptor& encoding); std::vector BatchAllocate(const std::vector& entries); @@ -220,6 +227,7 @@ class PeerDramAllocator : public OwnedLocationSource { struct DramCopyPin { std::vector> segments; size_t total_size = 0; + KvEncodingDescriptor encoding; uint64_t pin_token = 0; // release-time validation }; @@ -316,7 +324,8 @@ class PeerDramAllocator : public OwnedLocationSource { PageBitmapAllocator* AllocatorForLocked(TierType tier); const PageBitmapAllocator* AllocatorForLocked(TierType tier) const; - AllocateResult AllocateLocked(const std::string& key, uint64_t size, TierType tier); + AllocateResult AllocateLocked(const std::string& key, uint64_t size, TierType tier, + const KvEncodingDescriptor& encoding); bool CommitLocked(uint64_t slot_id, const std::string& key, uint64_t& bytes_committed); bool AbortLocked(uint64_t slot_id); diff --git a/src/umbp/include/umbp/distributed/peer/peer_ssd_manager.h b/src/umbp/include/umbp/distributed/peer/peer_ssd_manager.h index 852861dac..29f8c810d 100644 --- a/src/umbp/include/umbp/distributed/peer/peer_ssd_manager.h +++ b/src/umbp/include/umbp/distributed/peer/peer_ssd_manager.h @@ -80,6 +80,8 @@ class PeerSsdManager : public OwnedLocationSource { // (best-effort clean). bool Write(const std::string& key, const std::vector>& segments, size_t total_size); + bool Write(const std::string& key, const std::vector>& segments, + size_t total_size, const KvEncodingDescriptor& encoding); // Local eviction of a single key. Read priority: a key with an // in-flight PrepareRead (inflight_reads_ > 0) is NOT evicted (returns false). @@ -167,6 +169,7 @@ class PeerSsdManager : public OwnedLocationSource { // a touch is an O(1) splice and a victim lookup is an O(1) walk from the tail. struct OwnedEntry { uint64_t size = 0; + KvEncodingDescriptor encoding; std::list::iterator lru_it; // position of this key in lru_ }; diff --git a/src/umbp/include/umbp/distributed/pool_client.h b/src/umbp/include/umbp/distributed/pool_client.h index 603a6128e..512f54336 100644 --- a/src/umbp/include/umbp/distributed/pool_client.h +++ b/src/umbp/include/umbp/distributed/pool_client.h @@ -106,13 +106,25 @@ class PoolClient { // adds the failed node to the exclude set. bool Put(const std::string& key, const void* src, size_t size); bool Get(const std::string& key, void* dst, size_t size); + bool PutEncoded(const std::string& key, const void* src, size_t size, + const KvEncodingDescriptor& encoding); + bool GetEncoded(const std::string& key, void* dst, size_t size, + KvEncodingDescriptor* out_encoding); std::vector BatchPut(const std::vector& keys, const std::vector& srcs, const std::vector& sizes); + std::vector BatchPutEncoded(const std::vector& keys, + const std::vector& srcs, + const std::vector& sizes, + const std::vector& encodings); std::vector BatchGet(const std::vector& keys, const std::vector& dsts, const std::vector& sizes); + std::vector BatchGetEncoded(const std::vector& keys, + const std::vector& dsts, + const std::vector& sizes, + std::vector* out_encodings); // Cluster-wide existence check — issues a RouteGet and reports // whether master surfaced any replica. No RDMA, no lease bump. @@ -259,8 +271,9 @@ class PoolClient { enum class GetAttemptOutcome { kSuccess, kRetry, kFatal }; PutAttemptOutcome ExecuteLocalPut(const std::string& key, const void* src, size_t size, - TierType tier); - GetAttemptOutcome ExecuteLocalGet(const std::string& key, void* dst, size_t size); + TierType tier, const KvEncodingDescriptor& encoding); + GetAttemptOutcome ExecuteLocalGet(const std::string& key, void* dst, size_t size, + KvEncodingDescriptor* out_encoding); // Self-target SSD get: read straight from the local SSD tier into the user // buffer (no staging / RDMA / lease). GetAttemptOutcome ExecuteLocalSsdGet(const std::string& key, void* dst, size_t size); @@ -269,6 +282,7 @@ class PoolClient { const std::string* key; const void* src; size_t size; + const KvEncodingDescriptor* encoding; RoutePutResult route; }; struct BatchGetItem { @@ -276,6 +290,7 @@ class PoolClient { const std::string* key; void* dst; size_t size; + KvEncodingDescriptor* out_encoding; RouteGetResult route; }; @@ -289,7 +304,7 @@ class PoolClient { std::unordered_map> remote_dram_groups; std::unordered_map> remote_ssd_groups; std::vector local_dram_indices; - std::vector local_ssd_indices; + std::vector local_ssd_items; }; // Pure grouping for one BatchPut (no IO, no local put executed; mirrors @@ -305,6 +320,7 @@ class PoolClient { BatchPutPlan PartitionBatchPutTargets(const std::vector& keys, const std::vector& srcs, const std::vector& sizes, + const std::vector& encodings, const std::vector>& routes, std::vector* results); // Execute a BatchPutPlan. Zero-copy submits all peers (not waited), runs the @@ -319,6 +335,7 @@ class PoolClient { BatchGetPlan PartitionBatchGetTargets(const std::vector& keys, const std::vector& dsts, const std::vector& sizes, + std::vector* out_encodings, const std::vector>& routes); // Execute a BatchGetPlan: local DRAM/SSD, remote SSD, and the remote-DRAM // submit/wait arrangement. Zero-copy remote DRAM submits all peers, runs local @@ -327,6 +344,7 @@ class PoolClient { // (submit -> wait). Reads the plan; writes per-key outcomes into *results. void ExecuteBatchGetPlan(const BatchGetPlan& plan, const std::vector& keys, const std::vector& dsts, const std::vector& sizes, + std::vector* out_encodings, std::vector* results); struct TransferInstruction { @@ -363,6 +381,7 @@ class PoolClient { const BatchPutItem* item; SlotPlan plan; uint64_t slot_id; + KvEncodingDescriptor encoding; std::optional> zero_copy; bool use_staging = false; uint64_t staging_offset = 0; @@ -373,6 +392,7 @@ class PoolClient { size_t result_index; const BatchGetItem* item; SlotPlan plan; + KvEncodingDescriptor encoding; std::optional> zero_copy; bool use_staging = false; uint64_t staging_offset = 0; diff --git a/src/umbp/include/umbp/distributed/types.h b/src/umbp/include/umbp/distributed/types.h index 118dc232b..db225064c 100644 --- a/src/umbp/include/umbp/distributed/types.h +++ b/src/umbp/include/umbp/distributed/types.h @@ -37,6 +37,60 @@ enum class TierType : int { SSD = 3, }; +enum class KvEncodingKind : int { + UNKNOWN = 0, + RAW = 1, + TURBOQUANT = 2, +}; + +enum class KvDType : int { + UNKNOWN = 0, + FP16 = 1, + BF16 = 2, + FP8 = 3, + UINT8 = 4, +}; + +// Semantic descriptor for the bytes stored under one KV key. Mori-IO still moves +// opaque byte ranges; this descriptor is UMBP/client-layer metadata that lets the +// consumer know how to interpret those bytes. +struct KvEncodingDescriptor { + uint32_t schema_version = 1; + KvEncodingKind kind = KvEncodingKind::RAW; + KvDType original_dtype = KvDType::UNKNOWN; + std::string preset; + + uint32_t key_bits = 0; + uint32_t value_bits = 0; + uint32_t head_dim = 0; + uint32_t block_size = 0; + uint32_t num_layers = 0; + uint32_t num_tokens = 0; + uint32_t num_heads = 0; + uint32_t hidden_dim = 0; + uint32_t skip_first_layers = 0; + uint32_t skip_last_layers = 0; + uint32_t packing_version = 1; + + uint64_t stored_bytes = 0; + uint64_t logical_bytes = 0; + + bool operator==(const KvEncodingDescriptor& other) const { + return schema_version == other.schema_version && kind == other.kind && + original_dtype == other.original_dtype && preset == other.preset && + key_bits == other.key_bits && value_bits == other.value_bits && + head_dim == other.head_dim && block_size == other.block_size && + num_layers == other.num_layers && num_tokens == other.num_tokens && + num_heads == other.num_heads && hidden_dim == other.hidden_dim && + skip_first_layers == other.skip_first_layers && + skip_last_layers == other.skip_last_layers && + packing_version == other.packing_version && stored_bytes == other.stored_bytes && + logical_bytes == other.logical_bytes; + } + + bool operator!=(const KvEncodingDescriptor& other) const { return !(*this == other); } +}; + struct TierCapacity { uint64_t total_bytes = 0; uint64_t available_bytes = 0; @@ -55,9 +109,11 @@ struct Location { std::string node_id; uint64_t size = 0; TierType tier = TierType::UNKNOWN; + KvEncodingDescriptor encoding; bool operator==(const Location& other) const { - return node_id == other.node_id && size == other.size && tier == other.tier; + return node_id == other.node_id && size == other.size && tier == other.tier && + encoding == other.encoding; } }; @@ -111,6 +167,7 @@ struct KvEvent { std::string key; TierType tier = TierType::UNKNOWN; uint64_t size = 0; // ADD only; REMOVE leaves this 0 + KvEncodingDescriptor encoding; }; struct EventBundle { diff --git a/src/umbp/tests/CMakeLists.txt b/src/umbp/tests/CMakeLists.txt index 8886e7f5d..b0e1beb89 100644 --- a/src/umbp/tests/CMakeLists.txt +++ b/src/umbp/tests/CMakeLists.txt @@ -204,6 +204,17 @@ target_compile_features(test_ssd_read_lease_gating PRIVATE cxx_std_17) gtest_discover_tests(test_ssd_read_lease_gating) +# --------------------------------------------------------------------------- +# test_kv_encoding — raw/TurboQuant descriptor validation and byte sizing. +# --------------------------------------------------------------------------- +add_executable(test_kv_encoding test_kv_encoding.cpp) + +target_link_libraries(test_kv_encoding PRIVATE umbp_common GTest::gtest_main) + +target_compile_features(test_kv_encoding PRIVATE cxx_std_17) + +gtest_discover_tests(test_kv_encoding) + # --------------------------------------------------------------------------- # test_ssd_reliability — cross-component reliability: owned-source DRAM+SSD # merge, SSD evict -> REMOVE -> master index convergence, tier-priority over the diff --git a/src/umbp/tests/test_batch_resolve_codec.cpp b/src/umbp/tests/test_batch_resolve_codec.cpp index 29df4b893..f41ad6c39 100644 --- a/src/umbp/tests/test_batch_resolve_codec.cpp +++ b/src/umbp/tests/test_batch_resolve_codec.cpp @@ -49,6 +49,26 @@ BufferMemoryDescBytes MakeDesc(uint32_t buffer_index, const std::string& bytes) return d; } +KvEncodingDescriptor MakeTurboQuantEncoding(uint64_t stored_bytes) { + KvEncodingDescriptor d; + d.kind = KvEncodingKind::TURBOQUANT; + d.original_dtype = KvDType::BF16; + d.preset = "turboquant_k8v4"; + d.key_bits = 8; + d.value_bits = 4; + d.head_dim = 128; + d.block_size = 16; + d.num_layers = 32; + d.num_tokens = 16; + d.num_heads = 8; + d.hidden_dim = 1024; + d.skip_first_layers = 1; + d.skip_last_layers = 1; + d.stored_bytes = stored_bytes; + d.logical_bytes = stored_bytes; + return d; +} + // A representative batch: 4 keys mixing found / not-found, varying page counts // (including a found-but-zero-page key), and several tiers. std::vector SampleEntries() { @@ -57,6 +77,7 @@ std::vector SampleEntries() { entries[0].found = true; entries[0].tier = TierType::HBM; entries[0].size = 3 * 4096; + entries[0].encoding = MakeTurboQuantEncoding(entries[0].size); entries[0].pages = {{0, 10}, {1, 11}, {0, 12}}; entries[1].found = false; // miss: no tier / size / pages @@ -64,6 +85,7 @@ std::vector SampleEntries() { entries[2].found = true; entries[2].tier = TierType::DRAM; entries[2].size = 4096; + entries[2].encoding = RawKvEncoding(entries[2].size, KvDType::FP16); entries[2].pages = {{2, 7}}; entries[3].found = true; @@ -79,6 +101,7 @@ void ExpectKeyMatches(const ResolvedKeyEntry& src, const DecodedResolveKey& dec) if (!src.found) return; EXPECT_EQ(src.tier, dec.tier); EXPECT_EQ(src.size, dec.size); + EXPECT_EQ(src.encoding, dec.encoding); ASSERT_EQ(src.pages.size(), dec.pages.size()); for (size_t p = 0; p < src.pages.size(); ++p) { EXPECT_EQ(src.pages[p].buffer_index, dec.pages[p].buffer_index); diff --git a/src/umbp/tests/test_kv_encoding.cpp b/src/umbp/tests/test_kv_encoding.cpp new file mode 100644 index 000000000..7fd26e8b7 --- /dev/null +++ b/src/umbp/tests/test_kv_encoding.cpp @@ -0,0 +1,71 @@ +// Copyright © Advanced Micro Devices, Inc. All rights reserved. +// +// MIT License +#include + +#include "umbp/codec/kv_encoding.h" + +namespace mori::umbp { +namespace { + +TEST(KvEncoding, RawEncodingCarriesStoredByteSize) { + auto desc = RawKvEncoding(4096, KvDType::BF16); + EXPECT_TRUE(ValidateKvEncodingDescriptor(desc)); + EXPECT_EQ(desc.kind, KvEncodingKind::RAW); + EXPECT_EQ(desc.stored_bytes, 4096u); + EXPECT_EQ(desc.logical_bytes, 4096u); +} + +TEST(KvEncoding, TurboQuantK8V4ComputesSerializedBytes) { + KvEncodingDescriptor desc; + desc.kind = KvEncodingKind::TURBOQUANT; + desc.original_dtype = KvDType::BF16; + desc.preset = "turboquant_k8v4"; + desc.key_bits = 8; + desc.value_bits = 4; + desc.head_dim = 128; + desc.block_size = 16; + desc.num_layers = 4; + desc.num_tokens = 16; + desc.num_heads = 2; + desc.hidden_dim = 256; + desc.skip_first_layers = 1; + desc.skip_last_layers = 1; + + auto layout = ComputeTurboQuantLayout(desc); + ASSERT_TRUE(layout.has_value()); + EXPECT_EQ(layout->key_packed_size, 128u); + EXPECT_EQ(layout->value_packed_size, 68u); + EXPECT_EQ(layout->slot_size_aligned, 196u); + EXPECT_EQ(layout->compressed_layers, 2u); + EXPECT_EQ(layout->raw_layers, 2u); + EXPECT_EQ(layout->compressed_bytes, 2u * 1u * 16u * 2u * 196u); + EXPECT_EQ(layout->raw_bytes, 2u * 2u * 16u * 2u * 128u * 2u); + EXPECT_EQ(layout->stored_bytes, layout->compressed_bytes + layout->raw_bytes); + + desc.stored_bytes = layout->stored_bytes; + EXPECT_TRUE(ValidateKvEncodingDescriptor(desc)); +} + +TEST(KvEncoding, RejectsMismatchedTurboQuantStoredBytes) { + KvEncodingDescriptor desc; + desc.kind = KvEncodingKind::TURBOQUANT; + desc.original_dtype = KvDType::BF16; + desc.preset = "turboquant_4bit_nc"; + desc.key_bits = 4; + desc.value_bits = 4; + desc.head_dim = 64; + desc.block_size = 16; + desc.num_layers = 2; + desc.num_tokens = 16; + desc.num_heads = 1; + desc.hidden_dim = 64; + desc.stored_bytes = 1; + + std::string error; + EXPECT_FALSE(ValidateKvEncodingDescriptor(desc, &error)); + EXPECT_FALSE(error.empty()); +} + +} // namespace +} // namespace mori::umbp diff --git a/src/umbp/tests/test_peer_dram_allocator.cpp b/src/umbp/tests/test_peer_dram_allocator.cpp index dda9d4c82..37db73e79 100644 --- a/src/umbp/tests/test_peer_dram_allocator.cpp +++ b/src/umbp/tests/test_peer_dram_allocator.cpp @@ -29,6 +29,7 @@ #include #include +#include "umbp/codec/kv_encoding.h" #include "umbp/distributed/peer/peer_dram_allocator.h" namespace mori::umbp { @@ -62,6 +63,21 @@ std::optional AllocateOk(PeerDramAllocator& a, return a.Allocate(key, size, tier).slot; } +KvEncodingDescriptor TestEncoding(uint64_t stored_bytes) { + KvEncodingDescriptor d = RawKvEncoding(stored_bytes, KvDType::BF16); + d.kind = KvEncodingKind::TURBOQUANT; + d.preset = "turboquant_k8v4"; + d.key_bits = 8; + d.value_bits = 4; + d.head_dim = 128; + d.block_size = 16; + d.num_layers = 2; + d.num_tokens = 16; + d.num_heads = 1; + d.hidden_dim = 128; + return d; +} + } // namespace // ---- Allocate / Commit / Resolve happy path --------------------------------- @@ -90,6 +106,31 @@ TEST(PeerDramAllocator, CommitMakesKeyResolvable) { EXPECT_EQ(events[0].tier, TierType::DRAM); } +TEST(PeerDramAllocator, EncodingDescriptorSurvivesCommitResolveAndSnapshot) { + auto a = MakeAllocator(); + auto encoding = TestEncoding(kPageSize); + + auto result = a->Allocate("encoded", kPageSize, TierType::DRAM, encoding); + ASSERT_EQ(result.outcome, PeerDramAllocator::Outcome::kSuccessAllocated); + ASSERT_TRUE(result.slot.has_value()); + EXPECT_EQ(result.slot->encoding, encoding); + + uint64_t committed_bytes = 0; + ASSERT_TRUE(a->Commit(result.slot->slot_id, "encoded", committed_bytes)); + + auto resolved = a->Resolve("encoded"); + ASSERT_TRUE(resolved.found); + EXPECT_EQ(resolved.encoding, encoding); + + auto events = a->DrainPendingEvents(); + ASSERT_EQ(events.size(), 1u); + EXPECT_EQ(events[0].encoding, encoding); + + auto snapshot = a->SnapshotOwnedKeys(); + ASSERT_EQ(snapshot.size(), 1u); + EXPECT_EQ(snapshot[0].encoding, encoding); +} + // ---- Allocate-side dedup ---------------------------------------------------- // Defensive layer for master-index lag (primary dedup is at BatchRoutePut).