Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 39 additions & 1 deletion cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -600,10 +600,22 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus(
void CacheTransceiver::checkGenTransferStatus(std::optional<int> const& atLeastRequestNum)
{
bool blockAll = !atLeastRequestNum.has_value();
TLLM_LOG_INFO(
"[disagg-debug] C++ checkGenTransferStatus enter: atLeastRequestNum=%d blockAll=%d "
"requesterFutures=%zu",
atLeastRequestNum.value_or(-1), static_cast<int>(blockAll), mRequesterFutures.size());
std::vector<LlmRequest::RequestIdType> genTransferReadyRequestIds;
for (auto&& [request, future] : mRequesterFutures)
{
if (future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready)
auto const status = future.wait_for(std::chrono::milliseconds(0));
auto const ctxRequestId
= request->getContextPhaseParams().has_value() ? request->getContextPhaseParams().value().getReqId() : 0;
TLLM_LOG_INFO(
"[disagg-debug] C++ checkGenTransferStatus future probe: requestId=%zu ctxRequestId=%zu "
"ready=%d state=%d",
request->mRequestId, ctxRequestId, static_cast<int>(status == std::future_status::ready),
static_cast<int>(request->getState()));
if (status == std::future_status::ready)
{
genTransferReadyRequestIds.push_back(request->mRequestId);
}
Expand Down Expand Up @@ -710,13 +722,39 @@ void CacheTransceiver::checkGenTransferStatus(std::optional<int> const& atLeastR
" checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(),
atLeastRequestNum.value_or(0));
}
std::ostringstream selectedIds;
bool firstSelectedId = true;
for (auto const requestId : toCompleteIdSet)
{
if (!firstSelectedId)
{
selectedIds << ",";
}
selectedIds << requestId;
firstSelectedId = false;
}
TLLM_LOG_INFO(
"[disagg-debug] C++ checkGenTransferStatus selected requests: readyLocal=%zu freqVec=%zu "
"toComplete=%zu ids=[%s]",
genTransferReadyRequestIds.size(), freqVec.size(), toCompleteIdSet.size(), selectedIds.str().c_str());
for (auto it = mRequesterFutures.begin(); it != mRequesterFutures.end();)
{
if (blockAll || toCompleteIdSet.find(it->first->mRequestId) != toCompleteIdSet.end())
{
try
{
auto const ctxRequestId = it->first->getContextPhaseParams().has_value()
? it->first->getContextPhaseParams().value().getReqId()
: 0;
TLLM_LOG_INFO(
"[disagg-debug] C++ checkGenTransferStatus waiting on requester future: requestId=%zu "
"ctxRequestId=%zu blockAll=%d",
it->first->mRequestId, ctxRequestId, static_cast<int>(blockAll));
it->second.get();
TLLM_LOG_INFO(
"[disagg-debug] C++ checkGenTransferStatus requester future completed: requestId=%zu "
"ctxRequestId=%zu",
it->first->mRequestId, ctxRequestId);
it->first->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE);

// Gather the kv cache transfer time from all workers and update to leader rank
Expand Down
185 changes: 183 additions & 2 deletions cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp

Large diffs are not rendered by default.

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
@@ -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");
Expand Down Expand Up @@ -34,6 +34,7 @@
#include <nixl_types.h>
#include <numeric>
#include <set>
#include <sstream>
#include <sys/file.h>
#include <sys/stat.h>
#include <thread>
Expand Down Expand Up @@ -485,6 +486,7 @@ NixlTransferStatus::NixlTransferStatus(nixlAgent* agent, nixlXferReqH* handle)
TransferState NixlTransferStatus::wait(int64_t timeout_ms) const
{
auto startTime = std::chrono::steady_clock::now();
auto nextLogTime = startTime + std::chrono::seconds(30);

while (true)
{
Expand All @@ -498,6 +500,18 @@ TransferState NixlTransferStatus::wait(int64_t timeout_ms) const
return TransferState::kFAILURE;
}

auto const now = std::chrono::steady_clock::now();
auto const elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(now - startTime).count();
if (now >= nextLogTime)
{
TLLM_LOG_INFO(
"[disagg-debug] C++ NIXL transfer wait still in progress: handle=%p timeoutMs=%lld "
"elapsedMs=%lld status=%s",
static_cast<void*>(mHandle), static_cast<long long>(timeout_ms), static_cast<long long>(elapsed),
nixlEnumStrings::statusStr(status).c_str());
nextLogTime = now + std::chrono::seconds(30);
}

// If timeout_ms < 0, wait indefinitely until status is not NIXL_IN_PROG
if (timeout_ms < 0)
{
Expand All @@ -506,9 +520,6 @@ TransferState NixlTransferStatus::wait(int64_t timeout_ms) const
}

// Check if timeout has elapsed
auto elapsed
= std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now() - startTime)
.count();
if (elapsed >= timeout_ms)
{
return TransferState::kIN_PROGRESS;
Expand Down Expand Up @@ -818,17 +829,40 @@ void NixlTransferAgent::notifySyncMessage(std::string const& name, SyncMessage c
{

auto status = mRawAgent->genNotif(name, syncMessage);
TLLM_CHECK_WITH_INFO(
status == NIXL_SUCCESS, "genNotif failed with status: %s", nixlEnumStrings::statusStr(status).c_str());
auto const statusString = nixlEnumStrings::statusStr(status);
TLLM_LOG_INFO("[disagg-debug] C++ NIXL genNotif returned: selfAgent=%s remoteAgent=%s status=%s payloadBytes=%zu",
mName.c_str(), name.c_str(), statusString.c_str(), syncMessage.size());
TLLM_CHECK_WITH_INFO(status == NIXL_SUCCESS, "genNotif failed with status: %s", statusString.c_str());
}

[[nodiscard]] std::unordered_map<std::string, std::vector<SyncMessage>> NixlTransferAgent::getNotifiedSyncMessages()
{

nixl_notifs_t notifs;
auto status = mRawAgent->getNotifs(notifs);
TLLM_CHECK_WITH_INFO(
status == NIXL_SUCCESS, "getNotifs failed with status: %s", nixlEnumStrings::statusStr(status).c_str());
auto const statusString = nixlEnumStrings::statusStr(status);
TLLM_CHECK_WITH_INFO(status == NIXL_SUCCESS, "getNotifs failed with status: %s", statusString.c_str());
if (!notifs.empty())
{
size_t totalCount{0};
std::ostringstream sourceCounts;
bool firstSource{true};
for (auto const& [agent, messages] : notifs)
{
totalCount += messages.size();
if (!firstSource)
{
sourceCounts << ",";
}
firstSource = false;
sourceCounts << agent << ":" << messages.size();
}
auto const sourceCountsString = sourceCounts.str();
TLLM_LOG_INFO(
"[disagg-debug] C++ NIXL getNotifs returned: selfAgent=%s status=%s sources=%zu total=%zu "
"sourceCounts={%s}",
mName.c_str(), statusString.c_str(), notifs.size(), totalCount, sourceCountsString.c_str());
}

return notifs;
}
Expand Down
Loading