diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 735716e59e92..cf77a948b98c 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -28,11 +28,13 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/utils/pgUtils.h" #include +#include #include #include #include #include #include +#include #include #include #include @@ -305,6 +307,10 @@ class CacheTransceiver : public BaseCacheTransceiver void setContextState(LlmRequest* llmRequest); + // Append one row per completed request to the gen-side transfer summary CSV. Opens the file + // lazily on first use; expects timing to already be synced across ranks by the caller. + void writeGenTransferSummary(std::vector const& completedRequests); + std::unique_ptr mCacheSender; std::unique_ptr mCacheReceiver; // shared_ptr (not raw LlmRequest*) so the futures hold a strong reference for @@ -340,6 +346,13 @@ class CacheTransceiver : public BaseCacheTransceiver // TODO(shreyasm): update this to use same container as kv by using base trans buffers instead std::unique_ptr mRnnCacheTransBufferManager{nullptr}; + // Unique instance identifier for CSV file naming (avoids collisions across gen instances) + std::string mInstanceId; + + // Gen-side transfer summary CSV (written after timing sync) + std::ofstream mGenTransferSummaryFile; + std::mutex mGenTransferSummaryMutex; + // library handle to the communicator related features, // this is used to defer dependency resolution until needed. static std::mutex mDllMutex; diff --git a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h index 94c46c3fb25f..dea7bda60274 100644 --- a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h +++ b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h @@ -84,6 +84,15 @@ enum LlmRequestType class ContextProgress; +// Process-global offset between the local steady clock and the global steady +// clock (rank 0's steady clock). The storage lives in a single translation unit +// (llmRequest.cpp) and is reached through this accessor so that +// libtensorrt_llm.so and the nanobind extension module share one copy across +// .so boundaries. An inline-static member would instead give each shared object +// its own copy, so an offset calibrated on one side would be invisible to the +// other. +std::optional& globalSteadyClockOffset(); + template class GenericLlmRequest { @@ -1857,14 +1866,18 @@ class GenericLlmRequest mDecodingIter = iter; } + // Callers must pass a global-steady-clock time point (getSteadyClockNow(), + // or a value merged from such time points). Normalizing again here would + // apply the global steady clock offset twice, which corrupts cross-node + // min/max merging whenever the offset is non-zero. void setKvCacheTransferStart(TimePoint time) const { - mPerfMetrics.timingMetrics.kvCacheTransferStart = maybeToGlobalSteadyClock(time); + mPerfMetrics.timingMetrics.kvCacheTransferStart = time; } void setKvCacheTransferEnd(TimePoint time) const { - mPerfMetrics.timingMetrics.kvCacheTransferEnd = maybeToGlobalSteadyClock(time); + mPerfMetrics.timingMetrics.kvCacheTransferEnd = time; } TimePoint getKvCacheTransferStart() const @@ -2040,8 +2053,8 @@ class GenericLlmRequest return mUseDraftModel; } - // If sGlobalSteadyClockOffset is set, return a global steady clock time point, otherwise return local steady clock - // time point + // If the global steady clock offset is set, return a global steady clock time point, otherwise return local steady + // clock time point [[nodiscard]] static TimePoint getSteadyClockNow() { return maybeToGlobalSteadyClock(std::chrono::steady_clock::now()); @@ -2071,9 +2084,6 @@ class GenericLlmRequest // current position of the prompt tuning table (only used in chunked prefill mode) SizeType32 mPtableCurrentPosition{0}; - // The offset between local steady clock and global steady clock (at rank 0) - inline static std::optional sGlobalSteadyClockOffset{std::nullopt}; - protected: bool mIsStreaming; @@ -2385,9 +2395,10 @@ class GenericLlmRequest static TimePoint maybeToGlobalSteadyClock(TimePoint const& time_point) { - if (sGlobalSteadyClockOffset.has_value()) + auto const& offset = globalSteadyClockOffset(); + if (offset.has_value()) { - return time_point + *sGlobalSteadyClockOffset; + return time_point + *offset; } return time_point; } diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 69adf4246591..8a50bcc76d4f 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -57,7 +57,12 @@ #include #include #include +#include +#include +#include #include +#include +#include #include #include #include @@ -65,6 +70,28 @@ namespace tensorrt_llm::batch_manager { +namespace +{ + +/// Generate a UUID-like hex string (e.g. "a1b2c3d4-e5f6-7890-abcd-ef1234567890") +/// to uniquely identify a CacheTransceiver instance across gen instances. +std::string generateInstanceId() +{ + // The RNG state is comparatively expensive to construct/seed, so keep one + // per thread instead of building it on every call. + static thread_local std::mt19937_64 gen{std::random_device{}()}; + std::uniform_int_distribution dis; + uint64_t a = dis(gen); + uint64_t b = dis(gen); + std::ostringstream oss; + oss << std::hex << std::setfill('0') << std::setw(8) << (a >> 32) << "-" << std::setw(4) << ((a >> 16) & 0xFFFF) + << "-" << std::setw(4) << (a & 0xFFFF) << "-" << std::setw(4) << (b >> 48) << "-" << std::setw(12) + << (b & 0xFFFFFFFFFFFF); + return oss.str(); +} + +} // anonymous namespace + std::mutex CacheTransceiver::mDllMutex; namespace @@ -386,6 +413,89 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa mGroupComm = std::make_shared(tensorrt_llm::pg_utils::get_world_pg()); } + // Generate instance ID on rank 0 and broadcast to all ranks in the session + // so every rank in the same gen/ctx instance shares the same ID. + { + if (mGroupComm->getRank() == 0) + { + mInstanceId = generateInstanceId(); + } + if (useMPI()) + { + int len = static_cast(mInstanceId.size()); + tensorrt_llm::mpi::MpiComm::session().bcast(&len, 1, mpi::MpiType::kINT32, 0); + mInstanceId.resize(len); + tensorrt_llm::mpi::MpiComm::session().bcast(mInstanceId.data(), len, mpi::MpiType::kCHAR, 0); + } + else + { + // PG path: rank 0 sends via allgather, others receive. + constexpr int kUuidLen = 36; + std::vector sendBuf(kUuidLen, '\0'); + if (mGroupComm->getRank() == 0) + { + std::copy_n(mInstanceId.begin(), std::min(mInstanceId.size(), kUuidLen), sendBuf.begin()); + } + std::vector recvBuf(kUuidLen * mGroupComm->getSize(), '\0'); + mGroupComm->allgather(std::ref(sendBuf), std::ref(recvBuf), {}); + // Take rank 0's segment. + mInstanceId = std::string(recvBuf.begin(), recvBuf.begin() + kUuidLen); + } + } + + // Calibrate steady_clock across ranks so that cross-node allgather + // in batchUpdateKVCacheTransferBW can compare time points. + // globalSteadyClockOffset() reads a single process-global copy shared with + // the nanobind module, so if the Python runtime already calibrated the offset + // (PyExecutor::_set_global_steady_clock_offset) it is visible here and we skip; + // the pure-C++ path performs the calibration below. + // The check-and-set is guarded by a mutex so that CacheTransceiver instances + // constructed concurrently in the same process (e.g. multi-engine serving) do + // not race on the shared offset or issue mismatched collectives. + { + static std::mutex sSteadyClockCalibrationMutex; + std::lock_guard lock(sSteadyClockCalibrationMutex); + if (!globalSteadyClockOffset().has_value()) + { + using Duration = LlmRequest::Duration; + // Synchronize all ranks immediately before sampling the local clock so + // every rank measures from a consistent point. + if (useMPI()) + { + tensorrt_llm::mpi::MpiComm::session().barrier(); + } + else + { + // CacheTransceiverComm exposes no barrier primitive, so use a cheap + // allgather as a pseudo-barrier for the process-group path. + int64_t const dummy = 0; + std::vector dummyRecv(mGroupComm->getSize(), 0); + mGroupComm->allgather(dummy, std::ref(dummyRecv), {}); + } + auto localNow = std::chrono::steady_clock::now(); + auto localNs = std::chrono::duration_cast(localNow.time_since_epoch()).count(); + + // Allgather timestamps from all ranks + std::vector allNs(mGroupComm->getSize(), 0); + if (useMPI()) + { + tensorrt_llm::mpi::MpiComm::session().allgather(&localNs, allNs.data(), 1, mpi::MpiType::kINT64); + } + else + { + mGroupComm->allgather(localNs, std::ref(allNs), {}); + } + + // Offset = rank0's timestamp - my timestamp (same formula as Python) + auto offsetNs = allNs[0] - localNs; + globalSteadyClockOffset() = Duration(offsetNs); + + TLLM_LOG_INFO(mGroupComm->getRank(), + "CacheTransceiver: set global steady clock offset = %.6f sec for rank %d", + static_cast(offsetNs) / 1e9, mGroupComm->getRank()); + } + } + if (worldConfig.isTensorParallel() || worldConfig.isContextParallel()) { mGroupTensorParaComm = std::make_shared( @@ -567,8 +677,10 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa auto makeCacheTransferLayer = [&]() { return CacheTransferLayer(*mCacheState, makeFormatter(), makeRnnFormatter()); }; - mCacheSender = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); - mCacheReceiver = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer()); + mCacheSender + = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer(), mInstanceId); + mCacheReceiver + = std::make_unique(mManager.get(), worldConfig.getRank(), makeCacheTransferLayer(), mInstanceId); // Keep automatic enablement within the currently qualified C++ NIXL/UCX TP1/CP1 pipeline topology. bool const coordinatorTopologyEligible = worldConfig.getPipelineParallelism() > 1 && useMPI() @@ -762,67 +874,109 @@ std::vector gatherRequestIds( return retData; } -void updateKVCacheTransferBW(std::shared_ptr const& mComm, LlmRequest* request) +void batchUpdateKVCacheTransferBW( + std::shared_ptr const& comm, std::vector const& requests) { + // Key-based merge: each rank serializes (requestId, start, end, size) + // tuples and we use allgatherv so ranks may have different request counts. + // The merge matches by requestId, not by position — this tolerates + // ordering differences and count mismatches across ranks. + namespace su = executor::serialize_utils; - int worldSize = mComm->getSize(); + int const worldSize = comm->getSize(); + + // --- Serialize local entries keyed by requestId --- + std::size_t const numReqs = requests.size(); std::ostringstream oStream; - su::serialize(request->getKvCacheTransferStart(), oStream); - su::serialize(request->getKvCacheTransferEnd(), oStream); + su::serialize(numReqs, oStream); + for (auto* req : requests) + { + su::serialize(req->getContextPhaseParams().value().getReqId(), oStream); + su::serialize(req->getKvCacheTransferStart(), oStream); + su::serialize(req->getKvCacheTransferEnd(), oStream); + su::serialize(req->getKvCacheSize(), oStream); + } auto str = oStream.str(); std::vector sendBuffer(str.begin(), str.end()); - auto sendBufferSize = sendBuffer.size(); - auto recvBufferSize = sendBufferSize * worldSize; - std::vector recvBuffer(recvBufferSize); + int const sendSize = static_cast(sendBuffer.size()); + // --- Step 1: allgather per-rank buffer sizes --- + std::vector recvCounts(worldSize, 0); if (useMPI()) { - mComm->allgather(sendBuffer.data(), recvBuffer.data(), sendBufferSize, mpi::MpiType::kCHAR); + comm->allgather(&sendSize, recvCounts.data(), 1, mpi::MpiType::kINT32); } else { - mComm->allgather(std::ref(sendBuffer), std::ref(recvBuffer), {}); + comm->allgather(sendSize, std::ref(recvCounts), {}); } - su::VectorWrapBuf strbuf(recvBuffer); - std::istream is(&strbuf); - - auto minStartTime = executor::RequestPerfMetrics::TimePoint::max(); - auto maxEndTime = executor::RequestPerfMetrics::TimePoint::min(); - - for (int rank = 0; rank < worldSize; rank++) + // --- Step 2: allgatherv the serialized data --- + std::vector displs(worldSize, 0); + int totalRecvSize = 0; + for (int r = 0; r < worldSize; ++r) { - minStartTime = std::min(su::deserialize(is), minStartTime); - maxEndTime = std::max(su::deserialize(is), maxEndTime); + displs[r] = totalRecvSize; + totalRecvSize += recvCounts[r]; } - - // Handle KV cache size separately - gather all sizes to the leader rank - std::size_t localKVCacheSize = request->getKvCacheSize(); - std::vector allKVCacheSizes(worldSize, 0); + std::vector recvBuffer(totalRecvSize, 0); if (useMPI()) { - mComm->allgather(&localKVCacheSize, allKVCacheSizes.data(), 1, mpi::MpiType::kUINT64); + comm->allgatherv(sendBuffer.data(), sendSize, mpi::MpiType::kCHAR, recvBuffer.data(), recvCounts, displs, + mpi::MpiType::kCHAR); } else { - mComm->allgather(&localKVCacheSize, std::ref(allKVCacheSizes), {}); + comm->allgatherv(std::ref(sendBuffer), std::ref(recvBuffer), recvCounts, {}); } - std::size_t totalKVCacheSize = 0; - for (int rank = 0; rank < worldSize; rank++) + // --- Step 3: Deserialize and merge by requestId --- + using TimePoint = executor::RequestPerfMetrics::TimePoint; + using ReqIdType = LlmRequest::RequestIdType; + + struct MergedEntry + { + TimePoint minStart = TimePoint::max(); + TimePoint maxEnd = TimePoint::min(); + std::size_t totalSize = 0; + }; + + std::unordered_map merged; + + su::VectorWrapBuf strbuf(recvBuffer); + std::istream is(&strbuf); + + for (int rank = 0; rank < worldSize; ++rank) { - totalKVCacheSize += allKVCacheSizes[rank]; + auto rankNumReqs = su::deserialize(is); + for (std::size_t i = 0; i < rankNumReqs; ++i) + { + auto rid = su::deserialize(is); + auto start = su::deserialize(is); + auto end = su::deserialize(is); + auto size = su::deserialize(is); + + auto& entry = merged[rid]; + entry.minStart = std::min(entry.minStart, start); + entry.maxEnd = std::max(entry.maxEnd, end); + entry.totalSize += size; + } } - // Update the latest KV cache transfer time for leader rank - if (mComm->getRank() == 0) + // --- Step 4: Update local requests --- + for (auto* req : requests) { - request->setKvCacheTransferStart(minStartTime); - request->setKvCacheTransferEnd(maxEndTime); - request->setKvCacheSize(totalKVCacheSize); + auto reqId = req->getContextPhaseParams().value().getReqId(); + auto it = merged.find(reqId); + if (it != merged.end()) + { + req->setKvCacheTransferStart(it->second.minStart); + req->setKvCacheTransferEnd(it->second.maxEnd); + req->setKvCacheSize(it->second.totalSize); + } } } @@ -1340,6 +1494,10 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR eraseLocalTransferOutcome( requestId, mCompletedRequesterRequestIds, mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus); } + + // Collect consensus-completed requests so timing can be synced across ranks in a single + // batched allgather (instead of one collective per request). + std::vector completedRequests; for (auto const requestId : sortedRequestIds(consensusOutcome.completedRequestIds)) { auto const requestIt = mRequesterRequestsAwaitingConsensus.find(requestId); @@ -1348,17 +1506,44 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR continue; } requestIt->second->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); - - // Gather the kv cache transfer time from all workers and update to leader rank. - if (!common::getEnvKVCacheTimeOutputPath().empty()) - { - updateKVCacheTransferBW(syncComm, requestIt->second.get()); - } + completedRequests.push_back(requestIt->second.get()); mTimedOutRequesterIds.erase(requestId); mCancelRequestedRequesterIds.erase(requestId); eraseLocalTransferOutcome( requestId, mCompletedRequesterRequestIds, mFailedRequesterRequestIds, mRequesterRequestsAwaitingConsensus); } + + // Batch-sync timing across ranks in one allgather (instead of per-request), then write + // the gen-side transfer summary CSV. + if (!completedRequests.empty() && !common::getEnvKVCacheTimeOutputPath().empty()) + { + batchUpdateKVCacheTransferBW(syncComm, completedRequests); + writeGenTransferSummary(completedRequests); + } +} + +void CacheTransceiver::writeGenTransferSummary(std::vector const& completedRequests) +{ + std::lock_guard lock(mGenTransferSummaryMutex); + if (!mGenTransferSummaryFile.is_open()) + { + namespace fs = std::filesystem; + auto outputPath = fs::path(common::getEnvKVCacheTimeOutputPath()); + fs::create_directories(outputPath); + int rank = useMPI() ? mpi::MpiComm::world().getRank() : tensorrt_llm::pg_utils::get_world_pg()->getRank(); + auto filePath = outputPath / (mInstanceId + "_" + std::to_string(rank) + "_gen_transfer_summary.csv"); + mGenTransferSummaryFile.open(filePath); + TLLM_CHECK_WITH_INFO(mGenTransferSummaryFile.is_open(), "Failed to open gen transfer summary file: %s", + filePath.string().c_str()); + mGenTransferSummaryFile << "RequestID,gen_side_transfer_time(ms),kv_cache_size" << '\n'; + } + for (auto* req : completedRequests) + { + auto reqId = req->getContextPhaseParams().value().getReqId(); + mGenTransferSummaryFile << reqId << "," << req->getKvCacheTransferTimeMS() << "," << req->getKvCacheSize() + << '\n'; + } + mGenTransferSummaryFile << std::flush; } bool CacheTransceiver::checkGenTransferComplete() const diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 851f0abb04ef..52329b6fca28 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -229,7 +229,7 @@ int32_t tagFromRequestId(LlmRequest::RequestIdType requestId) return ((requestId & 0xFFF) << 8) | (kDATA_TAG & 0xFF); } -std::filesystem::path getTransferOutputPath(char const* tag) +std::filesystem::path getTransferOutputPath(char const* tag, std::string const& instanceId = "") { namespace fs = std::filesystem; auto outputPath = common::getEnvKVCacheTimeOutputPath(); @@ -238,7 +238,9 @@ std::filesystem::path getTransferOutputPath(char const* tag) auto rank = mpi::MpiComm::world().getRank(); auto path = fs::path(outputPath); fs::create_directories(path); - return path / ("rank_" + std::to_string(rank) + "_" + tag + ".csv"); + std::string prefix + = instanceId.empty() ? "rank_" + std::to_string(rank) : instanceId + "_" + std::to_string(rank); + return path / (prefix + "_" + tag + ".csv"); } return {}; } @@ -323,11 +325,13 @@ class CacheSender::Impl public: using RequestIdType = LlmRequest::RequestIdType; - Impl(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer) + Impl(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer, + std::string instanceId = "") : mManager{manager} , mSelfState{cacheLayer.getCacheState(), executor::kv_cache::CommState{manager->getCommState()}} , mCacheTransferLayer{std::move(cacheLayer)} , mBufferManager{std::make_shared()} + , mInstanceId{std::move(instanceId)} { TLLM_CHECK(mManager); TLLM_CHECK(mManager->getCommState().getSelfIdx() == selfIndex); @@ -405,7 +409,7 @@ class CacheSender::Impl { if (!mMeasuresFile.is_open()) { - auto outputPath = getTransferOutputPath("send"); + auto outputPath = getTransferOutputPath("send", mInstanceId); mMeasuresFile.open(outputPath); TLLM_CHECK_WITH_INFO(mMeasuresFile.is_open(), "Failed to open transfer output file: %s", outputPath.string().c_str()); @@ -979,16 +983,19 @@ class CacheSender::Impl std::ofstream mMeasuresFile; std::mutex mInFlightCancelMutex; std::unordered_map>> mInFlightCancelFlags; + std::string mInstanceId; }; class CacheReceiver::Impl { public: - Impl(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer) + Impl(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer, + std::string instanceId = "") : mManager{manager} , mSelfState{cacheLayer.getCacheState(), executor::kv_cache::CommState{manager->getCommState()}} , mCacheTransferLayer{std::move(cacheLayer)} , mBufferManager{std::make_shared()} + , mInstanceId{std::move(instanceId)} { TLLM_CHECK(mManager); TLLM_CHECK(mManager->getCommState().getSelfIdx() == selfIndex); @@ -1056,7 +1063,7 @@ class CacheReceiver::Impl std::unique_lock lock(mMeasuresFileMutex); if (!mMeasuresFile.is_open()) { - auto outputPath = getTransferOutputPath("recv"); + auto outputPath = getTransferOutputPath("recv", mInstanceId); mMeasuresFile.open(outputPath); TLLM_CHECK_WITH_INFO(mMeasuresFile.is_open(), "Failed to open transfer output file: %s", outputPath.string().c_str()); @@ -1720,6 +1727,7 @@ class CacheReceiver::Impl std::atomic mTerminate{false}; std::mutex mInFlightCancelMutex; std::unordered_map>> mInFlightCancelFlags; + std::string mInstanceId; }; void CacheSender::ImplDeleter::operator()(Impl* ptr) @@ -1732,9 +1740,10 @@ void CacheReceiver::ImplDeleter::operator()(Impl* ptr) delete ptr; } -CacheSender::CacheSender( - executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer) - : mImpl{std::unique_ptr(new Impl(manager, selfIndex, std::move(cacheLayer)))} +CacheSender::CacheSender(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, + CacheTransferLayer cacheLayer, std::string instanceId) + : mImpl{ + std::unique_ptr(new Impl(manager, selfIndex, std::move(cacheLayer), std::move(instanceId)))} { } @@ -1784,9 +1793,10 @@ void CacheSender::setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rn mImpl->setRnnConfig(std::move(rnnModelConfig), std::move(rnnLayerNumPerPP), convStateDataType, ssmStateDataType); } -CacheReceiver::CacheReceiver( - executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer) - : mImpl{std::unique_ptr(new Impl(manager, selfIndex, std::move(cacheLayer)))} +CacheReceiver::CacheReceiver(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, + CacheTransferLayer cacheLayer, std::string instanceId) + : mImpl{ + std::unique_ptr(new Impl(manager, selfIndex, std::move(cacheLayer), std::move(instanceId)))} { } diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h index c1da646916dc..07160f8a1d1d 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -268,7 +268,8 @@ class CacheSender /// @param manager The connection manager. /// @param selfIndex The sequential index of the current executor process. /// @param cacheLayer The cache layer bundling all cache states and formatters. - CacheSender(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer); + CacheSender(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer, + std::string instanceId = ""); CacheSender() = default; @@ -331,7 +332,8 @@ class CacheReceiver /// @param manager The connection manager. /// @param selfIndex The sequential index of the current executor process. /// @param cacheLayer The cache layer bundling all cache states and formatters. - CacheReceiver(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer); + CacheReceiver(executor::kv_cache::ConnectionManager* manager, SizeType32 selfIndex, CacheTransferLayer cacheLayer, + std::string instanceId = ""); CacheReceiver() = default; diff --git a/cpp/tensorrt_llm/batch_manager/llmRequest.cpp b/cpp/tensorrt_llm/batch_manager/llmRequest.cpp index e51fad8ba149..d5466b4a7539 100644 --- a/cpp/tensorrt_llm/batch_manager/llmRequest.cpp +++ b/cpp/tensorrt_llm/batch_manager/llmRequest.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -22,6 +22,16 @@ namespace tensorrt_llm::batch_manager { +// Single, process-global storage for the steady-clock offset. Keeping it in this +// translation unit (rather than as an inline-static member reachable from every +// shared object) guarantees that libtensorrt_llm.so and the nanobind extension +// module observe the same value once either side calibrates it. +std::optional& globalSteadyClockOffset() +{ + static std::optional offset{std::nullopt}; + return offset; +} + template runtime::SizeType32 GenericLlmRequest::getBeamWidthByIter(bool const forNextIteration) { diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp index 202228b394d8..3c1823e2a94a 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp @@ -196,6 +196,8 @@ void initBindings(nb::module_& m) .def_prop_ro("kv_cache_transfer_time_ms", &GenLlmReq::getKvCacheTransferTimeMS) .def_prop_ro("kv_cache_transfer_start", &GenLlmReq::getKvCacheTransferStart) .def_prop_ro("kv_cache_transfer_end", &GenLlmReq::getKvCacheTransferEnd) + .def("get_kv_cache_transfer_start", &GenLlmReq::getKvCacheTransferStart) + .def("get_kv_cache_transfer_end", &GenLlmReq::getKvCacheTransferEnd) .def_prop_ro("kv_cache_size", &GenLlmReq::getKvCacheSize) .def("set_kv_cache_transfer_start", &GenLlmReq::setKvCacheTransferStart, nb::arg("time")) .def("set_kv_cache_transfer_end", &GenLlmReq::setKvCacheTransferEnd, nb::arg("time")) @@ -479,7 +481,11 @@ void initBindings(nb::module_& m) .def("set_first_scheduled_time", &tb::LlmRequest::setFirstScheduledTime) .def("update_perf_metrics", &tb::LlmRequest::updatePerfMetrics, nb::arg("iter_counter")) .def("remove_lora_tensors", &tb::LlmRequest::removeLoraTensors) - .def_rw_static("global_steady_clock_offset", &tb::LlmRequest::sGlobalSteadyClockOffset); + // Bind to the single storage owned by libtensorrt_llm.so (reached through + // globalSteadyClockOffset()) instead of an inline-static member, so the + // offset is shared with the native library rather than living in this + // module's separate copy. + .def_rw_static("global_steady_clock_offset", &tb::globalSteadyClockOffset()); nb::class_(m, "SequenceSlotManager") .def(nb::init(), nb::arg("max_num_slots"), diff --git a/cpp/tensorrt_llm/nanobind/bindings.cpp b/cpp/tensorrt_llm/nanobind/bindings.cpp index 79dd147092bc..2c724adc279c 100644 --- a/cpp/tensorrt_llm/nanobind/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/bindings.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -544,4 +544,8 @@ NB_MODULE(TRTLLM_NB_MODULE, m) m.def("ipc_nvls_supported", &tr::ipcNvlsSupported); m.def("steady_clock_now", []() { return std::chrono::steady_clock::now(); }); + // Global (offset-normalized) steady clock, matching what + // LlmRequest::setKvCacheTransferStart/End expect. Reads the process-global + // steady clock offset, set by PyExecutor at startup. + m.def("global_steady_clock_now", []() { return tb::LlmRequest::getSteadyClockNow(); }); } diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index 5d1a958d37ce..2cf69983b34f 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -883,15 +883,18 @@ class AsymmetricalCacheTest : public ::testing::TestWithParam( - mConnectionManager.get(), mRankInInstance, CacheTransferLayer(*mCacheState, makeFormatter())); + mSender = std::make_unique(mConnectionManager.get(), mRankInInstance, + CacheTransferLayer(*mCacheState, makeFormatter()), instanceId); } else { - mRequester = std::make_unique( - mConnectionManager.get(), mRankInInstance, CacheTransferLayer(*mCacheState, makeFormatter())); + mRequester = std::make_unique(mConnectionManager.get(), mRankInInstance, + CacheTransferLayer(*mCacheState, makeFormatter()), instanceId); } TLLM_LOG_DEBUG("setUpCacheTransceiver mSender"); diff --git a/examples/disaggregated/slurm/cache_transceiver_test/README.md b/examples/disaggregated/slurm/cache_transceiver_test/README.md index 14619e4d387c..e99c57dc4fe3 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/README.md +++ b/examples/disaggregated/slurm/cache_transceiver_test/README.md @@ -64,7 +64,7 @@ resolved_config.json # the validated config the job ran with logs/ctt-.log # batch-level log (stdout+stderr merged) logs/install.log # TensorRT-LLM install log logs/sweep__rank*.log # per-rank logs: transfer START/DONE+verify, UCX_PROTO_INFO -csv//ctx|gen/rank_*_{send,recv}.csv # C++ transceiver timing (Bandwidth(Gbps)) +csv//ctx|gen/_*_{send,recv}.csv # C++ transceiver timing (Bandwidth(Gbps)); renamed to _*_{send,recv}__c.csv per combination csv//ctx/py_*_*.csv # Python transceiver perf log (throughput_mbs) status/sweep_.jsonl # PASS / MISMATCH / TRANSFER_ERROR / TIMEOUT results.json # full results, grouped per combination (longest req_len) @@ -96,7 +96,7 @@ deliverable for tuning your cluster. | Transceiver | Env enabling timing | File | Column (native) | |---|---|---|---| -| C++ (UCX/NIXL) | `TRTLLM_KVCACHE_TIME_OUTPUT_PATH` (set by the driver) | `rank_*_recv.csv` | `Bandwidth(Gbps)` | +| C++ (UCX/NIXL) | `TRTLLM_KVCACHE_TIME_OUTPUT_PATH` (set by the driver) | `_*_recv.csv` (renamed `_*_recv__c.csv` per combination) | `Bandwidth(Gbps)` | | Python (NIXL) | `TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO=1`, `TLLM_KV_TRANSFER_PERF_LOG_FILE` (set by the driver) | `py_*_*.csv` | `throughput_mbs` (MiB/s) | `report.py` normalizes both to **per-GPU GB/s** (bytes, ÷1e9): C++ diff --git a/examples/disaggregated/slurm/cache_transceiver_test/report.py b/examples/disaggregated/slurm/cache_transceiver_test/report.py index 70d2658eab6b..00ecab6b4e49 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/report.py +++ b/examples/disaggregated/slurm/cache_transceiver_test/report.py @@ -184,7 +184,7 @@ def _mean(values): def _parse_cpp_recv_csvs(csv_dir): - """Return {rid: [per_rank_GBps, ...]} from C++ rank_*_recv.csv files. + """Return {rid: [per_rank_GBps, ...]} from C++ *_recv.csv files. Each rank writes one row per request with a repeating Bandwidth(Gbps) column per transmission; we take the mean transmission bandwidth as that rank's @@ -194,9 +194,10 @@ def _parse_cpp_recv_csvs(csv_dir): per-rank rates with unequal durations would overstate the real throughput. """ per_rid = {} # rid -> list[per-rank mean bw] - # The driver renames each combination's rank_*_recv.csv to rank_*_recv__cl
  • .csv + # C++ writes "__recv.csv" (instanceId is a runtime UUID); + # the driver renames each combination's *_recv.csv to *_recv__cl
  • .csv # (and the un-renamed name may exist for the last iteration). Match both. - for path in glob.glob(os.path.join(csv_dir, "rank_*_recv*.csv")): + for path in glob.glob(os.path.join(csv_dir, "*_recv*.csv")): with open(path) as f: reader = csv.reader(f) header = next(reader, None) diff --git a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py index 6a6775d5c08e..138a43c96063 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py +++ b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py @@ -24,7 +24,7 @@ deterministic, rank-specific pattern, sends it, and the gen side verifies the received blocks regenerate to the same pattern. Bandwidth is emitted by the transceivers themselves into per-rank CSVs (parsed later by report.py): - C++ -> TRTLLM_KVCACHE_TIME_OUTPUT_PATH (rank_*_send.csv / rank_*_recv.csv) + C++ -> TRTLLM_KVCACHE_TIME_OUTPUT_PATH (_*_send.csv / _*_recv.csv) Py -> TLLM_KV_TRANSFER_PERF_LOG_FILE (py_*_*.csv, throughput_mbs) This driver mirrors the single-process test (tests/unittest/others/ @@ -437,15 +437,19 @@ def _preserve_cpp_csvs(csv_dir, ci, rank): request lengths of a combination and appends a row per request, so we move the whole combination's output aside (rid encodes req_len). - Each rank touches ONLY its own files: all ranks share `csv_dir`, so a glob - over `rank_*` would race -- multiple ranks renaming the same file, leaving - some with FileNotFoundError, crashing those ranks and deadlocking the rest on - the next case's collective KVCacheManager allreduce. + C++ names files "__.csv" (instanceId is a runtime + UUID). Each rank touches ONLY files carrying its own "__.csv" + suffix: all ranks share `csv_dir`, so matching a broader pattern would race + -- multiple ranks renaming the same file, leaving some with + FileNotFoundError, crashing those ranks and deadlocking the rest on the next + case's collective KVCacheManager allreduce. """ for tag in ("send", "recv"): - path = os.path.join(csv_dir, f"rank_{rank}_{tag}.csv") - if os.path.exists(path): - os.replace(path, os.path.join(csv_dir, f"rank_{rank}_{tag}__c{ci}.csv")) + suffix = f"_{rank}_{tag}.csv" + for name in os.listdir(csv_dir): + if name.endswith(suffix) and "__c" not in name: + base = name[: -len(".csv")] + os.replace(os.path.join(csv_dir, name), os.path.join(csv_dir, f"{base}__c{ci}.csv")) def main(): diff --git a/tensorrt_llm/_torch/disaggregation/native/perf_logger.py b/tensorrt_llm/_torch/disaggregation/native/perf_logger.py index 3c491df3495b..5528b8f1eb08 100644 --- a/tensorrt_llm/_torch/disaggregation/native/perf_logger.py +++ b/tensorrt_llm/_torch/disaggregation/native/perf_logger.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -130,11 +130,13 @@ def get_task_latency(self, peer_rank: int) -> float: class PerfLogManager: """Singleton manager for KV transfer performance logging. - Logic: - - TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO not set: no output - - TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO set, TLLM_KV_TRANSFER_PERF_LOG_FILE not set: - logger.info to stdout - - Both set: CSV output to {TLLM_KV_TRANSFER_PERF_LOG_FILE}_{instance_name}_{instance_rank}.csv + Logic (checked in priority order): + 1. TRTLLM_KVCACHE_TIME_OUTPUT_PATH set (C++ standard): enabled, CSV to + ``{path}/{instance_name}_{instance_rank}.csv`` + 2. TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO=1 with TLLM_KV_TRANSFER_PERF_LOG_FILE: + enabled, CSV to ``{base}_{instance_name}_{instance_rank}.csv`` + 3. TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO=1 alone: enabled, logger.info to stdout + 4. None of the above: disabled """ _instance = None @@ -154,8 +156,18 @@ def __init__(self): self._initialized = True self._file_loggers = {} # (instance_name, instance_rank) -> logger self._file_lock = threading.Lock() - self._perf_enabled = os.getenv("TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO", "0") == "1" - self._log_file_base = os.getenv("TLLM_KV_TRANSFER_PERF_LOG_FILE") + + # Primary: C++ standard env var (directory path) + cpp_output_path = os.getenv("TRTLLM_KVCACHE_TIME_OUTPUT_PATH") + if cpp_output_path: + self._perf_enabled = True + self._log_file_base = cpp_output_path + self._use_cpp_naming = True + else: + # Fallback: existing Python env vars + self._perf_enabled = os.getenv("TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO", "0") == "1" + self._log_file_base = os.getenv("TLLM_KV_TRANSFER_PERF_LOG_FILE") + self._use_cpp_naming = False @property def enabled(self) -> bool: @@ -175,8 +187,12 @@ def _get_or_create_file_logger(self, instance_name: str, instance_rank: int): if key in self._file_loggers: return self._file_loggers[key] - # Create file path: {base}_{instance_name}_{instance_rank}.csv - log_file = f"{self._log_file_base}_{instance_name}_{instance_rank}.csv" + if self._use_cpp_naming: + # C++ pattern: {dir}/{instance_name}_{instance_rank}.csv + log_file = os.path.join(self._log_file_base, f"{instance_name}_{instance_rank}.csv") + else: + # Legacy Python pattern: {base}_{instance_name}_{instance_rank}.csv + log_file = f"{self._log_file_base}_{instance_name}_{instance_rank}.csv" try: # Create directory if needed @@ -315,6 +331,73 @@ def log_recv_task_perf( ) self.log(instance_name, instance_rank, csv_line, info_msg) + def log_gen_transfer_summary( + self, + unique_rid: int, + instance_name: str, + instance_rank: int, + gen_side_transfer_time_ms: float, + kv_cache_size: int, + ) -> None: + """Log a gen-side transfer summary row to a separate CSV. + + Written after timing sync across ranks so values are globally + consistent. Only active when ``TRTLLM_KVCACHE_TIME_OUTPUT_PATH`` + is set. + + Args: + unique_rid: Unique request id. + instance_name: Instance name for file naming. + instance_rank: Instance rank for file naming. + gen_side_transfer_time_ms: Synced gen-side transfer time in ms. + kv_cache_size: Total KV cache size across ranks (bytes). + """ + cpp_output_path = os.getenv("TRTLLM_KVCACHE_TIME_OUTPUT_PATH") + if not cpp_output_path: + return + + _GEN_SUMMARY_HEADER = "timestamp,RequestID,gen_side_transfer_time(ms),kv_cache_size" + key = ("gen_summary", instance_name, instance_rank) + + if key not in self._file_loggers: + with self._file_lock: + if key not in self._file_loggers: + log_file = os.path.join( + cpp_output_path, + f"{instance_name}_{instance_rank}_gen_transfer_summary.csv", + ) + try: + log_dir = os.path.dirname(log_file) + if log_dir and not os.path.exists(log_dir): + os.makedirs(log_dir, exist_ok=True) + + write_header = not os.path.exists(log_file) + file_logger = logging.getLogger( + f"kv_gen_summary_{instance_name}_{instance_rank}" + ) + file_logger.setLevel(logging.INFO) + file_logger.propagate = False + file_handler = logging.FileHandler(log_file, mode="a", encoding="utf-8") + formatter = logging.Formatter( + fmt="%(asctime)s.%(msecs)03d,%(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + file_handler.setFormatter(formatter) + file_logger.addHandler(file_handler) + if write_header: + file_handler.stream.write(_GEN_SUMMARY_HEADER + "\n") + file_handler.stream.flush() + self._file_loggers[key] = file_logger + except OSError as e: + sys.stderr.write( + f"[KV Transfer] Warning: Failed to create gen summary log file {log_file}: {e}\n" + ) + return + + file_logger = self._file_loggers.get(key) + if file_logger: + file_logger.info(f"{unique_rid},{gen_side_transfer_time_ms:.3f},{kv_cache_size}") + # Singleton instance perf_log_manager = PerfLogManager() diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 8e0e004dda6c..0ee936f20251 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -176,16 +176,16 @@ class AgentResult(Enum): FAILED = "FAILED" -# KV_AGENT_RESULT prefix in one struct frame (was 5 ascii frames serialized/parsed under the -# GIL per slice per writer): instance_rank, unique_rid, slice_id, is_last, status. The optional -# bounce tail follows at message[2:]. -_KV_RESULT_PREFIX = struct.Struct(" SessionStatus: return SessionStatus.READY if self.receiver_ready else SessionStatus.INIT def send(self, slice: KVSlice) -> None: + if self.transfer_start_time is None: + self.transfer_start_time = tensorrt_llm.bindings.global_steady_clock_now() with self.lock: params = self._base_args.params slice_id = len(self.kv_tasks) @@ -1725,7 +1734,7 @@ def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): f"_process_kv_agent_result: unexpected msg_type={message[0]!r}, expected KV_AGENT_RESULT" ) return - peer_rank, unique_rid, sender_slice_id, is_last_slice, status_code = ( + peer_rank, unique_rid, sender_slice_id, is_last_slice, status_code, transfer_size = ( _KV_RESULT_PREFIX.unpack(message[1]) ) from .bounce import decode_result_tail @@ -1745,6 +1754,7 @@ def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): dst_ptrs=dst_ptrs, sizes=sizes, src_base=src_base, + transfer_size=transfer_size, ) def _process_aux_agent_result(self, _send_id: bytes, message: list[bytes]): @@ -1802,6 +1812,9 @@ def __init__( self._exception: Optional[Exception] = None self._closed = False self._terminal_status: Optional[SessionStatus] = None + self.transfer_start_time = None + self.transfer_end_time = None + self.kv_cache_size_bytes: int = 0 self._kv_tasks: list[KVRecvTask] = [] self._aux_count = 0 self._aux_status: TaskStatus = TaskStatus.INIT @@ -1842,6 +1855,8 @@ def mark_transferring(self, slice_id: int): self._kv_tasks[slice_id].status = TaskStatus.TRANSFERRING def receive(self, slice: KVSlice) -> None: + if self.transfer_start_time is None: + self.transfer_start_time = tensorrt_llm.bindings.global_steady_clock_now() params = self._base_args.params slice_id = len(self._kv_tasks) task = KVRecvTask( @@ -1863,8 +1878,10 @@ def process_kv_agent_result( dst_ptrs=None, sizes=None, src_base=None, + transfer_size: int = 0, ): with self.lock: + self.kv_cache_size_bytes += transfer_size assert sender_slice_id < len(self._kv_tasks), ( f"Receiver got slice_id={sender_slice_id} from sender but only has " f"{len(self._kv_tasks)} receive task(s) for request {self.request_id}. " @@ -1919,6 +1936,13 @@ def on_done( f"slice={sender_slice_id}: {e}" ) task.complete() + # Transfer end for perf/time-sync: only meaningful once every slice has + # landed. Plain attribute write (atomic under the GIL); on_done must stay + # lock-free, and consumers only read it after wait_complete succeeds. + if all(t.status == TaskStatus.TRANSFERRED for t in self._kv_tasks): + self.transfer_end_time = ( + tensorrt_llm.bindings.global_steady_clock_now() + ) logger.debug( f"KV transfer complete for request {request_id} " f"slice={sender_slice_id}" diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 4fe7fa28e50f..2c257a39f731 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -1,3 +1,4 @@ +import os import time import uuid from collections import defaultdict @@ -7,6 +8,7 @@ import numpy as np import torch +import tensorrt_llm.bindings from tensorrt_llm import logger from tensorrt_llm._torch.disaggregation.base.transfer import ( KVSlice, @@ -20,6 +22,7 @@ from tensorrt_llm._torch.disaggregation.native.bounce import ( config_from_size as bounce_config_from_size, ) +from tensorrt_llm._torch.disaggregation.native.perf_logger import perf_log_manager from tensorrt_llm._torch.disaggregation.native.transfer import TransferWorker, TransferWorkerConfig from tensorrt_llm._torch.disaggregation.resource.cache_reuse import ( CacheReuseAdapter, @@ -406,6 +409,57 @@ def _ctx_consensus_outcome(self, to_process, cancelled, failed, completed, timed c, f, d = self._consensus_outcome(to_process, c, f, d, pp_allgather, True) return c, f, d, timed_out + def _sync_transfer_timing(self, reqs: list): + """Allgather timing for a batch of completed requests in one collective. + + Matches C++ ``batchUpdateKVCacheTransferBW()`` in ``cacheTransceiver.cpp``. + Only runs when ``TRTLLM_KVCACHE_TIME_OUTPUT_PATH`` is set (same gate + as C++) and multi-rank sync is needed. All ranks that participate in + the allgather update their local request objects. + """ + if not reqs: + return + if not os.getenv("TRTLLM_KVCACHE_TIME_OUTPUT_PATH"): + return + if not self._gen_need_sync: + return + + # Pack local timing for all completed requests into one dict. + local_data = { + get_unique_rid(req): ( + req.get_kv_cache_transfer_start(), + req.get_kv_cache_transfer_end(), + req.kv_cache_size, + ) + for req in reqs + } + + # Single allgather for the whole batch. + all_data = self._gen_allgather(local_data) + + # Merge: per-rid min(start), max(end), sum(size) across ranks. + merged: dict = {} + for rank_data in all_data: + for rid, (start, end, size) in rank_data.items(): + if rid in merged: + prev = merged[rid] + merged[rid] = ( + min(prev[0], start), + max(prev[1], end), + prev[2] + size, + ) + else: + merged[rid] = (start, end, size) + + # Every rank updates its own local requests. + rid_to_req = {get_unique_rid(r): r for r in reqs} + for rid, (min_start, max_end, total_size) in merged.items(): + req = rid_to_req.get(rid) + if req is not None: + req.set_kv_cache_transfer_start(min_start) + req.set_kv_cache_transfer_end(max_end) + req.set_kv_cache_size(total_size) + def _collect_done(self, sessions: dict, reqs: dict): """Scan sessions and return (completed_rids, failed_rids).""" completed, failed = [], [] @@ -482,6 +536,7 @@ def _finalize_send(self, req: LlmRequest, session: TxSessionBase): @nvtx_range("KvCacheTransceiverV2.respond_and_send_async") def respond_and_send_async(self, req: LlmRequest): self._ever_had_send_session = True + req.set_kv_cache_transfer_start(tensorrt_llm.bindings.global_steady_clock_now()) session = self._get_or_create_send_session(req) req.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS session.send(self._create_kv_slice(req)) @@ -526,6 +581,7 @@ def request_and_receive_sync(self, req: LlmRequest): @nvtx_range("KvCacheTransceiverV2.request_and_receive_async") def request_and_receive_async(self, req: LlmRequest): self._ever_had_recv_session = True + req.set_kv_cache_transfer_start(tensorrt_llm.bindings.global_steady_clock_now()) rid = get_unique_rid(req) if rid in self._recv_sessions: logger.warning( @@ -633,6 +689,11 @@ def check_gen_transfer_status(self, at_least_request_num: Optional[int]): # distinguish the two cases and set the appropriate state. cancelled.append(rid) elif result == WaitResult.COMPLETED: + req = self._recv_reqs[rid] + if session.transfer_end_time is not None: + req.set_kv_cache_transfer_end(session.transfer_end_time) + if session.kv_cache_size_bytes > 0: + req.set_kv_cache_size(session.kv_cache_size_bytes) completed.append(rid) elif result == WaitResult.FAILED: failed.append(rid) @@ -650,6 +711,20 @@ def check_gen_transfer_status(self, at_least_request_num: Optional[int]): del self._recv_reqs[rid] del self._recv_sessions[rid] + # Log gen-side transfer summary after consensus. + if completed and os.getenv("TRTLLM_KVCACHE_TIME_OUTPUT_PATH"): + # Batch-sync timing for all completed requests in one allgather. + self._sync_transfer_timing([self._recv_reqs[rid] for rid in completed]) + for rid in completed: + req = self._recv_reqs[rid] + perf_log_manager.log_gen_transfer_summary( + unique_rid=rid, + instance_name=self._instance_name, + instance_rank=self._mapping.rank, + gen_side_transfer_time_ms=req.kv_cache_transfer_time_ms, + kv_cache_size=req.kv_cache_size, + ) + for rid in completed: session = self._recv_sessions[rid] req = self._recv_reqs[rid] diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index e4a2d1354498..feb02319a3e3 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1276,6 +1276,21 @@ def start_worker(self): def _set_global_steady_clock_offset(self): assert self.global_rank >= 0, "rank should be >= 0" + # First calibration wins (mirrors the C++ guard in CacheTransceiver). + # PyExecutor is constructed twice per process (memory-profiling dry run, + # then the real executor), so this method runs twice. Recalibration is + # idempotent as long as the measurement below reads the raw + # steady_clock, but skipping it avoids a redundant barrier+allgather + # and protects the offset if the measurement path ever becomes + # offset-aware (which would make a second pass observe ~zero skew and + # wipe the correct value). + if LlmRequest.global_steady_clock_offset is not None: + logger.info( + f"global_steady_clock_offset already set " + f"({LlmRequest.global_steady_clock_offset}); skipping recalibration " + f"for rank {self.global_rank}") + return + # Sync all ranks self.dist.barrier() # Immediately take the local steady clock timestamp diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 05c905db31bb..9347bc090c0b 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -1554,10 +1554,18 @@ def test_disaggregated_kv_cache_time_output(disaggregated_test_root, llm_venv, model_path=llama_model_root, cwd=llm_venv.get_working_directory()) assert os.path.isdir(output_path) - send_file = os.path.join(output_path, "rank_0_send.csv") - recv_file = os.path.join(output_path, "rank_0_recv.csv") - assert os.path.exists(send_file) - assert os.path.exists(recv_file) + # The C++ transceiver names timing files "__.csv" + # (instanceId is a runtime UUID that disambiguates instances sharing an + # output directory), so match by the "_.csv" suffix instead of a fixed + # "rank_0" prefix. + send_files = sorted(f for f in os.listdir(output_path) + if f.endswith("_send.csv")) + recv_files = sorted(f for f in os.listdir(output_path) + if f.endswith("_recv.csv")) + assert send_files, f"no *_send.csv in {output_path}: {os.listdir(output_path)}" + assert recv_files, f"no *_recv.csv in {output_path}: {os.listdir(output_path)}" + send_file = os.path.join(output_path, send_files[0]) + recv_file = os.path.join(output_path, recv_files[0]) with open(send_file, "r") as f: lines = f.readlines() assert len(lines) > 1 diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index f35cb99b80ff..133a34e15300 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -82,6 +82,7 @@ l0_h100: - unittest/disaggregated/region/test_region.py - unittest/disaggregated/test_disaggregated_params.py - unittest/disaggregated/test_perf_logger.py + - unittest/disaggregated/test_sync_transfer_timing.py - unittest/disaggregated/test_rank_info.py - unittest/disaggregated/test_request_id.py - unittest/disaggregated/test_kv_transfer.py diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index b58739d889a3..be94d873ed61 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -161,14 +161,16 @@ def test_encode_tail_handles_unset_base(self): def test_kv_result_prefix_roundtrip(): """The KV_AGENT_RESULT binary prefix (transfer.py) must round-trip exactly.""" tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") - for rank, rid, sl, last, status in [ - (7, 6925227277844486, 42, True, tfr.AgentResult.SUCCESS), - (0, 1, 0, False, tfr.AgentResult.FAILED), - (31, 2**62, 9999, True, tfr.AgentResult.SUCCESS), + for rank, rid, sl, last, status, size in [ + (7, 6925227277844486, 42, True, tfr.AgentResult.SUCCESS, 4096), + (0, 1, 0, False, tfr.AgentResult.FAILED, 0), + (31, 2**62, 9999, True, tfr.AgentResult.SUCCESS, 2**40), ]: - packed = tfr._KV_RESULT_PREFIX.pack(rank, rid, sl, last, tfr._AGENT_RESULT_CODE[status]) - r, i, s, last_out, c = tfr._KV_RESULT_PREFIX.unpack(packed) - assert (r, i, s, last_out) == (rank, rid, sl, last) + packed = tfr._KV_RESULT_PREFIX.pack( + rank, rid, sl, last, tfr._AGENT_RESULT_CODE[status], size + ) + r, i, s, last_out, c, sz = tfr._KV_RESULT_PREFIX.unpack(packed) + assert (r, i, s, last_out, sz) == (rank, rid, sl, last, size) assert tfr._AGENT_RESULT_BY_CODE[c] is status @@ -177,11 +179,11 @@ def test_make_kv_result_msg_uses_binary_frame(result_name): """Every KV result (success and failure) uses the binary frame so the receiver can decode it.""" tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") result = getattr(tfr.AgentResult, result_name) - msg = tfr._make_kv_result_msg(3, 12345, 7, True, result) + msg = tfr._make_kv_result_msg(3, 12345, 7, True, result, transfer_size=8192) assert msg[0] == tfr.MessageType.KV_AGENT_RESULT assert len(msg) == 2 # prefix only; no bounce tail when none is passed - r, rid, sl, last, code = tfr._KV_RESULT_PREFIX.unpack(msg[1]) - assert (r, rid, sl, last) == (3, 12345, 7, True) + r, rid, sl, last, code, size = tfr._KV_RESULT_PREFIX.unpack(msg[1]) + assert (r, rid, sl, last, size) == (3, 12345, 7, True, 8192) assert tfr._AGENT_RESULT_BY_CODE[code] is result diff --git a/tests/unittest/disaggregated/test_cache_transceiver_harness_report.py b/tests/unittest/disaggregated/test_cache_transceiver_harness_report.py index 24e3b2d72cfb..491eab1e0b67 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_harness_report.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_harness_report.py @@ -186,7 +186,8 @@ def test_emit_ucx_env_second_sweep(self, sample_cfg, capfd): # --------------------------------------------------------------------------- class TestParseCppRecvCsvs: def test_basic(self, tmp_path): - csv_path = tmp_path / "rank_0_recv.csv" + # C++ names files "__recv.csv" (instanceId is a UUID). + csv_path = tmp_path / "3c9f0e2a-1111-2222-3333-444455556666_0_recv.csv" with open(csv_path, "w", newline="") as f: w = csv.writer(f) w.writerow(["RequestID", "Bandwidth(Gbps)", "Bandwidth(Gbps)"]) @@ -199,6 +200,7 @@ def test_basic(self, tmp_path): assert abs(bws[0] - 15.0) < 0.01 def test_renamed_pattern(self, tmp_path): + # Legacy "rank_*" prefix must still parse (backward compatibility). csv_path = tmp_path / "rank_0_recv__c0.csv" with open(csv_path, "w", newline="") as f: w = csv.writer(f) @@ -466,7 +468,9 @@ def _setup_work_dir(self, tmp_path, sample_cfg): # Create CSV for gen side (C++ recv) gen_csv = work / "csv" / "0" / "gen" gen_csv.mkdir(parents=True) - with open(gen_csv / "rank_0_recv.csv", "w", newline="") as f: + with open( + gen_csv / "5d7b1f80-aaaa-bbbb-cccc-ddddeeeeffff_0_recv.csv", "w", newline="" + ) as f: w = csv.writer(f) w.writerow(["RequestID", "Bandwidth(Gbps)"]) # warmup rid (r=0) should be excluded diff --git a/tests/unittest/disaggregated/test_perf_logger.py b/tests/unittest/disaggregated/test_perf_logger.py index 682a69a13c15..73723a3e9b15 100644 --- a/tests/unittest/disaggregated/test_perf_logger.py +++ b/tests/unittest/disaggregated/test_perf_logger.py @@ -1,3 +1,5 @@ +import os +import tempfile import time from unittest.mock import patch @@ -45,9 +47,20 @@ def test_unrecorded_peer_returns_zero(self): # PerfLogManager tests # --------------------------------------------------------------------------- def _reset_singleton(): - """Reset PerfLogManager singleton so each test gets a fresh instance.""" + """Reset PerfLogManager singleton so each test gets a fresh instance. + + Closes any file handlers the previous instance attached to the global + logging registry so temporary directories can be cleaned up and no file + descriptors leak across tests. + """ from tensorrt_llm._torch.disaggregation.native import perf_logger + instance = perf_logger.PerfLogManager._instance + if instance is not None: + for file_logger in getattr(instance, "_file_loggers", {}).values(): + for handler in list(file_logger.handlers): + file_logger.removeHandler(handler) + handler.close() perf_logger.PerfLogManager._instance = None @@ -77,6 +90,78 @@ def test_disabled_when_env_unset(self): # log() should be a no-op, not crash mgr.log("inst", 0, "csv_line", "info_msg") + def test_log_gen_transfer_summary_writes_csv(self) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + try: + with patch.dict( + "os.environ", + {"TRTLLM_KVCACHE_TIME_OUTPUT_PATH": tmpdir}, + clear=False, + ): + from tensorrt_llm._torch.disaggregation.native.perf_logger import PerfLogManager + + mgr = PerfLogManager() + mgr.log_gen_transfer_summary( + unique_rid=42, + instance_name="test_inst", + instance_rank=0, + gen_side_transfer_time_ms=12.345, + kv_cache_size=1024, + ) + csv_path = os.path.join(tmpdir, "test_inst_0_gen_transfer_summary.csv") + assert os.path.exists(csv_path), f"CSV file not created at {csv_path}" + with open(csv_path) as f: + lines = f.readlines() + assert len(lines) >= 2 # header + at least 1 data row + assert "gen_side_transfer_time(ms)" in lines[0] + assert "42" in lines[1] + assert "12.345" in lines[1] + assert "1024" in lines[1] + finally: + # Close file handlers before the temporary directory exits. + _reset_singleton() + + def test_cpp_env_takes_priority_over_legacy_for_log(self) -> None: + """With both C++ and legacy env vars set, log() uses the C++ naming.""" + with tempfile.TemporaryDirectory() as tmpdir: + try: + legacy_base = os.path.join(tmpdir, "legacy") + with patch.dict( + "os.environ", + { + "TRTLLM_KVCACHE_TIME_OUTPUT_PATH": tmpdir, + "TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO": "1", + "TLLM_KV_TRANSFER_PERF_LOG_FILE": legacy_base, + }, + clear=False, + ): + from tensorrt_llm._torch.disaggregation.native.perf_logger import PerfLogManager + + mgr = PerfLogManager() + mgr.log("test_inst", 0, "csv_line", "info_msg") + cpp_path = os.path.join(tmpdir, "test_inst_0.csv") + legacy_path = f"{legacy_base}_test_inst_0.csv" + assert os.path.exists(cpp_path), f"C++-style CSV not created at {cpp_path}" + assert not os.path.exists(legacy_path), "legacy-style CSV should not be used" + finally: + # Close file handlers before the temporary directory exits. + _reset_singleton() + + @patch.dict("os.environ", {}, clear=False) + def test_log_gen_transfer_summary_disabled_without_env(self) -> None: + os.environ.pop("TRTLLM_KVCACHE_TIME_OUTPUT_PATH", None) + from tensorrt_llm._torch.disaggregation.native.perf_logger import PerfLogManager + + mgr = PerfLogManager() + # Should be a no-op, not crash + mgr.log_gen_transfer_summary( + unique_rid=1, + instance_name="test", + instance_rank=0, + gen_side_transfer_time_ms=0.0, + kv_cache_size=0, + ) + @patch.dict("os.environ", {"TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO": "1"}, clear=False) def test_log_task_perf_does_not_crash(self): from tensorrt_llm._torch.disaggregation.native.perf_logger import PerfLogManager, PerfTimer diff --git a/tests/unittest/disaggregated/test_sync_transfer_timing.py b/tests/unittest/disaggregated/test_sync_transfer_timing.py new file mode 100644 index 000000000000..51834429dd1d --- /dev/null +++ b/tests/unittest/disaggregated/test_sync_transfer_timing.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for KvCacheTransceiverV2._sync_transfer_timing.""" + +from typing import TYPE_CHECKING, Any +from unittest.mock import MagicMock, patch + +if TYPE_CHECKING: + from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 + + +class _FakeRequest: + """Minimal stand-in for LlmRequest with timing getters/setters.""" + + def __init__(self, rid: int, start: int, end: int, kv_size: int) -> None: + self.request_id = rid + self.py_disaggregated_params = None # forces get_unique_rid -> request_id + self._start = start + self._end = end + self._kv_cache_size = kv_size + + def get_kv_cache_transfer_start(self) -> int: + return self._start + + def get_kv_cache_transfer_end(self) -> int: + return self._end + + @property + def kv_cache_size(self) -> int: + return self._kv_cache_size + + def set_kv_cache_transfer_start(self, v: int) -> None: + self._start = v + + def set_kv_cache_transfer_end(self, v: int) -> None: + self._end = v + + def set_kv_cache_size(self, v: int) -> None: + self._kv_cache_size = v + + +def _make_transceiver(**overrides: Any) -> "KvCacheTransceiverV2": + """Create a minimal mock of KvCacheTransceiverV2 with _sync_transfer_timing.""" + from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 + + tc = object.__new__(KvCacheTransceiverV2) + tc._gen_need_sync = overrides.get("gen_need_sync", True) + tc._gen_allgather = overrides.get("gen_allgather", lambda x: [x]) + return tc + + +class TestSyncTransferTiming: + @patch.dict("os.environ", {"TRTLLM_KVCACHE_TIME_OUTPUT_PATH": "/tmp/test"}, clear=False) + def test_merges_correctly_two_ranks(self) -> None: + """Two simulated ranks — verify min(start), max(end), sum(size).""" + req = _FakeRequest(rid=1, start=10, end=20, kv_size=100) + + def fake_allgather(local_data): + # Simulate rank 0 = local_data, rank 1 = different timing + rank1_data = {1: (5, 25, 200)} + return [local_data, rank1_data] + + tc = _make_transceiver(gen_need_sync=True, gen_allgather=fake_allgather) + tc._sync_transfer_timing([req]) + + assert req._start == 5 # min(10, 5) + assert req._end == 25 # max(20, 25) + assert req._kv_cache_size == 300 # 100 + 200 + + @patch.dict("os.environ", {"TRTLLM_KVCACHE_TIME_OUTPUT_PATH": "/tmp/test"}, clear=False) + def test_multiple_requests_batched(self) -> None: + """Multiple requests in one allgather call.""" + req_a = _FakeRequest(rid=1, start=10, end=20, kv_size=100) + req_b = _FakeRequest(rid=2, start=30, end=40, kv_size=200) + + def fake_allgather(local_data): + rank1_data = { + 1: (8, 22, 150), + 2: (28, 45, 250), + } + return [local_data, rank1_data] + + tc = _make_transceiver(gen_need_sync=True, gen_allgather=fake_allgather) + tc._sync_transfer_timing([req_a, req_b]) + + assert req_a._start == 8 + assert req_a._end == 22 + assert req_a._kv_cache_size == 250 + + assert req_b._start == 28 + assert req_b._end == 45 + assert req_b._kv_cache_size == 450 + + @patch.dict("os.environ", {"TRTLLM_KVCACHE_TIME_OUTPUT_PATH": "/tmp/test"}, clear=False) + def test_all_ranks_updated(self) -> None: + """Every request object should be updated, not just rank-0.""" + req = _FakeRequest(rid=1, start=10, end=20, kv_size=100) + + def fake_allgather(local_data): + return [local_data, {1: (5, 25, 200)}] + + tc = _make_transceiver(gen_need_sync=True, gen_allgather=fake_allgather) + tc._sync_transfer_timing([req]) + + # All fields should reflect the merged values + assert req._start == 5 + assert req._end == 25 + assert req._kv_cache_size == 300 + + @patch.dict("os.environ", {}, clear=False) + def test_skips_when_no_env(self) -> None: + """Without TRTLLM_KVCACHE_TIME_OUTPUT_PATH, allgather is not called.""" + import os + + os.environ.pop("TRTLLM_KVCACHE_TIME_OUTPUT_PATH", None) + + allgather_mock = MagicMock() + tc = _make_transceiver(gen_need_sync=True, gen_allgather=allgather_mock) + + req = _FakeRequest(rid=1, start=10, end=20, kv_size=100) + tc._sync_transfer_timing([req]) + + allgather_mock.assert_not_called() + assert req._start == 10 # unchanged + + @patch.dict("os.environ", {"TRTLLM_KVCACHE_TIME_OUTPUT_PATH": "/tmp/test"}, clear=False) + def test_skips_when_single_rank(self) -> None: + """When _gen_need_sync is False, allgather is not called.""" + allgather_mock = MagicMock() + tc = _make_transceiver(gen_need_sync=False, gen_allgather=allgather_mock) + + req = _FakeRequest(rid=1, start=10, end=20, kv_size=100) + tc._sync_transfer_timing([req]) + + allgather_mock.assert_not_called() + assert req._start == 10 # unchanged + + @patch.dict("os.environ", {"TRTLLM_KVCACHE_TIME_OUTPUT_PATH": "/tmp/test"}, clear=False) + def test_empty_list(self) -> None: + """Empty request list should return immediately.""" + allgather_mock = MagicMock() + tc = _make_transceiver(gen_need_sync=True, gen_allgather=allgather_mock) + + tc._sync_transfer_timing([]) + + allgather_mock.assert_not_called()