Skip to content

[#12660][fix] Add error propagation for disagg KV cache transfer failures #12988

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
nv-yna wants to merge 1 commit into from
Closed

[#12660][fix] Add error propagation for disagg KV cache transfer failures #12988

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
98 changes: 95 additions & 3 deletions cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp
View file
Open in desktop
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 @@ -515,6 +515,95 @@ class CacheSender::Impl
}
}

/// @brief Notify the receiver that the sender encountered an error during KV cache transfer.
/// This unblocks the receiver so it can handle the failure instead of waiting indefinitely.
/// Must not throw -- called from noexcept context.
void sendErrorSignalToReceiver(RequestIdType id, std::string const& errorMessage) noexcept
{
try
{
auto* agentConnectionManager = dynamic_cast<executor::kv_cache::AgentConnectionManager*>(mManager);
if (!agentConnectionManager)
{
// Error signals are only supported for agent-based connections (NIXL/UCX).
return;
}
TransferSession* session = nullptr;
{
std::unique_lock<std::mutex> lock(mMtxForMap);
auto it = mRequestToSession.find(id);
if (it == mRequestToSession.end())
{
TLLM_LOG_WARNING(
"Cannot send error signal for request %ld: session not found", id);
return;
}
session = std::addressof(it->second);
}
auto const& connections = session->getConnections();
for (size_t i = 0; i < connections.size(); i++)
{
auto* agentConnection
= dynamic_cast<executor::kv_cache::AgentConnection const*>(connections.at(i));
if (agentConnection)
{
agentConnection->sendErrorSignal(
executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG}, id, errorMessage);
}
}
}
catch (std::exception const& signalErr)
{
TLLM_LOG_WARNING(
"Failed to send error signal to receiver for request %ld: %s", id, signalErr.what());
}
}

/// @brief Broadcast error signals to ALL in-flight requests on the receiver side.
/// When one transfer fails, the sender's executor may freeze, leaving other pending
/// transfers stuck. This ensures receivers for ALL pending requests are notified.
/// Must not throw -- called from noexcept context.
void broadcastErrorToAllPendingReceivers(std::string const& errorMessage) noexcept
{
try
{
std::unique_lock<std::mutex> lock(mMtxForMap);
for (auto& [reqId, session] : mRequestToSession)
{
sendErrorSignalToReceiverWithSession(reqId, session, errorMessage);
}
}
catch (std::exception const& e)
{
TLLM_LOG_WARNING("Failed to broadcast error signals: %s", e.what());
}
}

/// @brief Internal: send error signal using an already-locked session reference.
void sendErrorSignalToReceiverWithSession(
RequestIdType id, TransferSession& session, std::string const& errorMessage) noexcept
{
try
{
auto const& connections = session.getConnections();
for (size_t i = 0; i < connections.size(); i++)
{
auto* agentConnection
= dynamic_cast<executor::kv_cache::AgentConnection const*>(connections.at(i));
if (agentConnection)
{
agentConnection->sendErrorSignal(
executor::kv_cache::DataContext{TransceiverTag::kREADY_SIGNAL_TAG}, id, errorMessage);
}
}
}
catch (std::exception const& signalErr)
{
TLLM_LOG_WARNING(
"Failed to send error signal to receiver for request %ld: %s", id, signalErr.what());
}
}
Comment on lines +521 to +605

@coderabbitai coderabbitai Bot Apr 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

This helper block still needs clang-format.

Release Checks already flagged this section, so it needs a formatter pass before merge.

As per coding guidelines, "Use LLVM clang-format tool for formatting changes; maximum line length is 120 characters."

🧰 Tools
🪛 GitHub Actions: Release Checks

[error] 534-600: clang-format required formatting changes (TLLM_LOG_WARNING and dynamic_cast statements reflow).

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp` around lines 521 - 605,
Run LLVM clang-format on cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp and
reformat the helper block containing sendErrorSignalToReceiver,
broadcastErrorToAllPendingReceivers, and sendErrorSignalToReceiverWithSession so
it adheres to the project's style (LLVM profile) and maximum line length 120;
ensure consistent indentation, spacing around braces/parentheses, alignment of
long wrapped lines (e.g., TLLM_LOG_WARNING and agentConnection->sendErrorSignal
calls), and that noexcept and catch blocks are formatted per clang-format output
before committing.


void sendAndRemoveResponse(RequestIdType id, Response resp) noexcept
{
try
Expand All @@ -527,12 +616,16 @@ class CacheSender::Impl
catch (tensorrt_llm::common::RequestSpecificException const& e)
{
TLLM_LOG_ERROR("Exception in sendAndRemoveResponse: %s ", e.what());
sendErrorSignalToReceiver(id, e.what());
broadcastErrorToAllPendingReceivers(e.what());
Comment on lines 617 to +620

@coderabbitai coderabbitai Bot Apr 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Don't send the same request's error twice.

sendErrorSignalToReceiver(id, ...) already walks every connection in the failed session. Calling broadcastErrorToAllPendingReceivers(...) immediately afterward walks mRequestToSession again, so the same request gets a second ErrorSignalInfo. waitForNotification() consumes only one error before throwing, which means the duplicate can stay queued and fail a later unrelated wait on that agent.

Exclude id from the broadcast set, or remove that session from mRequestToSession before broadcasting.

Also applies to: 627-628

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp` around lines 617 - 620,
In sendAndRemoveResponse, avoid enqueueing the same ErrorSignalInfo twice by
excluding the failed session id from the subsequent broadcast: after calling
sendErrorSignalToReceiver(id, ...), call broadcastErrorToAllPendingReceivers
with logic that skips the provided id (or remove the session from
mRequestToSession prior to broadcasting) so the session that already received
the error does not get a duplicate ErrorSignalInfo; apply the same fix to the
other occurrence around lines 627-628 to prevent duplicated errors being left in
queues consumed by waitForNotification.

auto new_exception = TLLM_REQUEST_EXCEPTION(id, e.getErrorCode(), "%s", e.what());
resp.mPromise.set_exception(std::make_exception_ptr(new_exception));
}
catch (std::exception const& e)
{
TLLM_LOG_ERROR("Exception in sendAndRemoveResponse: %s request id: %ld", e.what(), id);
sendErrorSignalToReceiver(id, e.what());
broadcastErrorToAllPendingReceivers(e.what());
resp.mPromise.set_exception(std::current_exception());
}
}
Expand Down Expand Up @@ -1063,10 +1156,9 @@ class CacheReceiver::Impl
bool isReady = receiveReadySignal(session);
if (!isReady)
{
// Reuse the error state for the cancelled request.
llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR);
llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now());
return;
TLLM_THROW("Sender indicated transfer not ready for request %ld", llmRequest.mRequestId);
Comment on lines 1157 to +1161

@coderabbitai coderabbitai Bot Apr 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Keep this failure request-scoped.

Throwing TLLM_THROW here downgrades a sender-declared transfer abort into a generic exception, so the async path loses the request id and kNETWORK_ERROR classification that the rest of this PR is trying to preserve.

🛠️ Proposed fix 
         if (!isReady)
         {
             llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR);
             llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now());
-            TLLM_THROW("Sender indicated transfer not ready for request %ld", llmRequest.mRequestId);
+            throw TLLM_REQUEST_EXCEPTION(
+                llmRequest.mRequestId,
+                common::RequestErrorCode::kNETWORK_ERROR,
+                "Sender indicated transfer not ready for request %ld",
+                llmRequest.mRequestId);
         }
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp` around lines 1157 - 1161,
The code currently throws a global TLLM_THROW when isReady is false, which
converts a sender-declared transfer abort into a generic exception and loses the
request-scoped failure info; instead, remove the TLLM_THROW call inside the
isReady check, keep setting
llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR) and
llmRequest.setKvCacheTransferEnd(...), emit a request-scoped log that includes
llmRequest.mRequestId, and then return (or otherwise exit the current
request-processing path) so the failure remains tied to this llmRequest rather
than escalating as a global exception.

}
receiveSync(session);
llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now());
Expand Down
35 changes: 34 additions & 1 deletion cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp
View file
Open in desktop
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 All @@ -17,6 +17,8 @@

#include "connection.h"
#include "tensorrt_llm/common/envUtils.h"
#include "tensorrt_llm/common/logger.h"
#include "tensorrt_llm/common/tllmException.h"
#include "tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h"
#include <random>
#include <string>
Expand Down Expand Up @@ -275,6 +277,15 @@ std::optional<size_t> AgentConnection::getPreAssignedBufferId(uint8_t kind) cons
return std::nullopt;
}

void AgentConnection::sendErrorSignal(DataContext const& ctx, uint64_t requestId, std::string const& errorMessage) const
{
ErrorSignalInfo errorSignalInfo{mRemoteAgentName, ctx, requestId, errorMessage};
NotificationInfo notificationInfo{errorSignalInfo};
std::stringstream ss;
NotificationInfo::serialize(notificationInfo, ss);
mAgentConnectionManager->getAgent()->notifySyncMessage(mRemoteAgentName, ss.str());
}

AgentConnectionManager::AgentConnectionManager(
std::vector<batch_manager::BaseTransBufferManager*> cacheTransBufferManagers, CacheState cacheState,
std::string const& backendType, std::optional<CacheState::RnnCacheState> rnnCacheState)
Expand Down Expand Up @@ -648,6 +659,28 @@ void AgentConnectionManager::waitForNotification(
}
}

// Check for error signals from the remote agent regardless of what
// notification type we are waiting for. This unblocks the receiver
// when the sender encounters an error during KV cache transfer.
if (std::holds_alternative<ErrorSignalInfo>(notificationInfo.mInfo))
{
auto errorSignalData = std::get<ErrorSignalInfo>(notificationInfo.mInfo);
TLLM_LOG_ERROR(
"Received error signal from sender for request %ld: %s",
errorSignalData.mRequestId, errorSignalData.mErrorMessage.c_str());
erase = true;
notifIt = notifs.erase(notifIt);
if (notifs.empty())
{
it = mUnhandledNotifications.erase(it);
}
throw TLLM_REQUEST_EXCEPTION(
errorSignalData.mRequestId,
common::RequestErrorCode::kNETWORK_ERROR,
"Sender error for request %ld: %s",
errorSignalData.mRequestId, errorSignalData.mErrorMessage.c_str());
Comment on lines +665 to +681

@coderabbitai coderabbitai Bot Apr 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

This new error branch still needs clang-format.

CI is already flagging this section, so please reflow it with the repo formatter before merging.

As per coding guidelines, "Use LLVM clang-format tool for formatting changes; maximum line length is 120 characters."

🧰 Tools
🪛 GitHub Actions: Release Checks

[error] 665-665: clang-format required formatting changes (TLLM_LOG_ERROR and TLLM_REQUEST_EXCEPTION call wrapping/reflow).

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp`
around lines 665 - 681, The added ErrorSignalInfo branch is not formatted to
repo style; run the LLVM clang-format tool (max line length 120) on the block
around the ErrorSignalInfo handling (symbols: notificationInfo.mInfo,
ErrorSignalInfo, TLLM_LOG_ERROR, TLLM_REQUEST_EXCEPTION,
mUnhandledNotifications, notifs, notifIt, it) and reflow the long log/exception
lines so they wrap within 120 chars and adhere to the project's clang-format
rules before committing.

}

if (!erase)
{
notifIt++;
Expand Down
53 changes: 51 additions & 2 deletions cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h
View file
Open in desktop
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 @@ -174,10 +174,45 @@ struct NotificationSyncInfo
}
};

struct ErrorSignalInfo
{
std::string mAgentName;
DataContext mContext;
uint64_t mRequestId;
std::string mErrorMessage;

static void serialize(ErrorSignalInfo const& errorSignalInfo, std::ostream& os)
{
namespace su = executor::serialize_utils;
su::serialize(errorSignalInfo.mAgentName, os);
su::serialize(errorSignalInfo.mContext.getTag(), os);
su::serialize(errorSignalInfo.mRequestId, os);
su::serialize(errorSignalInfo.mErrorMessage, os);
}

static ErrorSignalInfo deserialize(std::istream& is)
{
namespace su = executor::serialize_utils;
auto agentName = su::deserialize<decltype(mAgentName)>(is);
auto contextTag = su::deserialize<decltype(mContext.getTag())>(is);
DataContext context{contextTag};
auto requestId = su::deserialize<decltype(mRequestId)>(is);
auto errorMessage = su::deserialize<decltype(mErrorMessage)>(is);
return ErrorSignalInfo{agentName, context, requestId, errorMessage};
}

static size_t serializedSize(ErrorSignalInfo const& errorSignalInfo)
{
namespace su = executor::serialize_utils;
return su::serializedSize(errorSignalInfo.mAgentName) + su::serializedSize(errorSignalInfo.mContext.getTag())
+ su::serializedSize(errorSignalInfo.mRequestId) + su::serializedSize(errorSignalInfo.mErrorMessage);
}
};

struct NotificationInfo
{

std::variant<RequestAndBufferInfo, NotificationSyncInfo, ReadySignalInfo> mInfo;
std::variant<RequestAndBufferInfo, NotificationSyncInfo, ReadySignalInfo, ErrorSignalInfo> mInfo;

static void serialize(NotificationInfo const& notificationInfo, std::ostream& os)
{
Expand All @@ -195,6 +230,10 @@ struct NotificationInfo
{
ReadySignalInfo::serialize(std::get<ReadySignalInfo>(notificationInfo.mInfo), os);
}
else if (std::holds_alternative<ErrorSignalInfo>(notificationInfo.mInfo))
{
ErrorSignalInfo::serialize(std::get<ErrorSignalInfo>(notificationInfo.mInfo), os);
}
else
{
TLLM_THROW("Unknown variant type");
Expand All @@ -208,6 +247,7 @@ struct NotificationInfo
constexpr std::size_t requestAndBufferInfoIdx{0};
constexpr std::size_t notificationSyncInfoIdx{1};
constexpr std::size_t readySignalInfoIdx{2};
constexpr std::size_t errorSignalInfoIdx{3};
if (variantIdx == requestAndBufferInfoIdx)
{
return NotificationInfo{RequestAndBufferInfo::deserialize(is)};
Expand All @@ -220,6 +260,10 @@ struct NotificationInfo
{
return NotificationInfo{ReadySignalInfo::deserialize(is)};
}
else if (variantIdx == errorSignalInfoIdx)
{
return NotificationInfo{ErrorSignalInfo::deserialize(is)};
}
else
{
TLLM_THROW("Unknown variant type");
Expand All @@ -243,6 +287,10 @@ struct NotificationInfo
{
totalSize += ReadySignalInfo::serializedSize(std::get<ReadySignalInfo>(notificationInfo.mInfo));
}
else if (std::holds_alternative<ErrorSignalInfo>(notificationInfo.mInfo))
{
totalSize += ErrorSignalInfo::serializedSize(std::get<ErrorSignalInfo>(notificationInfo.mInfo));
}
else
{
TLLM_THROW("Unknown variant type");
Expand All @@ -267,6 +315,7 @@ class AgentConnection : public Connection
[[nodiscard]] bool hasLoadRemoteAgent() const;
void sendReadySignal(DataContext const& ctx, bool isReady) const;
bool recvReadySignal(DataContext const& ctx) const;
void sendErrorSignal(DataContext const& ctx, uint64_t requestId, std::string const& errorMessage) const;

void activateBuffer(uint8_t kind) const override;
[[nodiscard]] std::optional<size_t> getPreAssignedBufferId(uint8_t kind) const override;
Expand Down
1 change: 1 addition & 0 deletions cpp/tests/unit_tests/batch_manager/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,3 +27,4 @@ add_gtest(microBatchSchedulerTest microBatchSchedulerTest.cpp)
add_gtest(peftCacheManagerTest peftCacheManagerTest.cpp)
add_gtest(staticThreadPoolTest staticThreadPoolTest.cpp)
add_gtest(rnnCacheFormatterTest rnnCacheFormatterTest.cpp)
add_gtest(errorPropagationTest errorPropagationTest.cpp)
Loading
Loading