[TRTLLM-17105][fix] Report a bounce failure only after the receiver has answered every sent chunk

cpp/include/tensorrt_llm/executor/transferAgent.h

@@ -374,9 +374,11 @@ class TransferStatus
 {
 public:
     virtual ~TransferStatus() = default;
-    /// NIXL and bounce statuses return true only on success; backends that cannot distinguish
-    /// failure from completion (mooncake) return true on any terminal state — use
-    /// wait()/getLastStatus for the outcome.
+    /// True once the transfer is over and will not read its source or write its destination memory
+    /// again. After a failure it may turn true when the result is set or later, or never if the backend
+    /// cannot prove it. It does not turn back to false, except that a NIXL status reads false after release() or
+    /// once its agent is gone. Mooncake cannot tell failure from completion and returns true on any
+    /// terminal state. The outcome comes from wait().
     [[nodiscard]] virtual bool isCompleted() const = 0;
     virtual TransferState wait(int64_t timeout_ms = -1) const = 0;
 

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h

@@ -166,7 +166,9 @@ struct BounceConfig
     int requestTimeoutMs{30000}; // request_timeout_ms; must be > 0 — the whole failure model
                                  // (abandoned-flow resolution, receiver lease, quarantine) hangs off
                                  // this timer, so applyParam rejects 0 (negatives already fail the
-                                 // strict unsigned parse) and keeps the current value.
+                                 // strict unsigned parse) and keeps the current value, and BounceSender
+                                 // rejects a non-positive one. It also bounds a failed request's answer
+                                 // wait (see FailedRequest).
     // Receiver-side lease on granted regions — DERIVED in deriveDependentTimeouts() as
     // 2 x requestTimeoutMs, not an independent knob (tests may still set the field directly). A
     // dead sender emits neither DATA nor a cancel, which is unobservable through the protocol

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.h

@@ -97,10 +97,11 @@ static_assert(sizeof(BounceCreditEntry) == 24, "BounceCreditEntry must be 24 byt
 static_assert(sizeof(BounceScatterRun) == 36, "BounceScatterRun must be 36 bytes");
 
 inline constexpr std::uint32_t kBounceMagic = 0x424E4332U; // 'B''N''C''2'
-// v2: DATA scatter entries became strided RUNS (BounceScatterRun). The out-of-band capability
-// handshake checks this version before bounce is selected; decodeHeader also rejects a mismatched
-// control message defensively.
-inline constexpr std::uint16_t kBounceVersion = 2U;
+// v2: DATA scatter entries became strided RUNS (BounceScatterRun). v3: a chunk NACK, like an ACK,
+// follows the chunk's last write into the destination, and a refused DATA is NACKed; the sender relies
+// on both. The out-of-band capability handshake checks this version before bounce is selected;
+// decodeHeader also rejects a mismatched control message defensively.
+inline constexpr std::uint16_t kBounceVersion = 3U;
 
 // ---- encode (each returns a self-contained blob: header + payload) ----
 /// WANT carries the per-chunk byte sizes the sender will write (the receiver allocates a region of
@@ -122,10 +123,11 @@ inline constexpr std::uint16_t kBounceVersion = 2U;
 [[nodiscard]] std::string encodeData(std::uint64_t requestId, std::uint32_t chunkIdx, std::uint32_t numChunks,
     std::uint64_t regionHandle, std::vector<BounceScatterRun> const& entries);
 [[nodiscard]] std::string encodeAck(std::uint64_t requestId, std::uint32_t chunkIdx, std::uint64_t regionHandle);
-/// NACK: the receiver KNOWS this chunk/request cannot complete (scatter failure, ungrantable chunk
-/// size, plan overflow), so the sender fails it now instead of waiting out requestTimeoutMs.
-/// Header-only like ACK. The sender fails the WHOLE request on any NACK, so chunkIdx/regionHandle are
-/// informational only (0 when the receiver has no chunk context, e.g. a rejected WANT).
+/// NACK: the receiver KNOWS this chunk/request cannot complete (scatter failure, a DATA it refused,
+/// ungrantable chunk size, plan overflow), so the sender fails it now instead of waiting out
+/// requestTimeoutMs. Header-only like ACK. A chunk NACK names (chunkIdx, regionHandle) like an ACK and
+/// promises, like an ACK, that the chunk will not write the destination again. A request-level NACK (a
+/// rejected WANT) carries 0/0; it precedes every grant, so it matches no sent chunk.
 [[nodiscard]] std::string encodeNack(std::uint64_t requestId, std::uint32_t chunkIdx, std::uint64_t regionHandle);
 
 // ---- decode ----

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceResult.cpp

@@ -0,0 +1,171 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 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.
+ */
+
+#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceResult.h"
+
+#include <algorithm>
+#include <utility>
+
+namespace tensorrt_llm::executor::kv_cache::bounce
+{
+
+char const* toString(BounceFailReason reason)
+{
+    // Prefixed with the module name: these strings travel up through TransferStatus::
+    // getLastStatusStr() into upper-layer error logs, where "RDMA write failed" alone would not
+    // say WHICH transport failed.
+    switch (reason)
+    {
+    case BounceFailReason::kNone: return "bounce: none";
+    case BounceFailReason::kPlanRejected: return "bounce: plan rejected (request did not fit a transfer plan)";
+    case BounceFailReason::kNoProgressTimeout: return "bounce: no GRANT/ACK progress within requestTimeoutMs";
+    case BounceFailReason::kPeerDropped: return "bounce: peer dropped (forgetPeer/invalidateRemoteAgent)";
+    case BounceFailReason::kGatherFailed: return "bounce: gather kernel failed (CUDA error)";
+    case BounceFailReason::kWriteFailed: return "bounce: RDMA write failed";
+    case BounceFailReason::kProtocolError: return "bounce: protocol error (GRANT mispair/plan overflow)";
+    case BounceFailReason::kShutdown: return "bounce: transport shut down while pending";
+    case BounceFailReason::kPeerNack: return "bounce: receiver reported it cannot complete the chunk/request (NACK)";
+    }
+    return "bounce: unknown";
+}
+
+std::chrono::milliseconds answerWaitFor(BounceFailReason trigger, std::chrono::milliseconds requestTimeout)
+{
+    return trigger == BounceFailReason::kNoProgressTimeout ? std::chrono::milliseconds{0} : requestTimeout;
+}
+
+OutstandingAccess::OutstandingAccess(std::vector<SentChunk> unansweredChunks)
+    : mUnansweredChunks(std::move(unansweredChunks))
+{
+}
+
+OutstandingAccess::SourceRead::SourceRead(std::shared_ptr<OutstandingAccess> const& access)
+    : mAccess(access)
+{
+    std::lock_guard<std::mutex> lk(access->mMu);
+    ++access->mSourceReads;
+}
+
+OutstandingAccess::SourceRead::~SourceRead()
+{
+    if (mAccess != nullptr) // null once moved from
+    {
+        std::lock_guard<std::mutex> lk(mAccess->mMu);
+        --mAccess->mSourceReads;
+    }
+}
+
+bool OutstandingAccess::answer(std::uint32_t chunkIdx, std::uint64_t regionHandle)
+{
+    std::lock_guard<std::mutex> lk(mMu);
+    auto const chunk = std::find_if(mUnansweredChunks.begin(), mUnansweredChunks.end(),
+        [&](SentChunk const& c) { return c.chunkIdx == chunkIdx && c.regionHandle == regionHandle; });
+    if (chunk == mUnansweredChunks.end())
+    {
+        return false;
+    }
+    mUnansweredChunks.erase(chunk);
+    return true;
+}
+
+std::size_t OutstandingAccess::unansweredChunks() const
+{
+    std::lock_guard<std::mutex> lk(mMu);
+    return mUnansweredChunks.size();
+}
+
+bool OutstandingAccess::quiesced() const
+{
+    std::lock_guard<std::mutex> lk(mMu);
+    return mSourceReads == 0 && mUnansweredChunks.empty();
+}
+
+bool BounceResult::destinationQuiesced() const
+{
+    return outstanding == nullptr || outstanding->unansweredChunks() == 0;
+}
+
+bool BounceResult::quiesced() const
+{
+    return outstanding == nullptr || outstanding->quiesced();
+}
+
+FailedRequest::FailedRequest(BounceFailReason reason, std::shared_ptr<std::promise<BounceResult>> promise,
+    std::shared_ptr<OutstandingAccess> outstanding, Clock::duration answerWait, Clock::time_point now)
+    : mReason(reason)
+    , mAnswerWait(answerWait)
+    , mRecord(outstanding)
+{
+    mAwaiting.emplace(AwaitingAnswers{std::move(promise), std::move(outstanding), now + answerWait});
+}
+
+bool FailedRequest::answer(std::uint32_t chunkIdx, std::uint64_t regionHandle, Clock::time_point now)
+{
+    auto const record = mRecord.lock();
+    if (record == nullptr || !record->answer(chunkIdx, regionHandle))
+    {
+        return false;
+    }
+    if (mAwaiting)
+    {
+        mAwaiting->deadline = now + mAnswerWait; // a live receiver gets a fresh wait for each remaining answer
+    }
+    return true;
+}
+
+std::optional<std::size_t> FailedRequest::resolveIfDue(Clock::time_point now)
+{
+    bool const due = mAwaiting && (mAwaiting->outstanding->unansweredChunks() == 0 || now >= mAwaiting->deadline);
+    return due ? resolve() : std::nullopt;
+}
+
+std::optional<std::size_t> FailedRequest::resolve()
+{
+    if (!mAwaiting)
+    {
+        return std::nullopt;
+    }
+    AwaitingAnswers awaiting = std::move(*mAwaiting);
+    mAwaiting.reset();
+    // Count before the record moves into the result.
+    auto const unanswered = awaiting.outstanding->unansweredChunks();
+    awaiting.promise->set_value({TransferState::kFAILURE, mReason, std::move(awaiting.outstanding)});
+    return unanswered;
+}
+
+BounceFailReason FailedRequest::reason() const
+{
+    return mReason;
+}
+
+bool FailedRequest::resolved() const
+{
+    return !mAwaiting;
+}
+
+bool FailedRequest::worthKeeping() const
+{
+    return !resolved() || unansweredChunks() > 0;
+}
+
+std::size_t FailedRequest::unansweredChunks() const
+{
+    auto const record = mRecord.lock();
+    return record == nullptr ? 0 : record->unansweredChunks();
+}
+
+} // namespace tensorrt_llm::executor::kv_cache::bounce

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceResult.h

@@ -0,0 +1,159 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 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.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/executor/transferAgent.h"
+
+#include <chrono>
+#include <cstddef>
+#include <cstdint>
+#include <future>
+#include <memory>
+#include <mutex>
+#include <optional>
+#include <vector>
+
+namespace tensorrt_llm::executor::kv_cache::bounce
+{
+
+/// Why a bounce request resolved kFAILURE — carried inside the request's result future and
+/// surfaced to the upper layer through TransferStatus::getLastStatusStr(); the bare kFAILURE
+/// enum alone tells an operator nothing.
+enum class BounceFailReason : std::uint8_t
+{
+    kNone = 0,
+    kPlanRejected,      // request did not fit a transfer plan (submit-time TLLM_CHECK)
+    kNoProgressTimeout, // no GRANT/ACK progress within requestTimeoutMs (peer unreachable/stalled)
+    kPeerDropped,       // forgetPeer()/invalidateRemoteAgent while the request was in flight
+    kGatherFailed,      // gather kernel launch / event-record / async poll failure (CUDA error)
+    kWriteFailed,       // the RDMA write reported failure (getXferStatus == failed)
+    kProtocolError,     // GRANT mispair or plan overflow — the flow was abandoned
+    kShutdown,          // transport shut down while the request was still pending
+    kPeerNack,          // receiver reported it cannot complete the chunk/request (NACK)
+};
+
+[[nodiscard]] char const* toString(BounceFailReason reason);
+
+/// How long a request failed by `trigger` holds its result back for the receiver's answers, restarted by
+/// each answer. None after kNoProgressTimeout, which already saw no progress for requestTimeoutMs: the
+/// result resolves at once, in-doubt if a sent chunk is unanswered (its DATA may be recent; a late answer
+/// still ends it).
+[[nodiscard]] std::chrono::milliseconds answerWaitFor(
+    BounceFailReason trigger, std::chrono::milliseconds requestTimeout);
+
+/// What a failed request may still do to caller memory after its result is set: gathers that may still
+/// read its source, and sent chunks the receiver has not answered — it answers a chunk only after the
+/// chunk's scatter has last written the destination. Once the result is set this only shrinks. Updated by
+/// the sender's IO thread, read from any thread.
+class OutstandingAccess
+{
+public:
+    /// A chunk whose DATA is out, named the way the receiver's ACK/NACK names it.
+    struct SentChunk
+    {
+        std::uint32_t chunkIdx{};
+        std::uint64_t regionHandle{};
+    };
+
+    /// Recorded once, before the result is set: afterwards answers and SourceReads can only shrink it.
+    explicit OutstandingAccess(std::vector<SentChunk> unansweredChunks = {});
+
+    /// One gather that may still read the source, until this is destroyed. It keeps the record alive.
+    class SourceRead
+    {
+    public:
+        explicit SourceRead(std::shared_ptr<OutstandingAccess> const& access);
+        SourceRead(SourceRead&& other) noexcept = default;
+        SourceRead(SourceRead const&) = delete;
+        SourceRead& operator=(SourceRead const&) = delete;
+        SourceRead& operator=(SourceRead&&) = delete;
+        ~SourceRead();
+
+    private:
+        std::shared_ptr<OutstandingAccess> mAccess;
+    };
+
+    /// The receiver answered (ACK or NACK) the chunk. False if it is not an unanswered chunk.
+    bool answer(std::uint32_t chunkIdx, std::uint64_t regionHandle);
+
+    [[nodiscard]] std::size_t unansweredChunks() const;
+    [[nodiscard]] bool quiesced() const;
+
+private:
+    mutable std::mutex mMu;
+    std::uint32_t mSourceReads{0};
+    std::vector<SentChunk> mUnansweredChunks;
+};
+
+/// A bounce request's single result: the terminal state, on kFAILURE its cause, and what the request may
+/// still access. The state and cause never change; the outstanding accesses only end.
+struct BounceResult
+{
+    TransferState state{TransferState::kFAILURE};
+    BounceFailReason reason{BounceFailReason::kNone};
+    std::shared_ptr<OutstandingAccess const> outstanding; // null: nothing outstanding
+
+    [[nodiscard]] bool destinationQuiesced() const;
+    [[nodiscard]] bool quiesced() const;
+};
+
+/// A failed request's result, held back until every sent chunk is answered or the answer wait is over (a
+/// caller that trusts wait() would otherwise reuse the destination while the receiver may still scatter into
+/// it). Answers keep shrinking the record after that, so isCompleted() can turn true after wait() returns.
+class FailedRequest
+{
+public:
+    using Clock = std::chrono::steady_clock;
+
+    /// `answerWait` is how long the receiver gets to answer, from now and again after each answer.
+    FailedRequest(BounceFailReason reason, std::shared_ptr<std::promise<BounceResult>> promise,
+        std::shared_ptr<OutstandingAccess> outstanding, Clock::duration answerWait, Clock::time_point now);
+
+    /// The receiver answered (ACK or NACK) a sent chunk. False if it is not an unanswered chunk.
+    bool answer(std::uint32_t chunkIdx, std::uint64_t regionHandle, Clock::time_point now);
+    /// Resolve the result once every sent chunk is answered or the answer wait is over. If this call resolved
+    /// it: the number of sent chunks still unanswered then.
+    std::optional<std::size_t> resolveIfDue(Clock::time_point now);
+    /// Resolve the result whatever is still unanswered (shutdown); returns like resolveIfDue.
+    std::optional<std::size_t> resolve();
+
+    [[nodiscard]] BounceFailReason reason() const;
+    [[nodiscard]] bool resolved() const;
+    /// The result is held back, or an answer can still change it (a sent chunk is unanswered and the record
+    /// is alive).
+    [[nodiscard]] bool worthKeeping() const;
+
+private:
+    struct AwaitingAnswers
+    {
+        std::shared_ptr<std::promise<BounceResult>> promise;
+        std::shared_ptr<OutstandingAccess> outstanding; // keeps mRecord alive until the result is set
+        Clock::time_point deadline;
+    };
+
+    /// 0 once the record has expired.
+    [[nodiscard]] std::size_t unansweredChunks() const;
+
+    BounceFailReason mReason;
+    Clock::duration mAnswerWait;
+    std::optional<AwaitingAnswers> mAwaiting; // empty once resolved
+    // Once resolved, the result's holders and running orphaned gathers keep the record alive.
+    std::weak_ptr<OutstandingAccess> mRecord;
+};
+
+} // namespace tensorrt_llm::executor::kv_cache::bounce