diff --git a/src/io/rdma/backend_impl.cpp b/src/io/rdma/backend_impl.cpp index 75daa3b8b..659a704a4 100644 --- a/src/io/rdma/backend_impl.cpp +++ b/src/io/rdma/backend_impl.cpp @@ -815,13 +815,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; @@ -1312,20 +1332,24 @@ 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); + RdmaTransferControl control{}; + control.chunkBytes = config.enableTransferChunking ? config.chunkBytes : 0; + control.maxChunks = config.maxChunksPerTransfer; + control.creditByWrCount = config.enableTransferChunking; + control.extraCompletionCredits = notifBatchSize; RdmaOpRet ret = RdmaBatchReadWrite(eps, localMrPerEp, remoteMrPerEp, {localOffset}, - {remoteOffset}, {size}, callbackMeta, id, isRead, 1, - config.enableTransferChunking ? config.chunkBytes : 0, - config.maxChunksPerTransfer, config.enableTransferChunking); + {remoteOffset}, {size}, callbackMeta, id, isRead, 1, control); assert(!ret.Init()); if (ret.Failed() || ret.Succeeded()) { 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); } @@ -1380,6 +1404,13 @@ void RdmaBackendSession::BatchReadWrite(const SizeVec& localOffsets, const SizeV totalCredit = static_cast(sizes.size()); } + const int notifBatchSize = config.enableNotification ? static_cast(eps.size()) : 0; + if (totalCredit > std::numeric_limits::max() - notifBatchSize) { + status->Update(StatusCode::ERR_INVALID_ARGS, "final WR count exceeds int range"); + return; + } + totalCredit += notifBatchSize; + auto callbackMeta = std::make_shared(status, id, totalCredit); internal::PublishCurrentIoCallDiagnostics(callbackMeta); RdmaOpRet ret; @@ -1416,6 +1447,7 @@ void RdmaBackendSession::BatchReadWrite(const SizeVec& localOffsets, const SizeV control.chunkBytes = chunk ? config.chunkBytes : 0; control.maxChunks = config.maxChunksPerTransfer; control.creditByWrCount = chunk; + control.extraCompletionCredits = notifBatchSize; ret = RdmaBatchReadWrite(eps, localMrPerEp, remoteMrPerEp, localOffsets, remoteOffsets, sizes, callbackMeta, id, isRead, config.postBatchSize, control); } @@ -1424,7 +1456,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 1dd8d0bea..4e9c7fef4 100644 --- a/src/io/rdma/common.cpp +++ b/src/io/rdma/common.cpp @@ -395,9 +395,9 @@ 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; @@ -413,6 +413,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())}; @@ -421,8 +427,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; @@ -431,8 +439,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 = @@ -705,7 +716,10 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, return {StatusCode::ERR_INVALID_ARGS, "final WR count exceeds int range"}; } for (size_t k = 0; k < mergedWrCount; ++k) mergedWrs[k].mergedRequests = 1; - if (control.ownsTotalBatchSize) callbackMeta->totalBatchSize = static_cast(mergedWrCount); + if (control.ownsTotalBatchSize) { + callbackMeta->totalBatchSize = + static_cast(mergedWrCount) + control.extraCompletionCredits; + } } size_t epNum = eps.size(); diff --git a/src/io/rdma/common.hpp b/src/io/rdma/common.hpp index 9712e195a..c968df5cb 100644 --- a/src/io/rdma/common.hpp +++ b/src/io/rdma/common.hpp @@ -253,7 +253,8 @@ 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); struct RdmaTransferControl { size_t chunkBytes{0}; @@ -261,6 +262,7 @@ struct RdmaTransferControl { bool creditByWrCount{false}; bool ownsTotalBatchSize{true}; bool disableMerge{false}; + int extraCompletionCredits{0}; }; RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, 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). diff --git a/tests/cpp/io/test_engine.cpp b/tests/cpp/io/test_engine.cpp index 79dbd2274..fee0678c9 100644 --- a/tests/cpp/io/test_engine.cpp +++ b/tests/cpp/io/test_engine.cpp @@ -247,6 +247,46 @@ 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 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); + (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"); @@ -1642,6 +1682,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},