diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index b1a55cd3b590..ad7e17255e87 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & +# SPDX-FileCopyrightText: Copyright (c) 1993-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 @@ -280,8 +280,15 @@ include_directories(${CMAKE_BINARY_DIR}/_deps/nanobind-src/include) FetchContent_MakeAvailable(cutlass cxxopts flashmla json msa xgrammar) +# cppzmq (header-only) is the bounce v2 control-plane dependency AND the UCX +# bootstrap dependency. Make it available whenever NIXL (-> bounce) or UCX is +# enabled, so bounce does NOT require UCX (its real deps are NIXL + zmq, +# orthogonal to UCX). ucxx stays UCX-only. +if(NIXL_ROOT OR ENABLE_UCX) + FetchContent_MakeAvailable(cppzmq) +endif() if(ENABLE_UCX) - FetchContent_MakeAvailable(cppzmq ucxx) + FetchContent_MakeAvailable(ucxx) endif() if(NOT NVTX_DISABLE) diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index acc0efe18966..9fd935d41864 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/cpp/include/tensorrt_llm/executor/executor.h @@ -1522,7 +1522,8 @@ class CacheTransceiverConfig explicit CacheTransceiverConfig(std::optional backendType = std::nullopt, std::optional maxNumTokens = std::nullopt, std::optional kvTransferTimeoutMs = std::nullopt, std::optional kvTransferSenderFutureTimeoutMs = std::nullopt, - std::optional kvTransferPollIntervalMs = kDefaultKvTransferPollIntervalMs); + std::optional kvTransferPollIntervalMs = kDefaultKvTransferPollIntervalMs, size_t agentBufferSizeMb = 0, + std::map agentBounceParams = {}); bool operator==(CacheTransceiverConfig const& other) const; void setBackendType(std::optional backendType); @@ -1530,12 +1531,16 @@ class CacheTransceiverConfig void setKvTransferTimeoutMs(std::optional kvTransferTimeoutMs); void setKvTransferSenderFutureTimeoutMs(std::optional kvTransferSenderFutureTimeoutMs); void setKvTransferPollIntervalMs(std::optional kvTransferPollIntervalMs); + void setAgentBufferSizeMb(size_t agentBufferSizeMb); + void setAgentBounceParams(std::map agentBounceParams); [[nodiscard]] std::optional getMaxTokensInBuffer() const; [[nodiscard]] std::optional getBackendType() const; [[nodiscard]] std::optional getKvTransferTimeoutMs() const; [[nodiscard]] std::optional getKvTransferSenderFutureTimeoutMs() const; [[nodiscard]] std::optional getKvTransferPollIntervalMs() const; + [[nodiscard]] size_t getAgentBufferSizeMb() const; + [[nodiscard]] std::map const& getAgentBounceParams() const; private: std::optional mBackendType; @@ -1550,6 +1555,15 @@ class CacheTransceiverConfig // @brief Bounded wait interval in milliseconds for polling KV transfer progress when active transfers block // disaggregated admission. std::optional mKvTransferPollIntervalMs; + // @brief Size in MiB of the transfer agent's staging-buffer (bounce) arena, currently + // implemented by the NIXL agent. 0 (default) disables the fast path; >0 enables it at that + // arena capacity (one shared arena for the sender and receiver roles). + size_t mAgentBufferSizeMb{0}; + // @brief Expert tuning knobs for the bounce pipeline, keyed by the TRTLLM_NIXL_BOUNCE_* + // environment-variable names without the prefix and the trailing _BYTES, lowercased (e.g. + // TRTLLM_NIXL_BOUNCE_MAX_CHUNK_SIZE_BYTES -> "max_chunk_size"). Precedence: this map > + // environment variable > built-in default. Ignored when mAgentBufferSizeMb == 0. + std::map mAgentBounceParams; }; /// @brief Configuration class for the model executor diff --git a/cpp/include/tensorrt_llm/executor/transferAgent.h b/cpp/include/tensorrt_llm/executor/transferAgent.h index 35e661b2aa06..1892bfb94b19 100644 --- a/cpp/include/tensorrt_llm/executor/transferAgent.h +++ b/cpp/include/tensorrt_llm/executor/transferAgent.h @@ -251,9 +251,8 @@ struct VramRegionMeta size_t chunkSize; ///< 0 = cudaMalloc (no split), >0 = VMM chunk size }; -// `AgentDesc` represents the unique identifier for reading and writing to the agent. -// By accessing this identifier, the backend can establish the correct connection. -// It also carries VMM region metadata so that remote agents can split at chunk boundaries. +// `AgentDesc` carries the backend metadata needed to connect to an agent, plus optional +// backend-independent metadata used by higher-level transfer paths. class AgentDesc final { public: @@ -262,9 +261,10 @@ class AgentDesc final { } - AgentDesc(std::string backendAgentDesc, std::vector vramRegions) + AgentDesc(std::string backendAgentDesc, std::vector vramRegions, std::string bounceHandshake = {}) : mBackendAgentDesc{std::move(backendAgentDesc)} , mVramRegions{std::move(vramRegions)} + , mBounceHandshake{std::move(bounceHandshake)} { } @@ -278,7 +278,18 @@ class AgentDesc final return mVramRegions; } - /// Serialize the entire AgentDesc (backend blob + VMM regions) into an opaque string. + /// Optional NIXL-bounce capability handshake advertised by this agent. + /// An opaque blob encoded/decoded by bounce::encodeHandshake/decodeHandshake: wire version, + /// control-channel kind + endpoint, and effective region-size limits. It is empty when bounce is + /// disabled. A caller must exchange the complete serialized AgentDesc, rather than only + /// getBackendAgentDesc(), for the peer to receive this handshake. + [[nodiscard]] std::string const& getBounceHandshake() const noexcept + { + return mBounceHandshake; + } + + /// Serialize the entire AgentDesc (backend blob + VMM regions + bounce handshake) into an + /// opaque string. [[nodiscard]] std::string serialize() const; /// Deserialize an opaque string back into an AgentDesc. @@ -287,6 +298,7 @@ class AgentDesc final private: std::string mBackendAgentDesc; std::vector mVramRegions; + std::string mBounceHandshake; }; // `TransferOp` is an enumeration that represents the types of transfer operations. @@ -371,6 +383,13 @@ class TransferStatus { return false; } + + /// Human-readable detail for the most recent terminal state (empty when unavailable). Backends + /// override this so a kFAILURE from wait() can carry its cause to the caller. + [[nodiscard]] virtual std::string getLastStatusStr() const + { + return {}; + } }; struct BaseAgentConfig @@ -383,6 +402,18 @@ struct BaseAgentConfig std::unordered_map backendParams; std::optional rank; std::optional worldSize; + /// Size in MiB of the agent's staging-buffer (bounce) arena, currently implemented by the + /// NIXL agent. 0 (default) disables the fast path; >0 enables it at that arena capacity. + /// Ignored by agents without bounce support (e.g. mooncake). + size_t agentBufferSizeMb{0}; + /// Expert tuning knobs for the bounce pipeline, keyed by the TRTLLM_NIXL_BOUNCE_* + /// environment-variable names without the prefix and the trailing _BYTES, lowercased (e.g. + /// TRTLLM_NIXL_BOUNCE_MAX_CHUNK_SIZE_BYTES -> "max_chunk_size"). + /// Kept SEPARATE from backendParams on purpose: backendParams is forwarded verbatim to the + /// backend plugin (e.g. NIXL createBackend), so bounce keys must not leak into it. + /// Precedence: this map > environment variable > built-in default. Ignored when + /// agentBufferSizeMb == 0. + std::unordered_map bounceParams; }; class BaseTransferAgent @@ -416,8 +447,8 @@ class BaseTransferAgent /// @return The descriptor of the local agent. virtual AgentDesc getLocalAgentDesc() = 0; - /// @brief Fetch the descriptor of the local agent. - /// @return The descriptor of the local agent. + /// @brief Fetch the backend-specific connection information for the local agent. + /// @return The local connection information. Optional AgentDesc metadata is not included. virtual ConnectionInfoType getLocalConnectionInfo() = 0; /// @brief Initiate the transfer by submitting the request. diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 97c079f917d8..a76747196494 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -664,7 +664,8 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa auto rnnState = mCacheState->hasRnnConfig() ? std::make_optional(mCacheState->getRnnCacheState()) : std::nullopt; mManager = std::make_unique( - mCacheTransBufferManagerPtrs, *mCacheState, "nixl", rnnState); + mCacheTransBufferManagerPtrs, *mCacheState, "nixl", rnnState, + mCacheTransceiverConfig->getAgentBufferSizeMb(), mCacheTransceiverConfig->getAgentBounceParams()); TLLM_LOG_INFO("NIXL Connection Manager created"); } else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::MOONCAKE) @@ -672,7 +673,8 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa auto rnnState = mCacheState->hasRnnConfig() ? std::make_optional(mCacheState->getRnnCacheState()) : std::nullopt; mManager = std::make_unique( - mCacheTransBufferManagerPtrs, *mCacheState, "mooncake", rnnState); + mCacheTransBufferManagerPtrs, *mCacheState, "mooncake", rnnState, + mCacheTransceiverConfig->getAgentBufferSizeMb(), mCacheTransceiverConfig->getAgentBounceParams()); TLLM_LOG_INFO("MOONCAKE Connection Manager created"); } else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::MPI) diff --git a/cpp/tensorrt_llm/executor/cacheTransceiverConfig.cpp b/cpp/tensorrt_llm/executor/cacheTransceiverConfig.cpp index 45d994a46d21..81c9e27cd355 100644 --- a/cpp/tensorrt_llm/executor/cacheTransceiverConfig.cpp +++ b/cpp/tensorrt_llm/executor/cacheTransceiverConfig.cpp @@ -23,11 +23,14 @@ namespace tensorrt_llm::executor CacheTransceiverConfig::CacheTransceiverConfig(std::optional backendType, std::optional maxNumTokens, std::optional kvTransferTimeoutMs, - std::optional kvTransferSenderFutureTimeoutMs, std::optional kvTransferPollIntervalMs) + std::optional kvTransferSenderFutureTimeoutMs, std::optional kvTransferPollIntervalMs, + size_t agentBufferSizeMb, std::map agentBounceParams) : mBackendType(backendType) , mMaxTokensInBuffer(maxNumTokens) , mKvTransferTimeoutMs(kvTransferTimeoutMs) , mKvTransferSenderFutureTimeoutMs(kvTransferSenderFutureTimeoutMs) + , mAgentBufferSizeMb(agentBufferSizeMb) + , mAgentBounceParams(std::move(agentBounceParams)) { setKvTransferPollIntervalMs(kvTransferPollIntervalMs); } @@ -37,7 +40,8 @@ bool CacheTransceiverConfig::operator==(CacheTransceiverConfig const& other) con return mMaxTokensInBuffer == other.mMaxTokensInBuffer && mBackendType == other.mBackendType && mKvTransferTimeoutMs == other.mKvTransferTimeoutMs && mKvTransferSenderFutureTimeoutMs == other.mKvTransferSenderFutureTimeoutMs - && mKvTransferPollIntervalMs == other.mKvTransferPollIntervalMs; + && mKvTransferPollIntervalMs == other.mKvTransferPollIntervalMs + && mAgentBufferSizeMb == other.mAgentBufferSizeMb && mAgentBounceParams == other.mAgentBounceParams; } void CacheTransceiverConfig::setBackendType(std::optional backendType) @@ -101,4 +105,24 @@ std::optional CacheTransceiverConfig::getKvTransferPollIntervalMs() const { return mKvTransferPollIntervalMs; } + +void CacheTransceiverConfig::setAgentBufferSizeMb(size_t agentBufferSizeMb) +{ + mAgentBufferSizeMb = agentBufferSizeMb; +} + +size_t CacheTransceiverConfig::getAgentBufferSizeMb() const +{ + return mAgentBufferSizeMb; +} + +void CacheTransceiverConfig::setAgentBounceParams(std::map agentBounceParams) +{ + mAgentBounceParams = std::move(agentBounceParams); +} + +std::map const& CacheTransceiverConfig::getAgentBounceParams() const +{ + return mAgentBounceParams; +} } // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp index b3bb08b57425..abcda675256b 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.cpp @@ -386,7 +386,8 @@ std::optional AgentConnection::getPreAssignedBufferId(uint8_t kind) cons AgentConnectionManager::AgentConnectionManager( std::vector cacheTransBufferManagers, CacheState cacheState, - std::string const& backendType, std::optional rnnCacheState) + std::string const& backendType, std::optional rnnCacheState, size_t agentBufferSizeMb, + std::map const& agentBounceParams) : mCacheState(std::move(cacheState)) , mRnnCacheState(std::move(rnnCacheState)) , mCacheTransBufferManagers(std::move(cacheTransBufferManagers)) @@ -425,6 +426,11 @@ AgentConnectionManager::AgentConnectionManager( BaseAgentConfig config{mAgentName, true, false, true}; config.rank = mRank; config.worldSize = mWorldSize; + config.agentBufferSizeMb = agentBufferSizeMb; + // BaseAgentConfig keeps bounce knobs separate from backendParams (which is forwarded verbatim + // to the backend plugin); convert the ordered map from CacheTransceiverConfig here. + config.bounceParams + = std::unordered_map(agentBounceParams.begin(), agentBounceParams.end()); m_Agent = makeTransferAgent(backendType, &config); TLLM_CHECK(!mCacheTransBufferManagers.empty()); mBufferKinds.reserve(mCacheTransBufferManagers.size()); diff --git a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h index 83f7df2541b9..2ecc6c657563 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/agent_utils/connection.h @@ -305,7 +305,8 @@ class AgentConnectionManager : public ConnectionManager public: AgentConnectionManager(std::vector cacheTransBufferManagers, CacheState cacheState, std::string const& backendType, - std::optional rnnCacheState = std::nullopt); + std::optional rnnCacheState = std::nullopt, size_t agentBufferSizeMb = 0, + std::map const& agentBounceParams = {}); ~AgentConnectionManager(); AgentConnection* recvConnect(DataContext const& ctx, void* data, size_t size) override; [[nodiscard]] std::vector getConnections(CommState const& state) override; diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt index ae3cf5e4d597..bd0fdcca5b2c 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt @@ -42,6 +42,42 @@ if(NIXL_ROOT) # Link against CUDA runtime (for cudaMemcpy in posix fallback) target_link_libraries(${NIXL_WRAPPER_TARGET} PRIVATE CUDA::cudart) + # ------------------------------------------------------------------------ + # Bounce v2 transport. Its real dependencies are NIXL (data plane, gated by + # the NIXL_ROOT block above) + zmq (control plane) — NOT UCX. So compile it + # whenever zmq is available, INDEPENDENT of ENABLE_UCX. cppzmq (header-only) + # is fetched at the top level when NIXL or UCX is on. Activates the + # TLLM_BOUNCE_V2 code paths; runtime use is still opt-in via + # CacheTransceiverConfig.agent_buffer_size_mb. + # ------------------------------------------------------------------------ + find_package(PkgConfig) + if(PKG_CONFIG_FOUND) + pkg_check_modules(ZMQ libzmq) # optional: bounce is skipped if libzmq is + # absent + endif() + if(ZMQ_FOUND) + target_sources( + ${NIXL_WRAPPER_TARGET} + PRIVATE bounce/BounceTransferPlan.cpp + bounce/BuddyAllocator.cpp + bounce/CreditScheduler.cpp + bounce/BounceMessage.cpp + bounce/BounceArena.cpp + bounce/ExecPool.cpp + bounce/GatherScatterKernel.cu + bounce/ZmqControlChannel.cpp + bounce/BounceTransport.cpp) + target_include_directories( + ${NIXL_WRAPPER_TARGET} PRIVATE ${CMAKE_BINARY_DIR}/_deps/cppzmq-src + ${ZMQ_INCLUDE_DIRS}) + target_link_libraries(${NIXL_WRAPPER_TARGET} PRIVATE ${ZMQ_LIBRARIES} + ${NIXL_BUILD_LIBRARY}) + target_compile_definitions(${NIXL_WRAPPER_TARGET} PRIVATE TLLM_BOUNCE_V2=1) + # The .cu device symbols must resolve into this shared lib. + set_target_properties(${NIXL_WRAPPER_TARGET} + PROPERTIES CUDA_RESOLVE_DEVICE_SYMBOLS ON) + endif() + set(NIXL_ENABLED TRUE) else() set(NIXL_ENABLED FALSE) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp index a963bd2e58e3..d0fa6862a5a2 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp @@ -189,7 +189,11 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m) // subclass type is not directly registered (e.g., agents created via factory). nb::class_(m, "TransferStatus") .def("is_completed", &kvc::TransferStatus::isCompleted, nb::call_guard()) - .def("wait", &kvc::TransferStatus::wait, nb::arg("timeout_ms") = -1, nb::call_guard()); + .def("wait", &kvc::TransferStatus::wait, nb::arg("timeout_ms") = -1, nb::call_guard()) + // Failure detail for the last terminal state (empty if unavailable). Named to match the + // lookup in BindingsNixlTransferStatus.last_status_str (nixl/_agent_cpp.py), which is what + // the Python transceiver's error log reads. + .def("get_last_status_str", &kvc::TransferStatus::getLastStatusStr); // BaseAgentConfig struct nb::class_(m, "BaseAgentConfig") @@ -199,15 +203,18 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m) [](kvc::BaseAgentConfig* self, std::string name, bool use_prog_thread, bool multi_thread, bool use_listen_thread, bool enable_telemetry, std::unordered_map backend_params, std::optional rank, - std::optional world_size) + std::optional world_size, size_t agent_buffer_size_mb, + std::unordered_map bounce_params) { new (self) kvc::BaseAgentConfig{std::move(name), use_prog_thread, multi_thread, use_listen_thread, - enable_telemetry, std::move(backend_params), rank, world_size}; + enable_telemetry, std::move(backend_params), rank, world_size, agent_buffer_size_mb, + std::move(bounce_params)}; }, nb::arg("name"), nb::arg("use_prog_thread") = true, nb::arg("multi_thread") = false, nb::arg("use_listen_thread") = false, nb::arg("enable_telemetry") = false, nb::arg("backend_params") = std::unordered_map{}, nb::arg("rank") = std::nullopt, - nb::arg("world_size") = std::nullopt) + nb::arg("world_size") = std::nullopt, nb::arg("agent_buffer_size_mb") = 0, + nb::arg("bounce_params") = std::unordered_map{}) .def_rw("name", &kvc::BaseAgentConfig::mName) .def_rw("use_prog_thread", &kvc::BaseAgentConfig::useProgThread) .def_rw("multi_thread", &kvc::BaseAgentConfig::multiThread) @@ -215,7 +222,9 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m) .def_rw("enable_telemetry", &kvc::BaseAgentConfig::enableTelemetry) .def_rw("backend_params", &kvc::BaseAgentConfig::backendParams) .def_rw("rank", &kvc::BaseAgentConfig::rank) - .def_rw("world_size", &kvc::BaseAgentConfig::worldSize); + .def_rw("world_size", &kvc::BaseAgentConfig::worldSize) + .def_rw("agent_buffer_size_mb", &kvc::BaseAgentConfig::agentBufferSizeMb) + .def_rw("bounce_params", &kvc::BaseAgentConfig::bounceParams); // BaseTransferAgent class (abstract base) // All transfer-engine operations release the GIL: they may block on NIXL / @@ -293,7 +302,10 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m) .def("get_notified_sync_messages", &kvc::NixlTransferAgent::getNotifiedSyncMessages, nb::call_guard()) .def("check_remote_descs", &kvc::NixlTransferAgent::checkRemoteDescs, nb::arg("name"), nb::arg("memory_descs"), - nb::call_guard()); + nb::call_guard()) + // Programmatic bounce v2 observability (deployment checks and tests; no log parsing). + .def_prop_ro("bounce_enabled", &kvc::NixlTransferAgent::isBounceEnabled) + .def_prop_ro("bounce_submit_count", &kvc::NixlTransferAgent::getBounceSubmitCount); #endif // NOTE: MooncakeTransferAgent/MooncakeTransferStatus class bindings are intentionally diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceArena.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceArena.cpp new file mode 100644 index 000000000000..0286704b4230 --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceArena.cpp @@ -0,0 +1,62 @@ +/* + * 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/BounceArena.h" + +#include "tensorrt_llm/batch_manager/cacheTransBuffer.h" // kv_cache_manager::FabricMemory (full def) +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/logger.h" + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +BounceArena::BounceArena(std::size_t bytes, int deviceId, bool allowFabric) + : mDeviceId(deviceId) + , mBytes(bytes) +{ + TLLM_CHECK_WITH_INFO(bytes > 0, "BounceArena: bytes must be > 0"); + TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); + + // On MNNVL parts (GH200/GB200) the arena (RDMA src/dst, NIXL-registered) must be fabric memory + // to be reachable over the NVLink fabric + GPUDirect-RDMA capable. Elsewhere (and CI, or when + // fabric is force-disabled) fall back to cudaMalloc. + if (allowFabric && tensorrt_llm::batch_manager::kv_cache_manager::FabricMemory::supportFabricMemory()) + { + mFabric = std::make_unique(bytes); + mBase = mFabric->getPtr(); + TLLM_LOG_DEBUG("BounceArena: %zuB backed by fabric memory", bytes); + } + else + { + TLLM_CUDA_CHECK(cudaMalloc(&mBase, bytes)); + } +} + +BounceArena::~BounceArena() +{ + if (!mFabric && mBase != nullptr) + { + // Select the owning device before freeing (multi-GPU: cudaFree otherwise targets the thread's + // current device). Fabric-backed arena is freed by ~FabricMemory (mFabric). A dtor can't + // throw, so use the project's warn-only cleanup check rather than discarding the result. + TLLM_CUDA_CHECK_WARN(cudaSetDevice(mDeviceId)); + TLLM_CUDA_CHECK_WARN(cudaFree(mBase)); + } +} + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceArena.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceArena.h new file mode 100644 index 000000000000..35aa63aa8e7b --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceArena.h @@ -0,0 +1,103 @@ +/* + * 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 +#include +#include + +// Reuse the KV-cache transfer buffers' fabric allocator (batch_manager) rather than a private copy: +// one MNNVL/GPUDirect-RDMA implementation, shared. Forward-declared here (the member is held by +// unique_ptr with an out-of-line dtor); the full definition is included in the .cpp. +namespace tensorrt_llm::batch_manager::kv_cache_manager +{ +class FabricMemory; +} // namespace tensorrt_llm::batch_manager::kv_cache_manager + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +// ============================================================================ +// BounceArena — the ONE shared bounce data buffer +// ---------------------------------------------------------------------------- +// Role +// Pure STORAGE: a single contiguous `bytes`-byte device buffer registered ONCE with NIXL and +// shared by BOTH roles. CreditScheduler (its BuddyAllocator) carves variable-size regions out of +// it by byte offset — receiver grants (remote senders' RDMA-write targets) and the local sender's +// gather staging both draw from the same arena. Unlike the old fixed-slot pool, each chunk requests +// only its packed extent; the scheduler rounds that extent to a buddy block. Small requests +// therefore use much less space than a full-size slot, while requests larger than the arena stream +// through chunk by chunk. +// +// Allocation +// On MNNVL parts (GH200/GB200) the buffer must be fabric memory (cuMemCreate + +// CU_MEM_HANDLE_TYPE_FABRIC) to be NVLink-fabric reachable + GPUDirect-RDMA capable, matching how +// the KV-cache transfer buffers are allocated. Elsewhere (and CI / x86, or when fabric is force- +// disabled) it falls back to cudaMalloc. The device-pointer surface is identical either way. +// +// Threading +// base()/baseAddr()/at() are const lookups into an immutable buffer; safe from the IO thread and +// the scatter workers concurrently. WHICH region a role may touch is arbitrated by CreditScheduler +// (a region is granted/acquired to exactly one user until it is freed). +// ============================================================================ +class BounceArena +{ +public: + /// Allocate one `bytes`-byte device buffer. When `allowFabric` and the device is fabric-capable, + /// it is fabric memory; otherwise cudaMalloc. Throws on CUDA allocation failure. + BounceArena(std::size_t bytes, int deviceId, bool allowFabric); + ~BounceArena(); + + BounceArena(BounceArena const&) = delete; + BounceArena& operator=(BounceArena const&) = delete; + + [[nodiscard]] void* base() const noexcept + { + return mBase; + } + + [[nodiscard]] std::uint64_t baseAddr() const noexcept + { + return reinterpret_cast(mBase); + } + + [[nodiscard]] std::size_t bytes() const noexcept + { + return mBytes; + } + + /// Device address of `offset` bytes into the arena (an arena region's start). + [[nodiscard]] void* at(std::uint64_t offset) const noexcept + { + return static_cast(mBase) + offset; + } + + [[nodiscard]] bool isFabric() const noexcept + { + return mFabric != nullptr; + } + +private: + int mDeviceId{0}; + std::size_t mBytes{0}; + void* mBase{nullptr}; // the registered RDMA src/dst buffer (offset 0) + // Non-null on MNNVL when fabric-backed; else null and mBase is cudaMalloc'd. + std::unique_ptr mFabric; +}; + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h new file mode 100644 index 000000000000..d36baf0fdc9c --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h @@ -0,0 +1,390 @@ +/* + * 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/common/logger.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +// --------------------------------------------------------------------------- +// Reusable string parsers, shared by the TRTLLM_NIXL_BOUNCE_* env fallback and +// the CacheTransceiverConfig agent_bounce_params dict. All are defensive: a +// value that fails to parse yields nullopt so the caller can keep its current +// (default or env-resolved) value and warn instead of aborting or silently +// producing 0 (a 0 would later trip TLLM_CHECKs, e.g. in BounceTransferPlan). +// --------------------------------------------------------------------------- + +/// Case-insensitive: 0/false/no/off -> false, 1/true/yes/on -> true, anything else -> nullopt. +[[nodiscard]] inline std::optional parseBoolValue(std::string const& raw) +{ + std::string s; + s.reserve(raw.size()); + for (char c : raw) + { + s.push_back(static_cast(std::tolower(static_cast(c)))); + } + if (s == "0" || s == "false" || s == "no" || s == "off") + { + return false; + } + if (s == "1" || s == "true" || s == "yes" || s == "on") + { + return true; + } + return std::nullopt; +} + +/// Strict unsigned decimal integer: garbage (typo like "abc"), trailing junk, or overflow -> nullopt. +[[nodiscard]] inline std::optional parseU64Value(std::string const& raw) +{ + if (raw.empty() || !std::isdigit(static_cast(raw[0]))) + { + return std::nullopt; + } + char* end = nullptr; + errno = 0; + std::uint64_t const parsed = std::strtoull(raw.c_str(), &end, 10); + if (errno == ERANGE || end == raw.c_str() || *end != '\0') + { + return std::nullopt; + } + return parsed; +} + +/// Byte size: unsigned integer with an optional binary suffix — K/KB/KiB, M/MB/MiB, G/GB/GiB +/// (case-insensitive, no space), e.g. "256MB", "1gb", "512kib". All suffixes are powers of two +/// (MB == MiB == 2^20). Bare numbers and a trailing "B" stay bytes. Unknown suffix, trailing +/// junk, or multiplication overflow -> nullopt. +[[nodiscard]] inline std::optional parseBytesValue(std::string const& raw) +{ + if (raw.empty() || !std::isdigit(static_cast(raw[0]))) + { + return std::nullopt; + } + char* end = nullptr; + errno = 0; + std::uint64_t const parsed = std::strtoull(raw.c_str(), &end, 10); + if (errno == ERANGE || end == raw.c_str()) + { + return std::nullopt; + } + std::string suffix; + for (char const* p = end; *p != '\0'; ++p) + { + suffix.push_back(static_cast(std::tolower(static_cast(*p)))); + } + std::uint64_t mult = 1; + if (suffix.empty() || suffix == "b") + { + mult = 1; + } + else if (suffix == "k" || suffix == "kb" || suffix == "kib") + { + mult = 1ULL << 10; + } + else if (suffix == "m" || suffix == "mb" || suffix == "mib") + { + mult = 1ULL << 20; + } + else if (suffix == "g" || suffix == "gb" || suffix == "gib") + { + mult = 1ULL << 30; + } + else + { + return std::nullopt; // unknown suffix + } + if (mult > 1 && parsed > std::numeric_limits::max() / mult) + { + return std::nullopt; // would overflow + } + return parsed * mult; +} + +/// POD config for the bounce v2 pipeline. There is no `enabled` field: the on/off switch is +/// CacheTransceiverConfig's agent_bounce_buffer_enable + kv_cache_bounce_size_mb (which the Python +/// frontend folds into the agent's arena size in MiB: 0 = off, >0 = arena size), and at runtime +/// "bounce on" simply means the owning agent built its bounce state (mBounce != nullptr). +/// `arenaSizeBytes` is NOT env-backed either — the owning agent sets it from the same knob. The +/// expert knobs resolve as: agent_bounce_params dict > TRTLLM_NIXL_BOUNCE_* env var > built-in +/// default (`fromEnv()` reads the env, `fromParams()` layers the dict on top). Byte-valued knobs +/// accept an optional case-insensitive binary suffix such as "256MB", "1gb", or "512KiB" +/// (K/M/G == KiB/MiB/GiB, powers of two). +struct BounceConfig +{ + // Overwritten with the agent arena size (MiB << 20, derived from CacheTransceiverConfig's + // kv_cache_bounce_size_mb + agent_bounce_buffer_enable) on every production path + // (maybeInitBounce); the + // default only serves unit tests that construct a BounceConfig{} directly. + std::size_t arenaSizeBytes{512ULL << 20}; + std::size_t arenaAllocationGranularityBytes{1ULL << 20}; // arena_allocation_granularity + std::size_t maxChunkSizeBytes{32ULL << 20}; // max_chunk_size + std::uint32_t maxInflightChunksPerRequest{8}; // max_inflight_chunks_per_request + std::uint32_t copyStreamCount{8}; // copy_stream_count + std::uint32_t scatterWorkerCount{4}; // scatter_worker_count + std::size_t minDescriptorCount{1024}; // min_descriptor_count + std::size_t maxAverageDescriptorSizeBytes{16ULL << 10}; // max_average_descriptor_size + int requestTimeoutMs{30000}; // request_timeout_ms; <=0 DISABLES the timeout + // (checkTimeouts no-ops; used by tests that intentionally wait) + // 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 + // alone — so a flow whose grants see no progress (no GRANT sent, no DATA received) for this + // long is reclaimed and its regions quarantined (below) before reuse. The lease must EXCEED + // the peers' requestTimeoutMs (a live sender abandons + cancels first, so only + // dead/unreachable peers ever hit this) — the 2x derivation assumes both ends run the SAME + // request_timeout_ms. + int receiverFlowTimeoutMs{60000}; + // How long a receiver-reclaimed, possibly-still-being-written region stays out of the arena + // before reuse — DERIVED in deriveDependentTimeouts() as requestTimeoutMs. A one-sided RDMA + // write cannot be aborted, so time is the only barrier against re-granting a region a gone + // peer's NIC may still be writing. + int quarantineMs{30000}; + bool disableFabricMemory{false}; // disable_fabric_memory + // enable_eager_gather: launch a chunk's gather at submit() time, before the receiver's GRANT + // arrives, overlapping the WANT->GRANT control round-trip with the gather kernel. Eager + // (credit-less) staging regions are capped at HALF the arena so that on a bidirectional + // deployment both sides can always still grant incoming regions (no mutual eager-starvation); + // the credit-backed path is unaffected by the cap. + bool enableEagerGather{true}; + // use_zero_copy_arguments: the copy kernel reads [srcs|dsts|sizes] directly from pinned host + // memory instead of staging them in device scratch first. Faster at every plan size (same + // bytes over the bus, but no H2D-then-kernel serialization), so on by default. + bool useZeroCopyArguments{true}; + + /// Re-derive the lease/quarantine values from the one user-visible timeout (see the field + /// comments): the lease must exceed the peers' request timeout, and time is the only write + /// barrier for a reclaimed region. Disabling the request timeout (<=0) disables both. Must be + /// re-run whenever requestTimeoutMs changes: fromEnv() always calls it, fromParams() only when + /// the dict actually provides request_timeout_ms — so a caller-tweaked receiverFlowTimeoutMs / + /// quarantineMs on the base config (tests set the fields directly) survives an unrelated dict. + /// The doubling is done in 64-bit and clamped: requestTimeoutMs admits values up to INT_MAX, + /// and a signed overflow here would wrap the lease negative — silently disabling the + /// dead-sender reclaim it drives. + void deriveDependentTimeouts() + { + receiverFlowTimeoutMs = requestTimeoutMs > 0 ? static_cast(std::min( + std::int64_t{2} * requestTimeoutMs, std::numeric_limits::max())) + : 0; + quarantineMs = requestTimeoutMs; + } + + /// Apply one knob given as a (key, value) string pair. Returns false when the key is unknown. + /// A value that fails to parse or is out of range keeps the current field value and warns + /// (`origin` names the source — the env var or the config dict — for the log line). + static bool applyParam(BounceConfig& cfg, std::string const& key, std::string const& value, char const* origin) + { + auto warnBad = [&] + { + TLLM_LOG_WARNING("BounceConfig: invalid value '%s' for '%s' (from %s) -> keeping the current value", + value.c_str(), key.c_str(), origin); + }; + auto setBool = [&](bool& field) + { + if (auto v = parseBoolValue(value)) + { + field = *v; + } + else + { + warnBad(); + } + }; + auto setSizeBytes = [&](std::size_t& field, bool allowZero) + { + auto const v = parseBytesValue(value); + // The size_t bound only bites on 32-bit size_t builds; it is trivially true on 64-bit. + if (v.has_value() && (allowZero || *v != 0) && *v <= std::numeric_limits::max()) + { + field = static_cast(*v); + } + else + { + warnBad(); + } + }; + auto setSize = [&](std::size_t& field, bool allowZero) + { + auto const v = parseU64Value(value); + // The size_t bound only bites on 32-bit size_t builds; it is trivially true on 64-bit. + if (v.has_value() && (allowZero || *v != 0) && *v <= std::numeric_limits::max()) + { + field = static_cast(*v); + } + else + { + warnBad(); + } + }; + auto setU32 = [&](std::uint32_t& field, bool allowZero) + { + auto const v = parseU64Value(value); + if (v.has_value() && (allowZero || *v != 0) && *v <= std::numeric_limits::max()) + { + field = static_cast(*v); + } + else + { + warnBad(); + } + }; + auto setInt = [&](int& field) + { + auto const v = parseU64Value(value); + if (v.has_value() && *v <= static_cast(std::numeric_limits::max())) + { + field = static_cast(*v); + } + else + { + warnBad(); + } + }; + + if (key == "max_chunk_size") + { + setSizeBytes(cfg.maxChunkSizeBytes, /*allowZero=*/false); + } + else if (key == "arena_allocation_granularity") + { + setSizeBytes(cfg.arenaAllocationGranularityBytes, /*allowZero=*/false); + } + else if (key == "max_average_descriptor_size") + { + setSizeBytes(cfg.maxAverageDescriptorSizeBytes, /*allowZero=*/true); + } + else if (key == "max_inflight_chunks_per_request") + { + setU32(cfg.maxInflightChunksPerRequest, /*allowZero=*/false); + } + else if (key == "copy_stream_count") + { + setU32(cfg.copyStreamCount, /*allowZero=*/false); + } + else if (key == "scatter_worker_count") + { + setU32(cfg.scatterWorkerCount, /*allowZero=*/false); + } + else if (key == "min_descriptor_count") + { + setSize(cfg.minDescriptorCount, /*allowZero=*/true); + } + else if (key == "request_timeout_ms") + { + setInt(cfg.requestTimeoutMs); + } + else if (key == "disable_fabric_memory") + { + setBool(cfg.disableFabricMemory); + } + else if (key == "enable_eager_gather") + { + setBool(cfg.enableEagerGather); + } + else if (key == "use_zero_copy_arguments") + { + setBool(cfg.useZeroCopyArguments); + } + else + { + return false; + } + return true; + } + + /// Defaults overridden by any set TRTLLM_NIXL_BOUNCE_* env var. Each call reads the current + /// environment. The on/off switch and the arena size are NOT read here — they only come from + /// CacheTransceiverConfig (agent_bounce_buffer_enable + kv_cache_bounce_size_mb). + [[nodiscard]] static BounceConfig fromEnv() + { + // paramKey -> env var name. Byte-valued knobs keep the historical _BYTES env suffix. + // KEEP IN SYNC with AGENT_BOUNCE_PARAM_KEYS in + // tensorrt_llm/_torch/disaggregation/nixl/bounce_knobs.py, which the llm_args validator + // uses to reject unknown agent_bounce_params keys upfront. + static constexpr std::array, 11> kEnvKnobs{{ + {"max_chunk_size", "TRTLLM_NIXL_BOUNCE_MAX_CHUNK_SIZE_BYTES"}, + {"arena_allocation_granularity", "TRTLLM_NIXL_BOUNCE_ARENA_ALLOCATION_GRANULARITY_BYTES"}, + {"max_average_descriptor_size", "TRTLLM_NIXL_BOUNCE_MAX_AVERAGE_DESCRIPTOR_SIZE_BYTES"}, + {"max_inflight_chunks_per_request", "TRTLLM_NIXL_BOUNCE_MAX_INFLIGHT_CHUNKS_PER_REQUEST"}, + {"copy_stream_count", "TRTLLM_NIXL_BOUNCE_COPY_STREAM_COUNT"}, + {"scatter_worker_count", "TRTLLM_NIXL_BOUNCE_SCATTER_WORKER_COUNT"}, + {"min_descriptor_count", "TRTLLM_NIXL_BOUNCE_MIN_DESCRIPTOR_COUNT"}, + {"request_timeout_ms", "TRTLLM_NIXL_BOUNCE_REQUEST_TIMEOUT_MS"}, + {"disable_fabric_memory", "TRTLLM_NIXL_BOUNCE_DISABLE_FABRIC_MEMORY"}, + {"enable_eager_gather", "TRTLLM_NIXL_BOUNCE_ENABLE_EAGER_GATHER"}, + {"use_zero_copy_arguments", "TRTLLM_NIXL_BOUNCE_USE_ZERO_COPY_ARGUMENTS"}, + }}; + BounceConfig cfg; + for (auto const& [paramKey, envName] : kEnvKnobs) + { + char const* v = std::getenv(envName); + if (v != nullptr && v[0] != '\0') // unset or empty -> keep the default + { + applyParam(cfg, paramKey, v, envName); + } + } + cfg.deriveDependentTimeouts(); + return cfg; + } + + /// Layer the CacheTransceiverConfig agent_bounce_params dict over `base` (normally the + /// fromEnv() result, so dict > env > default). Unknown keys warn and are skipped (defensive + /// backstop — the llm_args validator already rejects them at the Python boundary, but other + /// entry points can reach this directly). The lease/quarantine values are re-derived only + /// when the dict provides a request_timeout_ms that actually parses (a rejected value keeps + /// the base timeout, so re-deriving would only clobber directly-set values on `base`). + [[nodiscard]] static BounceConfig fromParams( + std::unordered_map const& params, BounceConfig base) + { + for (auto const& [key, value] : params) + { + if (!applyParam(base, key, value, "agent_bounce_params")) + { + TLLM_LOG_WARNING("BounceConfig: unknown agent_bounce_params key '%s' -> ignored", key.c_str()); + } + } + // Same validity check as applyParam's request_timeout_ms setter: only a value that was + // actually accepted triggers the re-derivation. + auto const it = params.find("request_timeout_ms"); + if (it != params.end()) + { + auto const parsed = parseU64Value(it->second); + if (parsed.has_value() && *parsed <= static_cast(std::numeric_limits::max())) + { + base.deriveDependentTimeouts(); + } + } + return base; + } +}; + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.cpp new file mode 100644 index 000000000000..e9937d7fc404 --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.cpp @@ -0,0 +1,254 @@ +/* + * 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/BounceMessage.h" + +#include "tensorrt_llm/executor/serializeUtils.h" + +#include +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +namespace +{ + +BounceMsgHeader makeHeader(BounceMsgType type, std::uint64_t requestId, std::uint32_t chunkIdx, std::uint32_t numChunks, + std::uint64_t regionHandle, std::uint32_t count, std::uint32_t payloadBytes) +{ + BounceMsgHeader header{}; + header.magic = kBounceMagic; + header.version = kBounceVersion; + header.msgType = static_cast(type); + header.requestId = requestId; + header.chunkIdx = chunkIdx; + header.numChunks = numChunks; + header.regionHandle = regionHandle; + header.count = count; + header.payloadBytes = payloadBytes; + return header; +} + +template +std::string encodeWithEntries(BounceMsgHeader const& header, std::vector const& entries) +{ + std::string blob; + blob.resize(sizeof(BounceMsgHeader) + header.payloadBytes); + std::memcpy(blob.data(), &header, sizeof(BounceMsgHeader)); + if (!entries.empty()) + { + std::memcpy(blob.data() + sizeof(BounceMsgHeader), entries.data(), entries.size() * sizeof(Entry)); + } + return blob; +} + +template +bool decodeEntries(std::string const& blob, BounceMsgHeader const& header, std::vector& out) +{ + std::size_t const expectBytes = static_cast(header.count) * sizeof(Entry); + if (expectBytes != header.payloadBytes) + { + return false; + } + if (blob.size() < sizeof(BounceMsgHeader) + expectBytes) + { + return false; + } + out.resize(header.count); + if (header.count > 0) + { + std::memcpy(out.data(), blob.data() + sizeof(BounceMsgHeader), expectBytes); + } + return true; +} + +std::string encodeHeaderOnly(BounceMsgHeader const& header) +{ + std::string blob; + blob.resize(sizeof(BounceMsgHeader)); + std::memcpy(blob.data(), &header, sizeof(BounceMsgHeader)); + return blob; +} + +} // namespace + +std::string encodeWant( + std::uint64_t requestId, std::vector const& chunkBytes, std::string const& endpoint) +{ + auto const n = static_cast(chunkBytes.size()); + auto const sizesBytes = static_cast(chunkBytes.size() * sizeof(std::uint32_t)); + auto const epLen = static_cast(endpoint.size()); + auto const payloadBytes = sizesBytes + static_cast(sizeof(std::uint32_t)) + epLen; + auto h = makeHeader(BounceMsgType::kWANT, requestId, 0, 0, 0, n, payloadBytes); + std::string blob; + blob.resize(sizeof(BounceMsgHeader) + payloadBytes); + std::memcpy(blob.data(), &h, sizeof(BounceMsgHeader)); + auto* payload = blob.data() + sizeof(BounceMsgHeader); + if (sizesBytes > 0) + { + std::memcpy(payload, chunkBytes.data(), sizesBytes); + } + std::memcpy(payload + sizesBytes, &epLen, sizeof(epLen)); + if (epLen > 0) + { + std::memcpy(payload + sizesBytes + sizeof(epLen), endpoint.data(), epLen); + } + return blob; +} + +std::string encodeCancel(std::uint64_t requestId, std::string const& endpoint) +{ + // A cancel IS an empty-chunk WANT (same wire form) — see the header. Keeping it a thin wrapper + // (rather than a new message type) reuses the receiver's onWant/reclaim path verbatim. + return encodeWant(requestId, {}, endpoint); +} + +std::string encodeGrant(std::uint64_t requestId, std::vector const& credits) +{ + auto const bytes = static_cast(credits.size() * sizeof(BounceCreditEntry)); + auto h = makeHeader(BounceMsgType::kGRANT, requestId, 0, 0, 0, static_cast(credits.size()), bytes); + return encodeWithEntries(h, credits); +} + +std::string encodeData(std::uint64_t requestId, std::uint32_t chunkIdx, std::uint32_t numChunks, + std::uint64_t regionHandle, std::vector const& entries) +{ + auto const bytes = static_cast(entries.size() * sizeof(BounceScatterRun)); + auto h = makeHeader(BounceMsgType::kDATA, requestId, chunkIdx, numChunks, regionHandle, + static_cast(entries.size()), bytes); + return encodeWithEntries(h, entries); +} + +std::string encodeAck(std::uint64_t requestId, std::uint32_t chunkIdx, std::uint64_t regionHandle) +{ + return encodeHeaderOnly(makeHeader(BounceMsgType::kACK, requestId, chunkIdx, 0, regionHandle, 0, 0)); +} + +bool decodeHeader(std::string const& blob, BounceMsgHeader& outHeader) +{ + if (blob.size() < sizeof(BounceMsgHeader)) + { + return false; + } + std::memcpy(&outHeader, blob.data(), sizeof(BounceMsgHeader)); + if (outHeader.magic != kBounceMagic || outHeader.version != kBounceVersion) + { + return false; + } + // Guard against a payloadBytes that the blob can't satisfy. + if (blob.size() < sizeof(BounceMsgHeader) + static_cast(outHeader.payloadBytes)) + { + return false; + } + return true; +} + +bool decodeCredits(std::string const& blob, BounceMsgHeader const& header, std::vector& out) +{ + return decodeEntries(blob, header, out); +} + +bool decodeScatter(std::string const& blob, BounceMsgHeader const& header, std::vector& out) +{ + return decodeEntries(blob, header, out); +} + +namespace +{ +// The handshake travels inside AgentDesc (whose fields use executor::serialize_utils), so it uses +// the SAME serialization utility rather than the raw-memcpy control-message codec above: a magic + +// format-version prefix, then each field via su::serialize. `formatVersion` versions THIS blob's +// layout (so the handshake itself can evolve); `wireVersion` inside is the control-message codec +// version being negotiated. +constexpr std::uint32_t kHandshakeMagic = 0x424E4853U; // 'B''N''H''S' +constexpr std::uint16_t kHandshakeFormatVersion = 1U; +} // namespace + +std::string encodeHandshake(BounceHandshake const& handshake) +{ + namespace su = tensorrt_llm::executor::serialize_utils; + std::ostringstream os; + su::serialize(kHandshakeMagic, os); + su::serialize(kHandshakeFormatVersion, os); + su::serialize(handshake.wireVersion, os); + su::serialize(static_cast(handshake.controlKind), os); + su::serialize(handshake.arenaUsableCapacityBytes, os); + su::serialize(handshake.maxChunkSizeBytes, os); + su::serialize(handshake.endpoint, os); + return os.str(); +} + +bool decodeHandshake(std::string const& blob, BounceHandshake& out) +{ + namespace su = tensorrt_llm::executor::serialize_utils; + std::istringstream is(blob); + auto const magic = su::deserialize(is); + auto const formatVersion = su::deserialize(is); + if (is.fail() || magic != kHandshakeMagic || formatVersion != kHandshakeFormatVersion) + { + return false; // absent / foreign / future-format blob -> peer treated as non-bounce + } + out.wireVersion = su::deserialize(is); + out.controlKind = static_cast(su::deserialize(is)); + out.arenaUsableCapacityBytes = su::deserialize(is); + out.maxChunkSizeBytes = su::deserialize(is); + auto const endpointLen = su::deserialize(is); + if (is.fail()) + { + return false; // truncated fixed fields (deserialize on a failed stream yields garbage) + } + // The endpoint is the final field, so its length must equal exactly the bytes left in the blob. + // Bounding it here keeps the "malformed blob -> false" contract: handing a corrupt 64-bit length + // to std::string would throw (length_error/bad_alloc) out of a peer-input path instead. + auto const consumed = static_cast(is.tellg()); + if (endpointLen != blob.size() - consumed) + { + return false; + } + out.endpoint = blob.substr(consumed); + return true; +} + +bool decodeWant(std::string const& blob, BounceMsgHeader const& header, std::vector& outChunkBytes, + std::string& outEndpoint) +{ + std::size_t const sizesBytes = static_cast(header.count) * sizeof(std::uint32_t); + std::size_t const off = sizeof(BounceMsgHeader); + // Need the chunk sizes plus the u32 endpoint-length prefix. + if (blob.size() < off + sizesBytes + sizeof(std::uint32_t)) + { + return false; + } + std::uint32_t epLen = 0; + std::memcpy(&epLen, blob.data() + off + sizesBytes, sizeof(epLen)); + std::size_t const expectPayload = sizesBytes + sizeof(std::uint32_t) + epLen; + // payloadBytes must match exactly, and the blob must hold the whole endpoint. + if (header.payloadBytes != expectPayload || blob.size() < off + expectPayload) + { + return false; + } + outChunkBytes.resize(header.count); + if (header.count > 0) + { + std::memcpy(outChunkBytes.data(), blob.data() + off, sizesBytes); + } + outEndpoint.assign(blob.data() + off + sizesBytes + sizeof(std::uint32_t), epLen); + return true; +} + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.h new file mode 100644 index 000000000000..4b58ebfa4c35 --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.h @@ -0,0 +1,187 @@ +/* + * 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 +#include +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +/// Bounce v2 control-plane wire format. Carried over the ControlChannel (zmq by default). +/// Pure encode/decode — no IO — so it is fully unit-testable. NOTE: this is NOT the NIXL +/// notifMsg (v2 does not use notifMsg at all); it is our own message set. +/// +/// ENDIANNESS: fields are memcpy'd raw (host byte order), so this assumes both peers are +/// little-endian — true for all NVIDIA GPU hosts (x86_64 / aarch64 LE). It is NOT a portable +/// cross-endian format; if a big-endian peer is ever introduced, add explicit byte-order +/// normalization here. `decodeHeader` validates magic + version + payload length; `dispatch` +/// treats any unknown `msgType` as a no-op (default branch), so a bogus type can't misroute. +enum class BounceMsgType : std::uint16_t +{ + kWANT = 1, // sender -> receiver: per-chunk byte sizes[] for requestId (empty list cancels) + kGRANT = 2, // receiver -> sender: credits[] for requestId + kDATA = 3, // sender -> receiver: region written; scatter entries[] (after getXferStatus SUCCESS) + kACK = 4, // receiver -> sender: chunk scattered; region freed +}; + +#pragma pack(push, 1) + +/// Fixed 40-byte header prefixing every message. Field meaning is per-type (documented at each +/// encode function); unused fields are 0. +struct BounceMsgHeader +{ + std::uint32_t magic; // kMagic + std::uint16_t version; // kBounceVersion + std::uint16_t msgType; // BounceMsgType + std::uint64_t requestId; // WANT/GRANT/DATA/ACK + std::uint64_t + regionHandle; // DATA/ACK: the arena region offset of this chunk; 0 elsewhere (64-bit -> arena may exceed 4 GiB) + std::uint32_t chunkIdx; // DATA/ACK + std::uint32_t numChunks; // DATA/ACK + std::uint32_t count; // number of trailing entries (credits / chunk sizes / scatter) + std::uint32_t payloadBytes; // bytes of trailing payload +}; + +/// Credit = permission to RDMA-write one chunk into a receiver arena allocation. `addr` is its +/// absolute remote address, `len` is the chunk's packed transfer length (the backing buddy block may +/// be larger), `devId` is the RECEIVER's GPU, and `regionHandle` is the receiver arena offset that +/// the sender echoes in DATA so the receiver can locate and release the allocation. +struct BounceCreditEntry +{ + std::uint64_t addr; + std::uint32_t len; + std::uint32_t devId; + std::uint64_t regionHandle; // 64-bit (was 32-bit + pad) -> arena offsets may exceed 4 GiB +}; + +/// One scatter RUN inside a DATA message: `count` equal pieces of `pieceSize` bytes; piece p copies +/// region[bounceOffset + p*bounceStride .. +pieceSize) -> dstAddr + p*dstStride. +/// count==1 is a single plain extent (strides unused, 0). Runs compress BOTH dst layouts a KV +/// transfer produces: a fully CONTIGUOUS dst (one count==1 run whose pieceSize covers the chunk — +/// e.g. ctx tp1 -> gen tp4, the gen rank's head-slice pool is dense) and a uniformly STRIDED dst +/// (one count==N run — e.g. ctx tp4 -> gen DP, each ctx rank's head slice lands every +/// bytes-per-layer in the gen rank's full-head pool). Either way a whole chunk's scatter plan is a +/// handful of runs instead of one entry per desc, keeping the DATA message (ACK critical path) tiny. +struct BounceScatterRun +{ + std::uint64_t bounceOffset; // first piece's offset within the granted region + std::uint64_t dstAddr; // first piece's destination address + std::uint64_t dstStride; // dst step between pieces (0 when count==1) + std::uint32_t bounceStride; // region-offset step between pieces (0 when count==1) + std::uint32_t pieceSize; // bytes per piece + std::uint32_t count; // number of pieces (>= 1) +}; + +#pragma pack(pop) + +static_assert(sizeof(BounceMsgHeader) == 40, "BounceMsgHeader must be 40 bytes"); +static_assert(sizeof(BounceCreditEntry) == 24, "BounceCreditEntry must be 24 bytes"); +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; + +// ---- 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 +/// each size as it grants) AND the sender's own bounce control endpoint. An EMPTY size list cancels. +/// The endpoint lets the receiver self-bootstrap the reverse control path (addPeer the sender) so +/// GRANT/ACK can flow back when disaggregated serving exchanges metadata only from this sender to +/// the receiver. Payload layout: +/// [count * u32 chunkBytes][u32 endpointLen][endpointLen bytes endpoint] +/// decode via decodeWant. +[[nodiscard]] std::string encodeWant( + std::uint64_t requestId, std::vector const& chunkBytes, std::string const& endpoint); +/// Cancel/retract a request: a WANT with an EMPTY chunk list (the receiver frees everything it +/// allocated/held for `requestId`). This is the ONLY meaning of an empty WANT — submit() never sends +/// a zero-chunk WANT (a 0-chunk transfer resolves SUCCESS without any WANT) — so an empty chunk list +/// is unambiguously a cancel. Thin named wrapper over encodeWant to make the intent explicit at call +/// sites; the wire form is identical (still a WANT), so it reuses the receiver's onWant/reclaim path. +[[nodiscard]] std::string encodeCancel(std::uint64_t requestId, std::string const& endpoint); +[[nodiscard]] std::string encodeGrant(std::uint64_t requestId, std::vector const& credits); +[[nodiscard]] std::string encodeData(std::uint64_t requestId, std::uint32_t chunkIdx, std::uint32_t numChunks, + std::uint64_t regionHandle, std::vector const& entries); +[[nodiscard]] std::string encodeAck(std::uint64_t requestId, std::uint32_t chunkIdx, std::uint64_t regionHandle); + +// ---- decode ---- +/// Decode the fixed header. Returns false on short blob, bad magic, or version mismatch. +[[nodiscard]] bool decodeHeader(std::string const& blob, BounceMsgHeader& outHeader); + +/// Decode credit entries (GRANT payload). `header.count` entries expected. +[[nodiscard]] bool decodeCredits( + std::string const& blob, BounceMsgHeader const& header, std::vector& out); + +/// Decode scatter entries (DATA payload). `header.count` entries expected. +[[nodiscard]] bool decodeScatter( + std::string const& blob, BounceMsgHeader const& header, std::vector& out); + +/// Decode a WANT: its per-chunk byte sizes (`header.count` entries; empty == cancel) and the +/// sender's bounce control endpoint. Returns false on a malformed/short blob. +[[nodiscard]] bool decodeWant(std::string const& blob, BounceMsgHeader const& header, + std::vector& outChunkBytes, std::string& outEndpoint); + +/// Does a decoded WANT mean "cancel/retract"? (An empty chunk list — see encodeCancel.) Makes the +/// receiver-side intent explicit instead of an inline `chunkBytes.empty()` check. +[[nodiscard]] inline bool isCancelWant(std::vector const& chunkBytes) +{ + return chunkBytes.empty(); +} + +// ---- handshake (out-of-band, travels in AgentDesc — NOT a control-channel message) ---- + +/// Which transport carries the bounce control messages. Peers MUST agree — the field stays on the +/// wire so a future alternative transport remains negotiable (an unknown value simply fails the +/// strict handshake equality and falls back to the standard NIXL path). +enum class BounceControlKind : std::uint8_t +{ + kZMQ = 0, +}; + +/// Bounce capability handshake, advertised in this agent's AgentDesc metadata. A peer that loads +/// the metadata validates compatibility (BounceTransport::registerPeerHandshake) BEFORE ever +/// engaging bounce toward us: incompatible (or absent) handshake -> that peer transparently uses +/// the standard per-desc NIXL path instead of waiting for requestTimeoutMs. +/// Compatibility is strict equality of wireVersion, controlKind and effective maxChunkSizeBytes. +/// Worker/stream counts, timeouts, the zero-copy-argument path, allocation granularity and +/// configured arena size are local settings and are not compared. The usable arena capacity is +/// carried for diagnostics. +struct BounceHandshake +{ + std::uint16_t wireVersion{kBounceVersion}; // control-message codec version; must match exactly + BounceControlKind controlKind{BounceControlKind::kZMQ}; // must match exactly + // The largest single allocation this agent's buddy arena can ever grant. Informational only: + // maxChunkSizeBytes below is already clamped to it before the handshake is generated. + std::uint64_t arenaUsableCapacityBytes{0}; + // This agent's effective per-chunk cap after the arena-capacity clamp. Peers compare this field + // exactly, which also keeps their scatter-plan capacities consistent. + std::uint64_t maxChunkSizeBytes{0}; + // Control-channel address: a zmq endpoint (kZMQ). + std::string endpoint; +}; + +/// Encode/decode the handshake blob carried in AgentDesc. decode returns false on a short blob, +/// bad magic, or unknown handshake-format version — the caller then treats the peer as +/// non-bounce-capable (fallback), never as an error. +[[nodiscard]] std::string encodeHandshake(BounceHandshake const& handshake); +[[nodiscard]] bool decodeHandshake(std::string const& blob, BounceHandshake& out); + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceNvtx.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceNvtx.h new file mode 100644 index 000000000000..60a1992488dc --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceNvtx.h @@ -0,0 +1,162 @@ +/* + * 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/common/nvtxUtils.h" + +#include +#include +#include + +#ifndef NVTX_DISABLE +#include +#include +#endif + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +// NVTX instrumentation for the bounce pipeline (perf analysis with nsys). Everything here +// compiles to a no-op when the build defines NVTX_DISABLE (the default; build with +// -DNVTX_DISABLE=OFF to profile). +// +// Two kinds of spans: +// - BounceNvtxScope: same-thread RAII start/end span, for synchronous sections (buildPlan, the +// gather launch, the scatter kernel + sync in a worker). +// - bounceRangeStart()/bounceRangeEnd(): process-wide start/end ranges for the ASYNC legs — +// started on one thread and ended on another (gather in flight: IO-thread launch -> IO-thread +// event poll later; RDMA write: post -> poll Done; ACK wait: DATA sent -> ACK received; +// scatter queueing: IO-thread enqueue -> worker dequeue). The handle is a plain uint64 so the +// reactor structs that carry it across threads need no NVTX include. + +/// Dedicated domain: bounce ranges get their own row group in nsys instead of mixing with the +/// global TRT-LLM ranges. +struct BounceNvtxDomain +{ + static constexpr char const* name{"trtllm.disagg.bounce"}; +}; + +// Span colors (ARGB), one per pipeline stage so the nsys timeline reads at a glance. +inline constexpr std::uint32_t kNvtxBuildPlan = 0xFF9E9E9EU; // gray +inline constexpr std::uint32_t kNvtxRequest = 0xFF2196F3U; // blue: submit -> resolve +inline constexpr std::uint32_t kNvtxGrantWait = 0xFFFF9800U; // orange: WANT sent -> first GRANT +inline constexpr std::uint32_t kNvtxGatherLaunch = 0xFF8BC34AU; // light green +inline constexpr std::uint32_t kNvtxGather = 0xFF4CAF50U; // green: gather launched -> event done +inline constexpr std::uint32_t kNvtxNixlWrite = 0xFF3F51B5U; // indigo: postWrite -> poll Done +inline constexpr std::uint32_t kNvtxAckWait = 0xFFE91E63U; // pink: DATA sent -> ACK +inline constexpr std::uint32_t kNvtxScatterQueue = 0xFFFFC107U; // dark yellow: enqueue -> worker dequeue +inline constexpr std::uint32_t kNvtxScatter = 0xFFFFEB3BU; // yellow: scatter kernel + sync +inline constexpr std::uint32_t kNvtxCreditStarved = 0xFFFF5722U; // deep orange: out of credits -> next GRANT +inline constexpr std::uint32_t kNvtxArenaStarved = 0xFF795548U; // brown: credits parked on local arena/exec + +// Fine-grained control-path spans decomposing ackWait (DATA sent -> ACK received). Together with +// the coarse spans these localize where the ACK round-trip actually goes: sender-side DATA +// build/enqueue, wire+peer time (the ackWait residual), receiver decode, scatter prep vs the real +// GPU wait, ACK enqueue, and the IO-thread bookkeeping drain. +inline constexpr std::uint32_t kNvtxDataSend = 0xFF00BCD4U; // cyan: build entries + encode + zmq enqueue of DATA +inline constexpr std::uint32_t kNvtxOnData = 0xFF009688U; // teal: DATA decode + scatter-job enqueue (receiver IO) +inline constexpr std::uint32_t kNvtxScatterPrep = 0xFFCDDC39U; // lime: scatter plan-array build + kernel launch +inline constexpr std::uint32_t kNvtxScatterSync = 0xFFB2A429U; // dark lime: cudaStreamSynchronize (the GPU wait) +inline constexpr std::uint32_t kNvtxAckSend = 0xFFF06292U; // light pink: encode + zmq enqueue of ACK (worker) +inline constexpr std::uint32_t kNvtxOnAck = 0xFFAD1457U; // dark pink: ACK dispatch incl. mReqMu wait (sender IO) +inline constexpr std::uint32_t kNvtxDoneDrain = 0xFF607D8BU; // blue gray: drainScatterDone region bookkeeping +inline constexpr std::uint32_t kNvtxZmqSend = 0xFF9C27B0U; // purple: zmq sendTo (lock + msg copy + enqueue) +inline constexpr std::uint32_t kNvtxZmqRecv = 0xFF673AB7U; // deep purple: zmq frame reads + blob copies + +/// Scoped range with a printf-formatted message, e.g. +/// `BounceNvtxScope s(kNvtxScatter, "scatter rid=%llu chunk=%u", rid, chunk);` +/// RAII (begins on construction, ends on destruction — early-exit safe), but implemented with +/// start/end ranges rather than push/pop: start/end is PROCESS-scoped, so these spans land in the +/// single `trtllm.disagg.bounce` domain row in nsys-ui together with the async spans, instead of +/// being scattered across per-thread rows (push/pop is thread-scoped and only shows under the +/// pushing thread, which for scatter workers is easy to miss). +class BounceNvtxScope +{ +public: + BounceNvtxScope(BounceNvtxScope const&) = delete; + BounceNvtxScope& operator=(BounceNvtxScope const&) = delete; + +#ifndef NVTX_DISABLE + BounceNvtxScope(std::uint32_t argb, char const* fmt, ...) + { + char msg[128]; + std::va_list args; + va_start(args, fmt); + std::vsnprintf(msg, sizeof(msg), fmt, args); + va_end(args); + mHandle = ::nvtx3::start_range_in(msg, ::nvtx3::color{argb}); + } + + ~BounceNvtxScope() + { + ::nvtx3::end_range_in(mHandle); + } + +private: + ::nvtx3::range_handle mHandle{}; +#else + BounceNvtxScope(std::uint32_t /*argb*/, char const* /*fmt*/, ...) {} +#endif +}; + +/// Start a cross-thread span. Returns an opaque handle (0 == no range, e.g. NVTX disabled); +/// end it — possibly on a different thread — with bounceRangeEnd(). +#ifndef NVTX_DISABLE +inline std::uint64_t bounceRangeStart(std::uint32_t argb, char const* fmt, ...) +{ + char msg[128]; + std::va_list args; + va_start(args, fmt); + std::vsnprintf(msg, sizeof(msg), fmt, args); + va_end(args); + return ::nvtx3::start_range_in(msg, ::nvtx3::color{argb}).get_value(); +} +#else +inline std::uint64_t bounceRangeStart(std::uint32_t /*argb*/, char const* /*fmt*/, ...) +{ + return 0; +} +#endif + +/// Label the calling thread in profiler timelines (e.g. "bounceIO", "bounceScatter"), so the +/// bounce threads are identifiable rows in nsys-ui instead of anonymous worker threads (the +/// names label the bounceIO/bounceScatter thread rows even though the spans themselves land in +/// the process-scoped domain row). +inline void bounceNameThread(char const* name) +{ +#ifndef NVTX_DISABLE + nvtxNameOsThreadA(static_cast(::syscall(SYS_gettid)), name); +#else + (void) name; +#endif +} + +/// End a span started by bounceRangeStart() and zero the handle. Safe on a 0 handle (no-op), +/// so failure paths can end every handle unconditionally. +inline void bounceRangeEnd(std::uint64_t& handle) +{ +#ifndef NVTX_DISABLE + if (handle != 0) + { + ::nvtx3::end_range_in(::nvtx3::range_handle{handle}); + } +#endif + handle = 0; +} + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferPlan.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferPlan.cpp new file mode 100644 index 000000000000..50d160b56aa6 --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferPlan.cpp @@ -0,0 +1,201 @@ +/* + * 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/BounceTransferPlan.h" + +#include "tensorrt_llm/common/assert.h" + +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +namespace +{ +// 32-byte alignment is enough for memory-coalesced vectorized copies and imposes no stricter +// requirement than the underlying registered memory already has. +constexpr std::uint64_t kAlignment = 32ULL; + +constexpr std::uint64_t alignUp(std::uint64_t value, std::uint64_t align) noexcept +{ + return (value + align - 1ULL) / align * align; +} + +// Build the chunk's coalesced scatter view (see BounceScatterRun). Greedy single pass; per desc, +// try to extend the last run in one of three ways before opening a new one: +// (a) contiguous growth (count==1): bounce AND dst both continue exactly where the run ends -> +// grow pieceSize in place. Captures a fully-dense dst (whole chunk -> ONE run). +// (b) stride latch (count==1, same size): the second desc fixes (dstStride, bounceStride) and the +// run becomes count=2. Only forward, u32-representable bounce steps latch. +// (c) stride extension (count>=2): the desc lands exactly one stride past the run's last piece. +// Captures a uniformly-strided dst (e.g. a head slice into a wider pool -> ONE run). +// Irregular layouts simply break runs (worst case: one count==1 run per desc == the old per-desc +// plan). Correctness never depends on merging. +void buildScatterRuns(BounceChunk& chunk) +{ + auto const n = chunk.dstPtrs.size(); + chunk.scatterRuns.clear(); + for (std::size_t i = 0; i < n; ++i) + { + std::uint64_t const dst = chunk.dstPtrs[i]; + std::uint64_t const bounce = chunk.bounceOffsets[i]; + std::uint32_t const size = chunk.sizes[i]; + if (!chunk.scatterRuns.empty()) + { + auto& r = chunk.scatterRuns.back(); + if (r.count == 1) + { + // (a) contiguous growth. Piece size stays within u32 (packedBytes <= 4 GiB - 1 is + // enforced in build(), but guard the sum explicitly anyway). + if (dst == r.dstAddr + r.pieceSize && bounce == r.bounceOffset + r.pieceSize + && static_cast(r.pieceSize) + size <= std::numeric_limits::max()) + { + r.pieceSize += size; + continue; + } + // (b) stride latch. + if (size == r.pieceSize && dst > r.dstAddr && bounce > r.bounceOffset + && bounce - r.bounceOffset <= std::numeric_limits::max()) + { + r.dstStride = dst - r.dstAddr; + r.bounceStride = static_cast(bounce - r.bounceOffset); + r.count = 2; + continue; + } + } + else if (size == r.pieceSize && dst == r.dstAddr + static_cast(r.count) * r.dstStride + && bounce == r.bounceOffset + static_cast(r.count) * r.bounceStride) + { + // (c) stride extension. + r.count += 1; + continue; + } + } + chunk.scatterRuns.push_back(BounceScatterRun{bounce, dst, 0, 0, size, 1}); + } +} +} // namespace + +BounceTransferPlan BounceTransferPlan::build(TransferDescs const& srcDescs, TransferDescs const& dstDescs, + std::size_t maxChunkSizeBytes, std::size_t maxDescsPerChunk) +{ + BounceTransferPlan plan; + + auto const& srcVec = srcDescs.getDescs(); + auto const& dstVec = dstDescs.getDescs(); + TLLM_CHECK_WITH_INFO(srcVec.size() == dstVec.size(), "BounceTransferPlan: src/dst desc count mismatch (%zu vs %zu)", + srcVec.size(), dstVec.size()); + TLLM_CHECK_WITH_INFO(maxChunkSizeBytes > 0 && maxDescsPerChunk > 0, + "BounceTransferPlan: maxChunkSizeBytes/maxDescsPerChunk must be > 0"); + // A chunk's packed size flows through 32-bit fields on the wire (Grant.len, WANT chunk sizes, + // scatter entry size, Posted.writeBytes), so a chunk must fit in 32 bits. Arena offsets are + // 64-bit (arena may exceed 4 GiB) but a single staging chunk above 4 GiB is nonsensical. + TLLM_CHECK_WITH_INFO(maxChunkSizeBytes <= std::numeric_limits::max(), + "BounceTransferPlan: maxChunkSizeBytes (%zu) must be <= 4 GiB (chunk size is 32-bit on the wire)", + maxChunkSizeBytes); + + if (srcVec.empty()) + { + return plan; // 0 descs -> 0 chunks + } + + auto const sourceDeviceId = srcVec.front().getDeviceId(); + auto const destinationDeviceId = dstVec.front().getDeviceId(); + BounceChunk current; + current.dstDeviceId = destinationDeviceId; + std::uint64_t cursor = 0; // running write offset within the current chunk region (aligned) + + auto flush = [&]() + { + if (!current.srcPtrs.empty()) + { + buildScatterRuns(current); + plan.mChunks.emplace_back(std::move(current)); + current = BounceChunk{}; + cursor = 0; + } + }; + + for (std::size_t i = 0; i < srcVec.size(); ++i) + { + auto const& src = srcVec[i]; + auto const& dst = dstVec[i]; + std::size_t const len = src.getLen(); + TLLM_CHECK_WITH_INFO(len == dst.getLen(), "BounceTransferPlan: src/dst len mismatch at idx %zu", i); + TLLM_CHECK_WITH_INFO(src.getDeviceId() == sourceDeviceId, + "BounceTransferPlan: mixed source device ids are unsupported (idx %zu has %u, expected %u)", i, + src.getDeviceId(), sourceDeviceId); + TLLM_CHECK_WITH_INFO(dst.getDeviceId() == destinationDeviceId, + "BounceTransferPlan: mixed destination device ids are unsupported (idx %zu has %u, expected %u)", i, + dst.getDeviceId(), destinationDeviceId); + TLLM_CHECK_WITH_INFO(len <= maxChunkSizeBytes, + "BounceTransferPlan: single desc (%zu B) exceeds maxChunkSizeBytes (%zu B)", len, maxChunkSizeBytes); + TLLM_CHECK_WITH_INFO(len < (1ULL << 32U), "BounceTransferPlan: single desc (%zu B) exceeds 4 GiB", len); + + // A zero-length desc carries no data; skip it so it never forces an empty chunk. + if (len == 0) + { + plan.mTotalDescs += 1; + continue; + } + + bool const overflow = (cursor + len > maxChunkSizeBytes); + bool const tooManyDescs = (current.srcPtrs.size() >= maxDescsPerChunk); + + // Extend the previous desc in place when src, dst AND the bounce cursor all advance + // contiguously (the aligned cursor left no gap): one desc instead of two shrinks the gather + // plan, the scatter runs and the wire messages. Only within the current chunk (`!overflow`) + // and staying within the u32 per-desc size field. + bool const srcDstContig = !current.srcPtrs.empty() && !overflow + && src.getAddr() == current.srcPtrs.back() + current.sizes.back() + && dst.getAddr() == current.dstPtrs.back() + current.sizes.back() + && cursor == current.bounceOffsets.back() + current.sizes.back() + && static_cast(current.sizes.back()) + len <= std::numeric_limits::max(); + if (srcDstContig) + { + current.sizes.back() += static_cast(len); + current.totalBytes += len; + current.packedBytes = current.bounceOffsets.back() + current.sizes.back(); + cursor = alignUp(current.packedBytes, kAlignment); + plan.mTotalBytes += len; + plan.mTotalDescs += 1; + continue; + } + + if (overflow || tooManyDescs) + { + flush(); + current.dstDeviceId = destinationDeviceId; + } + + current.srcPtrs.push_back(static_cast(src.getAddr())); + current.dstPtrs.push_back(static_cast(dst.getAddr())); + current.sizes.push_back(static_cast(len)); + current.bounceOffsets.push_back(cursor); + current.totalBytes += len; + current.packedBytes = cursor + len; // extent to transfer (this desc is the furthest so far) + cursor = alignUp(cursor + len, kAlignment); + + plan.mTotalBytes += len; + plan.mTotalDescs += 1; + } + flush(); + + return plan; +} + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferPlan.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferPlan.h new file mode 100644 index 000000000000..824efae77758 --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferPlan.h @@ -0,0 +1,96 @@ +/* + * 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/cache_transmission/nixl_utils/bounce/BounceMessage.h" +#include "tensorrt_llm/executor/transferAgent.h" + +#include +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +/// One chunk = the set of (src,dst) descriptors that get packed into a single bounce region +/// and moved with one RDMA write. Offsets are byte offsets within that region. +struct BounceChunk +{ + std::vector srcPtrs; // sender-local source addresses + std::vector dstPtrs; // receiver-local final destination addresses + std::vector sizes; // per-desc byte counts + std::vector bounceOffsets; // per-desc offset within the region + // Coalesced scatter view — what goes on the wire as the DATA scatter plan. Adjacent descs merge + // into one BounceScatterRun when (bounceOffset, dstPtr) advance contiguously (run grows its + // pieceSize; the fully-dense dst, e.g. ctx tp1 -> gen tp4) OR by a uniform stride (run grows its + // count; the strided dst, e.g. ctx tp-slice -> gen DP full-head pool). Either way thousands of + // per-desc entries collapse to a handful of runs, shrinking the DATA control message (which sits + // on the ACK critical path) by orders of magnitude. The per-desc arrays above stay as-is for the + // gather (src layout is independent of dst). + std::vector scatterRuns; + std::uint64_t totalBytes{0}; // sum of desc sizes (payload only, excludes padding) + std::uint64_t packedBytes{0}; // region extent to RDMA-write: last bounceOffset + its size + std::uint32_t dstDeviceId{0}; // receiver device id (uniform across the entire request) +}; + +/// Pure bin-packing of a TransferRequest's (src,dst) descriptor pairs into chunks that each +/// fit in one bounce region (<= maxChunkSizeBytes). No CUDA / NIXL / threads — trivially +/// unit-testable. +/// +/// Source descriptors must all use one device id, and destination descriptors must all use one +/// device id. Mixed-device requests are unsupported. +/// +/// Packing rules (a chunk is flushed when either holds): +/// - adding the next desc would exceed `maxChunkSizeBytes`, +/// - the chunk already holds `maxDescsPerChunk` descs. +class BounceTransferPlan +{ +public: + /// @param maxChunkSizeBytes Per-chunk byte cap (one transfer writes at most this many bytes). + /// @param maxDescsPerChunk Upper bound on descriptors per chunk (bounds scatter-plan size). + /// Throws (TLLM_CHECK) on src/dst count, length, or device-id mismatch, or when a single desc is + /// larger than maxChunkSizeBytes. + [[nodiscard]] static BounceTransferPlan build(TransferDescs const& srcDescs, TransferDescs const& dstDescs, + std::size_t maxChunkSizeBytes, std::size_t maxDescsPerChunk); + + [[nodiscard]] std::vector const& chunks() const noexcept + { + return mChunks; + } + + [[nodiscard]] std::size_t numChunks() const noexcept + { + return mChunks.size(); + } + + [[nodiscard]] std::uint64_t totalBytes() const noexcept + { + return mTotalBytes; + } + + [[nodiscard]] std::size_t totalDescs() const noexcept + { + return mTotalDescs; + } + +private: + std::vector mChunks; + std::uint64_t mTotalBytes{0}; + std::size_t mTotalDescs{0}; +}; + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.cpp new file mode 100644 index 000000000000..ce7b76e12d4e --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.cpp @@ -0,0 +1,1822 @@ +/* + * 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/BounceTransport.h" + +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceNvtx.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/GatherScatterKernel.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h" + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmException.h" + +#include +#include +#include +#include +#include + +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"; + } + return "bounce: unknown"; +} + +namespace +{ +constexpr char kSep = '\x1f'; // unit separator: agent names won't contain it + +// Split large copy runs into pieces of this size when building the batched-copy plan arrays. The +// copy kernel assigns ONE thread block per plan entry, so a plan of a few huge coalesced runs would +// use only a few SMs; splitting restores the grid-level parallelism the pre-coalescing per-desc plan +// had, without giving up the small wire messages. +constexpr std::uint32_t kCopySplitBytes = 64U << 10; + +// This entry's split budget: the scratch slots still free after reserving ONE slot for every raw +// entry not yet appended (each needs at least one). Never below 1. +std::size_t splitBudget(std::size_t appended, std::size_t rawRemaining, std::size_t maxEntries) +{ + std::size_t const reserved = appended + rawRemaining; + return maxEntries > reserved ? maxEntries - reserved : 1; +} + +// Number of pieces appendSplitInto() will emit for one raw entry under `maxPieces`: one per full +// kCopySplitBytes piece up to the budget, the remainder as a single (possibly oversized) entry. +// Used for the exact-count pass that lets callers write the plan arrays straight into the pinned +// buffer (the [srcs|dsts|sizes] layout needs the total BEFORE the first write). +std::size_t piecesFor(std::uint32_t size, std::size_t maxPieces) +{ + std::size_t const want = (static_cast(size) + kCopySplitBytes - 1) / kCopySplitBytes; + return std::min(std::max(want, 1), std::max(maxPieces, 1)); +} + +// Plan-array views into an exec context's pinned buffer: [srcs(n) | dsts(n) | sizes(n)] — the +// layout launchPrepared() consumes. Callers fill these IN PLACE: one write pass replaces the old +// build-std::vectors-then-memcpy-into-pinned flow (two full passes over the plan arrays per chunk). +struct PlanBufs +{ + std::uint64_t* srcs; + std::uint64_t* dsts; + std::uint32_t* sizes; +}; + +PlanBufs planBufs(ExecCtx* ctx, std::size_t n) +{ + auto* host = static_cast(ctx->hostPinned); + std::size_t const b64 = n * sizeof(std::uint64_t); + return {reinterpret_cast(host), reinterpret_cast(host + b64), + reinterpret_cast(host + 2 * b64)}; +} + +// Write (src, dst, size) into the plan buffers at `idx`, split into <= kCopySplitBytes pieces but +// at most `maxPieces` entries (>= 1) — when the budget runs out the remainder goes in as ONE +// oversized entry (the kernel's strided loop handles any size, so an unsplit entry only costs +// parallelism, never correctness). Emits exactly piecesFor(size, maxPieces) entries. +void appendSplitInto(PlanBufs const& bufs, std::size_t& idx, std::uint64_t src, std::uint64_t dst, std::uint32_t size, + std::size_t maxPieces) +{ + while (size > kCopySplitBytes && maxPieces > 1) + { + bufs.srcs[idx] = src; + bufs.dsts[idx] = dst; + bufs.sizes[idx] = kCopySplitBytes; + ++idx; + src += kCopySplitBytes; + dst += kCopySplitBytes; + size -= kCopySplitBytes; + --maxPieces; + } + bufs.srcs[idx] = src; + bufs.dsts[idx] = dst; + bufs.sizes[idx] = size; + ++idx; +} + +// Number of plan entries an exec context's scratch can hold ([srcs|dsts|sizes] packed). +std::size_t maxPlanEntries(ExecCtx const* ctx) +{ + return ctx->scratchBytes / (2 * sizeof(std::uint64_t) + sizeof(std::uint32_t)); +} + +std::string makeKey(std::string const& peer, std::uint64_t rid) +{ + return peer + kSep + std::to_string(rid); +} + +std::pair splitKey(std::string const& key) +{ + auto pos = key.rfind(kSep); + return {key.substr(0, pos), std::strtoull(key.c_str() + pos + 1, nullptr, 10)}; +} + +// Launch the batched copy over plan arrays the CALLER already wrote into the exec context's pinned +// host buffer (via planBufs()/appendSplitInto(): [srcs(n)|dsts(n)|sizes(n)]). Direction-agnostic +// (gather or scatter); the data region (arena offset) is encoded into the srcs/dsts by the caller. +// zeroCopy (default on) skips the H2D of the plan arrays — the kernel reads them straight from +// pinned host via ctx->hostPinnedDev, removing the H2D-then-kernel serialization. Off stages them +// in device scratch first (faster on machines where mapped-host reads are slow). +cudaError_t launchPrepared(ExecCtx* ctx, std::size_t n, bool zeroCopy) +{ + if (n == 0) + { + return cudaSuccess; + } + std::size_t const b64 = n * sizeof(std::uint64_t); + std::size_t const b32 = n * sizeof(std::uint32_t); + // The plan arrays must fit this context's scratch/hostPinned (both sized for maxDescsPerChunk). + // Callers bound n before writing (planBufs is capacity-unchecked), so this is defense in depth. + if (2 * b64 + b32 > ctx->scratchBytes) + { + return cudaErrorInvalidValue; + } + auto* host = static_cast(ctx->hostPinned); + // Pick the DEVICE-accessible base for the plan arrays: the pinned buffer's device alias (zeroCopy) + // or the device scratch we H2D-copy into. + std::uint8_t* base = nullptr; + if (zeroCopy && ctx->hostPinnedDev != nullptr) + { + base = static_cast(ctx->hostPinnedDev); + } + else + { + base = static_cast(ctx->scratch); + cudaError_t const st = cudaMemcpyAsync(base, host, 2 * b64 + b32, cudaMemcpyHostToDevice, ctx->stream); + if (st != cudaSuccess) + { + return st; + } + } + auto* dsrcs = reinterpret_cast(base); + auto* ddsts = reinterpret_cast(base + b64); + auto* dsizes = reinterpret_cast(base + 2 * b64); + auto const n32 = static_cast(n); // n <= scratchBytes/20, far below u32 max + return launchBatchedCopy(dsrcs, ddsts, dsizes, n32, ctx->stream); +} +} // namespace + +// ============================================================================ +// BounceContext +// ============================================================================ + +void BounceContext::sendGrants(std::vector const& grants) +{ + if (grants.empty()) + { + return; + } + // grants are keyed by flow id ("peerrid"); split back to (agent name, rid) to address them. + std::unordered_map> byFlow; + for (auto const& g : grants) + { + // Carry OUR (receiver) device id so the sender writes the remote desc to the right GPU + // even if the two agents don't share a device index. `regionHandle` (our arena offset) is + // echoed back in DATA so we can locate + free the region. + byFlow[g.flow].push_back(BounceCreditEntry{g.addr, g.len, static_cast(deviceId), g.offset}); + } + for (auto const& [flow, creds] : byFlow) + { + auto [peer, rid] = splitKey(flow); + channel->sendTo(peer, encodeGrant(rid, creds)); + } +} + +// ============================================================================ +// BounceReceiver — [R] role +// ============================================================================ + +BounceReceiver::BounceReceiver(BounceContext& ctx) + : mCtx(ctx) +{ +} + +void BounceReceiver::startWorkers() +{ + std::uint32_t const workers = mCtx.cfg.scatterWorkerCount > 0 ? mCtx.cfg.scatterWorkerCount : 1; + mWorkers.reserve(workers); + for (std::uint32_t i = 0; i < workers; ++i) + { + mWorkers.emplace_back(&BounceReceiver::scatterWorkerLoop, this); + } +} + +void BounceReceiver::wake() +{ + // Take mJobMu so the notify is ordered against a worker's predicate check. The wait predicate + // reads mCtx.stop, which shutdown() sets WITHOUT holding mJobMu — a naked notify_all can then + // fire in the window between a worker evaluating the predicate false and parking, and is lost + // forever (no later notifier exists once the IO thread is joined), hanging joinWorkers(). + std::lock_guard lk(mJobMu); + mJobCv.notify_all(); +} + +void BounceReceiver::joinWorkers() +{ + for (auto& t : mWorkers) + { + if (t.joinable()) + { + t.join(); + } + } + // Workers are gone; any never-dequeued job still holds an open queue-wait span — close them so + // shutdown doesn't leave dangling NVTX ranges (their regions die with the arena right after). + std::lock_guard lk(mJobMu); + for (auto& j : mJobs) + { + bounceRangeEnd(j.nvtxQueue); + } +} + +void BounceReceiver::onWant(std::string const& peer, BounceMsgHeader const& h, std::string const& blob) +{ + std::vector chunkBytes; + std::string endpoint; + if (!decodeWant(blob, h, chunkBytes, endpoint)) + { + return; + } + // Self-bootstrap the reverse control path. Disaggregated serving supports one-directional + // metadata exchange: the KV sender may load our agent metadata without us loading the sender's. + // In that case we have no reverse route for GRANT/ACK, so register the sender endpoint carried + // by WANT here (addPeer is idempotent). A cancel is still honored when registration fails so it + // can reclaim any flow state left by an earlier, valid WANT. + bool reversePathReady = false; + try + { + reversePathReady = mCtx.channel->addPeer(peer, endpoint); + } + catch (tensorrt_llm::common::TllmException const& e) + { + // A malformed WANT is peer input on the reactor thread. Reject it without allowing an + // exception to escape the thread; handshake registration still propagates this error to + // the caller because a ZMQ endpoint is mandatory there. + TLLM_LOG_WARNING( + "BounceTransport(%s): rejected WANT from peer %s: %s", mCtx.selfName.c_str(), peer.c_str(), e.what()); + } + auto const key = makeKey(peer, h.requestId); + if (isCancelWant(chunkBytes)) + { + // Explicit cancel/abort (the sender failed or retracted): precisely free this flow's + // granted-but-unwritten regions now — otherwise they stay held until peer loss and a + // long-running receiver leaks up to one request's in-flight allocation cap per failed rid. + // Any region whose scatter is still running is deferred (flagged orphaned in mScattering, + // freed on completion). + // Immediate free (no quarantine) is safe HERE because a cancel is sender-initiated: the + // sender defers it until its last in-flight RDMA write reached a terminal state + // (mPendingCancel / drainOrphanLocal), so nothing can still be writing these regions. + reclaimAndFlagDeferred( + [&](auto const& busy, auto& deferred) { return mCtx.scheduler.reclaimFlow(key, busy, deferred); }); + return; + } + if (!reversePathReady) + { + return; + } + // Chunk sizes are untrusted peer input. The capability handshake pins both sides to the same + // effective maxChunkSizeBytes (already clamped to usable arena capacity in the ctor), so a + // compliant sender never emits a chunk of 0 or above that cap. An ungrantable size must not + // reach the scheduler: BuddyAllocator::alloc() can never satisfy it, the flow would stay + // blocked forever, and once maybeActivateDrain() latches it as the drain flow every grant to + // every peer stops — with no reclaim path, because a pending-only flow holds no regions and + // the lease sweep skips it. + for (auto const bytes : chunkBytes) + { + if (bytes == 0 || bytes > mCtx.cfg.maxChunkSizeBytes) + { + TLLM_LOG_WARNING("BounceTransport(%s): rejected WANT from peer %s rid=%llu: chunk size %u outside (0, %zu]", + mCtx.selfName.c_str(), peer.c_str(), static_cast(h.requestId), bytes, + static_cast(mCtx.cfg.maxChunkSizeBytes)); + return; + } + } + // A non-empty WANT has no retransmission path (submit() sends it exactly once per fresh rid), + // so one for an already-tracked flow is a replay or rid collision. Re-queueing would re-grant + // over the still-held regions — the sender never writes the extras, so they leak — and the + // lease refresh inside onWant would keep the flow forever off staleFlows(), defeating the + // reclaim path that exists for exactly this state. Drop it, mirroring the duplicate-DATA and + // stale-ACK handling. (Check-then-act is safe: all flow-state mutation happens on this IO + // thread; app threads only acquireLocal().) + if (mCtx.scheduler.knowsFlow(key)) + { + TLLM_LOG_WARNING("BounceTransport(%s): dropped duplicate WANT from peer %s rid=%llu", mCtx.selfName.c_str(), + peer.c_str(), static_cast(h.requestId)); + return; + } + mCtx.sendGrants(mCtx.scheduler.onWant(key, chunkBytes)); +} + +void BounceReceiver::onData(std::string const& peer, BounceMsgHeader const& h, std::string const& blob) +{ + // Covers DATA decode + scatter-job enqueue on the receiver IO thread — the receiver-side + // software leg of the sender's ackWait (the message copy cost scales with the entry count). + BounceNvtxScope onDataScope(kNvtxOnData, "onData rid=%llu chunk=%u bytes=%zu", + static_cast(h.requestId), h.chunkIdx, blob.size()); + std::vector entries; + if (!decodeScatter(blob, h, entries)) + { + return; + } + auto const key = makeKey(peer, h.requestId); + // Drop a DATA whose region this flow no longer holds — it was cancelled/reclaimed (e.g. an empty + // WANT raced ahead of this DATA), and the region may have been re-granted to another flow. + // Scattering it would read a freed/re-owned region and corrupt that other flow's data. + if (!mCtx.scheduler.heldByFlow(key, h.regionHandle)) + { + return; + } + ScatterJob job; + job.key = key; + job.peer = peer; + job.rid = h.requestId; + job.chunkIdx = h.chunkIdx; + job.offset = h.regionHandle; + // Capture the granted region's byte size HERE (IO thread) — the scatter worker cannot query the + // IO-thread-only scheduler. Used to bound scatter reads to THIS flow's region (below). + job.regionBytes = mCtx.scheduler.regionBytes(h.regionHandle); + job.entries = std::move(entries); + job.nvtxQueue = bounceRangeStart( + kNvtxScatterQueue, "scatterQueue rid=%llu chunk=%u", static_cast(h.requestId), h.chunkIdx); + // DATA is not retransmitted by this protocol. An occupied region therefore indicates stale or + // malformed input; drop it instead of adding replay state or risking two readers of one grant. + if (!mScattering.emplace(job.offset, false).second) + { + TLLM_LOG_WARNING("BounceTransport(%s): dropping duplicate DATA peer=%s rid=%llu chunk=%u region=%llu", + mCtx.selfName.c_str(), peer.c_str(), static_cast(h.requestId), h.chunkIdx, + static_cast(h.regionHandle)); + bounceRangeEnd(job.nvtxQueue); + return; + } + { + std::lock_guard lk(mJobMu); + mJobs.emplace_back(std::move(job)); + } + mJobCv.notify_one(); +} + +std::unordered_set BounceReceiver::scatteringRegions() const +{ + std::unordered_set busy; + busy.reserve(mScattering.size()); + for (auto const& [off, orphaned] : mScattering) + { + busy.insert(off); + } + return busy; +} + +void BounceReceiver::reclaimAndFlagDeferred( + std::function(std::unordered_set const&, std::vector&)> const& + reclaim) +{ + std::vector deferred; + mCtx.sendGrants(reclaim(scatteringRegions(), deferred)); + for (auto off : deferred) + { + mScattering[off] = true; + } +} + +bool BounceReceiver::drainScatterDone() +{ + std::deque done; + { + std::lock_guard lk(mDoneMu); + done.swap(mDone); + } + if (done.empty()) + { + return false; + } + // Bookkeeping only (the ACK itself was sent by the worker): region frees + re-grants. + BounceNvtxScope drainScope(kNvtxDoneDrain, "doneDrain n=%zu", done.size()); + for (auto& d : done) + { + // The ACK was already sent by the scatter worker itself (latency: it is on the sender's + // ackWait critical path); only the region bookkeeping happens here, on the IO thread. + // Worker finished reading this region. Was its flow reclaimed (peer gone / cancel) mid-scatter? + auto it = mScattering.find(d.offset); + bool const orphaned = (it != mScattering.end()) && it->second; + if (it != mScattering.end()) + { + mScattering.erase(it); + } + if (orphaned) + { + // Flow was reclaimed while this scatter ran; the region was kept out of the arena so it + // couldn't be re-granted under the worker. Now it's safe to free. + mCtx.sendGrants(mCtx.scheduler.freeOrphanRegion(d.offset)); + } + else + { + mCtx.sendGrants(mCtx.scheduler.onScatterDone(d.key, d.offset)); + } + } + return true; +} + +void BounceReceiver::forget(std::string const& peer) +{ + // Reclaim every flow this peer was granted ("peer\x1f rid"). + // (1) Drop this peer's not-yet-started scatter jobs — no point scattering for a gone peer; their + // incoming regions are no longer busy and get freed by the reclaim below. + { + std::lock_guard lk(mJobMu); + std::deque keep; + for (auto& j : mJobs) + { + if (j.peer == peer) + { + bounceRangeEnd(j.nvtxQueue); // job dropped, close its queue-wait span + mScattering.erase(j.offset); + } + else + { + keep.push_back(std::move(j)); + } + } + mJobs.swap(keep); + } + // (2) Regions of this peer still scattering are reads already RUNNING in a worker — they must not + // be re-granted until the worker finishes. reclaimByPrefix defers those; we flag them orphaned + // in mScattering so drainScatterDone frees them via freeOrphanRegion on completion. + // (3) Every OTHER held region may STILL be RDMA-written by the peer: forget() is + // receiver-initiated (invalidateRemoteAgent), so — unlike an explicit cancel — no + // sender-side drain guarantees those writes ended, and a one-sided write cannot be aborted. + // quarantineFor > 0 keeps them out of the arena until checkTimeouts()'s reapQuarantine. + reclaimAndFlagDeferred( + [&](auto const& busy, auto& deferred) + { + return mCtx.scheduler.reclaimByPrefix( + peer + kSep, busy, deferred, std::chrono::milliseconds(std::max(0, mCtx.cfg.quarantineMs))); + }); +} + +void BounceReceiver::checkTimeouts() +{ + auto const now = std::chrono::steady_clock::now(); + if (now < mNextSweep) + { + return; // throttled: no need to scan on every ~1ms tick + } + // Sweep granularity follows the smallest ENABLED timeout (a tenth of it, clamped to + // [50ms, 1s]) instead of a fixed constant: the 60s/30s defaults yield a 1s sweep, while a + // config that shrinks the timeouts to sub-second (tests, debugging) automatically gets a + // proportionally finer sweep — no hidden floor under the configured values. + int smallestMs = mCtx.cfg.receiverFlowTimeoutMs > 0 ? mCtx.cfg.receiverFlowTimeoutMs : 0; + if (mCtx.cfg.quarantineMs > 0) + { + smallestMs = smallestMs > 0 ? std::min(smallestMs, mCtx.cfg.quarantineMs) : mCtx.cfg.quarantineMs; + } + if (smallestMs <= 0) + { + // Lease disabled AND quarantine disabled: nothing is ever quarantined (reclaims free + // immediately) and no flow can go stale — nothing to sweep for. + mNextSweep = now + std::chrono::seconds(1); + return; + } + mNextSweep = now + + std::clamp( + std::chrono::milliseconds(smallestMs / 10), std::chrono::milliseconds(50), std::chrono::milliseconds(1000)); + // Quarantine over for any expired region: no write posted before its flow's lease expired can + // plausibly still be in flight, so it may re-enter circulation (and may grant a waiter). + mCtx.sendGrants(mCtx.scheduler.reapQuarantine()); + if (mCtx.cfg.receiverFlowTimeoutMs <= 0) + { + return; // lease disabled (mirror of requestTimeoutMs <= 0 on the sender) + } + for (auto const& flow : mCtx.scheduler.staleFlows(std::chrono::milliseconds(mCtx.cfg.receiverFlowTimeoutMs))) + { + // The flow's sender went silent after taking grants: it is dead or unreachable (a LIVE + // sender either makes progress or abandons via requestTimeoutMs + cancel well before this + // lease expires). Reclaim the whole flow — regions a worker still reads defer as orphans, + // all others quarantine (the silent peer's NIC may still be writing them). + auto const [peer, rid] = splitKey(flow); + TLLM_LOG_WARNING( + "BounceTransport(%s): flow lease expired (no progress within %d ms) peer=%s rid=%llu -> reclaiming " + "(regions quarantined for %d ms before reuse)", + mCtx.selfName.c_str(), mCtx.cfg.receiverFlowTimeoutMs, peer.c_str(), static_cast(rid), + mCtx.cfg.quarantineMs); + reclaimAndFlagDeferred( + [&](auto const& busy, auto& deferred) + { + return mCtx.scheduler.reclaimFlow( + flow, busy, deferred, std::chrono::milliseconds(std::max(0, mCtx.cfg.quarantineMs))); + }); + } +} + +void BounceReceiver::scatterWorkerLoop() +{ + bounceNameThread("bounceScatter"); + // Pin this worker to our device. Can't throw out of a thread fn -> warn-only (the loop's CUDA ops + // would then target the wrong device, so this is a real fault, just non-recoverable here). + TLLM_CUDA_CHECK_WARN(cudaSetDevice(mCtx.deviceId)); + while (true) + { + ScatterJob job; + { + std::unique_lock lk(mJobMu); + mJobCv.wait(lk, [this] { return !mJobs.empty() || mCtx.stop.load(std::memory_order_acquire); }); + if (mJobs.empty()) + { + if (mCtx.stop.load(std::memory_order_acquire)) + { + break; + } + continue; + } + job = std::move(mJobs.front()); + mJobs.pop_front(); + } + bounceRangeEnd(job.nvtxQueue); // dequeued: the queue-wait leg ends here + // Covers exec-context acquire + scatter launch + stream sync (the scatter's real GPU wait). + BounceNvtxScope scatterScope(kNvtxScatter, "scatter rid=%llu chunk=%u n=%zu", + static_cast(job.rid), job.chunkIdx, job.entries.size()); + // Borrow an exec context (stream/scratch) for this scatter. The arena region (job.offset) is + // held by the scheduler until ACK; the exec context is needed only while the kernel runs, so + // it comes from the small shared pool. If all are busy, briefly retry (backpressure, never + // deadlock — senders release contexts independently of local scatter progress). + ExecCtx* ctx = nullptr; + while ((ctx = mCtx.exec->tryAcquire()) == nullptr) + { + if (mCtx.stop.load(std::memory_order_acquire)) + { + break; + } + std::this_thread::sleep_for(std::chrono::microseconds(100)); + } + if (ctx == nullptr) + { + break; // shutting down + } + auto const n = static_cast(job.entries.size()); + // Validate every scatter SOURCE stays inside THIS flow's granted region before launching. + // bounceOffset/size come from the peer's DATA message; a buggy/hostile peer (or a reordered + // GRANT) could point them past the region and, if we only bounded against the whole arena, + // read from an ADJACENT flow's region and copy its bytes into our KV — silent cross-flow + // corruption. The region is one buddy block [regionBase, regionBase+regionBytes) owned solely + // by this flow, so bounding to it prevents any cross-flow read. (dstAddr is the caller's own KV + // target by design, so it isn't bounded here.) Any bad entry -> skip launch, no ACK (sender + // times out). regionBytes==0 means the region wasn't allocated (stale) -> reject the whole job. + std::uint64_t const arenaLo = mCtx.arena->baseAddr(); + auto const regionBase = arenaLo + job.offset; + bool srcInBounds = (job.regionBytes > 0 && regionBase + job.regionBytes <= arenaLo + mCtx.arena->bytes()); + // Scatter into the final dst, then wait so the data is at dst before we ACK. A bad source, a + // failed launch, OR a stream error must NOT produce an ACK — an ACK tells the sender its KV + // data landed, so a false ACK here is silent corruption. On error the worker sends no ACK, + // but the done queue still releases the region; the sender then times out -> FAILURE. + cudaError_t launchErr = cudaErrorInvalidValue; + { + // Prep leg: bounds-check + plan-array build + kernel launch (host-side cost). + BounceNvtxScope prepScope(kNvtxScatterPrep, "scatterPrep rid=%llu chunk=%u n=%u", + static_cast(job.rid), job.chunkIdx, n); + // Entries arrive as COALESCED runs (contiguous or strided; see BounceScatterRun); expand + // them back to <= kCopySplitBytes pieces so the copy kernel keeps its grid-level + // parallelism. Exact-count pass first (it also validates every run stays inside THIS + // flow's granted region), then fill the pinned plan buffers DIRECTLY (no intermediate + // vectors). Piece counts come from the peer's DATA message — a peer built with a larger + // maxChunkSizeBytes can carry more pieces than our pinned/scratch hold; reject rather than + // overflow (no launch -> no ACK -> the sender times out, never a false ACK). + std::size_t const maxEntries = maxPlanEntries(ctx); + std::uint64_t rawPieces = 0; + for (std::uint32_t i = 0; i < n; ++i) + { + auto const& e = job.entries[i]; + // Run-level source bounds: every piece p reads region[bounceOffset + p*bounceStride + // .. +pieceSize). count-1 and bounceStride are both u32 so the span product cannot + // overflow u64. A count of 0 is malformed (a run always carries >= 1 piece). + std::uint64_t const span = static_cast(e.count - 1) * e.bounceStride + e.pieceSize; + srcInBounds = srcInBounds && e.count >= 1 && e.bounceOffset <= job.regionBytes + && span <= job.regionBytes - e.bounceOffset; + rawPieces += std::max(e.count, 1); + } + // Reject an oversized run list BEFORE the exact-count pass below: that pass iterates + // once per PIECE, so a hostile/corrupt DATA (per-run count near 2^32 passes the span + // check when bounceStride is 0) could otherwise pin this worker — and its region — + // for days. piecesFor() emits at least one entry per piece, so rawPieces > maxEntries + // implies the nTotal <= maxEntries check would reject the job anyway. + if (srcInBounds && rawPieces > maxEntries) + { + TLLM_LOG_WARNING("BounceTransport(%s): rejected scatter with %llu pieces (max %zu) rid=%llu chunk=%u", + mCtx.selfName.c_str(), static_cast(rawPieces), maxEntries, + static_cast(job.rid), job.chunkIdx); + srcInBounds = false; + } + std::size_t nTotal = 0; + std::uint64_t seen = 0; + for (std::uint32_t i = 0; i < n && srcInBounds; ++i) + { + auto const& e = job.entries[i]; + for (std::uint32_t p = 0; p < e.count; ++p) + { + ++seen; + nTotal += piecesFor( + e.pieceSize, splitBudget(nTotal, static_cast(rawPieces - seen), maxEntries)); + } + } + if (srcInBounds && nTotal > 0 && nTotal <= maxEntries) + { + auto const bufs = planBufs(ctx, nTotal); + std::size_t idx = 0; + seen = 0; + for (std::uint32_t i = 0; i < n; ++i) + { + auto const& e = job.entries[i]; + for (std::uint32_t p = 0; p < e.count; ++p) + { + ++seen; + appendSplitInto(bufs, idx, + regionBase + e.bounceOffset + static_cast(p) * e.bounceStride, + e.dstAddr + static_cast(p) * e.dstStride, e.pieceSize, + splitBudget(idx, static_cast(rawPieces - seen), maxEntries)); + } + } + launchErr = launchPrepared(ctx, nTotal, mCtx.cfg.useZeroCopyArguments); + } + else if (srcInBounds && nTotal == 0) + { + launchErr = cudaSuccess; // empty plan: nothing to scatter (0 runs) -> vacuous success + } + } + cudaError_t syncErr = cudaSuccess; + if (launchErr == cudaSuccess) + { + // Sync leg: the actual GPU wait (kernel queueing + run time on the exec stream). + BounceNvtxScope syncScope(kNvtxScatterSync, "scatterSync rid=%llu chunk=%u", + static_cast(job.rid), job.chunkIdx); + syncErr = cudaStreamSynchronize(ctx->stream); + } + bool const ok = srcInBounds && launchErr == cudaSuccess && syncErr == cudaSuccess; + if (!ok) + { + (void) cudaGetLastError(); // clear sticky error so the reused context isn't poisoned + TLLM_LOG_WARNING( + "BounceTransport(%s): scatter failed (srcInBounds=%d launch=%d sync=%d) rid=%llu chunk=%u -> no ACK", + mCtx.selfName.c_str(), static_cast(srcInBounds), static_cast(launchErr), + static_cast(syncErr), static_cast(job.rid), job.chunkIdx); + } + mCtx.exec->release(ctx); // kernel done (or failed) -> return the context + // ACK straight from the worker (ControlChannel::sendTo is thread-safe): the data IS at its + // final dst here, and skipping the done-queue -> IO-thread hop shaves its drain latency off + // the sender's ackWait critical path. Region bookkeeping (scheduler free / re-grant) still + // goes through drainScatterDone on the IO thread. A failed scatter sends NO ACK — a false + // ACK would tell the sender corrupt/absent data landed; it must time out instead. + if (ok) + { + BounceNvtxScope ackScope( + kNvtxAckSend, "ackSend rid=%llu chunk=%u", static_cast(job.rid), job.chunkIdx); + mCtx.channel->sendTo(job.peer, encodeAck(job.rid, job.chunkIdx, job.offset)); + } + { + std::lock_guard lk(mDoneMu); + mDone.push_back(ScatterDone{job.key, job.offset}); + } + } +} + +// ============================================================================ +// BounceSender — [S] role +// ============================================================================ + +BounceSender::BounceSender(BounceContext& ctx) + : mCtx(ctx) +{ +} + +std::shared_future BounceSender::submit( + TransferDescs const& srcDescs, TransferDescs const& dstDescs, std::string const& peer) +{ + auto promise = std::make_shared>(); + auto fut = promise->get_future().share(); + BounceTransferPlan plan; + try + { + BounceNvtxScope planScope(kNvtxBuildPlan, "buildPlan nDesc=%zu", srcDescs.getDescs().size()); + plan = BounceTransferPlan::build(srcDescs, dstDescs, mCtx.cfg.maxChunkSizeBytes, + std::max(1024ULL, mCtx.cfg.maxChunkSizeBytes / 256ULL)); + } + catch (std::exception const& e) + { + // The eligibility gate (shouldUseBounce) screens the plan's preconditions, so this is a + // should-not-happen defense: resolve the future to FAILURE instead of letting the exception + // unwind out of submit(). Bounce admission is final for a request — a failure here (or later + // in the protocol) fails the transfer; there is deliberately no automatic re-submit on the + // standard NIXL path, so a bounce-layer fault surfaces to the caller instead of being + // silently absorbed as a slow success. Whether to retry is the caller's decision. + TLLM_LOG_WARNING( + "BounceTransport(%s): plan build for peer %s rejected: %s", mCtx.selfName.c_str(), peer.c_str(), e.what()); + promise->set_value({TransferState::kFAILURE, BounceFailReason::kPlanRejected}); + return fut; + } + auto const numChunks = static_cast(plan.numChunks()); + if (numChunks == 0) + { + promise->set_value({TransferState::kSUCCESS, BounceFailReason::kNone}); + return fut; + } + + // Per-chunk packed byte sizes: the receiver allocates a region of each size as it grants, so the + // WANT both announces how many chunks we have and how big each one's bounce region must be. + std::vector chunkBytes(numChunks); + for (std::uint32_t i = 0; i < numChunks; ++i) + { + chunkBytes[i] = static_cast(plan.chunks()[i].packedBytes); + } + + std::uint64_t const rid = mNextRid.fetch_add(1, std::memory_order_relaxed); + std::uint64_t const planBytes = plan.totalBytes(); + { + std::lock_guard lk(mReqMu); + Request req; + req.peer = peer; + req.numChunks = numChunks; + req.plan = std::move(plan); + req.promise = promise; + req.lastProgress = std::chrono::steady_clock::now(); + req.nvtxReq = bounceRangeStart(kNvtxRequest, "req rid=%llu chunks=%u bytes=%llu", + static_cast(rid), numChunks, static_cast(planBytes)); + // Ends at the FIRST GRANT (onGrant) — the credit-wait leg of the request. + req.nvtxGrantWait + = bounceRangeStart(kNvtxGrantWait, "grantWait rid=%llu", static_cast(rid)); + mRequests.emplace(rid, std::move(req)); + } + // Ask the receiver to grant a region for each chunk of this request flow. The WANT carries our + // own control endpoint so the receiver can addPeer us and send GRANT/ACK back (self-bootstrap). + mCtx.channel->sendTo(peer, encodeWant(rid, chunkBytes, mCtx.channel->localEndpoint())); + if (mCtx.cfg.enableEagerGather) + { + // Overlap the WANT->GRANT control round-trip with the gather: launch this request's first + // chunks NOW instead of waiting for the GRANT (they were measured back-to-back at roughly + // the same duration, so eager gather hides one of the two). Running the launch on the + // caller's thread also keeps its prep cost off the IO thread. The exec streams live on our + // device but the caller thread's current device is not guaranteed — pin it first (warn-only: + // on failure the pump's CUDA calls fail and the request degrades to the classic GRANT path + // or a deterministic failure, never a hang). + TLLM_CUDA_CHECK_WARN(cudaSetDevice(mCtx.deviceId)); + std::lock_guard lk(mReqMu); + auto it = mRequests.find(rid); + if (it != mRequests.end()) // a racing GRANT may have already pumped (or failed) the request + { + pumpRequest(rid, it->second); + } + } + return fut; +} + +void BounceSender::onGrant(std::string const& peer, BounceMsgHeader const& h, std::string const& blob) +{ + std::vector credits; + if (!decodeCredits(blob, h, credits)) + { + return; + } + std::lock_guard lk(mReqMu); + auto it = mRequests.find(h.requestId); + if (it == mRequests.end()) + { + return; // late grant for a finished/cancelled request + } + Request& req = it->second; + // A credit contains a peer-owned address. Never let an unrelated peer redirect this request's + // RDMA write, even if it guesses the request id. + if (peer != req.peer) + { + TLLM_LOG_WARNING("BounceTransport(%s): dropping wrong-peer GRANT peer=%s expected=%s rid=%llu", + mCtx.selfName.c_str(), peer.c_str(), req.peer.c_str(), static_cast(h.requestId)); + return; + } + bounceRangeEnd(req.nvtxGrantWait); // first GRANT ends the credit-wait span (no-op on later GRANTs) + bounceRangeEnd(req.nvtxCreditStarved); // a GRANT ends the current starvation period (pump may reopen one) + for (auto const& credit : credits) + { + req.pendingCredits.push_back(credit); + } + pumpRequest(h.requestId, req); +} + +bool BounceSender::abandonOnCreditMispair( + std::uint64_t rid, Request& req, std::uint32_t chunkIdx, std::uint64_t packedBytes, std::uint32_t creditLen) +{ + // Credits pair with chunks by FIFO order, and the receiver sizes each granted region to + // chunkBytes[chunkIdx] in WANT, so packedBytes always fits when the channel honors its FIFO + // contract. A malformed or misordered GRANT would make us RDMA-write packedBytes into a smaller + // region, overflowing into an adjacent flow's region on the peer. Detect that protocol violation + // and abandon the flow (it then fails via checkTimeouts) rather than corrupt the peer. + if (packedBytes <= creditLen) + { + return false; + } + TLLM_LOG_WARNING( + "BounceTransport(%s): rid=%llu chunk=%u packedBytes=%zu > granted region len=%u (GRANT " + "mispair/reorder); abandoning flow", + mCtx.selfName.c_str(), static_cast(rid), chunkIdx, static_cast(packedBytes), + static_cast(creditLen)); + req.abandonReason = BounceFailReason::kProtocolError; + req.pendingCredits.clear(); + return true; +} + +void BounceSender::attachCredits(std::uint64_t rid, Request& req) +{ + // Credits pair with chunks strictly in order (the receiver serves the WANT size list FIFO), so + // pendingCredits.front() is always chunk `nextCredit`'s credit. Attach parked credits to chunks + // that eager gather already posted credit-less; chunks not yet posted keep their credit parked + // for pumpRequest to consume at gather-launch time. + while (!req.pendingCredits.empty() && req.nextCredit < req.nextPost) + { + BounceCreditEntry const& credit = req.pendingCredits.front(); + Posted* target = nullptr; + for (auto& p : req.posted) + { + if (p.chunkIdx == req.nextCredit) + { + target = &p; + break; + } + } + if (target == nullptr || target->hasCredit) + { + // A posted chunk is only erased on ACK, and it can only be ACKed after consuming its + // credit (nextCredit already passed it) — so a missing or already-credited target is a + // protocol anomaly (dup GRANT after reconnect). Drop the credit, never mispair it. + TLLM_LOG_WARNING("BounceTransport(%s): rid=%llu chunk=%u unexpected credit (dup GRANT?); dropping", + mCtx.selfName.c_str(), static_cast(rid), req.nextCredit); + req.pendingCredits.pop_front(); + req.nextCredit += 1; + continue; + } + auto const& chunk = req.plan.chunks()[target->chunkIdx]; + if (abandonOnCreditMispair(rid, req, target->chunkIdx, chunk.packedBytes, credit.len)) + { + return; + } + target->remoteHandle = credit.regionHandle; + target->remoteAddr = credit.addr; + target->remoteDevId = credit.devId; + target->hasCredit = true; + // Now credit-backed: its staging region stops counting against the eager budget. + mCtx.scheduler.promoteLocal(target->localOffset); + req.pendingCredits.pop_front(); + req.nextCredit += 1; + req.lastProgress = std::chrono::steady_clock::now(); // forward progress: a chunk got its credit + } + if (req.nextCredit >= req.numChunks && !req.pendingCredits.empty()) + { + // Over-grant: receiver handed more credits than we have chunks. Shouldn't happen under + // the protocol (receiver grants at most numChunks). Log it — silently dropping would + // mask an upstream bug and leak the receiver-side regions backing these credits. + TLLM_LOG_WARNING("BounceTransport(%s): rid=%llu over-grant, dropping %zu extra credit(s)", + mCtx.selfName.c_str(), static_cast(rid), req.pendingCredits.size()); + req.pendingCredits.clear(); + } +} + +void BounceSender::pumpRequest(std::uint64_t rid, Request& req) +{ + attachCredits(rid, req); + while (req.nextPost < req.numChunks) + { + std::uint32_t const chunkIdx = req.nextPost; + auto const& chunk = req.plan.chunks()[chunkIdx]; + // After attachCredits, a non-empty pendingCredits implies nextCredit == nextPost, i.e. the + // front credit is exactly THIS chunk's credit. + bool const haveCredit = !req.pendingCredits.empty(); + if (!haveCredit) + { + if (!mCtx.cfg.enableEagerGather) + { + break; // classic path: a chunk's gather starts only once its GRANT arrived + } + if (req.posted.size() >= mCtx.cfg.maxInflightChunksPerRequest) + { + break; // cap eager gathers by this request's configured in-flight chunk limit + } + } + // Defensive pairing check BEFORE committing resources (see abandonOnCreditMispair). + if (haveCredit && abandonOnCreditMispair(rid, req, chunkIdx, chunk.packedBytes, req.pendingCredits.front().len)) + { + break; + } + // Non-blocking: borrow an exec context (cheap to return), then a gather-staging region sized + // to this chunk from the SHARED arena. If either is unavailable, leave the credit parked and + // bail; the IO loop retries via drainPendingPosts() once an ACK frees a region / context — + // never blocks here, so an oversubscribed arena (many peers, or both roles) degrades to + // backpressure, not deadlock. Credit-less (eager) staging is additionally capped by the + // scheduler's eager budget (half the arena) so it can never starve incoming grants. + ExecCtx* ctx = mCtx.exec->tryAcquire(); + if (ctx == nullptr) + { + break; + } + auto localOff = mCtx.scheduler.acquireLocal(chunk.packedBytes, /*eager=*/!haveCredit); + if (!localOff) + { + mCtx.exec->release(ctx); + break; + } + BounceCreditEntry credit{}; + if (haveCredit) + { + credit = req.pendingCredits.front(); + req.pendingCredits.pop_front(); + req.nextCredit += 1; + } + auto const nDesc = static_cast(chunk.srcPtrs.size()); + // Covers plan-array prep + gather launch + event record (the synchronous launch cost; + // the gather's GPU time is the async `gather` span ended in drainGatherReady). + BounceNvtxScope gatherLaunchScope(kNvtxGatherLaunch, "gatherLaunch rid=%llu chunk=%u n=%u bytes=%llu", + static_cast(rid), chunkIdx, nDesc, static_cast(chunk.packedBytes)); + auto const regionBase = mCtx.arena->baseAddr() + *localOff; + // Coalescing in the plan can leave very large per-desc runs; split them so the copy kernel + // keeps its one-thread-block-per-entry parallelism (bounded by the scratch capacity). + // Two passes: an exact piece count first (the packed [srcs|dsts|sizes] pinned layout needs + // the total before the first write), then fill the pinned buffer DIRECTLY — one write pass + // instead of building std::vectors and memcpy'ing them in. + std::size_t const maxEntries = maxPlanEntries(ctx); + std::size_t nTotal = 0; + for (std::uint32_t i = 0; i < nDesc; ++i) + { + nTotal += piecesFor(chunk.sizes[i], splitBudget(nTotal, nDesc - 1 - i, maxEntries)); + } + // Same capacity guard as the scatter path: planBufs is capacity-unchecked, so writing more + // than maxEntries would overflow the pinned/scratch plan buffers. The bound holds by + // construction today (submit's maxDescsPerChunk and the ExecPool's maxDescs derive from the + // same expression), but a divergence must fail the flow, not corrupt the heap. + if (nTotal > maxEntries) + { + TLLM_LOG_WARNING( + "BounceTransport(%s): rid=%llu chunk=%u plan entries %zu > exec capacity %zu; abandoning flow", + mCtx.selfName.c_str(), static_cast(rid), chunkIdx, nTotal, maxEntries); + mCtx.exec->release(ctx); + mCtx.sendGrants(mCtx.scheduler.releaseLocal(*localOff)); + req.abandonReason = BounceFailReason::kProtocolError; + req.pendingCredits.clear(); // abandon: the request then fails via checkTimeouts + break; + } + auto const bufs = planBufs(ctx, nTotal); + std::size_t idx = 0; + for (std::uint32_t i = 0; i < nDesc; ++i) + { + appendSplitInto(bufs, idx, chunk.srcPtrs[i], regionBase + chunk.bounceOffsets[i], chunk.sizes[i], + splitBudget(idx, nDesc - 1 - i, maxEntries)); + } + // gather into the region (cfg.useZeroCopyArguments selects the H2D-vs-zero-copy plan-argument path) + cudaError_t const gatherErr = launchPrepared(ctx, nTotal, mCtx.cfg.useZeroCopyArguments); + // Record an event for gather completion and DEFER the write. The gather must finish before + // NIXL reads the region, but on a shared GPU the gather can be delayed behind model kernels + // — blocking the IO thread on cudaStreamSynchronize here would stall the whole reactor + // (no recv/poll/ACK) for that delay. Instead drainGatherReady() polls this event and posts + // the write only once the gather is done (NIXL's postXferReq is not stream-ordered anyway). + cudaError_t const recordErr = cudaEventRecord(ctx->event, ctx->stream); + bool const gatherFailed = (gatherErr != cudaSuccess || recordErr != cudaSuccess); + if (gatherFailed) + { + // If the launch or the event-record failed we must NOT trust the event: an event never + // successfully recorded queries as "complete", which would make drainGatherReady post a + // write of an UN-gathered region (garbage). Clear the sticky error and flag the Posted so + // drainGatherReady fails the request deterministically (region/ctx released there). + (void) cudaGetLastError(); + TLLM_LOG_WARNING("BounceTransport(%s): gather launch/record failed (launch=%d record=%d) rid=%llu chunk=%u", + mCtx.selfName.c_str(), static_cast(gatherErr), static_cast(recordErr), + static_cast(rid), chunkIdx); + } + Posted p; + p.chunkIdx = chunkIdx; + p.localOffset = *localOff; + p.ctx = ctx; + p.hasCredit = haveCredit; + if (haveCredit) + { + p.remoteHandle = credit.regionHandle; + p.remoteAddr = credit.addr; + p.remoteDevId = credit.devId; + } + p.writeBytes = static_cast(chunk.packedBytes); + // Gather in flight; the write is issued later by drainGatherReady once the event signals. If + // the gather launch/record failed, go straight to GatherFailed so drainGatherReady fails the + // request without ever trusting the (un)recorded event. + p.state = gatherFailed ? PostState::GatherFailed : PostState::Gathering; + if (!gatherFailed) + { + p.nvtxGather = bounceRangeStart(kNvtxGather, "gather rid=%llu chunk=%u bytes=%llu", + static_cast(rid), chunkIdx, static_cast(chunk.packedBytes)); + } + req.posted.push_back(std::move(p)); + req.nextPost += 1; + req.lastProgress = std::chrono::steady_clock::now(); // forward progress: a chunk's gather launched + } + // Reconcile the pipeline-starvation NVTX spans after every pump pass (perf visibility only): + // - creditStarved: every granted credit is consumed but chunks still lack one -> the flow is + // waiting on the receiver's next re-GRANT (with eager gather a chunk may be posted yet still + // credit-less). Ended in onGrant, so each wait period is its own range. + // - arenaStarved: exited the loop with credits still parked -> blocked on LOCAL resources (gather + // region / exec ctx). One continuous range per park period, ended here once the park drains. + // Both are idempotent across repeated pump attempts (drainPendingPosts retries every IO pass). + if (req.pendingCredits.empty()) + { + bounceRangeEnd(req.nvtxArenaStarved); + if (req.nextCredit < req.numChunks && req.nvtxCreditStarved == 0) + { + req.nvtxCreditStarved = bounceRangeStart(kNvtxCreditStarved, "creditStarved rid=%llu posted=%u/%u", + static_cast(rid), req.nextPost, req.numChunks); + } + } + else if (req.nvtxArenaStarved == 0) + { + req.nvtxArenaStarved = bounceRangeStart(kNvtxArenaStarved, "arenaStarved rid=%llu parked=%zu", + static_cast(rid), req.pendingCredits.size()); + } +} + +void BounceSender::drainPendingPosts() +{ + std::lock_guard lk(mReqMu); + for (auto& [rid, req] : mRequests) + { + // Retry parked credits, and (with eager gather) chunks that couldn't start earlier because + // the arena/ExecPool/eager budget was exhausted — ACKs may have freed resources since. + if (!req.pendingCredits.empty() || (mCtx.cfg.enableEagerGather && req.nextPost < req.numChunks)) + { + pumpRequest(rid, req); + } + } +} + +bool BounceSender::drainGatherReady() +{ + std::lock_guard lk(mReqMu); + bool didWork = false; + std::vector toFail; + for (auto& [rid, req] : mRequests) + { + for (auto& p : req.posted) + { + if (p.state == PostState::Writing || p.state == PostState::Sent) + { + continue; // gather already done + write issued + } + if (p.state == PostState::GatherFailed) + { + // Gather launch / event-record failed in pumpRequest -> never trust the event; fail. + toFail.push_back(rid); + break; + } + if (p.state == PostState::Gathering) + { + // Poll the gather event (non-blocking). + cudaError_t const ev = cudaEventQuery(p.ctx->event); + if (ev == cudaErrorNotReady) + { + continue; // gather still running (possibly delayed behind other GPU work) — no block + } + if (ev != cudaSuccess) + { + // Gather kernel / stream error -> fail the request deterministically (never hang). + (void) cudaGetLastError(); + toFail.push_back(rid); + break; + } + // Gather done — return the exec context immediately for another chunk to reuse (the + // write path never needs it: postXferReq is not stream-ordered). The region stays + // held until ACK. + bounceRangeEnd(p.nvtxGather); + mCtx.exec->release(p.ctx); + p.ctx = nullptr; + p.state = PostState::Gathered; + didWork = true; + } + // state == Gathered: issue the RDMA write once the credit is here (an eagerly-gathered + // chunk may finish before its GRANT arrives; it then waits in place until attachCredits + // fills in the remote target). + if (!p.hasCredit) + { + continue; + } + p.nvtxWrite = bounceRangeStart(kNvtxNixlWrite, "nixlWrite rid=%llu chunk=%u bytes=%u", + static_cast(rid), p.chunkIdx, p.writeBytes); + { + // One already-final (src, dst) pair — the remote address came from the credit, so + // this goes through the agent's below-the-splitter primitive, not the public path. + // A nullptr result (submission failure, logged by the agent) polls as kFAILURE. + TransferDescs const src{MemoryType::kVRAM, + {MemoryDesc{reinterpret_cast(mCtx.arena->at(p.localOffset)), p.writeBytes, + static_cast(mCtx.deviceId)}}}; + TransferDescs const dst{MemoryType::kVRAM, {MemoryDesc{p.remoteAddr, p.writeBytes, p.remoteDevId}}}; + p.xfer = mCtx.agent.postXferRequest(TransferOp::kWRITE, src, dst, req.peer, std::nullopt); + } + p.state = PostState::Writing; + didWork = true; + req.lastProgress = std::chrono::steady_clock::now(); // forward progress: a chunk was posted + } + } + for (auto rid : toFail) + { + auto it = mRequests.find(rid); + if (it != mRequests.end()) + { + failRequest(rid, it->second, BounceFailReason::kGatherFailed); + didWork = true; + } + } + return didWork; +} + +bool BounceSender::pollSenderHandles() +{ + std::lock_guard lk(mReqMu); + bool didWork = false; + std::vector toFail; + for (auto& [rid, req] : mRequests) + { + for (auto& p : req.posted) + { + if (p.state != PostState::Writing) + { + continue; // still gathering (drainGatherReady handles it) or DATA already sent + } + // wait(0) is one non-blocking status query: IN_PROGRESS / SUCCESS / FAILURE. + // A failed post left p.xfer null -> report FAILURE. + TransferState const st = p.xfer != nullptr ? p.xfer->wait(0) : TransferState::kFAILURE; + if (st == TransferState::kSUCCESS) + { + // End the write span BEFORE building the DATA message so nixlWrite measures only the + // RDMA in-flight time; the DATA build/encode/enqueue cost gets its own span (it used + // to hide inside nixlWrite). + bounceRangeEnd(p.nvtxWrite); + auto const& chunk = req.plan.chunks()[p.chunkIdx]; + // The DATA scatter plan is the chunk's COALESCED run list (built once at plan time): + // same bytes, but typically orders of magnitude fewer entries than the per-desc view + // — this message sits on the ACK critical path. Sent as-is, no per-send rebuild. + auto const nRuns = static_cast(chunk.scatterRuns.size()); + { + BounceNvtxScope dataScope(kNvtxDataSend, "dataSend rid=%llu chunk=%u n=%u bytes=%zu", + static_cast(rid), p.chunkIdx, nRuns, + static_cast(nRuns) * sizeof(BounceScatterRun)); + mCtx.channel->sendTo( + req.peer, encodeData(rid, p.chunkIdx, req.numChunks, p.remoteHandle, chunk.scatterRuns)); + } + if (!p.xfer->release()) + { + // The write is terminal, so progress is safe; the status object keeps the + // handle and its destructor retries the backend release. + TLLM_LOG_WARNING("BounceTransport(%s): terminal write handle release deferred rid=%llu chunk=%u", + mCtx.selfName.c_str(), static_cast(rid), p.chunkIdx); + } + p.xfer.reset(); + p.state = PostState::Sent; + p.nvtxAckWait = bounceRangeStart( + kNvtxAckWait, "ackWait rid=%llu chunk=%u", static_cast(rid), p.chunkIdx); + didWork = true; + } + else if (st == TransferState::kFAILURE) + { + toFail.push_back(rid); + break; + } + } + } + for (auto rid : toFail) + { + auto it = mRequests.find(rid); + if (it != mRequests.end()) + { + failRequest(rid, it->second, BounceFailReason::kWriteFailed); + didWork = true; + } + } + return didWork; +} + +void BounceSender::onAck(std::string const& peer, BounceMsgHeader const& h) +{ + // Starts BEFORE taking mReqMu: the span exposes ACK-processing latency INCLUDING lock wait — + // pumpRequest holds mReqMu during gather launches, so a long onAck here means the ACK stalled + // behind another flow's launch prep. + BounceNvtxScope ackScope( + kNvtxOnAck, "onAck rid=%llu chunk=%u", static_cast(h.requestId), h.chunkIdx); + std::lock_guard lk(mReqMu); + auto it = mRequests.find(h.requestId); + if (it == mRequests.end()) + { + return; + } + Request& req = it->second; + if (peer != req.peer) + { + TLLM_LOG_WARNING("BounceTransport(%s): dropping wrong-peer ACK peer=%s expected=%s rid=%llu chunk=%u", + mCtx.selfName.c_str(), peer.c_str(), req.peer.c_str(), static_cast(h.requestId), + h.chunkIdx); + return; + } + bool found = false; + for (auto pit = req.posted.begin(); pit != req.posted.end(); ++pit) + { + // Only a Sent chunk has finished reading its local staging region and may recycle it on ACK. + if (pit->chunkIdx == h.chunkIdx && pit->remoteHandle == h.regionHandle && pit->state == PostState::Sent) + { + bounceRangeEnd(pit->nvtxAckWait); + // Return the gather-staging region to the shared arena; re-schedule may hand the freed + // bytes to a waiting remote flow. + mCtx.sendGrants(mCtx.scheduler.releaseLocal(pit->localOffset)); + req.posted.erase(pit); + found = true; + break; + } + } + if (!found) + { + TLLM_LOG_WARNING("BounceTransport(%s): dropping stale/invalid ACK peer=%s rid=%llu chunk=%u region=%llu", + mCtx.selfName.c_str(), peer.c_str(), static_cast(h.requestId), h.chunkIdx, + static_cast(h.regionHandle)); + // Duplicate / unknown ACK (zmq reconnect, retransmit). Do NOT count it — an over-count + // could push acked past numChunks and resolve SUCCESS before all chunks actually landed. + return; + } + req.acked += 1; + req.lastProgress = std::chrono::steady_clock::now(); // forward progress: a chunk was ACKed + if (req.acked >= req.numChunks) + { + bounceRangeEnd(req.nvtxGrantWait); + bounceRangeEnd(req.nvtxCreditStarved); + bounceRangeEnd(req.nvtxArenaStarved); + bounceRangeEnd(req.nvtxReq); + try + { + req.promise->set_value({TransferState::kSUCCESS, BounceFailReason::kNone}); + } + catch (...) + { + // set_value throws std::future_error ONLY if the promise is already satisfied — a benign + // double-resolve (a request resolves exactly once); intentionally ignored, not a failure. + } + mRequests.erase(it); + } +} + +void BounceSender::checkTimeouts() +{ + if (mCtx.cfg.requestTimeoutMs <= 0) + { + return; // timeout disabled (e.g. tests that intentionally wait forever) + } + auto const now = std::chrono::steady_clock::now(); + auto const limit = std::chrono::milliseconds(mCtx.cfg.requestTimeoutMs); + std::vector stuck; + { + std::lock_guard lk(mReqMu); + for (auto& [rid, req] : mRequests) + { + if (now - req.lastProgress > limit) + { + stuck.push_back(rid); + } + } + for (auto rid : stuck) + { + auto it = mRequests.find(rid); + if (it != mRequests.end()) + { + // Peer never granted or stopped making progress (unreachable / not bounce-ready / + // congested). Fail the request so wait() returns FAILURE instead of hanging. + failRequest(rid, it->second, BounceFailReason::kNoProgressTimeout); + } + } + } +} + +bool BounceSender::drainOrphanLocal() +{ + if (mOrphanLocal.empty()) + { + return false; + } + bool didWork = false; + std::vector keep; + keep.reserve(mOrphanLocal.size()); + for (auto& o : mOrphanLocal) + { + if (o.xfer != nullptr && o.xfer->wait(0) == TransferState::kIN_PROGRESS) + { + keep.push_back(std::move(o)); // write still in flight -> the NIC may still read the region; wait + continue; + } + // Terminal (Done or Failed): the NIC is finished with the region (source AND the receiver's + // destination) -> recycle the local source now. + if (o.xfer != nullptr && !o.xfer->release()) + { + // The write is terminal, so recycling is safe; only the backend handle remains retained + // (the status object's destructor retries the release). + TLLM_LOG_WARNING("BounceTransport(%s): terminal orphan handle release deferred rid=%llu", + mCtx.selfName.c_str(), static_cast(o.rid)); + } + mCtx.sendGrants(mCtx.scheduler.releaseLocal(o.offset)); + didWork = true; + } + mOrphanLocal.swap(keep); + // Send any deferred cancel whose flow now has NO in-flight write left: the receiver may safely + // reclaim its regions (the writes have landed/failed, no more DMA targets them). + for (auto it = mPendingCancel.begin(); it != mPendingCancel.end();) + { + std::uint64_t const rid = it->first; + bool const stillInFlight = std::any_of( + mOrphanLocal.begin(), mOrphanLocal.end(), [rid](OrphanLocal const& o) { return o.rid == rid; }); + if (stillInFlight) + { + ++it; + continue; + } + mCtx.channel->sendTo(it->second, encodeCancel(rid, mCtx.channel->localEndpoint())); + it = mPendingCancel.erase(it); + didWork = true; + } + return didWork; +} + +void BounceSender::failRequest(std::uint64_t rid, Request& req, BounceFailReason reason) +{ + // An abandoned flow (GRANT mispair / plan overflow) reaches here through the timeout path; + // report the specific abandon cause, not the generic timeout. + if (req.abandonReason != BounceFailReason::kNone) + { + reason = req.abandonReason; + } + // The one log line every sender-side failure passes through (timeout, peer drop, write/gather + // failure) — keep it WARNING and keep the progress context. + TLLM_LOG_WARNING("BounceTransport(%s): request FAILED rid=%llu peer=%s reason=\"%s\" chunks acked=%u posted=%u/%u", + mCtx.selfName.c_str(), static_cast(rid), req.peer.c_str(), toString(reason), req.acked, + req.nextPost, req.numChunks); + // Release in-flight transfer handles and return each gather-staging region to the shared arena — + // but only once nothing is still touching the region's memory, else recycling races a live DMA: + // - Writing: the RDMA write may still be reading the region as its source. Defer to mOrphanLocal; + // drainOrphanLocal() releases the xfer + region once poll() is terminal. + // - Gathering / GatherFailed: our gather kernel may still be WRITING the region; sync its stream + // before recycling (else an abandoned gather scribbles a re-granted region — the write is not + // ordered against the new owner). Rare failure path, so a sync here is fine. + // - Sent: write landed (poll==kDone), xfer already released in pollSenderHandles, NIC done + // reading -> recycle now. + bool deferredWrite = false; + for (auto& p : req.posted) + { + // Close this chunk's NVTX spans (whichever leg it died in); 0 handles are no-ops. + bounceRangeEnd(p.nvtxGather); + bounceRangeEnd(p.nvtxWrite); + bounceRangeEnd(p.nvtxAckWait); + if (p.state == PostState::Writing) + { + mOrphanLocal.push_back(OrphanLocal{std::move(p.xfer), p.localOffset, rid}); + deferredWrite = true; + continue; // do NOT release xfer or recycle the region yet + } + if ((p.state == PostState::Gathering || p.state == PostState::GatherFailed) && p.ctx != nullptr) + { + cudaError_t const se = cudaStreamSynchronize(p.ctx->stream); + if (se != cudaSuccess) + { + TLLM_LOG_WARNING("BounceTransport(%s): failRequest stream sync error rid=%llu chunk=%u: %s", + mCtx.selfName.c_str(), static_cast(rid), p.chunkIdx, cudaGetErrorString(se)); + (void) cudaGetLastError(); + } + mCtx.exec->release(p.ctx); + p.ctx = nullptr; + } + mCtx.sendGrants(mCtx.scheduler.releaseLocal(p.localOffset)); + } + // Retract the credit request so the receiver stops holding/granting for it. If any RDMA write is + // still in flight it is landing on the receiver's region; sending the cancel NOW would let the + // receiver reclaim+re-grant that region under the write -> corruption. Defer the cancel until the + // flow's writes drain (drainOrphanLocal sends it). With no in-flight write, send it immediately. + if (deferredWrite) + { + mPendingCancel[rid] = req.peer; + } + else + { + mCtx.channel->sendTo(req.peer, encodeCancel(rid, mCtx.channel->localEndpoint())); + } + bounceRangeEnd(req.nvtxGrantWait); + bounceRangeEnd(req.nvtxCreditStarved); + bounceRangeEnd(req.nvtxArenaStarved); + bounceRangeEnd(req.nvtxReq); + try + { + req.promise->set_value({TransferState::kFAILURE, reason}); + } + catch (...) + { + // set_value throws std::future_error ONLY if the promise is already satisfied — a benign + // double-resolve (a request resolves exactly once); intentionally ignored, not a failure. + } + mRequests.erase(rid); +} + +void BounceSender::forget(std::string const& peer) +{ + // Fail any in-flight request targeting the gone peer so its wait() returns. + std::lock_guard lk(mReqMu); + std::vector toFail; + for (auto const& [rid, req] : mRequests) + { + if (req.peer == peer) + { + toFail.push_back(rid); + } + } + for (auto rid : toFail) + { + auto it = mRequests.find(rid); + if (it != mRequests.end()) + { + failRequest(rid, it->second, BounceFailReason::kPeerDropped); + } + } +} + +void BounceSender::failAll() +{ + // Fail any still-pending requests so no submit() future hangs, releasing their in-flight + // transfer handles first (same handle-leak fix as failRequest). Called after the device has been + // synced and the IO thread joined, so no lock contention — but keep mReqMu for consistency. + // cudaDeviceSynchronize covers only local kernels. NIXL_SUCCESS from releaseXferReq means an + // active write was canceled/released; retain failures for retry instead of polling forever and + // letting a failed peer/backend hang teardown. + std::lock_guard lk(mReqMu); + std::vector> releaseRetry; + for (auto& [rid, req] : mRequests) + { + for (auto& p : req.posted) + { + bounceRangeEnd(p.nvtxGather); + bounceRangeEnd(p.nvtxWrite); + bounceRangeEnd(p.nvtxAckWait); + if (p.state == PostState::Writing && p.xfer != nullptr) + { + if (!p.xfer->release()) + { + releaseRetry.push_back(std::move(p.xfer)); + } + } + if (p.ctx != nullptr) + { + mCtx.exec->release(p.ctx); // still-gathering chunk: return its borrowed exec context + p.ctx = nullptr; + } + } + bounceRangeEnd(req.nvtxGrantWait); + bounceRangeEnd(req.nvtxCreditStarved); + bounceRangeEnd(req.nvtxArenaStarved); + bounceRangeEnd(req.nvtxReq); + try + { + req.promise->set_value({TransferState::kFAILURE, BounceFailReason::kShutdown}); + } + catch (...) + { + // set_value throws std::future_error ONLY if the promise is already satisfied — a benign + // double-resolve (a request resolves exactly once); intentionally ignored, not a failure. + } + } + mRequests.clear(); + // Cancel/release deferred writes without recycling their regions or sending deferred control + // messages: no producer remains, and shutdown must not grant work against an arena being torn down. + for (auto& o : mOrphanLocal) + { + if (o.xfer != nullptr && !o.xfer->release()) + { + releaseRetry.push_back(std::move(o.xfer)); + } + } + mOrphanLocal.clear(); + // Retry failed cancellations once after all producers/futures have been stopped. A persistent + // backend failure gets a final bounded attempt from the status object's destructor. + for (auto& status : releaseRetry) + { + if (!status->release()) + { + TLLM_LOG_WARNING( + "BounceTransport(%s): NIXL handle still retained after shutdown retry", mCtx.selfName.c_str()); + } + } + // Deferred cancels aren't sent at shutdown: in-flight RDMA writes aren't drained by the device + // sync (they're NIXL, not CUDA-stream), so a cancel could still race a write. The receiver + // reclaims those regions via its own teardown / peer-loss path. + mPendingCancel.clear(); +} + +// ============================================================================ +// BounceTransport — the reactor +// ============================================================================ + +BounceTransport::BounceTransport(std::string selfName, BounceConfig cfg, int deviceId, ControlChannel* channel, + NixlTransferAgent& agent, BounceArena* arena, ExecPool* exec) + : mCtx(std::move(selfName), cfg, deviceId, channel, agent, arena, exec) + , mReceiver(mCtx) + , mSender(mCtx) +{ + // A bounce chunk must fit a fully-drained arena, or its GRANT can never succeed and the flow + // stalls until requestTimeoutMs. The buddy allocator rounds usable capacity down and each allocation up + // to a power of two, so comparing maxChunkSizeBytes directly with arenaSizeBytes is insufficient + // (e.g. a 96 MiB arena has only 64 MiB usable, so a 65 MiB chunk never fits). Clamp to the largest + // block the drained arena can actually hand out. + std::size_t const cap = mCtx.scheduler.arenaCapacity(); + if (mCtx.cfg.maxChunkSizeBytes > cap) + { + TLLM_LOG_WARNING( + "BounceTransport(%s): maxChunkSizeBytes=%zu exceeds usable arena capacity=%zu; clamping to %zu", + mCtx.selfName.c_str(), static_cast(mCtx.cfg.maxChunkSizeBytes), cap, cap); + mCtx.cfg.maxChunkSizeBytes = cap; + } + try + { + mReceiver.startWorkers(); + mIoThread = std::thread(&BounceTransport::ioLoop, this); + } + catch (...) + { + // Thread creation failed partway (resource exhaustion). The destructor won't run — the + // constructor is throwing — so join any scatter workers already spawned here, or their + // std::thread destructors call std::terminate. + mCtx.stop.store(true, std::memory_order_release); + mReceiver.wake(); + mReceiver.joinWorkers(); + throw; + } +} + +BounceTransport::~BounceTransport() +{ + shutdown(); +} + +void BounceTransport::shutdown() +{ + bool expected = false; + if (!mCtx.stop.compare_exchange_strong(expected, true)) + { + return; // already shut down + } + mReceiver.wake(); // wake scatter workers so they observe stop + if (mIoThread.joinable()) + { + mIoThread.join(); + } + mReceiver.joinWorkers(); + // Threads are joined, but in-flight gather/scatter kernels may still be queued on ExecPool + // streams referencing the arena. Drain the device BEFORE we release contexts and let the caller + // tear down ExecPool/BounceArena — otherwise ~ExecPool's cudaStreamDestroy + ~BounceArena's + // cudaFree could race a kernel still reading the arena. (Teardown-only; a full-device sync is + // acceptable here.) Warn-only (we're in teardown / a dtor path -> must not throw); the user + // should still see a GPU fault. shutdown() may run on a thread whose current device isn't ours. + TLLM_CUDA_CHECK_WARN(cudaSetDevice(mCtx.deviceId)); + TLLM_CUDA_CHECK_WARN(cudaDeviceSynchronize()); + mSender.failAll(); +} + +bool BounceTransport::addPeer(std::string const& peer, std::string const& endpoint) +{ + return mCtx.channel->addPeer(peer, endpoint); +} + +std::string BounceTransport::localHandshakeBlob() const +{ + auto endpoint = mCtx.channel->localEndpoint(); + if (endpoint.empty()) + { + TLLM_THROW("BounceTransport(%s): ZMQ control requires a non-empty local endpoint", mCtx.selfName.c_str()); + } + BounceHandshake handshake; + handshake.wireVersion = kBounceVersion; + handshake.controlKind = BounceControlKind::kZMQ; + handshake.arenaUsableCapacityBytes = mCtx.scheduler.arenaCapacity(); + handshake.maxChunkSizeBytes = mCtx.cfg.maxChunkSizeBytes; // post-ctor-clamp effective value + handshake.endpoint = std::move(endpoint); + return encodeHandshake(handshake); +} + +bool BounceTransport::registerPeerHandshake(std::string const& peer, std::string const& blob) +{ + // Treat registration as replacement, not an additive update. A peer may be reloaded under the + // same name with bounce disabled or with changed settings; clear the previously validated route + // first so any missing/malformed/incompatible replacement immediately falls back to NIXL. + mCtx.channel->removePeer(peer); + { + std::lock_guard lk(mPeerMu); + mHandshakedPeers.erase(peer); + } + if (blob.empty()) + { + return false; // bounce not advertised by this peer + } + + BounceHandshake handshake; + if (!decodeHandshake(blob, handshake)) + { + TLLM_LOG_WARNING( + "BounceTransport(%s): peer %s advertised an unparsable bounce handshake -> bounce disabled for " + "this peer (NIXL fallback)", + mCtx.selfName.c_str(), peer.c_str()); + return false; + } + auto const localControlKind = BounceControlKind::kZMQ; + // STRICT equality. maxChunkSizeBytes is compared on the effective (post-clamp) values both + // sides advertise; each side already clamped its own value to its usable arena capacity, so + // equality also guarantees our chunks always fit the peer's arena and its scatter scratch + // (sized for its own maxChunkSizeBytes). Local-only knobs (worker/stream counts, timeouts, + // the zero-copy-argument path, granularity, arena size) intentionally do NOT have to match. + if (handshake.wireVersion != kBounceVersion || handshake.controlKind != localControlKind + || handshake.maxChunkSizeBytes != mCtx.cfg.maxChunkSizeBytes) + { + TLLM_LOG_WARNING( + "BounceTransport(%s): peer %s bounce handshake incompatible (wireVersion %u vs %u, controlKind " + "%u vs %u, maxChunkSizeBytes %llu vs %zu, peer arenaUsableCapacityBytes %llu vs %zu) -> bounce " + "disabled for this peer (NIXL fallback)", + mCtx.selfName.c_str(), peer.c_str(), static_cast(handshake.wireVersion), + static_cast(kBounceVersion), static_cast(handshake.controlKind), + static_cast(localControlKind), static_cast(handshake.maxChunkSizeBytes), + static_cast(mCtx.cfg.maxChunkSizeBytes), + static_cast(handshake.arenaUsableCapacityBytes), mCtx.scheduler.arenaCapacity()); + return false; + } + if (!mCtx.channel->addPeer(peer, handshake.endpoint)) + { + TLLM_LOG_WARNING( + "BounceTransport(%s): peer %s bounce endpoint could not be registered -> bounce disabled for " + "this peer (NIXL fallback)", + mCtx.selfName.c_str(), peer.c_str()); + return false; + } + std::lock_guard lk(mPeerMu); + mHandshakedPeers.insert(peer); + return true; +} + +bool BounceTransport::hasPeerHandshake(std::string const& peer) const +{ + std::lock_guard lk(mPeerMu); + return mHandshakedPeers.count(peer) > 0; +} + +void BounceTransport::forgetPeer(std::string const& peer) +{ + // Drop the control-channel DEALER to this peer SYNCHRONOUSLY here (ControlChannel::removePeer is + // thread-safe). Doing it on the caller thread — rather than on the IO thread in drainForgets — + // gives a deterministic happens-before for any addPeer() the caller issues after forgetPeer() + // returns (e.g. re-establishing a peer that came back): the dealer is already gone, so that + // addPeer() rebuilds it instead of racing an async removePeer that would otherwise erase the + // freshly re-added dealer. A pending send to the now-removed peer is dropped (it is being + // invalidated), which degrades any in-flight request to a FAILURE — never corruption. + mCtx.channel->removePeer(peer); + { + // Also drop it from the bounce-capable set so the fast path stops engaging it immediately; + // a fresh loadRemoteAgent must re-validate its handshake. + std::lock_guard lk(mPeerMu); + mHandshakedPeers.erase(peer); + } + // The scheduler / request-table reclaim still runs on the IO thread (drainForgets) so that state + // stays owned by a single thread. Safe to call from invalidateRemoteAgent. + std::lock_guard lk(mForgetMu); + mForgetPeers.push_back(peer); +} + +void BounceTransport::drainForgets() +{ + std::vector peers; + { + std::lock_guard lk(mForgetMu); + if (mForgetPeers.empty()) + { + return; + } + peers.swap(mForgetPeers); + } + for (auto const& peer : peers) + { + mReceiver.forget(peer); // reclaim receiver-side credits/jobs of the gone peer + mSender.forget(peer); // fail in-flight sender requests to the gone peer + // NOTE: the control-channel DEALER for this peer was already dropped synchronously in + // forgetPeer() (see there); we only reclaim scheduler/request state on this (IO) thread. + } +} + +void BounceTransport::ioLoop() +{ + bounceNameThread("bounceIO"); + // Pin this thread to our device up front; if it fails, every CUDA op in the loop targets the wrong + // device (the transport is effectively broken). Can't throw out of a thread fn -> warn-only. + TLLM_CUDA_CHECK_WARN(cudaSetDevice(mCtx.deviceId)); + std::string peer; + std::string blob; + while (!mCtx.stop.load(std::memory_order_acquire)) + { + // Exception boundary: this is a thread entry, so anything escaping here would + // std::terminate the process. recv() can throw zmq::error_t (e.g. EINTR/ETERM from + // zmq_poll) and the handlers may throw on peer input; log, back off briefly (so a + // persistent error can't hot-spin the core), and keep serving — in-flight requests still + // resolve via checkTimeouts, never a hang (R5). + try + { + tick(peer, blob); + } + catch (std::exception const& e) + { + TLLM_LOG_WARNING("BounceTransport(%s): IO tick failed: %s", mCtx.selfName.c_str(), e.what()); + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + } +} + +void BounceTransport::tick(std::string& peer, std::string& blob) +{ + // Adaptive poll wait. When there is GPU/RDMA work to poll — gather/write in flight on the + // sender (localHeldCount>0) or scatter in flight on the receiver (mScattering) — use a 0ms + // timeout so cudaEventQuery / getXferStatus / scatter completions are detected ASAP (low + // latency on the critical path; also keeps UCX progress driven). When fully idle (only + // waiting on a control message), sleep up to 1ms so the IO thread doesn't busy-spin a core. + // A request merely waiting for a GRANT (nothing posted yet) is NOT "busy", so a stalled / + // unreachable peer won't spin a core until requestTimeoutMs. Both checks are IO-thread-only. + bool const busy = mSender.busy() || mReceiver.busy(); + int const timeoutMs = busy ? 0 : 1; + bool work = false; + if (mCtx.channel->recv(peer, blob, timeoutMs)) + { + dispatch(peer, blob); + work = true; + } + work |= mSender.drainGatherReady(); // post writes for chunks whose gather kernel just finished + work |= mSender.pollSenderHandles(); + work |= mReceiver.drainScatterDone(); + work |= mSender.drainOrphanLocal(); // recycle failed-but-in-flight write regions once their write ends + drainForgets(); + mSender.drainPendingPosts(); // retry credits parked when the arena was full (onAck freed regions) + mSender.checkTimeouts(); + mReceiver.checkTimeouts(); // expire grant leases of silent senders + free post-quarantine regions + // Idle backoff: when there IS in-flight work (busy → 0ms poll) but nothing actually advanced + // this pass — the classic case being a gather stalled behind unrelated model kernels (gather + // event stays NotReady) — keep latency low for the first few spins, then sleep briefly so we + // don't peg a core at 100%. Any control message or forward progress resets the counter. + if (work) + { + mIdleSpins = 0; + } + else if (busy) + { + constexpr std::uint32_t kSpinBeforeBackoff = 64; + if (mIdleSpins < kSpinBeforeBackoff) + { + ++mIdleSpins; + } + else + { + std::this_thread::sleep_for(std::chrono::microseconds(50)); + } + } +} + +void BounceTransport::dispatch(std::string const& peer, std::string const& blob) +{ + BounceMsgHeader h{}; + if (!decodeHeader(blob, h)) + { + return; + } + switch (static_cast(h.msgType)) + { + case BounceMsgType::kWANT: mReceiver.onWant(peer, h, blob); break; // [R] + case BounceMsgType::kGRANT: mSender.onGrant(peer, h, blob); break; // [S] + case BounceMsgType::kDATA: mReceiver.onData(peer, h, blob); break; // [R] + case BounceMsgType::kACK: mSender.onAck(peer, h); break; // [S] + default: break; + } +} + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h new file mode 100644 index 000000000000..0534fa62e45d --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h @@ -0,0 +1,527 @@ +/* + * 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/cache_transmission/nixl_utils/bounce/BounceArena.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransferPlan.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ControlChannel.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ExecPool.h" +#include "tensorrt_llm/executor/transferAgent.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::executor::kv_cache +{ +// The data plane: bounce posts its per-chunk RDMA writes through the agent's low-level +// postXferRequest / registerRegionImpl primitives (below the VMM splitter). Forward-declared so +// bounce headers stay independent of the NIXL headers; BounceTransport.cpp includes the real one. +class NixlTransferAgent; +} // namespace tensorrt_llm::executor::kv_cache + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +// ============================================================================ +// Bounce v2 reactor +// ---------------------------------------------------------------------------- +// One BounceTransport per agent. Protocol progress is driven by ONE IO thread plus M scatter +// workers. submit() may also prepare eager gathers on application threads, so the sender request +// table, scheduler and shared resource pools provide their own synchronization. +// +// BounceContext — shared dependencies used by both roles: the injected +// channel/agent/arena/exec, one CreditScheduler, and sendGrants(). +// BounceSender — the [S] role: submit() -> WANT, GRANT -> gather+write, ACK +// -> resolve. Owns the request table + sender-side orphan state. +// BounceReceiver — the [R] role: WANT -> grant regions, DATA -> scatter, ACK. +// Owns the scatter workers + scatter job/done queues. +// BounceTransport— the thin reactor: owns the three above + the IO thread, and +// routes control messages / drains to sender vs receiver. +// +// IO thread loop (primary owner of scheduler + sender request table; submit() may access both +// through their locks): +// 1. recv() one control message (short timeout) and dispatch it: +// [S] GRANT(rid, credits) -> attach credits to eager-gathered chunks or acquire a local +// arena region + ExecCtx and launch gather. postWrite is deferred +// to drainGatherReady() until both gather completion and credit +// availability have been observed. +// [S] ACK(rid, chunk, h) -> releaseLocal() that chunk's region; acked++; +// all acked -> resolve the request's future SUCCESS. +// [R] WANT(rid, sizes[]) -> scheduler.onWant(peer:rid, sizes) -> send GRANT(s). +// [R] DATA(rid, chunk, h, plan) -> enqueue a scatter job over region `h`. +// 2. poll every in-flight write (TransferStatus::wait(0)); on Done send DATA (data has +// landed at the remote) and mark on-wire. THIS is why no notifMsg is needed. +// 3. drain scatter completions posted by workers: release receiver regions and re-GRANT. +// +// Scatter workers (receiver): pop a job, scatter region->dst via the kernel, synchronize the +// stream, send ACK immediately on success, then post bookkeeping completion to the IO thread. +// +// Sender requests are keyed by a monotonic rid; the receiver-side scheduler is keyed by +// "peer\x1f rid" so multiple concurrent requests from one peer are independent flows. +// +// Lifetime: submit() returns a shared_future; it resolves SUCCESS on full ACK and +// FAILURE on transfer error, peer invalidation, shutdown, or request timeout. Setting +// requestTimeoutMs to 0 intentionally disables timeout-based failure. +// +// The agent (data plane) & ControlChannel are injected (not owned) so the same reactor runs +// unchanged in tests and production (both over the real NIXL/zmq stack; failure tests subclass +// NixlTransferAgent and override postXferRequest to inject deterministic transfer faults). +// ============================================================================ + +/// 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 +}; + +[[nodiscard]] char const* toString(BounceFailReason reason); + +/// A bounce request's single result: the terminal state plus, on kFAILURE, its cause. The future +/// is the ONE source of truth — no side channel to keep in sync. +struct BounceResult +{ + TransferState state{TransferState::kFAILURE}; + BounceFailReason reason{BounceFailReason::kNone}; +}; + +/// Shared dependencies used by both roles. Holds the injected channel/agent/arena/exec (borrowed, +/// not owned), the one CreditScheduler that carves the shared arena for both roles, and +/// sendGrants(). Owned by value by BounceTransport and referenced by the sender and receiver. +/// Individual mutable collaborators provide their own synchronization; the immutable arena lookup +/// helpers require none. +class BounceContext +{ +public: + BounceContext(std::string selfName, BounceConfig cfg, int deviceId, ControlChannel* channel, + NixlTransferAgent& agent, BounceArena* arena, ExecPool* exec) + : selfName(std::move(selfName)) + , cfg(cfg) + , deviceId(deviceId) + , channel(channel) + , agent(agent) + , arena(arena) + , exec(exec) + , scheduler(arena->baseAddr(), cfg.arenaSizeBytes, cfg.arenaAllocationGranularityBytes, + cfg.maxInflightChunksPerRequest) + { + } + + BounceContext(BounceContext const&) = delete; + BounceContext& operator=(BounceContext const&) = delete; + + /// Split a scheduler GRANT batch (keyed by flow id "peerrid") back to (peer, rid) and send a + /// GRANT control message to each peer. Used by BOTH roles (the receiver grants incoming regions; + /// the sender's releaseLocal/reclaim can free arena bytes that re-grant a waiting remote flow). + void sendGrants(std::vector const& grants); + + std::string selfName; + BounceConfig cfg; + int deviceId{}; + ControlChannel* channel{}; + NixlTransferAgent& agent; // data plane: postXferRequest (per-chunk RDMA writes) + BounceArena* arena{}; // ONE shared data buffer: receiver grants + local gather both carve regions from it + ExecPool* exec{}; // gather/scatter exec contexts (streams/scratch), borrowed per kernel + CreditScheduler scheduler; // shared region allocator; internally synchronized + std::atomic stop{false}; +}; + +/// Receiver role ([R]): WANT -> grant regions, DATA -> scatter into the caller's KV, then ACK. Owns +/// the M scatter workers and the IO<->worker job/done queues. Protocol handlers run on the IO +/// thread; start/wake/join are lifecycle calls made by the owner. +class BounceReceiver +{ +public: + explicit BounceReceiver(BounceContext& ctx); + + /// Launch the scatter workers (call once after construction, before the IO loop starts). + void startWorkers(); + /// Wake workers so they observe ctx.stop and exit (notify; call before joinWorkers()). + void wake(); + /// Join the scatter workers (call during shutdown after wake()). + void joinWorkers(); + + void onWant(std::string const& peer, BounceMsgHeader const& h, std::string const& blob); + void onData(std::string const& peer, BounceMsgHeader const& h, std::string const& blob); + /// Drain scatter bookkeeping completions posted by workers, release regions, and re-grant newly + /// available space. Successful workers have already sent ACK. Returns true if any were drained. + bool drainScatterDone(); + /// A peer is gone: reclaim every receiver-side flow of that peer (deferring regions a worker is + /// still reading, quarantining regions the peer may still be RDMA-writing) and drop its + /// not-yet-started scatter jobs. + void forget(std::string const& peer); + + /// Time-driven reclamation (IO thread, called from tick()): reclaim flows the scheduler reports + /// idle beyond cfg.receiverFlowTimeoutMs — a dead/unreachable sender emits neither DATA nor a + /// cancel, so without this its granted regions leak forever — and return quarantined regions to + /// the arena once their cfg.quarantineMs deadline passes (scheduler.reapQuarantine). + void checkTimeouts(); + + /// True while any scatter is enqueued-or-running (drives the IO loop's 0ms busy-poll). + [[nodiscard]] bool busy() const + { + return !mScattering.empty(); + } + +private: + struct ScatterJob + { + std::string key; // "peer\x1f rid" + std::string peer; + std::uint64_t rid{}; + std::uint32_t chunkIdx{}; + std::uint64_t offset{}; // receiver arena region offset (the granted region handle) to scatter from + std::uint64_t regionBytes{}; // byte size of the buddy block backing `offset` (bounds scatter reads) + std::vector entries; + std::uint64_t nvtxQueue{0}; // NVTX span: enqueued on the IO thread -> dequeued by a worker + }; + + struct ScatterDone + { + std::string key; + std::uint64_t offset{}; // receiver arena region offset freed once scattered + }; + + /// The set of incoming regions currently scattering — the `busy` set passed to the scheduler's + /// reclaim calls so a region a worker is still reading is deferred, not re-granted. + std::unordered_set scatteringRegions() const; + /// Shared reclaim bookkeeping (cancel / peer loss / lease expiry): run `reclaim` with the + /// currently-scattering regions as the busy set, send the resulting grants, and flag every + /// deferred region orphaned in mScattering so drainScatterDone frees it on completion. + void reclaimAndFlagDeferred( + std::function(std::unordered_set const&, std::vector&)> const& + reclaim); + void scatterWorkerLoop(); + + BounceContext& mCtx; + std::vector mWorkers; + + // scatter job queue (IO thread -> workers) + std::mutex mJobMu; + std::condition_variable mJobCv; + std::deque mJobs; + + // scatter completion queue (workers -> IO thread) + std::mutex mDoneMu; + std::deque mDone; + + // Every incoming region with a scatter job enqueued-or-running, mapped to whether its flow was + // reclaimed (peer gone / cancel) mid-scatter (orphaned == true). Touched only by the IO thread + // (onData / drainScatterDone / forget / onWant) so it needs no lock. Membership stops forget() + // from re-granting an incoming region a worker is still reading; the orphaned flag decides whether + // drainScatterDone frees the region via freeOrphanRegion (flow already gone) or onScatterDone + // (normal completion). + std::unordered_map mScattering; + + std::chrono::steady_clock::time_point mNextSweep{}; // checkTimeouts() throttle (IO-thread-only) +}; + +/// Sender role ([S]): submit() announces a WRITE via WANT; GRANT -> gather into a local arena region +/// + RDMA-write; ACK -> resolve the chunk; all chunks ACKed -> resolve the request SUCCESS. Owns the +/// request table + the sender-side deferred-cleanup state. All methods except submit() run on the IO +/// thread; submit() is called from app threads and is mutex-guarded against the IO thread. +class BounceSender +{ +public: + explicit BounceSender(BounceContext& ctx); + + /// Submit a WRITE of (src -> dst) descriptors to `peer`. Returns a future that resolves + /// {kSUCCESS} once every chunk is scattered+ACKed, or {kFAILURE, reason} on error/shutdown. + /// Chunks are cut to cfg.maxChunkSizeBytes. NixlTransferAgent validates peer compatibility + /// before calling this; direct users of BounceTransport must establish the same invariant + /// themselves. + [[nodiscard]] std::shared_future submit( + TransferDescs const& srcDescs, TransferDescs const& dstDescs, std::string const& peer); + + void onGrant(std::string const& peer, BounceMsgHeader const& h, std::string const& blob); + void onAck(std::string const& peer, BounceMsgHeader const& h); + /// Issue the RDMA write for any chunk whose gather kernel has now completed (poll ctx->event, + /// non-blocking). Decouples the IO thread from gather latency on a shared GPU. Returns true if it + /// posted a write (or failed a request) this pass — used to drive the IO loop's idle backoff. + bool drainGatherReady(); + bool pollSenderHandles(); + /// Free local regions whose failed request still had an RDMA write in flight, once that write + /// reaches a terminal state (so the NIC is done reading the source). Returns true if any freed. + bool drainOrphanLocal(); + /// Retry parked credits for every request (called on the IO loop after ACKs free regions). + void drainPendingPosts(); + /// Fail requests that have made no progress within requestTimeoutMs (e.g. peer never granted). + void checkTimeouts(); + /// A peer is gone: fail any in-flight request targeting it so its wait() returns. + void forget(std::string const& peer); + /// Shutdown: fail every still-pending request and release each active write's TransferStatus + /// handle (a bounded cancel/release, with one retry pass over any release that failed). Does + /// not recycle regions or send new grants. Called after the IO/workers are joined and local + /// CUDA work has been synced. + void failAll(); + + /// True while a local gather region is held or an orphan-local write is in flight (drives the IO + /// loop's 0ms busy-poll). Called on the IO thread. + [[nodiscard]] bool busy() const + { + return mCtx.scheduler.localHeldCount() > 0 || !mOrphanLocal.empty(); + } + +private: + // A chunk's transfer state — the per-chunk view of the sender Request state machine. + // Linear progression Gathering -> Gathered -> Writing -> Sent, with GatherFailed as the one + // off-ramp (a gather whose launch/event-record failed). Replaces a former trio of bools whose + // illegal combinations (e.g. !writePosted && dataSent) were only ruled out by convention. + // With eager gather a chunk may sit in Gathered while its credit (GRANT) is still in flight; + // the write posts once BOTH the gather event has signalled AND the credit is attached. + enum class PostState + { + Gathering, // gather kernel launched + event recorded; write deferred until the event signals + GatherFailed, // gather launch / event-record failed -> drainGatherReady fails the request + // (an unrecorded event queries as "complete", so its event must NOT be trusted) + Gathered, // gather event signalled, exec ctx returned; waiting for the credit (eager gather) + Writing, // gather done; RDMA write issued (xfer valid); polling getXferStatus + Sent, // getXferStatus==Done -> DATA emitted; xfer released; awaiting ACK + }; + + struct Posted + { + std::uint32_t chunkIdx{}; + std::uint64_t localOffset{}; // shared-arena region held for gather/write until ACK (OUTGOING_HELD) + ExecCtx* ctx{nullptr}; // gather exec context, borrowed while the gather runs (until Gathered) + std::unique_ptr xfer; // in-flight RDMA write (set once state == Writing; + // nullptr after release, or if the post itself failed) + std::uint64_t remoteHandle{}; // receiver's region handle (its arena offset); echoed in DATA + // Remote write target (from the GRANT), kept so postWrite can be issued LATER — after the + // gather kernel completes — instead of blocking the IO thread on cudaStreamSynchronize. + // The gather may be delayed behind unrelated GPU work (shared device), so we never sync; + // drainGatherReady() polls ctx->event and posts the write only when the gather is done. + // With eager gather these fields are filled in later, by attachCredits() (hasCredit==true). + std::uint64_t remoteAddr{}; + std::uint32_t remoteDevId{}; + std::uint32_t writeBytes{}; + bool hasCredit{false}; // remote fields valid; a Gathered chunk without it waits in place + PostState state{PostState::Gathering}; + // Cross-thread NVTX spans for this chunk's async legs (0 == none; ended on failure paths too). + std::uint64_t nvtxGather{0}; // gather launched -> event signaled (in-flight incl. GPU queueing) + std::uint64_t nvtxWrite{0}; // RDMA write posted -> poll terminal + std::uint64_t nvtxAckWait{0}; // DATA sent -> ACK received + }; + + struct Request + { + std::string peer; + std::uint32_t numChunks{}; + BounceTransferPlan plan; + std::uint32_t nextPost{0}; + // Credits are granted strictly in chunk order (the receiver serves the WANT size list + // FIFO), so credit k always belongs to chunk k. Chunks [0, nextCredit) have consumed their + // credit (attached to a Posted, or popped at gather-launch time); pendingCredits.front() + // is chunk `nextCredit`'s credit. + std::uint32_t nextCredit{0}; + std::uint32_t acked{0}; + std::vector posted; + // Credits granted by the receiver but not yet consumed: either their chunk's gather couldn't + // start (no gather region / exec context free — the shared arena is used across peers and can + // be oversubscribed) or, with eager gather, the chunk posted before its GRANT arrived and the + // credit is about to be attached. The IO thread NEVER blocks on acquire; it parks the credit + // here and retries each loop iteration as ACKs free regions. FIFO: paired with chunks in order. + std::deque pendingCredits; + std::shared_ptr> promise; + // Set by the abandon sites (GRANT mispair / plan overflow): the request then dies through + // the timeout path, but failRequest reports THIS more specific cause instead. + BounceFailReason abandonReason{BounceFailReason::kNone}; + // Last time this request made forward progress (granted+posted a chunk, or got an ACK). + // A request stuck with no progress for requestTimeoutMs (e.g. the peer never GRANTs because + // it is unreachable / not bounce-ready) is failed rather than hanging wait() forever. + std::chrono::steady_clock::time_point lastProgress; + // Cross-thread NVTX spans (0 == none; ended on failure paths too). + std::uint64_t nvtxReq{0}; // submit -> promise resolved + std::uint64_t nvtxGrantWait{0}; // WANT sent -> first GRANT arrives + // Pipeline-starvation spans (can open/close several times per request, one range per period): + std::uint64_t nvtxCreditStarved{0}; // every credit consumed, chunks left -> next GRANT (onGrant ends) + std::uint64_t nvtxArenaStarved{0}; // credits parked: no local region / exec ctx -> parked credits drain + }; + + // Regions of a FAILED request whose RDMA write was still in flight (state == Writing). + // The NIC may still be reading the source region, so recycling is deferred until the write reaches + // a terminal state — drainOrphanLocal() polls and only then releases the xfer handle + returns the + // region. IO-thread-only (no lock). (xfer handle, arena offset). + struct OrphanLocal + { + std::unique_ptr xfer; + std::uint64_t offset{}; + std::uint64_t rid{}; + }; + + /// GRANT-mispair guard shared by attachCredits()/pumpRequest(): a credit smaller than its chunk + /// would make the RDMA write overflow the granted region into an adjacent flow's region on the + /// peer. On a mispair it abandons the flow (kProtocolError; the request then fails via + /// checkTimeouts) and returns true — the caller stops consuming credits. Caller MUST hold mReqMu. + [[nodiscard]] bool abandonOnCreditMispair( + std::uint64_t rid, Request& req, std::uint32_t chunkIdx, std::uint64_t packedBytes, std::uint32_t creditLen); + /// Attach parked credits to already-posted chunks (eager gather posted them credit-less), in + /// strict chunk order, promoting their staging regions out of the eager budget. Also flags the + /// protocol anomalies (credit/chunk size mispair -> abandon flow; over-grant -> drop extras). + /// Caller MUST hold mReqMu. + void attachCredits(std::uint64_t rid, Request& req); + /// Launch gathers for as many not-yet-posted chunks as resources allow (non-blocking): with a + /// parked credit when available, else eagerly (cfg.enableEagerGather, capped by the per-request + /// in-flight limit and the scheduler's eager budget). Stops when the arena/ExecPool is empty + /// (rest parked/retried). + /// Caller MUST hold mReqMu. Runs attachCredits() first so credit pairing stays in chunk order. + void pumpRequest(std::uint64_t rid, Request& req); + /// Resolve `rid` to kFAILURE, recording `reason` (unless a more specific one was tagged earlier) + /// and logging ONE warning with the request's progress context — the single choke point that + /// keeps every failure mode observable. See the body for the region-recycling rules. + void failRequest(std::uint64_t rid, Request& req, BounceFailReason reason); + + BounceContext& mCtx; + + // request table (IO thread + submit() from app threads) + std::mutex mReqMu; + std::unordered_map mRequests; + std::atomic mNextRid{1}; + + std::vector mOrphanLocal; + // A failed request's cancel (empty WANT) must be DEFERRED while any of its RDMA writes are still + // in flight: those writes are landing on the RECEIVER's regions, and an early cancel would let the + // receiver reclaim + re-grant those regions under the in-flight write -> cross-node corruption. + // drainOrphanLocal() sends the cancel once the flow's last in-flight write reaches terminal. + // rid -> peer. IO-thread-only. + std::unordered_map mPendingCancel; +}; + +/// The thin reactor: owns the shared context, the sender + receiver, and the IO thread. Routes each +/// control message to the right role and drives both roles' per-loop drains. One per agent. +class BounceTransport +{ +public: + /// @param arena ONE shared data buffer serving BOTH roles: receiver (remote senders' RDMA-write + /// targets, granted as variable regions by the scheduler) and sender (local gather staging, via + /// acquireLocal). The caller must already have NIXL-registered it with `agent`. + /// @param exec the gather/scatter execution contexts (streams/scratch) borrowed for the + /// duration of one kernel. `channel`/`agent`/`arena`/`exec` are borrowed for the transport's + /// lifetime. (Most disagg agents are sender-only OR receiver-only, so a single arena avoids + /// wasting a second one.) + BounceTransport(std::string selfName, BounceConfig cfg, int deviceId, ControlChannel* channel, + NixlTransferAgent& agent, BounceArena* arena, ExecPool* exec); + ~BounceTransport(); + + BounceTransport(BounceTransport const&) = delete; + BounceTransport& operator=(BounceTransport const&) = delete; + + /// Register where to reach `peer` (its ControlChannel endpoint). Returns false when the + /// underlying channel cannot load an optional endpoint; a mandatory invalid endpoint may throw. + bool addPeer(std::string const& peer, std::string const& endpoint); + + /// This agent's capability handshake, encoded for AgentDesc: control endpoint plus the fields a + /// peer validates before engaging bounce (wire version, control kind and effective chunk cap). + /// Usable arena capacity is included for diagnostics. Thread-safe because post-construction + /// state is immutable. + [[nodiscard]] std::string localHandshakeBlob() const; + + /// Validate a peer's handshake blob (from its AgentDesc) and, when compatible, register the + /// peer: control channel (addPeer) + the bounce-capable peer set. Compatibility is STRICT + /// equality of wireVersion, controlKind and the effective maxChunkSizeBytes — both sides + /// advertise post-clamp values, so equality also guarantees every chunk fits the peer's arena + /// and scatter scratch. Registration replaces any previous result for the same peer; an empty, + /// incompatible or unparsable replacement clears the old route/capability and returns false, so + /// the caller keeps that peer on the standard per-desc NIXL path. Thread-safe. + bool registerPeerHandshake(std::string const& peer, std::string const& blob); + + /// Did `peer` pass registerPeerHandshake (and not get forgotten since)? Gates the agent's + /// bounce fast path so we never WANT a peer that cannot answer. Thread-safe. + [[nodiscard]] bool hasPeerHandshake(std::string const& peer) const; + + /// A peer is gone (NixlTransferAgent::invalidateRemoteAgent). Remove its channel endpoint and + /// compatibility record synchronously, then queue receiver-credit reclamation and sender-request + /// failure for the IO thread. A fresh loadRemoteAgent must validate a new handshake. + void forgetPeer(std::string const& peer); + + /// Submit a WRITE of (src -> dst) descriptors to `peer`. Returns a future that resolves + /// {kSUCCESS} once every chunk is scattered+ACKed, or {kFAILURE, reason} on error/shutdown. + [[nodiscard]] std::shared_future submit( + TransferDescs const& srcDescs, TransferDescs const& dstDescs, std::string const& peer) + { + return mSender.submit(srcDescs, dstDescs, peer); + } + + /// This side's EFFECTIVE per-chunk cap: cfg.maxChunkSizeBytes AFTER the constructor's clamp to + /// the usable arena capacity. Advertised in the capability handshake; bounce engages toward a + /// peer only when both sides advertise the SAME value (see NixlTransferAgent::loadRemoteAgent), + /// which also guarantees our chunks fit the peer's arena and scatter scratch. + [[nodiscard]] std::size_t maxChunkSizeBytes() const + { + return mCtx.cfg.maxChunkSizeBytes; + } + + /// Stop threads, fail any in-flight requests, and join. Idempotent. + void shutdown(); + +private: + void ioLoop(); + /// One reactor pass (recv+dispatch, drains, timeouts). `peer`/`blob` are caller-owned scratch + /// buffers reused across ticks. Runs inside ioLoop()'s exception boundary. + void tick(std::string& peer, std::string& blob); + void dispatch(std::string const& peer, std::string const& blob); + /// Apply queued forgetPeer() requests on the IO thread (reclaim receiver credits + fail + /// sender requests to the gone peer). + void drainForgets(); + + BounceContext mCtx; // declared first: sender/receiver hold a reference to it + BounceReceiver mReceiver; + BounceSender mSender; + + std::thread mIoThread; + // IO-loop idle backoff: consecutive "busy poll but nothing happened" iterations. When in-flight + // work exists (so the poll timeout is 0ms) but a gather is stalled behind unrelated GPU kernels, + // the loop would otherwise spin a core at 100% on cudaEventQuery. After a threshold of no-progress + // spins we sleep briefly. Reset on any control message or forward progress. IO-thread-only. + std::uint32_t mIdleSpins{0}; + + // forgetPeer() requests queued from app threads -> applied on the IO thread + std::mutex mForgetMu; + std::vector mForgetPeers; + + // Peers whose capability handshake matched ours (registerPeerHandshake); the agent engages + // bounce only toward these. Touched from app threads (loadRemoteAgent / shouldUseBounce / + // invalidateRemoteAgent) -> own mutex. + mutable std::mutex mPeerMu; + std::unordered_set mHandshakedPeers; +}; + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BuddyAllocator.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BuddyAllocator.cpp new file mode 100644 index 000000000000..63da2a2ce0c8 --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BuddyAllocator.cpp @@ -0,0 +1,156 @@ +/* + * 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/BuddyAllocator.h" + +#include "tensorrt_llm/common/assert.h" + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +std::size_t BuddyAllocator::roundUpPow2(std::size_t v) +{ + std::size_t p = 1; + while (p < v) + { + p <<= 1; + } + return p; +} + +BuddyAllocator::BuddyAllocator(std::size_t capacity, std::size_t minBlock) +{ + TLLM_CHECK_WITH_INFO(capacity > 0 && minBlock > 0, "BuddyAllocator: capacity/minBlock must be > 0"); + mMinBlock = roundUpPow2(minBlock); + TLLM_CHECK_WITH_INFO(capacity >= mMinBlock, "BuddyAllocator: capacity < minBlock"); + + // Largest order L with minBlock< BuddyAllocator::alloc(std::size_t bytes) +{ + // Reject empty and anything larger than the usable arena BEFORE orderForBytes/roundUpPow2 — + // a `bytes` near SIZE_MAX would overflow roundUpPow2's doubling (p wraps to 0 -> infinite loop). + // `mUsable` is a power of two well under 2^63, so this bound is also the over-large guard. + if (bytes == 0 || bytes > mUsable) + { + return std::nullopt; + } + std::size_t const want = orderForBytes(bytes); + if (want > mMaxOrder) + { + return std::nullopt; // larger than the whole arena + } + // Find the smallest order >= want that has a free block. + std::uint32_t cur = static_cast(want); + while (cur <= mMaxOrder && mFree[cur].empty()) + { + ++cur; + } + if (cur > mMaxOrder) + { + return std::nullopt; // no block big enough is free (fragmented / full) + } + // Take a block at `cur` and split down to `want`, returning the lower buddy each time. + std::uint64_t block = *mFree[cur].begin(); + mFree[cur].erase(mFree[cur].begin()); + while (cur > want) + { + --cur; + std::uint64_t const buddy = block + (mMinBlock << cur); + mFree[cur].insert(buddy); // keep the upper half free, descend into the lower half + } + mAllocOrder.emplace(block, static_cast(want)); + return block; +} + +void BuddyAllocator::free(std::uint64_t offset) +{ + auto it = mAllocOrder.find(offset); + if (it == mAllocOrder.end()) + { + return; // not a live allocation (double free / bad offset) -> ignore, stay robust + } + std::uint32_t order = it->second; + mAllocOrder.erase(it); + + // Coalesce with the buddy while it is free at the same order. + while (order < mMaxOrder) + { + std::uint64_t const buddy = offset ^ (mMinBlock << order); + auto bit = mFree[order].find(buddy); + if (bit == mFree[order].end()) + { + break; // buddy not free -> stop merging + } + mFree[order].erase(bit); + offset = offset < buddy ? offset : buddy; // merged block starts at the lower address + ++order; + } + mFree[order].insert(offset); +} + +std::size_t BuddyAllocator::blockBytes(std::uint64_t offset) const noexcept +{ + auto it = mAllocOrder.find(offset); + return it == mAllocOrder.end() ? 0 : (mMinBlock << it->second); +} + +std::size_t BuddyAllocator::freeBytes() const noexcept +{ + std::size_t total = 0; + for (std::uint32_t o = 0; o <= mMaxOrder; ++o) + { + total += mFree[o].size() * (mMinBlock << o); + } + return total; +} + +std::size_t BuddyAllocator::maxAllocBytes() const noexcept +{ + for (std::uint32_t o = mMaxOrder + 1; o-- > 0;) + { + if (!mFree[o].empty()) + { + return mMinBlock << o; + } + } + return 0; +} + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BuddyAllocator.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BuddyAllocator.h new file mode 100644 index 000000000000..cc2cc54821de --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BuddyAllocator.h @@ -0,0 +1,136 @@ +/* + * 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 +#include +#include +#include +#include +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +// ============================================================================ +// BuddyAllocator — power-of-two buddy allocator over a logical byte space +// ---------------------------------------------------------------------------- +// Pure logic, no GPU / threads / IO — fully unit-testable. It is the data-region allocator for the +// bounce v2 arena (one shared registered buffer): instead of fixed full slots, each chunk gets a +// region whose buddy block is sized to its actual bytes, so MANY small transfers fit with less +// internal waste while a transfer larger than the whole buffer still streams through (each chunk is +// bounded by maxChunkSizeBytes and allocations are recycled after scatter completion). +// +// alloc(bytes) rounds up to a power-of-two multiple of minBlock and returns the byte offset of a +// free block of that order (splitting a larger one if needed). free(offset) coalesces with the +// buddy block whenever it is also free. Properties: +// - Structured external fragmentation: only free buddies of the same order coalesce, so total +// free bytes may exceed a request while no sufficiently large block is currently available. +// - Internal fragmentation is less than 2x because requests round up to a power-of-two block. +// - alloc returns nullopt when no block of the needed order is free (caller applies backpressure; +// a large alloc may transiently fail under fragmentation but a later free+coalesce frees a +// high-order block — no deadlock, since frees come from independently-completing flows). +// +// The allocator manages OFFSETS in [0, capacity); the actual GPU buffer (base ptr) lives elsewhere +// (BounceArena); CreditScheduler wraps this allocator to hand out region offsets over that buffer. +// +// SIZING +// capacity is rounded DOWN to the largest minBlock * 2^L that fits (the remainder is unused); +// minBlock is rounded UP to a power of two. So capacity()/minBlock() may differ from the ctor args +// — always query them, never assume the raw inputs. +// +// INTERNALS +// Block sizes are exactly minBlock< order for every live block (so free() knows its size to coalesce). +// alloc(bytes): want = the order of the smallest power-of-two block >= max(bytes, minBlock); find +// the LOWEST order >= want that has a free block; SPLIT down to `want`, pushing the upper half of +// each split back to mFree[order-1] and descending into the lower half; return the block offset. +// free(offset): look up its order; while the buddy (offset XOR (minBlock< orders 0..2; mFree shown as {order:[offsets]}): +// start mFree={2:[0]} one 1024-block at offset 0 +// alloc(200) split 1024 -> upper 512@512, 256@256; return 0 mFree={0:[256], 1:[512]} +// alloc(200) take the order-0 block at 256; return 256 mFree={1:[512]} +// free(0) buddy(0)=256 is LIVE -> no merge mFree={0:[0], 1:[512]} +// free(256) buddy(256)=0 free -> merge to 512@0; buddy=512 free -> merge to 1024@0 mFree={2:[0]} +// ============================================================================ +class BuddyAllocator +{ +public: + /// @param capacity total bytes (rounded DOWN to the largest minBlock * 2^L that fits). + /// @param minBlock smallest allocatable block (rounded UP to a power of two); the order-0 size. + BuddyAllocator(std::size_t capacity, std::size_t minBlock); + + /// Allocate at least `bytes` (>0). Returns the block's OFFSET only (not its size), or nullopt if + /// no free block of the required order exists right now. + /// + /// Why offset-only is sufficient: free() is offset-keyed (the order/size is looked up internally + /// via mAllocOrder), so the caller never needs to remember the block size to release it. And the + /// caller must use its OWN requested length for the actual transfer — NOT the rounded-up block + /// size: the slack between `bytes` and minBlock< alloc(std::size_t bytes); + + /// Free a block previously returned by alloc(). Coalesces with its buddy when possible. + void free(std::uint64_t offset); + + /// The actual (rounded-up, power-of-two) size of the live block at `offset`, or 0 if `offset` is + /// not a live allocation. For metrics/inspection — normal callers use their own requested length. + [[nodiscard]] std::size_t blockBytes(std::uint64_t offset) const noexcept; + + [[nodiscard]] std::size_t capacity() const noexcept + { + return mUsable; + } + + [[nodiscard]] std::size_t minBlock() const noexcept + { + return mMinBlock; + } + + /// Sum of all free block sizes (bytes). + [[nodiscard]] std::size_t freeBytes() const noexcept; + + /// Largest single alloc that can succeed right now (0 if none) — i.e. minBlock << (highest + /// order with a free block). Lets callers detect temporary allocation backpressure. + [[nodiscard]] std::size_t maxAllocBytes() const noexcept; + + /// Number of currently-allocated blocks (for tests / metrics / leak checks). + [[nodiscard]] std::size_t liveBlocks() const noexcept + { + return mAllocOrder.size(); + } + +private: + static std::size_t roundUpPow2(std::size_t v); + [[nodiscard]] std::size_t orderForBytes(std::size_t bytes) const; // smallest order fitting bytes + + std::size_t mMinBlock{}; // order-0 block size (power of two) + std::size_t mUsable{}; // minBlock << mMaxOrder + std::uint32_t mMaxOrder{}; // top order (one block of this size = whole arena) + std::vector> mFree; // mFree[order] = free block offsets at that order + std::unordered_map mAllocOrder; // allocated offset -> order +}; + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ControlChannel.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ControlChannel.h new file mode 100644 index 000000000000..fdd13c098749 --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ControlChannel.h @@ -0,0 +1,94 @@ +/* + * 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 +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +// ============================================================================ +// ControlChannel — bounce v2 control-plane transport +// ---------------------------------------------------------------------------- +// Role +// Carries the small control messages of the credit protocol between agents: +// WANT / GRANT / DATA / ACK (encoded by BounceMessage). It is a named-peer, +// asynchronous message channel: sendTo(peer, blob) is fire-and-forget and recv() +// pulls the next inbound (peer, blob). The capability handshake is exchanged +// out-of-band in AgentDesc, not as a ControlChannel message. WANT still carries +// the sender's endpoint so the receiver can self-bootstrap the reverse path; +// an empty WANT cancels a request (there is no separate RETURN message). +// +// Why it is abstract +// The data plane (bulk RDMA) is always NIXL. The control plane is pluggable so +// an alternative transport can be added without touching the reactor / credit +// logic; ZmqControlChannel (dedicated, event-driven zmq sockets) is the only +// production implementation. +// Crucially, it does NOT use NIXL's postXferReq notifMsg — v2 abandons it entirely +// (data-landed is detected by the sender polling getXferStatus, which already +// includes NIXL's ucp_ep_flush_nbx). +// +// Ordering / correctness contract +// The channel need NOT order messages w.r.t. the RDMA data plane: the only +// message with a data dependency is DATA ("go scatter"), and the sender emits +// DATA *after* observing getXferStatus==SUCCESS, so the data is already landed +// at the remote regardless of which transport carries DATA. Messages sent from +// one peer MUST be delivered FIFO: GRANT entries are associated with chunks by +// order and do not carry an explicit chunk index. +// +// Threading contract +// recv() is called only by the single IO-reactor thread. sendTo()/addPeer() may +// be called concurrently (reactor, app threads issuing WANT, and scatter workers +// issuing ACK). Implementations must make sendTo()/addPeer() mutually thread-safe. +// Only the reactor calls recv(), so implementations need not support concurrent +// receivers. +// ============================================================================ +class ControlChannel +{ +public: + virtual ~ControlChannel() = default; + + /// The endpoint other agents need to reach us: a ZMQ address or serialized NIXL metadata, + /// depending on the implementation. Advertised in the AgentDesc capability handshake and in + /// WANT for reverse-path bootstrap. + [[nodiscard]] virtual std::string localEndpoint() const = 0; + + /// Register where to reach `peer` (its localEndpoint()). Must be called before + /// the first sendTo(peer, ...). Idempotent. Returns false when the underlying transport + /// cannot load an optional endpoint; implementations whose endpoint is mandatory may throw + /// when it is empty or malformed. + virtual bool addPeer(std::string const& peer, std::string const& endpoint) = 0; + + /// Forget `peer`: drop its send-side socket/endpoint (e.g. when the peer is lost, so the + /// per-peer resource isn't kept alive until shutdown). Idempotent; thread-safe with + /// sendTo()/addPeer(). A sendTo() to a removed peer is dropped until it is re-added — either + /// explicitly via addPeer() or implicitly by the WANT self-bootstrap in the receiver's onWant. + virtual void removePeer(std::string const& peer) = 0; + + /// Fire-and-forget send of an encoded BounceMessage `blob` to `peer`. + /// Thread-safe with other sendTo()/addPeer() calls. + virtual void sendTo(std::string const& peer, std::string const& blob) = 0; + + /// Block up to `timeoutMs` for one inbound message. On arrival, fills `outPeer` + /// (sender agent name) and `outBlob` (encoded message) and returns true; + /// returns false on timeout. Reactor-thread only. + [[nodiscard]] virtual bool recv(std::string& outPeer, std::string& outBlob, int timeoutMs) = 0; +}; + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.cpp new file mode 100644 index 000000000000..72cae057868d --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.cpp @@ -0,0 +1,468 @@ +/* + * 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/CreditScheduler.h" + +#include "tensorrt_llm/common/assert.h" + +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +CreditScheduler::CreditScheduler(std::uint64_t baseAddr, std::size_t arenaSizeBytes, + std::size_t arenaAllocationGranularityBytes, std::uint32_t maxInflightChunksPerRequest, + std::function clock) + : mArena(arenaSizeBytes, arenaAllocationGranularityBytes) + , mBaseAddr(baseAddr) + , mMaxInflightChunksPerRequest(maxInflightChunksPerRequest == 0 ? 1 : maxInflightChunksPerRequest) + , mEagerBudgetBytes(mArena.capacity() / 2) + , mClock(clock ? std::move(clock) : [] { return std::chrono::steady_clock::now(); }) +{ +} + +std::size_t CreditScheduler::activeFlows() const +{ + std::lock_guard lk(mMu); + std::size_t n = 0; + for (auto const& [_, st] : mFlows) + { + if (!st.pending.empty()) + { + ++n; + } + } + return n; +} + +std::size_t CreditScheduler::heldCount(std::string const& flow) const +{ + std::lock_guard lk(mMu); + auto it = mFlows.find(flow); + return it == mFlows.end() ? 0 : it->second.held.size(); +} + +void CreditScheduler::ensureInRing(std::string const& flow) +{ + if (std::find(mRing.begin(), mRing.end(), flow) == mRing.end()) + { + mRing.push_back(flow); + } +} + +void CreditScheduler::dropFromRing(std::string const& flow) +{ + if (mDrainFlow && *mDrainFlow == flow) + { + mDrainFlow.reset(); + } + auto it = std::find(mRing.begin(), mRing.end(), flow); + if (it != mRing.end()) + { + auto const idx = static_cast(it - mRing.begin()); + mRing.erase(it); + if (mRing.empty()) + { + // Erased the LAST flow: the modulo below would divide by zero (SIGFPE), so reset and bail. + mCursor = 0; + return; + } + if (idx < mCursor) + { + --mCursor; + } + mCursor %= mRing.size(); + } +} + +void CreditScheduler::maybeActivateDrain() +{ + if (mDrainFlow || mRing.empty()) + { + return; + } + + std::uint64_t const ringSize = static_cast(mRing.size()); + std::uint64_t const bypassThreshold = std::max(kMinimumBypassGrants, kBypassRounds * ringSize); + std::optional oldestBlockedSequence; + + // Scan from the round-robin cursor so equal-age candidates retain the scheduler's normal order. + for (std::size_t k = 0; k < mRing.size(); ++k) + { + std::size_t const idx = (mCursor + k) % mRing.size(); + auto const fit = mFlows.find(mRing[idx]); + TLLM_CHECK_DEBUG(fit != mFlows.end()); + if (fit == mFlows.end()) + { + continue; + } + auto const& st = fit->second; + if (st.pending.empty() || st.held.size() >= mMaxInflightChunksPerRequest || !st.blockedAtGrantSequence) + { + continue; + } + std::uint64_t const bypassed = mGrantSequence - *st.blockedAtGrantSequence; + if (bypassed < bypassThreshold) + { + continue; + } + if (!oldestBlockedSequence || *st.blockedAtGrantSequence < *oldestBlockedSequence) + { + oldestBlockedSequence = st.blockedAtGrantSequence; + mDrainFlow = mRing[idx]; + } + } +} + +void CreditScheduler::issueGrant( + std::string const& flow, FlowState& st, std::size_t ringIdx, std::uint64_t offset, std::vector& grants) +{ + std::uint32_t const want = st.pending.front(); + st.pending.pop_front(); + st.held.insert(offset); + st.blockedAtGrantSequence.reset(); + st.lastProgress = mClock(); // issuing a grant renews the flow's lease + grants.push_back(Grant{flow, offset, mBaseAddr + offset, want}); + mCursor = (ringIdx + 1) % mRing.size(); // next sweep starts AFTER this flow -> round-robin + ++mGrantSequence; +} + +// Hand out as many region grants as possible RIGHT NOW, fairly and bounded. Called whenever space +// frees up or new demand arrives (onWant / onScatterDone / releaseLocal / reclaim*). Three rules: +// 1. Fair: rotate over flows round-robin (mRing + mCursor). A head chunk that repeatedly fails to +// fit while other grants pass it eventually enters receiver drain mode: no NEW remote grants +// are issued until that head fits, so sustained smaller traffic cannot starve it. +// 2. Per-request cap: a flow may hold at most mMaxInflightChunksPerRequest allocations; more must +// wait for scatter completion to free one. +// 3. Arena capacity: all flows' regions share one buddy arena; if the next chunk doesn't fit, skip +// this flow (a smaller chunk elsewhere may still fit) -> backpressure, never an error. +// Shape: each inner sweep grants AT MOST ONE region then breaks, advancing the cursor past the flow +// just served; the outer loop re-sweeps from there. That one-at-a-time + advance gives strict +// rotation (A,B,A,B,...) instead of filling one request to its in-flight cap first. It stops when a full +// sweep grants nothing (every flow is done / at its in-flight cap / can't fit). Returns the GRANTs. +std::vector CreditScheduler::schedule() +{ + std::vector grants; + // Re-sweep as long as the previous sweep granted something; stop when a whole sweep makes no + // progress or the ring is empty. + while (!mRing.empty()) + { + maybeActivateDrain(); + if (mDrainFlow) + { + auto fit = mFlows.find(*mDrainFlow); + TLLM_CHECK_DEBUG(fit != mFlows.end()); + if (fit == mFlows.end() || fit->second.pending.empty() + || fit->second.held.size() >= mMaxInflightChunksPerRequest) + { + mDrainFlow.reset(); + continue; + } + + auto& st = fit->second; + std::uint32_t const want = st.pending.front(); + auto off = mArena.alloc(want); + if (!off) + { + // Existing remote regions keep progressing and freeing space, but do not refill them + // with smaller chunks. acquireLocal() is deliberately unaffected: this is a + // receiver-only admission barrier and cannot introduce a bidirectional circular wait. + return grants; + } + + auto const ringIt = std::find(mRing.begin(), mRing.end(), *mDrainFlow); + TLLM_CHECK_DEBUG(ringIt != mRing.end()); + std::size_t const idx = static_cast(ringIt - mRing.begin()); + issueGrant(*mDrainFlow, st, idx, *off, grants); + mDrainFlow.reset(); + continue; + } + + bool progress = false; + // One round-robin sweep from the cursor; grant the first eligible flow whose next chunk fits. + for (std::size_t k = 0; k < mRing.size(); ++k) + { + std::size_t const idx = (mCursor + k) % mRing.size(); + // find() (not operator[]) — the ring/flows invariant says the key exists, but [] would + // fabricate an un-erasable empty FlowState tombstone if it ever didn't (reintroducing the + // very leak the design fights). Skip loudly instead. + auto fit = mFlows.find(mRing[idx]); + TLLM_CHECK_DEBUG(fit != mFlows.end()); + if (fit == mFlows.end()) + { + continue; + } + auto& st = fit->second; + if (st.pending.empty() || st.held.size() >= mMaxInflightChunksPerRequest) + { + continue; // nothing more wanted, or at the per-request in-flight cap + } + std::uint32_t const want = st.pending.front(); + auto off = mArena.alloc(want); + if (!off) + { + if (!st.blockedAtGrantSequence) + { + st.blockedAtGrantSequence = mGrantSequence; + } + continue; // arena can't fit this chunk now -> try another flow (smaller may fit) + } + issueGrant(mRing[idx], st, idx, *off, grants); + progress = true; + // One grant per sweep: break out and let the outer loop re-sweep from the advanced cursor, + // so grants alternate across flows (strict rotation) rather than filling one flow first. + break; + } + if (!progress) + { + break; + } + } + return grants; +} + +void CreditScheduler::eraseIfDone(std::string const& flow) +{ + auto it = mFlows.find(flow); + if (it != mFlows.end() && it->second.pending.empty() && it->second.held.empty()) + { + mFlows.erase(it); + dropFromRing(flow); + } +} + +std::vector CreditScheduler::onWant(std::string const& flow, std::vector const& chunkBytes) +{ + std::lock_guard lk(mMu); + auto& st = mFlows[flow]; + st.pending.assign(chunkBytes.begin(), chunkBytes.end()); + st.blockedAtGrantSequence.reset(); + st.lastProgress = mClock(); // a fresh WANT renews the flow's lease + if (!chunkBytes.empty()) + { + ensureInRing(flow); + } + else + { + // cancel: drop now if nothing is still in flight; otherwise reclaimed when held drains. + eraseIfDone(flow); + } + return schedule(); +} + +std::vector CreditScheduler::onScatterDone(std::string const& flow, std::uint64_t offset) +{ + std::lock_guard lk(mMu); + auto it = mFlows.find(flow); + if (it != mFlows.end() && it->second.held.erase(offset) > 0) + { + mArena.free(offset); + it->second.lastProgress = mClock(); // a completed scatter is flow progress: renew the lease + } + // else: not held by this flow (dup ACK / already reclaimed) -> ignore, stay idempotent. + eraseIfDone(flow); + return schedule(); +} + +void CreditScheduler::dropFlow(std::string const& flow, std::unordered_set const& busy, + std::vector& deferredOut, std::chrono::milliseconds quarantineFor) +{ + auto it = mFlows.find(flow); + if (it == mFlows.end()) + { + return; + } + for (auto off : it->second.held) + { + if (busy.count(off) > 0) + { + // A scatter is still reading this region -> caller frees it later via freeOrphanRegion; + // track it as an orphan so that call only frees a genuine deferred region. + deferredOut.push_back(off); + mOrphans.insert(off); + } + else if (quarantineFor.count() > 0) + { + // Receiver-initiated reclaim: the peer may still be RDMA-writing this granted region + // (a one-sided write cannot be aborted, and no sender-side drain preceded this call). + // Keep it allocated (never re-grantable) until reapQuarantine() passes the deadline. + mQuarantined.emplace(off, mClock() + quarantineFor); + } + else + { + mArena.free(off); + } + } + mFlows.erase(it); + dropFromRing(flow); +} + +std::vector CreditScheduler::reclaimByPrefix(std::string const& prefix, + std::unordered_set const& busy, std::vector& deferredOut, + std::chrono::milliseconds quarantineFor) +{ + // Guard the degenerate empty prefix: compare(0,0,"")==0 for EVERY key, so it would reclaim all + // flows of all peers. Callers always pass a real "peer" prefix; refuse empty defensively. + if (prefix.empty()) + { + return {}; + } + std::lock_guard lk(mMu); + std::vector victims; + for (auto const& [key, st] : mFlows) + { + if (key.size() >= prefix.size() && key.compare(0, prefix.size(), prefix) == 0) + { + victims.push_back(key); + } + } + for (auto const& key : victims) + { + dropFlow(key, busy, deferredOut, quarantineFor); + } + return schedule(); +} + +std::vector CreditScheduler::reclaimFlow(std::string const& flow, std::unordered_set const& busy, + std::vector& deferredOut, std::chrono::milliseconds quarantineFor) +{ + std::lock_guard lk(mMu); + dropFlow(flow, busy, deferredOut, quarantineFor); + return schedule(); +} + +std::vector CreditScheduler::staleFlows(std::chrono::milliseconds idleLimit) const +{ + std::lock_guard lk(mMu); + auto const now = mClock(); + std::vector stale; + for (auto const& [key, st] : mFlows) + { + // The lease protects arena REGIONS, so only flows holding one can go stale. A pending-only + // flow ties up no memory and may legitimately be idle for long: it is simply queued behind a + // full arena (its own sender's requestTimeoutMs bounds that wait). If its sender is dead it + // lingers as a few bytes of bookkeeping until a grant finally starts its lease. + if (!st.held.empty() && now - st.lastProgress > idleLimit) + { + stale.push_back(key); + } + } + return stale; +} + +std::vector CreditScheduler::reapQuarantine() +{ + std::lock_guard lk(mMu); + auto const now = mClock(); + bool freed = false; + for (auto it = mQuarantined.begin(); it != mQuarantined.end();) + { + if (now >= it->second) + { + mArena.free(it->first); + it = mQuarantined.erase(it); + freed = true; + } + else + { + ++it; + } + } + // Nothing freed -> nothing changed; skip the schedule() ring sweep this polling tick runs into. + return freed ? schedule() : std::vector{}; +} + +bool CreditScheduler::heldByFlow(std::string const& flow, std::uint64_t offset) const +{ + std::lock_guard lk(mMu); + auto it = mFlows.find(flow); + return it != mFlows.end() && it->second.held.count(offset) > 0; +} + +bool CreditScheduler::knowsFlow(std::string const& flow) const +{ + std::lock_guard lk(mMu); + return mFlows.find(flow) != mFlows.end(); +} + +std::vector CreditScheduler::freeOrphanRegion(std::uint64_t offset) +{ + std::lock_guard lk(mMu); + // Only free a region we actually deferred as an orphan. Guards against a stray/duplicate call + // freeing an offset that may since have been re-allocated to a live flow (defense in depth; the + // transport already gates on its own mScattering orphaned-flag map). + if (mOrphans.erase(offset) > 0) + { + mArena.free(offset); + } + return schedule(); +} + +std::optional CreditScheduler::acquireLocal(std::size_t bytes, bool eager) +{ + std::lock_guard lk(mMu); + auto off = mArena.alloc(bytes); + if (!off) + { + return std::nullopt; // arena full/fragmented -> caller parks + retries (never blocks) + } + if (eager) + { + // Cap all eager (credit-less) staging at half the arena so incoming grants can always + // progress (see header). Budget accounting uses the ROUNDED buddy-block size — that is what + // the arena actually loses. + std::size_t const rounded = mArena.blockBytes(*off); + if (mEagerHeldBytes + rounded > mEagerBudgetBytes) + { + mArena.free(*off); + return std::nullopt; // over the eager budget -> caller parks; credit path unaffected + } + mEagerHeld.emplace(*off, rounded); + mEagerHeldBytes += rounded; + } + mLocalHeld.insert(*off); + return *off; +} + +void CreditScheduler::promoteLocal(std::uint64_t offset) +{ + std::lock_guard lk(mMu); + auto it = mEagerHeld.find(offset); + if (it != mEagerHeld.end()) + { + mEagerHeldBytes -= it->second; + mEagerHeld.erase(it); + } +} + +std::vector CreditScheduler::releaseLocal(std::uint64_t offset) +{ + std::lock_guard lk(mMu); + if (mLocalHeld.erase(offset) > 0) + { + auto it = mEagerHeld.find(offset); + if (it != mEagerHeld.end()) + { + mEagerHeldBytes -= it->second; + mEagerHeld.erase(it); + } + mArena.free(offset); + } + return schedule(); // freed bytes may now let a waiting remote flow alloc its next chunk +} + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.h new file mode 100644 index 000000000000..ea9accc21dbb --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/CreditScheduler.h @@ -0,0 +1,263 @@ +/* + * 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/cache_transmission/nixl_utils/bounce/BuddyAllocator.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +/// A credit handed to a flow: exclusive write permission for one receiver arena allocation. +/// `offset` is the region's arena offset (its opaque handle), `addr = baseAddr + offset` the +/// absolute device address, and `len` the chunk's packed transfer length. The buddy allocation +/// backing the region may be larger than `len`. +struct Grant +{ + std::string flow; // flow id ("peerrid") the grant belongs to (NOT a bare agent name) + std::uint64_t offset{}; // arena offset (region handle) + std::uint64_t addr{}; // baseAddr + offset (what the sender RDMA-writes to) + std::uint32_t len{}; // transfer length = the chunk's packed bytes +}; + +/// Receiver-side credit allocator + fair scheduler over a single shared arena — pure logic, +/// no threads / IO / CUDA (the GPU buffer lives in the transport; this owns only a BuddyAllocator +/// over byte offsets + the base address for computing absolute addrs). +/// +/// VARIABLE REGIONS: instead of fixed full slots (every region maxChunkSizeBytes regardless of +/// need), each chunk requests only its packed byte extent; the buddy allocator rounds that up to a +/// power of two. This reduces waste vs fixed slots for small requests, while a request larger than +/// the whole arena streams through chunk by chunk. The per-flow cap is in allocation count; arena +/// capacity bounds aggregate concurrency. +/// +/// WHY BUDDY (vs an exact-fit / first-fit allocator): the up-to-2x internal fragmentation is the +/// accepted price for deterministic O(1) coalescing under this workload, and it rarely bites — +/// full chunks are exactly maxChunkSizeBytes (a power of two, zero rounding), only tail/small +/// chunks round up, and regions live for milliseconds so the waste costs transient concurrency, +/// not resident memory. Exact-fit would erase that rounding but frees here arrive in arbitrary +/// order (scatters of independent flows), the worst case for unstructured external fragmentation +/// in a long-running arena; buddy's fragmentation is structured and self-healing (freed buddies +/// always re-coalesce). It also keeps backpressure detection trivial (maxAllocBytes = highest +/// non-empty order) and the metadata host-side-only (XOR buddy lookup needs no boundary tags in +/// the GPU buffer). +/// +/// TERMINOLOGY: the string identifying a client is an opaque **flow id** ("peerNamerid"), NOT +/// an agent name (cf. `peer` in ControlChannel/BounceTransport). reclaimByPrefix() is the only +/// peer-level op. The local sender shares this same arena via acquireLocal() (gather staging). +/// +/// Each event method mutates state and returns the GRANTs the caller should now send. +/// +/// THREADING CONTRACT: every public method takes the internal mutex, so calls are safe from any +/// thread. The IO thread remains the primary owner (all flow-state events happen there); the one +/// cross-thread caller is acquireLocal() from submit() app threads (eager gather staging). +class CreditScheduler +{ +public: + /// @param baseAddr Device address of arena offset 0 (Grant.addr = baseAddr + offset). + /// @param arenaSizeBytes Total arena size. + /// @param arenaAllocationGranularityBytes Smallest requested buddy allocation size. + /// @param maxInflightChunksPerRequest Per-request in-flight allocation cap. + /// @param clock Injectable time source for the flow leases / quarantine deadlines (tests pass a + /// fake clock and advance it instead of sleeping); default is steady_clock::now. + CreditScheduler(std::uint64_t baseAddr, std::size_t arenaSizeBytes, std::size_t arenaAllocationGranularityBytes, + std::uint32_t maxInflightChunksPerRequest, std::function clock = {}); + + /// Flow announces the per-chunk byte sizes it wants to write (in order). EMPTY = cancel. + [[nodiscard]] std::vector onWant(std::string const& flow, std::vector const& chunkBytes); + + /// A region finished scattering on the receiver -> free it and re-schedule. Idempotent. + [[nodiscard]] std::vector onScatterDone(std::string const& flow, std::uint64_t offset); + + /// Reclaim every flow whose id starts with `prefix` (drop all flows of a gone peer). DEFERS + /// freeing any held region whose offset is in `busy` (a scatter is still reading it): those are + /// appended to `deferredOut` instead of freed, and the caller MUST later call freeOrphanRegion() + /// for each once its scatter completes. `quarantineFor` > 0 marks a RECEIVER-initiated reclaim + /// (peer loss / lease expiry): unlike a sender cancel there is no sender-side drain guaranteeing + /// the RDMA writes into the granted regions have ended, and a one-sided write cannot be aborted + /// — so every non-busy held region is QUARANTINED for that duration (kept out of the arena; + /// reapQuarantine() frees it once the deadline passes) instead of freed immediately. + [[nodiscard]] std::vector reclaimByPrefix(std::string const& prefix, + std::unordered_set const& busy, std::vector& deferredOut, + std::chrono::milliseconds quarantineFor = std::chrono::milliseconds{0}); + + /// Free a region deferred by reclaimByPrefix (its in-flight scatter has finished) + re-schedule. + [[nodiscard]] std::vector freeOrphanRegion(std::uint64_t offset); + + /// Cancel ONE flow (explicit abort / empty WANT): free its held regions and drop it. Any held + /// region in `busy` (a scatter is still reading it) is deferred to `deferredOut` instead of freed + /// (caller later calls freeOrphanRegion). Frees the granted-but-unwritten regions a failed sender + /// would otherwise leak until peer loss. The immediate-free default is safe ONLY when the SENDER + /// initiated the reclaim (its cancel is sent after draining in-flight writes); receiver-initiated + /// reclaims pass `quarantineFor` > 0 (same semantics as reclaimByPrefix). + [[nodiscard]] std::vector reclaimFlow(std::string const& flow, std::unordered_set const& busy, + std::vector& deferredOut, + std::chrono::milliseconds quarantineFor = std::chrono::milliseconds{0}); + + /// Flows HOLDING at least one region with no progress (no WANT refresh, no grant issued, no + /// scatter completed) for longer than `idleLimit`. A dead sender emits neither DATA nor a cancel + /// — unobservable through the protocol alone — so the receiver reclaims these via + /// reclaimFlow(quarantineFor > 0). Pending-only flows tie up no memory and are never reported: + /// they may legitimately queue behind a full arena for a long time. + [[nodiscard]] std::vector staleFlows(std::chrono::milliseconds idleLimit) const; + + /// Free every quarantined region whose deadline has passed (no write posted before its flow's + /// lease expired can plausibly still be in flight) + re-schedule. Returns the resulting grants. + [[nodiscard]] std::vector reapQuarantine(); + + /// True if `flow` currently holds region `offset`. Lets the transport drop a late DATA for a + /// region this flow no longer owns (cancelled/reclaimed) — scattering a freed/re-granted region + /// would corrupt another flow's data. + [[nodiscard]] bool heldByFlow(std::string const& flow, std::uint64_t offset) const; + + /// True if `flow` is currently tracked (has pending chunks and/or held regions). A non-empty + /// WANT has no retransmission path, so the transport drops a repeat one for a tracked flow — + /// re-queueing would re-grant over the still-held regions (the sender never writes the extras, + /// leaking them) and the lastProgress refresh would keep renewing the very lease that exists + /// to reclaim that state. + [[nodiscard]] bool knowsFlow(std::string const& flow) const; + + // ---- local (sender) role: gather staging from the SAME arena ---- + /// Allocate a region of `bytes` for local gather staging (non-blocking). Returns its offset, or + /// nullopt if the arena can't fit it right now (caller parks and retries). With `eager` the + /// allocation is additionally capped so that all eager (credit-less) local regions together stay + /// under HALF the arena capacity — on a bidirectional deployment this guarantees each side can + /// always still grant incoming regions, so two eager senders can never starve each other into a + /// circular wait (chunk waits for a GRANT the peer can't give because ITS arena is full of eager + /// staging, and vice versa). Credit-backed (non-eager) acquisitions are not capped. + [[nodiscard]] std::optional acquireLocal(std::size_t bytes, bool eager = false); + /// An eager region's credit arrived: it is now credit-backed, so stop counting it against the + /// eager budget (otherwise steady-state pipelining would be throttled to the eager cap even + /// though every in-flight chunk has its credit). No-op for non-eager offsets. + void promoteLocal(std::uint64_t offset); + /// Return a locally-held region (its chunk was ACKed / failed) to the arena + re-schedule. + [[nodiscard]] std::vector releaseLocal(std::uint64_t offset); + + [[nodiscard]] std::size_t localHeldCount() const + { + std::lock_guard lk(mMu); + return mLocalHeld.size(); + } + + // ---- inspectors (for tests / metrics) ---- + [[nodiscard]] std::size_t freeBytes() + { + std::lock_guard lk(mMu); + return mArena.freeBytes(); + } + + /// Largest region a fully-drained arena can ever hand out (the buddy allocator's usable capacity, + /// rounded down to one buddy block). A chunk larger than this can never be granted, so callers + /// clamp maxChunkSizeBytes to it. + [[nodiscard]] std::size_t arenaCapacity() const noexcept + { + return mArena.capacity(); + } + + /// Byte size of the buddy block backing a granted region offset (0 if not allocated). The whole + /// block belongs to one flow, so it bounds how far a scatter may read without touching another + /// flow's region. + [[nodiscard]] std::size_t regionBytes(std::uint64_t offset) const + { + std::lock_guard lk(mMu); + return mArena.blockBytes(offset); + } + + [[nodiscard]] std::size_t heldCount(std::string const& flow) const; + [[nodiscard]] std::size_t activeFlows() const; + + [[nodiscard]] std::size_t trackedFlows() const + { + std::lock_guard lk(mMu); + return mFlows.size(); + } + +private: + struct FlowState + { + std::deque pending; // per-chunk byte sizes still wanting a grant (FIFO) + std::unordered_set held; // region offsets currently granted to this flow + std::optional blockedAtGrantSequence; // first grant sequence where head did not fit + // Lease stamp: last WANT / grant issued / scatter completed. staleFlows() reports flows idle + // beyond the receiver's lease so their (possibly dead) sender can't leak regions forever. + std::chrono::steady_clock::time_point lastProgress; + }; + + std::vector schedule(); // grant while the arena has room and eligible flows exist + // Book one issued grant (shared by schedule()'s drain branch and round-robin sweep): pop the + // flow's head chunk, hold `offset`, renew the lease, append the Grant, and advance the cursor + // past ring slot `ringIdx`. Caller must hold mMu and have alloc'd `offset` for the head chunk. + void issueGrant( + std::string const& flow, FlowState& st, std::size_t ringIdx, std::uint64_t offset, std::vector& grants); + void maybeActivateDrain(); // age a repeatedly bypassed flow into receiver drain mode + void ensureInRing(std::string const& flow); // add flow to the round-robin ring if absent + void dropFromRing(std::string const& flow); // remove flow from the round-robin ring + void eraseIfDone(std::string const& flow); // pending empty && held empty -> drop the flow + // Free one flow's held regions (busy ones deferred to deferredOut + tracked as orphans; with + // quarantineFor > 0 the non-busy ones are quarantined instead of freed), then erase the flow + + // drop it from the ring. Shared by reclaimByPrefix and reclaimFlow. + void dropFlow(std::string const& flow, std::unordered_set const& busy, + std::vector& deferredOut, std::chrono::milliseconds quarantineFor); + + // Guards ALL mutable state below. Uncontended in practice: everything runs on the IO thread + // except acquireLocal() from submit() app threads (eager gather staging). + mutable std::mutex mMu; + + BuddyAllocator mArena; // the single shared region allocator (byte offsets) + std::uint64_t mBaseAddr{}; // device addr of offset 0 + std::uint32_t mMaxInflightChunksPerRequest{}; // per-request in-flight allocation cap + + std::unordered_map mFlows; // per-flow state (opaque flow id, NOT agent name) + std::unordered_set mLocalHeld; // regions taken for local gather staging + // Rounded (buddy-block) bytes of local regions acquired EAGERLY (before their credit arrived), + // and the cap they must stay under (half the arena) — see acquireLocal(). Offsets still tracked + // in mEagerHeld so releaseLocal() knows how much to subtract. + std::unordered_map mEagerHeld; + std::size_t mEagerHeldBytes{0}; + std::size_t mEagerBudgetBytes{0}; + std::unordered_set + mOrphans; // regions deferred by flow/peer reclamation (busy scatter), awaiting freeOrphanRegion + // Regions reclaimed by a RECEIVER-initiated teardown while possibly still being RDMA-written by + // the peer (granted, DATA never arrived), mapped to their reuse deadline. They stay allocated in + // the arena (so schedule() can never re-grant them) until reapQuarantine() passes the deadline. + std::unordered_map mQuarantined; + std::function mClock; // injectable for tests + // Receiver-only anti-starvation barrier. Every successful remote grant advances mGrantSequence. + // A flow whose head allocation keeps failing while enough other grants pass it becomes mDrainFlow; + // schedule() then pauses NEW remote grants until that one head fits. Local acquireLocal() remains + // untouched, avoiding a cross-direction circular wait in the shared-arena case. + static constexpr std::uint64_t kMinimumBypassGrants{8}; + static constexpr std::uint64_t kBypassRounds{2}; + std::uint64_t mGrantSequence{0}; + std::optional mDrainFlow; + // Round-robin ring of active flow keys (insertion order). NOTE: "ring" not "order" — distinct + // from BuddyAllocator's size `order` (mArena), which is the power-of-two block exponent. + std::vector mRing; + std::size_t mCursor{0}; // round-robin cursor into mRing +}; + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ExecPool.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ExecPool.cpp new file mode 100644 index 000000000000..634bef303afa --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ExecPool.cpp @@ -0,0 +1,132 @@ +/* + * 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/ExecPool.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/logger.h" + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +ExecPool::ExecPool(std::uint32_t count, std::size_t maxDescsPerChunk, int deviceId, bool useZeroCopyArguments) + : mDeviceId(deviceId) +{ + TLLM_CHECK_WITH_INFO(count > 0, "ExecPool: count must be > 0"); + TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); + std::size_t const scratchBytes = maxDescsPerChunk * (2 * sizeof(std::uint64_t) + sizeof(std::uint32_t)); + // Gather/scatter copies sit on the KV-transfer critical path but share the GPU with model + // kernels (prefill/decode). At default priority a copy kernel queues behind those (measured + // ~260us avg, >1ms tail, vs ~65us of actual copy time). Greatest priority lets its blocks be + // scheduled as soon as any SM frees up, without preempting running work. + int leastPriority = 0; + int greatestPriority = 0; + TLLM_CUDA_CHECK(cudaDeviceGetStreamPriorityRange(&leastPriority, &greatestPriority)); + mCtxs.resize(count); + // A mid-loop TLLM_CUDA_CHECK throw would skip the destructor (the object is not fully + // constructed), leaking every stream/event/buffer allocated so far -> clean up and rethrow. + try + { + for (std::uint32_t i = 0; i < count; ++i) + { + auto& c = mCtxs[i]; + c.id = i; + c.scratchBytes = scratchBytes; + TLLM_CUDA_CHECK(cudaMalloc(&c.scratch, scratchBytes)); + // Map the pinned buffer when kernels should read plan arrays in place without an H2D copy. + TLLM_CUDA_CHECK(cudaHostAlloc( + &c.hostPinned, scratchBytes, useZeroCopyArguments ? cudaHostAllocMapped : cudaHostAllocDefault)); + if (useZeroCopyArguments) + { + TLLM_CUDA_CHECK(cudaHostGetDevicePointer(&c.hostPinnedDev, c.hostPinned, 0)); + } + TLLM_CUDA_CHECK(cudaStreamCreateWithPriority(&c.stream, cudaStreamNonBlocking, greatestPriority)); + TLLM_CUDA_CHECK(cudaEventCreateWithFlags(&c.event, cudaEventDisableTiming)); + mFree.push_back(i); + } + } + catch (...) + { + destroyContexts(); + throw; + } +} + +ExecPool::~ExecPool() +{ + destroyContexts(); +} + +void ExecPool::destroyContexts() +{ + // Select the owning device before freeing; otherwise CUDA teardown targets the current device of + // this thread. Must not throw (runs from the destructor and the constructor's failure path), so + // every cleanup uses the warn-only check. Fields of never-initialized contexts are nullptr. + TLLM_CUDA_CHECK_WARN(cudaSetDevice(mDeviceId)); + for (auto& c : mCtxs) + { + if (c.scratch != nullptr) + { + TLLM_CUDA_CHECK_WARN(cudaFree(c.scratch)); + } + if (c.hostPinned != nullptr) + { + TLLM_CUDA_CHECK_WARN(cudaFreeHost(c.hostPinned)); // also frees its hostPinnedDev alias + } + if (c.stream != nullptr) + { + TLLM_CUDA_CHECK_WARN(cudaStreamDestroy(c.stream)); + } + if (c.event != nullptr) + { + TLLM_CUDA_CHECK_WARN(cudaEventDestroy(c.event)); + } + } + mCtxs.clear(); + mFree.clear(); +} + +ExecCtx* ExecPool::tryAcquire() +{ + std::lock_guard lk(mMu); + if (mFree.empty()) + { + return nullptr; + } + std::uint32_t const id = mFree.front(); + mFree.pop_front(); + return &mCtxs[id]; +} + +void ExecPool::release(ExecCtx* ctx) +{ + if (ctx == nullptr) + { + return; + } + std::lock_guard lk(mMu); + mFree.push_back(ctx->id); +} + +std::size_t ExecPool::freeCount() +{ + std::lock_guard lk(mMu); + return mFree.size(); +} + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ExecPool.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ExecPool.h new file mode 100644 index 000000000000..a66be5706824 --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ExecPool.h @@ -0,0 +1,94 @@ +/* + * 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 +#include +#include +#include +#include + +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +// ============================================================================ +// ExecPool — gather/scatter execution contexts, decoupled from the data region +// ---------------------------------------------------------------------------- +// With the variable-region arena, the data buffer a chunk uses is just a region (offset,len) held +// until ACK — a LONG lifetime. The CUDA resources needed to RUN one gather/scatter kernel +// (a stream, a completion event, device `scratch` for the plan arrays, and pinned `hostPinned` +// H2D staging) are only needed DURING the kernel — a SHORT lifetime. Bundling them per-region (as +// the old per-slot model did) would force one stream/scratch per concurrent region, which doesn't +// scale when the arena holds many small regions. So they live in a small fixed pool of E contexts, +// borrowed for one gather/scatter and returned when the kernel completes. E bounds GPU kernel +// concurrency (independent of region count); the sender's gather (IO thread) and the receiver's +// scatter (worker threads) both borrow here, so acquire/release are thread-safe. +// ============================================================================ + +struct ExecCtx +{ + std::uint32_t id{}; + cudaStream_t stream{nullptr}; + cudaEvent_t event{nullptr}; // gather-completion event (cudaEventRecord/Query, no blocking sync) + void* scratch{nullptr}; // device plan arrays; unused when zero-copy arguments are enabled + void* hostPinned{nullptr}; // pinned staging for the plan arrays (H2D source, or read in-kernel) + void* hostPinnedDev{nullptr}; // device alias of hostPinned when zero-copy arguments are enabled + std::size_t scratchBytes{0}; +}; + +class ExecPool +{ +public: + /// Allocate `count` contexts, each with a stream/event + scratch/hostPinned sized for + /// `maxDescsPerChunk` plan entries. Throws on CUDA allocation failure. + /// @param useZeroCopyArguments Map hostPinned into the device address space so a kernel can read + /// plan arrays directly. + ExecPool(std::uint32_t count, std::size_t maxDescsPerChunk, int deviceId, bool useZeroCopyArguments = false); + ~ExecPool(); + + ExecPool(ExecPool const&) = delete; + ExecPool& operator=(ExecPool const&) = delete; + + /// Borrow a free context, or nullptr if all are in use (non-blocking — caller parks/retries, + /// never blocks the IO thread). Thread-safe. + [[nodiscard]] ExecCtx* tryAcquire(); + + /// Return a context borrowed via tryAcquire(). Thread-safe. + void release(ExecCtx* ctx); + + [[nodiscard]] std::uint32_t size() const noexcept + { + return static_cast(mCtxs.size()); + } + + [[nodiscard]] std::size_t freeCount(); + +private: + /// Free every allocated context resource (never throws; used by the destructor and the + /// constructor's failure path). + void destroyContexts(); + + int mDeviceId{0}; + std::vector mCtxs; + std::mutex mMu; + std::deque mFree; +}; + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/GatherScatterKernel.cu b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/GatherScatterKernel.cu new file mode 100644 index 000000000000..92445d51d41f --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/GatherScatterKernel.cu @@ -0,0 +1,75 @@ +/* + * 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/GatherScatterKernel.h" + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +namespace +{ + +// One thread block per buffer. Within a block, threads stride over the buffer. When src, dst +// and size are all 16-byte aligned, a vectorized uint4 path gives coalesced 16-byte stores. +__global__ void batchedCopyKernel( + std::uint64_t const* srcs, std::uint64_t const* dsts, std::uint32_t const* sizes, std::uint32_t n) +{ + std::uint32_t const buf = blockIdx.x; + if (buf >= n) + { + return; + } + auto const* src = reinterpret_cast(srcs[buf]); + auto* dst = reinterpret_cast(dsts[buf]); + std::uint32_t const size = sizes[buf]; + + auto const addrMask = reinterpret_cast(src) | reinterpret_cast(dst) + | static_cast(size); + if ((addrMask & 0xFU) == 0U) + { + auto const* s4 = reinterpret_cast(src); + auto* d4 = reinterpret_cast(dst); + std::uint32_t const n4 = size >> 4U; + for (std::uint32_t i = threadIdx.x; i < n4; i += blockDim.x) + { + d4[i] = s4[i]; + } + } + else + { + for (std::uint32_t i = threadIdx.x; i < size; i += blockDim.x) + { + dst[i] = src[i]; + } + } +} + +} // namespace + +cudaError_t launchBatchedCopy(std::uint64_t const* srcs, std::uint64_t const* dsts, std::uint32_t const* sizes, + std::uint32_t n, cudaStream_t stream) +{ + if (n == 0) + { + return cudaSuccess; + } + constexpr int kThreads = 256; + batchedCopyKernel<<>>(srcs, dsts, sizes, n); + return cudaGetLastError(); +} + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/GatherScatterKernel.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/GatherScatterKernel.h new file mode 100644 index 000000000000..0b7524eda4f0 --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/GatherScatterKernel.h @@ -0,0 +1,36 @@ +/* + * 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 +#include +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +/// Batched device-to-device copy: for each of `n` buffers, copy `sizes[i]` bytes from +/// `srcs[i]` to `dsts[i]`. `srcs`, `dsts`, `sizes` are DEVICE pointers to arrays of length `n`. +/// +/// Direction-agnostic: used for gather (srcs=scattered, dsts=arena-region+offset) and scatter +/// (srcs=arena-region+offset, dsts=scattered) — only the pointer arrays differ. 16-byte-aligned +/// buffers take a vectorized uint4 path; otherwise a byte path. +[[nodiscard]] cudaError_t launchBatchedCopy(std::uint64_t const* srcs, std::uint64_t const* dsts, + std::uint32_t const* sizes, std::uint32_t n, cudaStream_t stream); + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ZmqControlChannel.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ZmqControlChannel.cpp new file mode 100644 index 000000000000..4a699e9ef476 --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ZmqControlChannel.cpp @@ -0,0 +1,191 @@ +/* + * 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/ZmqControlChannel.h" + +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceNvtx.h" + +#include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/tllmException.h" + +#include +#include +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +namespace +{ +// Per-peer outbound queue cap (messages). With the non-blocking send in sendTo(), hitting this cap +// DROPS the message instead of blocking the caller. Control messages are tiny (tens to a few hundred +// bytes), so this bounds a stalled peer's queue to a few MB while sitting far above any legitimate +// in-flight burst (per-request chunk cap times concurrent flows), so drops indicate peer stall. +constexpr int kSendHwm = 1 << 16; +} // namespace + +ZmqControlChannel::ZmqControlChannel(std::string selfName, std::string const& bindAddr) + : mSelfName(std::move(selfName)) + , mCtx(/*io_threads=*/1) + , mRouter(mCtx, zmq::socket_type::router) +{ + // Identify ourselves to peers' ROUTERs and fail fast rather than queue forever if a peer + // is unreachable (ROUTER drops messages to unknown/again-full peers by default). + mRouter.set(zmq::sockopt::routing_id, mSelfName); + mRouter.set(zmq::sockopt::linger, 0); + // Accept a reconnecting peer that reuses an existing routing id. Our peers' DEALERs identify by a + // fixed agent name, so when a peer is forgotten (removePeer drops its DEALER) and later comes back + // — same agent name = same identity — it reconnects with the SAME routing id. Without handover the + // ROUTER REJECTS the new connection while the old one is still being reaped and SILENTLY DROPS its + // messages (a forgotten-then-readded peer's WANTs vanish until request timeout). HANDOVER hands the + // identity to the new connection instead. (Loopback reconnect makes this race easy to hit.) + mRouter.set(zmq::sockopt::router_handover, 1); + // zmq disables IPv6 on a socket by default, so binding an IPv6 address would fail. Enable it when + // the bind address is IPv6 (brackets, e.g. "tcp://[::1]:*"). Mirrors ucx_utils. (Default ctor arg + // is the IPv4 loopback, so tests are unaffected.) + if (bindAddr.find('[') != std::string::npos) + { + mRouter.set(zmq::sockopt::ipv6, 1); + } + mRouter.bind(bindAddr); + mEndpoint = mRouter.get(zmq::sockopt::last_endpoint); +} + +ZmqControlChannel::~ZmqControlChannel() = default; + +std::string ZmqControlChannel::localEndpoint() const +{ + return mEndpoint; +} + +bool ZmqControlChannel::addPeer(std::string const& peer, std::string const& endpoint) +{ + std::lock_guard lk(mMu); + if (mDealers.find(peer) != mDealers.end()) + { + return true; // idempotent + } + if (endpoint.empty()) + { + TLLM_THROW("ZmqControlChannel(%s): peer %s requires a non-empty endpoint", mSelfName.c_str(), peer.c_str()); + } + try + { + zmq::socket_t dealer(mCtx, zmq::socket_type::dealer); + dealer.set(zmq::sockopt::routing_id, mSelfName); // so peer's ROUTER sees us by name + dealer.set(zmq::sockopt::linger, 0); + dealer.set(zmq::sockopt::sndhwm, kSendHwm); // bound the queue; full -> sendTo drops (never blocks) + // Enable IPv6 unconditionally (off by default in zmq) so a DEALER can connect to an IPv6 peer + // endpoint; harmless when the endpoint is IPv4. Mirrors ucx_utils' connect socket. + dealer.set(zmq::sockopt::ipv6, 1); + dealer.connect(endpoint); + mDealers.emplace(peer, std::move(dealer)); + return true; + } + catch (zmq::error_t const& e) + { + TLLM_THROW( + "ZmqControlChannel(%s): peer %s has an invalid endpoint: %s", mSelfName.c_str(), peer.c_str(), e.what()); + } +} + +void ZmqControlChannel::removePeer(std::string const& peer) +{ + std::lock_guard lk(mMu); + // Erasing closes the DEALER (linger=0 -> no block on close); idempotent if `peer` is unknown. + // sendTo() holds mMu while looking up the dealer, so the close can never race a concurrent send. + mDealers.erase(peer); +} + +void ZmqControlChannel::sendTo(std::string const& peer, std::string const& blob) +{ + // Starts BEFORE the lock: exposes cross-thread contention (IO thread vs scatter-worker ACKs) + // plus the message copy + enqueue cost. Large spans here mean the control channel itself is + // the bottleneck (message too big, or a stalled peer backing up the DEALER queue). + BounceNvtxScope sendScope(kNvtxZmqSend, "zmqSend bytes=%zu", blob.size()); + std::lock_guard lk(mMu); + auto it = mDealers.find(peer); + if (it == mDealers.end()) + { + TLLM_LOG_WARNING( + "ZmqControlChannel(%s): sendTo unknown peer %s (call addPeer first)", mSelfName.c_str(), peer.c_str()); + return; + } + zmq::message_t msg(blob.data(), blob.size()); + try + { + // DEALER -> peer ROUTER; the peer receives [our routing id, blob]. NON-BLOCKING: this runs on + // the IO thread (and under mMu, which also gates submit()), so a blocking send to a stalled / + // unreachable peer whose queue is full (kSendHwm + TCP buffers) would wedge the whole reactor + // — exactly the hang the design forbids. With dontwait a full queue returns an empty result + // (EAGAIN) instead; we DROP the message. A dropped control message degrades the affected + // request to a request-timeout FAILURE rather than blocking this thread. + auto const sent = it->second.send(msg, zmq::send_flags::dontwait); + if (!sent.has_value()) + { + TLLM_LOG_WARNING("ZmqControlChannel(%s): send to %s dropped (queue full / peer stalled)", mSelfName.c_str(), + peer.c_str()); + } + } + catch (zmq::error_t const& e) + { + TLLM_LOG_WARNING("ZmqControlChannel(%s): send to %s failed: %s", mSelfName.c_str(), peer.c_str(), e.what()); + } +} + +bool ZmqControlChannel::recv(std::string& outPeer, std::string& outBlob, int timeoutMs) +{ + std::array items{{{mRouter.handle(), 0, ZMQ_POLLIN, 0}}}; + zmq::poll(items.data(), items.size(), std::chrono::milliseconds(timeoutMs)); + if ((items[0].revents & ZMQ_POLLIN) == 0) + { + return false; + } + // A DEALER->ROUTER message arrives as [identity, body]. The span excludes the poll wait above + // (that's idle time, not work): it measures the frame reads + the blob copy out to the caller. + BounceNvtxScope recvScope(kNvtxZmqRecv, "zmqRecv"); + zmq::message_t idFrame; + auto r1 = mRouter.recv(idFrame, zmq::recv_flags::none); + if (!r1.has_value()) + { + return false; + } + zmq::message_t bodyFrame; + auto r2 = mRouter.recv(bodyFrame, zmq::recv_flags::none); + if (!r2.has_value()) + { + return false; + } + outPeer.assign(static_cast(idFrame.data()), idFrame.size()); + outBlob.assign(static_cast(bodyFrame.data()), bodyFrame.size()); + // A well-formed message is exactly [identity, body]. If a malformed peer sent extra frames, + // drain them so they don't desync the NEXT recv() (which would then read this message's leftover + // frame as an identity). We accept [identity, body] and discard any trailing parts. + while (bodyFrame.more()) + { + zmq::message_t extra; + auto re = mRouter.recv(extra, zmq::recv_flags::none); + if (!re.has_value()) + { + break; + } + bodyFrame.swap(extra); // advance the "more" flag to the just-read frame + } + return true; +} + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ZmqControlChannel.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ZmqControlChannel.h new file mode 100644 index 000000000000..28a14eae8191 --- /dev/null +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ZmqControlChannel.h @@ -0,0 +1,83 @@ +/* + * 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/cache_transmission/nixl_utils/bounce/ControlChannel.h" + +#include + +#include +#include +#include + +namespace tensorrt_llm::executor::kv_cache::bounce +{ + +// ============================================================================ +// ZmqControlChannel — zmq implementation of ControlChannel +// ---------------------------------------------------------------------------- +// Topology (one-way, symmetric peers) +// Each agent owns: +// - ONE ROUTER socket, bound to a tcp endpoint, ZMQ_ROUTING_ID = self name. +// It only RECEIVES. A peer's DEALER connects here; on recv the ROUTER yields +// [senderName, blob] (senderName = the peer DEALER's routing id). +// - ONE DEALER per peer, ZMQ_ROUTING_ID = self name, connected to that peer's +// ROUTER endpoint. It only SENDS. The peer's ROUTER then sees us by name. +// So "A -> B" travels A's DEALER(B) -> B's ROUTER; "B -> A" travels B's DEALER(A) +// -> A's ROUTER. No request/reply coupling — purely async one-way, which is all +// the credit protocol needs. +// +// Bootstrap +// localEndpoint() returns the ROUTER's bound address (auto-port via "tcp://host:*" +// + ZMQ_LAST_ENDPOINT). Endpoints are exchanged out-of-band (NIXL agent metadata) +// and registered with addPeer() before the first sendTo(). +// +// Threading +// The ROUTER is touched only by the reactor thread (recv). DEALERs (+ the dealer +// map) are guarded by mMu so app threads and the reactor can sendTo()/addPeer() +// concurrently. zmq sockets are not thread-safe; the mutex provides the required +// serialization + memory fence for the dealer sockets, and the ROUTER is never +// shared. recv() polls the ROUTER with a timeout so the reactor can interleave +// getXferStatus polling. +// ============================================================================ +class ZmqControlChannel : public ControlChannel +{ +public: + /// @param selfName this agent's name (becomes the zmq routing id). + /// @param bindAddr address pattern to bind the ROUTER ("tcp://127.0.0.1:*" by default, + /// "*" picks a free port; query the chosen one via localEndpoint()). + explicit ZmqControlChannel(std::string selfName, std::string const& bindAddr = "tcp://127.0.0.1:*"); + ~ZmqControlChannel() override; + + [[nodiscard]] std::string localEndpoint() const override; + bool addPeer(std::string const& peer, std::string const& endpoint) override; + void removePeer(std::string const& peer) override; + void sendTo(std::string const& peer, std::string const& blob) override; + [[nodiscard]] bool recv(std::string& outPeer, std::string& outBlob, int timeoutMs) override; + +private: + std::string mSelfName; + zmq::context_t mCtx; + zmq::socket_t mRouter; // receive-only (reactor thread) + std::string mEndpoint; // resolved bound endpoint + + mutable std::mutex mMu; // guards mDealers + dealer sends + std::unordered_map mDealers; // peer -> send-only DEALER +}; + +} // namespace tensorrt_llm::executor::kv_cache::bounce diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp index adfd6d1deffe..e5096dbec130 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -26,9 +26,11 @@ #include #include #include +#include #include #include #include +#include #include #include #include @@ -42,9 +44,303 @@ #include #include +#ifdef TLLM_BOUNCE_V2 +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceArena.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ExecPool.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ZmqControlChannel.h" +#include +#include +#endif + namespace tensorrt_llm::executor::kv_cache { +// ============================================================================ +// Bounce v2 integration (opt-in via CacheTransceiverConfig: agent_bounce_buffer_enable + +// kv_cache_bounce_size_mb). When disabled, transfers remain on the standard NIXL path. +// ============================================================================ +#ifdef TLLM_BOUNCE_V2 +namespace bounce +{ +struct NixlBounceState +{ + BounceConfig cfg; + int deviceId{}; + NixlTransferAgent* owner{nullptr}; // the agent this state belongs to (owns this state) + bool arenaRegistered{false}; // arena was NIXL-registered -> dtor must deregister it + std::unique_ptr channel; // ZmqControlChannel + std::unique_ptr arena; // ONE shared buffer: receiver targets + local gather staging + std::unique_ptr exec; // gather/scatter exec contexts (streams/scratch) + // Declared last -> destroyed first: BounceTransport::~ joins its threads (which use the + // agent/channel/arena/exec) before those are torn down. + std::unique_ptr transport; + + ~NixlBounceState() + { + transport.reset(); // join the IO/worker threads before anything they use goes away + if (owner != nullptr && arenaRegistered && arena) + { + // Deregister the arena from NIXL while its memory is still alive (cudaFree follows). + owner->deregisterRegionImpl(arena->base(), arena->bytes(), deviceId); + } + } +}; +} // namespace bounce + +namespace +{ +/// TransferStatus over the bounce transport's completion future (the value Python's wait() reads). +class BounceTransferStatus final : public TransferStatus +{ +public: + explicit BounceTransferStatus(std::shared_future fut) + : mFut(std::move(fut)) + { + } + + [[nodiscard]] bool isCompleted() const override + { + return mFut.valid() && mFut.wait_for(std::chrono::seconds(0)) == std::future_status::ready; + } + + [[nodiscard]] TransferState wait(int64_t timeoutMs) const override + { + if (!mFut.valid()) + { + return TransferState::kFAILURE; + } + if (timeoutMs < 0) + { + return mFut.get().state; + } + if (mFut.wait_for(std::chrono::milliseconds(timeoutMs)) == std::future_status::ready) + { + return mFut.get().state; + } + return TransferState::kIN_PROGRESS; + } + + /// Failure cause recorded by the bounce transport (empty while in flight / on success). + [[nodiscard]] std::string getLastStatusStr() const override + { + if (!mFut.valid() || mFut.wait_for(std::chrono::seconds(0)) != std::future_status::ready) + { + return {}; + } + auto const& result = mFut.get(); + return result.reason == bounce::BounceFailReason::kNone ? std::string{} : bounce::toString(result.reason); + } + +private: + std::shared_future mFut; +}; +} // namespace + +void NixlTransferAgent::maybeInitBounce( + std::size_t agentBufferSizeMb, std::unordered_map const& bounceParams) +{ + // The size doubles as the on/off switch (CacheTransceiverConfig.kv_cache_bounce_size_mb when + // agent_bounce_buffer_enable is set, 0 otherwise): 0 keeps + // bounce disabled, >0 enables it at that arena capacity. The expert knobs resolve as + // agent_bounce_params dict > TRTLLM_NIXL_BOUNCE_* env var > built-in default. + // The retired on/off & arena-size env vars no longer do anything — call that out instead of + // silently ignoring a deployment that still sets them. + for (char const* legacyEnv : {"TRTLLM_NIXL_BOUNCE_ENABLE", "TRTLLM_NIXL_BOUNCE_ARENA_SIZE_BYTES"}) + { + if (std::getenv(legacyEnv) != nullptr) + { + TLLM_LOG_WARNING( + "NixlTransferAgent(%s): %s is deprecated and has NO effect; configure bounce via " + "CacheTransceiverConfig.agent_bounce_buffer_enable + kv_cache_bounce_size_mb instead", + mName.c_str(), legacyEnv); + } + } + if (agentBufferSizeMb == 0) + { + if (!bounceParams.empty()) + { + TLLM_LOG_WARNING( + "NixlTransferAgent(%s): agent_bounce_params set but the bounce arena size is 0 " + "(agent_bounce_buffer_enable off or kv_cache_bounce_size_mb 0) -> bounce stays " + "disabled and the params are ignored", + mName.c_str()); + } + return; + } + auto cfg = bounce::BounceConfig::fromParams(bounceParams, bounce::BounceConfig::fromEnv()); + cfg.arenaSizeBytes = agentBufferSizeMb << 20; + // A single chunk must fit a fresh arena; otherwise its request cannot make progress before + // requestTimeoutMs. Clamp the per-chunk cap to the configured arena size first. BounceTransport + // applies a second clamp to the buddy allocator's smaller usable capacity when necessary. + if (cfg.maxChunkSizeBytes > cfg.arenaSizeBytes) + { + TLLM_LOG_WARNING("NixlTransferAgent(%s): maxChunkSizeBytes (%zu) > arenaSizeBytes (%zu) -> clamping to arena", + mName.c_str(), cfg.maxChunkSizeBytes, cfg.arenaSizeBytes); + cfg.maxChunkSizeBytes = cfg.arenaSizeBytes; + } + // A chunk's packed size travels in 32-bit wire fields, so it must fit in 32 bits even though the + // arena (and thus region offsets) may exceed 4 GiB. Clamp rather than abort on a large config. + if (cfg.maxChunkSizeBytes > std::numeric_limits::max()) + { + TLLM_LOG_WARNING( + "NixlTransferAgent(%s): maxChunkSizeBytes (%zu) > 4 GiB -> clamping " + "(chunk size is 32-bit)", + mName.c_str(), cfg.maxChunkSizeBytes); + cfg.maxChunkSizeBytes = std::numeric_limits::max(); + } + // Any setup failure (e.g. the arena/ExecPool cudaMalloc fails on a busy GPU, or fabric alloc + // throws) must NOT take down agent construction — bounce is an opt-in fast path. Catch, warn, + // and leave mBounce null so the agent runs the standard per-desc NIXL path unchanged. + try + { + int dev = 0; + if (cudaGetDevice(&dev) != cudaSuccess) + { + // Don't silently bind the arena/exec to device 0 on a multi-GPU host — warn and clear the + // sticky error; the enclosing try still proceeds (best-effort) with dev=0. + (void) cudaGetLastError(); + TLLM_LOG_WARNING("NixlTransferAgent(%s): cudaGetDevice failed; bounce assuming device 0", mName.c_str()); + } + auto st = std::make_unique(); + st->cfg = cfg; + st->deviceId = dev; + // Bind the control channel to a ROUTABLE interface (not loopback) so peers on OTHER nodes can + // reach it — the bounce endpoint is advertised cross-node via AgentDesc and the receiver + // self-bootstraps a DEALER to it from WANT. Pick the IP the same way the NIXL agent does + // (TRTLLM_NIXL_INTERFACE NIC if set, else auto-detect via outbound route / hostname; the + // shared common::getLocalIp util). IPv6 needs brackets in a zmq tcp endpoint. + std::string const localIp = common::getLocalIp(common::getEnvNixlInterface(), mpi::MpiComm::world().getRank()); + std::string const bindAddr + = (localIp.find(':') != std::string::npos) ? "tcp://[" + localIp + "]:*" : "tcp://" + localIp + ":*"; + st->channel = std::make_unique(mName, bindAddr); + std::string const controlDesc = st->channel->localEndpoint(); + st->owner = this; + std::size_t const maxDescs = std::max(1024ULL, cfg.maxChunkSizeBytes / 256ULL); + // ONE shared arena for both roles (receiver RDMA-write targets + local gather staging), + // carved into variable-size regions by the scheduler. Register it ONCE NOW (before any + // metadata exchange) so peers' loaded MD includes it. Exec contexts (streams/scratch) are a + // separate small pool borrowed per gather/scatter kernel. + st->arena = std::make_unique(cfg.arenaSizeBytes, dev, !cfg.disableFabricMemory); + if (!registerRegionImpl(st->arena->base(), st->arena->bytes(), dev)) + { + // Arena couldn't be NIXL-registered -> bounce can't move data. Leave mBounce null so the + // agent falls back transparently to the standard per-desc NIXL path (no partial enable). + TLLM_LOG_WARNING( + "NixlTransferAgent(%s): bounce arena registerMem failed -> bounce disabled " + "(NIXL fallback)", + mName.c_str()); + return; + } + st->arenaRegistered = true; + st->exec = std::make_unique(cfg.copyStreamCount, maxDescs, dev, cfg.useZeroCopyArguments); + st->transport = std::make_unique( + mName, cfg, dev, st->channel.get(), *this, st->arena.get(), st->exec.get()); + mBounce = std::move(st); + // Log EVERY resolved knob so a mis-tuned deployment can tell which of dict/env/default won. + TLLM_LOG_INFO( + "NixlTransferAgent(%s): bounce v2 enabled (arenaSizeBytes=%zu arenaAllocationGranularityBytes=%zu " + "maxChunkSizeBytes=%zu maxInflightChunksPerRequest=%u copyStreamCount=%u scatterWorkerCount=%u " + "minDescriptorCount=%zu maxAverageDescriptorSizeBytes=%zu requestTimeoutMs=%d receiverFlowTimeoutMs=%d " + "quarantineMs=%d disableFabricMemory=%d enableEagerGather=%d useZeroCopyArguments=%d control=%s)", + mName.c_str(), cfg.arenaSizeBytes, cfg.arenaAllocationGranularityBytes, cfg.maxChunkSizeBytes, + static_cast(cfg.maxInflightChunksPerRequest), static_cast(cfg.copyStreamCount), + static_cast(cfg.scatterWorkerCount), cfg.minDescriptorCount, cfg.maxAverageDescriptorSizeBytes, + cfg.requestTimeoutMs, cfg.receiverFlowTimeoutMs, cfg.quarantineMs, + static_cast(cfg.disableFabricMemory), static_cast(cfg.enableEagerGather), + static_cast(cfg.useZeroCopyArguments), controlDesc.c_str()); + } + catch (std::exception const& e) + { + mBounce.reset(); + TLLM_LOG_WARNING("NixlTransferAgent(%s): bounce init failed (%s) -> bounce disabled (NIXL fallback)", + mName.c_str(), e.what()); + } +} + +bool NixlTransferAgent::shouldUseBounce(TransferRequest const& request) const +{ + if (!mBounce) + { + return false; + } + auto const& cfg = mBounce->cfg; + if (request.getOp() != TransferOp::kWRITE) + { + return false; + } + if (request.getSrcDescs().getType() != MemoryType::kVRAM || request.getDstDescs().getType() != MemoryType::kVRAM) + { + return false; + } + if (request.getSyncMessage().has_value()) + { + return false; // sync message rides the standard notif path + } + if (!mBounce->transport->hasPeerHandshake(request.getRemoteName())) + { + // The peer did not advertise a compatible bounce handshake (bounce disabled, missing or + // malformed handshake, or version/control-kind/chunk-size mismatch). Without a compatible + // receiver, WANT would remain unanswered until requestTimeoutMs, so use standard NIXL. + return false; + } + auto const& srcs = request.getSrcDescs().getDescs(); + auto const& dsts = request.getDstDescs().getDescs(); + if (srcs.empty() || srcs.size() != dsts.size() || srcs.size() < cfg.minDescriptorCount) + { + return false; + } + auto const sourceDeviceId = srcs.front().getDeviceId(); + auto const destinationDeviceId = dsts.front().getDeviceId(); + if (sourceDeviceId != static_cast(mBounce->deviceId)) + { + return false; + } + // Screen every precondition BounceTransferPlan::build enforces with TLLM_CHECK: an admitted + // request must never throw out of submitTransferRequests — it either runs on bounce or is + // routed to the standard per-descriptor NIXL path. + std::uint64_t totalBytes = 0; + // BounceTransport may clamp the configured cap further to the buddy allocator's usable capacity. + auto const maxChunkSizeBytes = mBounce->transport->maxChunkSizeBytes(); + for (std::size_t i = 0; i < srcs.size(); ++i) + { + auto const len = srcs[i].getLen(); + if (len != dsts[i].getLen() || len > maxChunkSizeBytes || srcs[i].getDeviceId() != sourceDeviceId + || dsts[i].getDeviceId() != destinationDeviceId) + { + return false; + } + totalBytes += len; + } + std::uint64_t const avg = totalBytes / srcs.size(); + return avg <= cfg.maxAverageDescriptorSizeBytes; +} +#else // !TLLM_BOUNCE_V2 — bounce not built; the member stays null and these are no-ops. +namespace bounce +{ +struct NixlBounceState +{ +}; +} // namespace bounce + +void NixlTransferAgent::maybeInitBounce( + std::size_t agentBufferSizeMb, std::unordered_map const& /*bounceParams*/) +{ + if (agentBufferSizeMb > 0) + { + TLLM_LOG_WARNING( + "agent_bounce_buffer_enable requested (kv_cache_bounce_size_mb > 0) but bounce support is not built; " + "ignoring"); + } +} + +bool NixlTransferAgent::shouldUseBounce(TransferRequest const&) const +{ + return false; +} +#endif // TLLM_BOUNCE_V2 + class FileLock { private: @@ -429,8 +725,8 @@ NixlTransferAgent::NixlTransferAgent(BaseAgentConfig const& config) uint32_t numWorker = config.backendParams.find("num_workers") != config.backendParams.end() ? std::stoi(config.backendParams.at("num_workers")) : 1; - nixlAgentConfig nixlConfig{config.useProgThread, true, port, nixl_thread_sync_t::NIXL_THREAD_SYNC_DEFAULT, - numWorker, 0, 10000, config.enableTelemetry}; + nixlAgentConfig nixlConfig{config.useProgThread, true, port, nixl_thread_sync_t::NIXL_THREAD_SYNC_RW, numWorker, + 0, 10000, config.enableTelemetry}; std::string localIp = common::getLocalIp(common::getEnvNixlInterface(), mRank); // Bracket IPv6 literals (RFC 3986) so the last ':' always separates the port; // IPv4 keeps the legacy "ip:port" format for cross-version compatibility. @@ -447,8 +743,8 @@ NixlTransferAgent::NixlTransferAgent(BaseAgentConfig const& config) ? std::stoi(config.backendParams.at("num_workers")) : 1; mAddress.clear(); - nixlAgentConfig nixlConfig{config.useProgThread, false, 0, nixl_thread_sync_t::NIXL_THREAD_SYNC_DEFAULT, - numWorker, 0, 10000, config.enableTelemetry}; + nixlAgentConfig nixlConfig{config.useProgThread, false, 0, nixl_thread_sync_t::NIXL_THREAD_SYNC_RW, numWorker, + 0, 10000, config.enableTelemetry}; mRawAgent = std::make_shared(config.mName, std::move(nixlConfig)); } @@ -485,6 +781,10 @@ NixlTransferAgent::NixlTransferAgent(BaseAgentConfig const& config) } mExtraParams.backends.push_back(mRawBackend); TLLM_LOG_INFO("NixlTransferAgent::NixlTransferAgent mAddress: %s", mAddress.c_str()); + + // Bring up bounce v2 now, if enabled, so its shared arena is registered before any peer fetches + // our metadata. This is a no-op when bounce is disabled or not built. + maybeInitBounce(config.agentBufferSizeMb, config.bounceParams); } void NixlTransferAgent::registerMemory(RegisterDescs const& descs) @@ -540,6 +840,15 @@ void NixlTransferAgent::loadRemoteAgent(std::string const& name, AgentDesc const TLLM_CHECK(status == NIXL_SUCCESS); TLLM_CHECK_WITH_INFO( name == remoteName, "loadRemoteAgent gets error agent name: %s != %s", name.c_str(), remoteName.c_str()); +#ifdef TLLM_BOUNCE_V2 + // Validate the peer's bounce capability handshake from AgentDesc and register its control + // channel. Only peers with the same wire version, control kind and effective maxChunkSizeBytes + // pass; missing or incompatible handshakes leave the peer on standard NIXL. + if (mBounce) + { + mBounce->transport->registerPeerHandshake(name, agentDesc.getBounceHandshake()); + } +#endif // Store remote VMM region info for chunk boundary calculations in // VmmDescSplitter::splitAndCoalesceTransferDescs. Per-agent map because different remote agents may have @@ -572,7 +881,16 @@ AgentDesc NixlTransferAgent::getLocalAgentDesc() regions.push_back({base, info.totalLen, info.chunkSize}); } - return AgentDesc{nixlBlob, std::move(regions)}; + std::string bounceHandshake; +#ifdef TLLM_BOUNCE_V2 + // Include the capability handshake in the structured AgentDesc. Callers must exchange the full + // AgentDesc serialization for the peer to receive and validate it. + if (mBounce) + { + bounceHandshake = mBounce->transport->localHandshakeBlob(); + } +#endif + return AgentDesc{nixlBlob, std::move(regions), std::move(bounceHandshake)}; } void NixlTransferAgent::invalidateRemoteAgent(std::string const& name) @@ -585,6 +903,15 @@ void NixlTransferAgent::invalidateRemoteAgent(std::string const& name) } // Clean up remote VMM region info before invalidating the remote agent. mRemoteVramRegionInfo.erase(name); +#ifdef TLLM_BOUNCE_V2 + // Drop bounce credits and in-flight requests tied to this peer so it cannot retain receiver + // allocations or leave a sender request pending. forgetPeer also removes the handshake record; + // bounce re-engages only after a fresh loadRemoteAgent validates new metadata. + if (mBounce && mBounce->transport) + { + mBounce->transport->forgetPeer(name); + } +#endif mRawAgent->invalidateRemoteMD(name); } @@ -592,21 +919,28 @@ void NixlTransferAgent::invalidateRemoteAgent(std::string const& name) { std::shared_lock lock(mLock); TLLM_CHECK_WITH_INFO(!mShutdown.load(), "NixlTransferAgent::submitTransferRequests called after shutdown"); - nixl_status_t status; - nixlXferReqH* handle; - // Local per-request copy: hasNotif / notifMsg vary per call; a shared mExtraParams - // would race between concurrent submits even under shared_lock. - nixl_opt_args_t reqParams = mExtraParams; - if (request.getSyncMessage().has_value()) - { - reqParams.hasNotif = true; - reqParams.notifMsg = request.getSyncMessage().value(); - } - else +#ifdef TLLM_BOUNCE_V2 + // Bounce fast path for many-small-desc VRAM writes (opt-in). Falls through to the standard + // NIXL path for everything else. The returned future resolves once every chunk is + // scattered+ACKed at the peer (or FAILURE) — never hangs. Admission is final: once a request + // enters bounce, a failure (peer loss / requestTimeoutMs) fails the transfer with no automatic + // fallback to the standard path — the dominant failure causes (dead peer, partition) would fail + // there too, and silent fallback would mask bounce-layer faults. Fallbacks happen only at + // admission time (this gate, unhandshaked peer, disabled/failed arena setup). + if (shouldUseBounce(request)) { - reqParams.hasNotif = false; + mBounceSubmitCount.fetch_add(1, std::memory_order_relaxed); + // Debug-level so it's silent by default but lets ops confirm the bounce fast path actually + // engaged for a given write (enable via TLLM_LOG_LEVEL_BY_MODULE "debug:executor"). + // Programmatic check: getBounceSubmitCount() / bounce_submit_count on the Python binding. + TLLM_LOG_DEBUG("NixlTransferAgent(%s): bounce path engaged for write to %s (%zu descs)", mName.c_str(), + request.getRemoteName().c_str(), request.getSrcDescs().getDescs().size()); + auto fut = mBounce->transport->submit(request.getSrcDescs(), request.getDstDescs(), request.getRemoteName()); + return std::make_unique(std::move(fut)); } +#endif + // Split transfer descriptors at VMM chunk boundaries to match registered memory, then coalesce // contiguous pieces. A coalesced descriptor never crosses a chunk boundary or a registered // region boundary on either side, so every descriptor still falls within a single registered @@ -622,22 +956,95 @@ void NixlTransferAgent::invalidateRemoteAgent(std::string const& name) auto [xferSrc, xferDst] = VmmDescSplitter::splitAndCoalesceTransferDescs(request.getSrcDescs(), request.getDstDescs(), mLocalVramRegionInfo, remoteRegionMap, !common::getEnvNixlDisableCoalesce()); + auto status = postXferRequest(request.getOp(), xferSrc, xferDst, request.getRemoteName(), request.getSyncMessage()); + TLLM_CHECK_WITH_INFO(status != nullptr, + " rank: %d createXferReq failed (see warning above) selfname: %s remoteAgent name: %s", mRank, mName.c_str(), + request.getRemoteName().c_str()); + return status; +} + +std::unique_ptr NixlTransferAgent::postXferRequest(TransferOp op, TransferDescs const& srcDescs, + TransferDescs const& dstDescs, std::string const& remoteName, std::optional const& syncMessage) +{ + // No mLock and no shutdown gate here: the public path already holds the shared lock, and the + // bounce IO thread (the other caller) is joined before agent teardown. mExtraParams is set once + // at construction and only read afterwards — its hasNotif is never set, so the no-notif case + // (every bounce chunk write, on the sub-ms pipeline) passes it directly without a copy. Only a + // notif-carrying call takes a local copy (hasNotif / notifMsg vary per call; mutating the shared + // mExtraParams would race between concurrent submits even under shared_lock). + nixl_opt_args_t notifParams; + nixl_opt_args_t const* reqParams = &mExtraParams; + if (syncMessage.has_value()) { - NVTX3_SCOPED_RANGE(createXferReq); - status = mRawAgent->createXferReq(NixlHelper::convert(request.getOp()), NixlHelper::convertXferDist(xferSrc), - NixlHelper::convertXferDist(xferDst), request.getRemoteName(), handle, &reqParams); + notifParams = mExtraParams; + notifParams.hasNotif = true; + notifParams.notifMsg = syncMessage.value(); + reqParams = ¬ifParams; } - TLLM_CHECK_WITH_INFO(status == NIXL_SUCCESS, - " rank: %d createXferReq failed with status: %s selfname: %s remoteAgent name: %s", mRank, - nixlEnumStrings::statusStr(status).c_str(), mName.c_str(), request.getRemoteName().c_str()); + nixl_status_t status; + nixlXferReqH* handle = nullptr; + { + NVTX3_SCOPED_RANGE(createXferReq); + status = mRawAgent->createXferReq(NixlHelper::convert(op), NixlHelper::convertXferDist(srcDescs), + NixlHelper::convertXferDist(dstDescs), remoteName, handle, reqParams); + } + if (status != NIXL_SUCCESS) + { + TLLM_LOG_WARNING("NixlTransferAgent(%s): createXferReq to %s failed: %s", mName.c_str(), remoteName.c_str(), + nixlEnumStrings::statusStr(status).c_str()); + if (handle != nullptr) + { + (void) mRawAgent->releaseXferReq(handle); + } + return nullptr; + } { NVTX3_SCOPED_RANGE(postXferReq); - status = mRawAgent->postXferReq(handle, &reqParams); + status = mRawAgent->postXferReq(handle, reqParams); + } + if (status != NIXL_SUCCESS && status != NIXL_IN_PROG) + { + TLLM_LOG_WARNING("NixlTransferAgent(%s): postXferReq to %s failed: %s", mName.c_str(), remoteName.c_str(), + nixlEnumStrings::statusStr(status).c_str()); + // The status object still owns the handle; its release()/dtor retries the backend release. } return std::make_unique(std::weak_ptr(mRawAgent), handle); } +bool NixlTransferAgent::registerRegionImpl(void* base, std::size_t bytes, int deviceId) +{ + nixl_reg_dlist_t list{VRAM_SEG}; + list.addDesc(nixlBlobDesc{reinterpret_cast(base), bytes, static_cast(deviceId)}); + nixl_status_t const st = mRawAgent->registerMem(list); + if (st != NIXL_SUCCESS) + { + TLLM_LOG_WARNING( + "NixlTransferAgent(%s): registerMem failed: %s", mName.c_str(), nixlEnumStrings::statusStr(st).c_str()); + return false; + } + return true; +} + +void NixlTransferAgent::deregisterRegionImpl(void* base, std::size_t bytes, int deviceId) +{ + try + { + nixl_reg_dlist_t list{VRAM_SEG}; + list.addDesc(nixlBlobDesc{reinterpret_cast(base), bytes, static_cast(deviceId)}); + nixl_status_t const st = mRawAgent->deregisterMem(list); + if (st != NIXL_SUCCESS) + { + TLLM_LOG_WARNING("NixlTransferAgent(%s): deregisterMem failed: %s", mName.c_str(), + nixlEnumStrings::statusStr(st).c_str()); + } + } + catch (std::exception const& e) + { + TLLM_LOG_WARNING("NixlTransferAgent(%s): deregisterMem threw: %s", mName.c_str(), e.what()); + } +} + void NixlTransferAgent::notifySyncMessage(std::string const& name, SyncMessage const& syncMessage) { std::shared_lock lock(mLock); @@ -661,7 +1068,8 @@ void NixlTransferAgent::notifySyncMessage(std::string const& name, SyncMessage c ConnectionInfoType NixlTransferAgent::getLocalConnectionInfo() { - // mAddress is set in ctor and never mutated; no lock needed. + // mAddress is set in ctor and never mutated; no lock needed. This connection-info string does + // not contain the bounce handshake; callers that need bounce must also exchange AgentDesc. return mAddress; } @@ -669,6 +1077,7 @@ void NixlTransferAgent::loadRemoteAgent(std::string const& name, ConnectionInfoT { std::unique_lock lock(mLock); TLLM_CHECK_WITH_INFO(!mShutdown.load(), "NixlTransferAgent::loadRemoteAgent called after shutdown"); + // Bounce bootstrap is handled on the AgentDesc overload, which is the production path. auto const separator = connectionInfo.rfind(':'); TLLM_CHECK_WITH_INFO(separator != std::string::npos, "Invalid NIXL connection info, expected 'ip:port' or '[ipv6]:port': %s", connectionInfo.c_str()); @@ -730,6 +1139,17 @@ void NixlTransferAgent::shutdown() noexcept } TLLM_LOG_DEBUG("NixlTransferAgent::shutdown"); +#ifdef TLLM_BOUNCE_V2 + // Stop the bounce transport first: its IO/scatter threads poll mRawAgent (getXferStatus), + // so they must be joined before the agent is torn down. Failing pending futures here means + // no submit() waiter hangs across shutdown. + if (mBounce && mBounce->transport) + { + mBounce->transport->shutdown(); + } + mBounce.reset(); +#endif + if (mRawAgent) { // Inline invalidate: invalidateRemoteAgent() would re-enter the non-recursive lock. diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h index 20ed7ff83ab5..07986d72ee4e 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h @@ -28,6 +28,15 @@ namespace tensorrt_llm::executor::kv_cache { +namespace bounce +{ +// Pimpl holding the bounce v2 transport + its arena/exec pool/control channel. The full definition +// lives in the .cpp: the real one under TLLM_BOUNCE_V2 (set when libzmq is found, independent of +// ENABLE_UCX), an empty struct otherwise; the member below is always present so the +// NixlTransferAgent layout is identical across all translation units. +struct NixlBounceState; +} // namespace bounce + struct NixlHelper { [[nodiscard]] static nixl_mem_t convert(MemoryType type); @@ -57,7 +66,7 @@ class NixlTransferStatus final : public TransferStatus [[nodiscard]] TransferState wait(int64_t timeout_ms = -1) const override; [[nodiscard]] int getLastStatus() const noexcept; - [[nodiscard]] std::string getLastStatusStr() const; + [[nodiscard]] std::string getLastStatusStr() const override; [[nodiscard]] bool release() override; @@ -72,7 +81,9 @@ class NixlTransferStatus final : public TransferStatus mutable std::mutex mHandleMutex; }; -class NixlTransferAgent final : public BaseTransferAgent +// Not `final`: the low-level virtuals below (postXferRequest / registerRegionImpl) are the +// fault-injection seam — bounce failure tests subclass this agent and override them. +class NixlTransferAgent : public BaseTransferAgent { public: NixlTransferAgent(BaseAgentConfig const& config); @@ -93,10 +104,22 @@ class NixlTransferAgent final : public BaseTransferAgent [[nodiscard]] std::unique_ptr submitTransferRequests(TransferRequest const& request) override; - [[nodiscard]] nixlAgent* getRawAgent() const noexcept - { - return mRawAgent.get(); - } + // ---- Low-level transfer primitives (below the VMM splitter) ------------------------------- + // submitTransferRequests() = bounce fork + VMM split/coalesce + postXferRequest(). The bounce + // transport calls these directly: its remote address comes from a credit (already final, no + // VMM resolution) and its per-chunk cadence cannot afford the full public path. Virtual so + // bounce failure tests can inject deterministic transfer faults. + + /// Post one transfer whose descriptors are already FINAL device addresses (no VMM splitting, + /// no bounce fork). Returns nullptr on submission failure (logged) instead of aborting. + /// Takes no agent lock: safe from the bounce IO thread, which is joined before agent teardown. + [[nodiscard]] virtual std::unique_ptr postXferRequest(TransferOp op, TransferDescs const& srcDescs, + TransferDescs const& dstDescs, std::string const& remoteName, std::optional const& syncMessage); + + /// Register/deregister one raw device range with NIXL, WITHOUT the VMM split or the AgentDesc + /// VRAM-region bookkeeping (the bounce arena must not enter the splitter's region maps). + [[nodiscard]] virtual bool registerRegionImpl(void* base, std::size_t bytes, int deviceId); + virtual void deregisterRegionImpl(void* base, std::size_t bytes, int deviceId); nixl_opt_args_t* getExtraParams() noexcept { @@ -113,7 +136,23 @@ class NixlTransferAgent final : public BaseTransferAgent bool checkRemoteDescs(std::string const& name, MemoryDescs const& memoryDescs) override; + /// Whether the bounce v2 transport is active on this agent (built + enabled + init succeeded). + /// Programmatic alternative to grepping logs (deployment checks and tests). + [[nodiscard]] bool isBounceEnabled() const noexcept + { + return mBounce != nullptr; + } + + /// Number of transfer requests routed to the bounce fast path so far. + [[nodiscard]] std::uint64_t getBounceSubmitCount() const noexcept + { + return mBounceSubmitCount.load(std::memory_order_relaxed); + } + private: + /// Counts requests admitted by shouldUseBounce (see getBounceSubmitCount). + std::atomic mBounceSubmitCount{0}; + // shared_ptr so outstanding NixlTransferStatus (via weak_ptr) can detect agent reset. std::shared_ptr mRawAgent; nixlBackendH* mRawBackend{}; @@ -135,6 +174,21 @@ class NixlTransferAgent final : public BaseTransferAgent /// Remote VMM region info (from loadRemoteAgent). Keyed by {agentName → {addr → info}}. /// Per-agent maps because different remote agents may have overlapping virtual addresses. std::unordered_map mRemoteVramRegionInfo; + + /// Bounce v2 transport (opt-in via CacheTransceiverConfig.agent_bounce_buffer_enable + + /// kv_cache_bounce_size_mb). Null unless + /// enabled & built; when null the agent behaves exactly as before. See the design overview at + /// the top of bounce/BounceTransport.h. + std::unique_ptr mBounce; + + /// Lazily create the bounce transport (ctor, before any metadata exchange) when enabled. + /// @param agentBufferSizeMb arena size in MiB from BaseAgentConfig; 0 keeps bounce disabled. + /// @param bounceParams expert knobs from BaseAgentConfig, layered over the + /// TRTLLM_NIXL_BOUNCE_* env fallback (dict > env > default). + void maybeInitBounce( + std::size_t agentBufferSizeMb, std::unordered_map const& bounceParams); + /// Heuristic gate: is this request eligible for the bounce fast path? + [[nodiscard]] bool shouldUseBounce(TransferRequest const& request) const; }; class NixlLoopbackAgent final : public BaseLoopbackAgent diff --git a/cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp index 9c5aaaac78ae..c0b9f5138194 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp @@ -106,6 +106,7 @@ std::string AgentDesc::serialize() const su::serialize(r.totalLen, os); su::serialize(r.chunkSize, os); } + su::serialize(mBounceHandshake, os); // NIXL-bounce capability handshake (empty when bounce disabled) return os.str(); } @@ -124,10 +125,11 @@ AgentDesc AgentDesc::deserialize(std::string const& data) auto chunkSize = su::deserialize(is); regions.push_back({baseAddr, totalLen, chunkSize}); } + auto bounceHandshake = su::deserialize(is); TLLM_CHECK_WITH_INFO(!is.fail(), "AgentDesc::deserialize failed: stream error after reading %zu/%zu regions (data size=%zu)", regions.size(), numRegions, data.size()); - return AgentDesc{std::move(backendAgentDesc), std::move(regions)}; + return AgentDesc{std::move(backendAgentDesc), std::move(regions), std::move(bounceHandshake)}; } // ── VmmDescSplitter utilities (backend-agnostic, no NIXL dependency) ── diff --git a/cpp/tensorrt_llm/executor/serialization.cpp b/cpp/tensorrt_llm/executor/serialization.cpp index eff5c62b92ec..47f79293ba96 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -1466,8 +1466,12 @@ CacheTransceiverConfig Serialization::deserializeCacheTransceiverConfig(std::ist auto kvTransferTimeoutMs = su::deserialize>(is); auto kvTransferSenderFutureTimeoutMs = su::deserialize>(is); auto kvTransferPollIntervalMs = su::deserialize>(is); - return CacheTransceiverConfig{ - backendType, maxTokensInBuffer, kvTransferTimeoutMs, kvTransferSenderFutureTimeoutMs, kvTransferPollIntervalMs}; + auto agentBufferSizeMb = su::deserialize(is); + // serializeUtils has no native map support; the params travel as a sorted vector of pairs. + auto agentBounceParamsVec = su::deserialize>>(is); + std::map agentBounceParams(agentBounceParamsVec.begin(), agentBounceParamsVec.end()); + return CacheTransceiverConfig{backendType, maxTokensInBuffer, kvTransferTimeoutMs, kvTransferSenderFutureTimeoutMs, + kvTransferPollIntervalMs, agentBufferSizeMb, std::move(agentBounceParams)}; } void Serialization::serialize(CacheTransceiverConfig const& cacheTransceiverConfig, std::ostream& os) @@ -1477,6 +1481,11 @@ void Serialization::serialize(CacheTransceiverConfig const& cacheTransceiverConf su::serialize(cacheTransceiverConfig.getKvTransferTimeoutMs(), os); su::serialize(cacheTransceiverConfig.getKvTransferSenderFutureTimeoutMs(), os); su::serialize(cacheTransceiverConfig.getKvTransferPollIntervalMs(), os); + su::serialize(cacheTransceiverConfig.getAgentBufferSizeMb(), os); + auto const& agentBounceParams = cacheTransceiverConfig.getAgentBounceParams(); + std::vector> const agentBounceParamsVec( + agentBounceParams.begin(), agentBounceParams.end()); + su::serialize(agentBounceParamsVec, os); } size_t Serialization::serializedSize(CacheTransceiverConfig const& cacheTransceiverConfig) @@ -1487,6 +1496,11 @@ size_t Serialization::serializedSize(CacheTransceiverConfig const& cacheTranscei totalSize += su::serializedSize(cacheTransceiverConfig.getKvTransferTimeoutMs()); totalSize += su::serializedSize(cacheTransceiverConfig.getKvTransferSenderFutureTimeoutMs()); totalSize += su::serializedSize(cacheTransceiverConfig.getKvTransferPollIntervalMs()); + totalSize += su::serializedSize(cacheTransceiverConfig.getAgentBufferSizeMb()); + auto const& agentBounceParams = cacheTransceiverConfig.getAgentBounceParams(); + std::vector> const agentBounceParamsVec( + agentBounceParams.begin(), agentBounceParams.end()); + totalSize += su::serializedSize(agentBounceParamsVec); return totalSize; } diff --git a/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp b/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp index acd33c0df769..6a02c1abd511 100644 --- a/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp @@ -438,11 +438,12 @@ void initConfigBindings(nb::module_& m) auto cacheTransceiverConfigGetstate = [](tle::CacheTransceiverConfig const& self) { return nb::make_tuple(self.getBackendType(), self.getMaxTokensInBuffer(), self.getKvTransferTimeoutMs(), - self.getKvTransferSenderFutureTimeoutMs(), self.getKvTransferPollIntervalMs()); + self.getKvTransferSenderFutureTimeoutMs(), self.getKvTransferPollIntervalMs(), self.getAgentBufferSizeMb(), + self.getAgentBounceParams()); }; auto cacheTransceiverConfigSetstate = [](tle::CacheTransceiverConfig& self, nb::tuple const& state) { - if (state.size() < 3 || state.size() > 5) + if (state.size() < 3 || state.size() > 7) { throw std::runtime_error("Invalid CacheTransceiverConfig state!"); } @@ -451,9 +452,21 @@ void initConfigBindings(nb::module_& m) auto kvTransferPollIntervalMs = state.size() >= 5 ? nb::cast>(state[4]) : std::optional{tle::CacheTransceiverConfig::kDefaultKvTransferPollIntervalMs}; + // Older pickles stored the retired Optional[bool] agent_buffer_enable at index 5. + // convert=false accepts only an exact int (the new format): a legacy True would otherwise + // be number-converted to 1 and silently enable bounce with a 1 MiB arena. Legacy + // True/False/None all fall back to 0 (bounce off) instead of throwing. + size_t agentBufferSizeMb = 0; + if (state.size() >= 6 && !nb::try_cast(state[5], agentBufferSizeMb, /*convert=*/false)) + { + agentBufferSizeMb = 0; + } + auto agentBounceParams = state.size() >= 7 ? nb::cast>(state[6]) + : std::map{}; auto backendType = nb::cast>(state[0]); new (&self) tle::CacheTransceiverConfig(backendType, nb::cast>(state[1]), - nb::cast>(state[2]), kvTransferSenderFutureTimeoutMs, kvTransferPollIntervalMs); + nb::cast>(state[2]), kvTransferSenderFutureTimeoutMs, kvTransferPollIntervalMs, + agentBufferSizeMb, std::move(agentBounceParams)); }; nb::enum_(m, "CacheTransceiverBackendType") @@ -479,12 +492,14 @@ void initConfigBindings(nb::module_& m) }); nb::class_(m, "CacheTransceiverConfig") - .def(nb::init, std::optional, - std::optional, std::optional, std::optional>(), + .def( + nb::init, std::optional, std::optional, + std::optional, std::optional, size_t, std::map>(), nb::arg("backend") = std::nullopt, nb::arg("max_tokens_in_buffer") = std::nullopt, nb::arg("kv_transfer_timeout_ms") = std::nullopt, nb::arg("kv_transfer_sender_future_timeout_ms") = std::nullopt, - nb::arg("kv_transfer_poll_interval_ms") = tle::CacheTransceiverConfig::kDefaultKvTransferPollIntervalMs) + nb::arg("kv_transfer_poll_interval_ms") = tle::CacheTransceiverConfig::kDefaultKvTransferPollIntervalMs, + nb::arg("agent_buffer_size_mb") = 0, nb::arg("agent_bounce_params") = std::map{}) .def_prop_rw( "backend", &tle::CacheTransceiverConfig::getBackendType, &tle::CacheTransceiverConfig::setBackendType) .def_prop_rw("max_tokens_in_buffer", &tle::CacheTransceiverConfig::getMaxTokensInBuffer, @@ -496,6 +511,10 @@ void initConfigBindings(nb::module_& m) &tle::CacheTransceiverConfig::setKvTransferSenderFutureTimeoutMs) .def_prop_rw("kv_transfer_poll_interval_ms", &tle::CacheTransceiverConfig::getKvTransferPollIntervalMs, &tle::CacheTransceiverConfig::setKvTransferPollIntervalMs) + .def_prop_rw("agent_buffer_size_mb", &tle::CacheTransceiverConfig::getAgentBufferSizeMb, + &tle::CacheTransceiverConfig::setAgentBufferSizeMb) + .def_prop_rw("agent_bounce_params", &tle::CacheTransceiverConfig::getAgentBounceParams, + &tle::CacheTransceiverConfig::setAgentBounceParams) .def("__getstate__", cacheTransceiverConfigGetstate) .def("__setstate__", cacheTransceiverConfigSetstate); diff --git a/cpp/tests/unit_tests/executor/CMakeLists.txt b/cpp/tests/unit_tests/executor/CMakeLists.txt index 6d851ebd4544..58d53b542b58 100644 --- a/cpp/tests/unit_tests/executor/CMakeLists.txt +++ b/cpp/tests/unit_tests/executor/CMakeLists.txt @@ -31,6 +31,11 @@ add_gtest(ucxCommTest ucxCommTest.cpp TIMEOUT 60) target_link_libraries(ucxCommTest PRIVATE ${Python3_LIBRARIES}) target_link_libraries(serializeUtilsTest PRIVATE ${Python3_LIBRARIES}) +# NIXL bounce-buffer v2 KV-cache transfer tests live in their own subdirectory +# (mirrors the source layout under nixl_utils/bounce/). It self-gates on zmq and +# NIXL_ROOT, so it is safe to descend unconditionally. +add_subdirectory(bounce) + # Skip MOONCAKE related tests on Rocky8 set(IS_ROCKY8 FALSE) if(EXISTS "/etc/redhat-release") @@ -47,6 +52,9 @@ if(NIXL_ROOT OR (MOONCAKE_ROOT AND NOT IS_ROCKY8)) ${Python3_LIBRARIES}) target_compile_definitions(transferAgentTest PRIVATE TEST_NIXL_BACKEND=1) target_compile_definitions(agentCommTest PRIVATE TEST_NIXL_BACKEND=1) + + # NIXL bounce v2 tests (bounceTransportTest / bounceTransportFailureTest / + # bounceAgentE2ETest) live in bounce/CMakeLists.txt. endif() if(MOONCAKE_ROOT) diff --git a/cpp/tests/unit_tests/executor/bounce/CMakeLists.txt b/cpp/tests/unit_tests/executor/bounce/CMakeLists.txt new file mode 100644 index 000000000000..76edc38937d9 --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/CMakeLists.txt @@ -0,0 +1,147 @@ +# 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. + +# Unit tests for the NIXL bounce-buffer v2 KV-cache transfer. The code under +# test lives in tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/; the +# bounce core .cpp files are compiled directly into each test (they are also +# built into the NIXL wrapper for production). Three tiers, each independently +# gated: - pure-logic / GPU tests (no NIXL, no zmq, no network) - control-plane +# tests (need zmq + the fetched cppzmq header) - real-NIXL e2e tests (need +# NIXL_ROOT + zmq; two real agents, real RDMA) +set(BOUNCE_DIR + ${PROJECT_SOURCE_DIR}/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce +) + +# ---- Pure-logic + GPU unit tests (no NIXL / zmq / network) ---- +add_gtest(creditSchedulerTest creditSchedulerTest.cpp) +target_sources(creditSchedulerTest PRIVATE ${BOUNCE_DIR}/CreditScheduler.cpp + ${BOUNCE_DIR}/BuddyAllocator.cpp) +add_gtest(bounceTransferPlanTest bounceTransferPlanTest.cpp) +target_sources(bounceTransferPlanTest + PRIVATE ${BOUNCE_DIR}/BounceTransferPlan.cpp) +add_gtest(bounceMessageCodecTest bounceMessageCodecTest.cpp) +target_sources(bounceMessageCodecTest PRIVATE ${BOUNCE_DIR}/BounceMessage.cpp) +# Header-only: knob resolution (agent_bounce_params dict > TRTLLM_NIXL_BOUNCE_* +# env > default), byte-size suffixes, garbage fallback, derived-timeout +# invariants. +add_gtest(bounceConfigTest bounceConfigTest.cpp) +add_gtest(buddyAllocatorTest buddyAllocatorTest.cpp) +target_sources(buddyAllocatorTest PRIVATE ${BOUNCE_DIR}/BuddyAllocator.cpp) +# GPU test: execution-context pool (stream/event/scratch/hostPinned) +# acquire/release. +add_gtest(execPoolTest execPoolTest.cpp) +target_sources(execPoolTest PRIVATE ${BOUNCE_DIR}/ExecPool.cpp) +# GPU test: gather/scatter CUDA kernel round-trip (needs a CUDA device to run). +add_gtest(gatherScatterKernelTest gatherScatterKernelTest.cpp) +target_sources(gatherScatterKernelTest + PRIVATE ${BOUNCE_DIR}/GatherScatterKernel.cu) +# GPU test: shared data-arena alloc / address mapping / shutdown. +add_gtest(bounceArenaTest bounceArenaTest.cpp) +target_sources(bounceArenaTest PRIVATE ${BOUNCE_DIR}/BounceArena.cpp) + +# ---- Control-plane unit test — needs only zmq (+ the fetched cppzmq header), +# no NIXL. ZMQ_FOUND is reused by the real-NIXL block below. ---- +find_package(PkgConfig) +if(PKG_CONFIG_FOUND) + pkg_check_modules(ZMQ libzmq) +endif() +# Both zmq tiers need libzmq AND the fetched cppzmq header: without the latter +# the sources compile-fail on zmq.hpp even though libzmq itself was found. +if(ZMQ_FOUND AND EXISTS ${CMAKE_BINARY_DIR}/_deps/cppzmq-src) + set(BOUNCE_CPPZMQ_AVAILABLE TRUE) +else() + set(BOUNCE_CPPZMQ_AVAILABLE FALSE) +endif() +if(BOUNCE_CPPZMQ_AVAILABLE) + add_gtest(zmqControlChannelTest zmqControlChannelTest.cpp) + target_sources( + zmqControlChannelTest PRIVATE ${BOUNCE_DIR}/ZmqControlChannel.cpp + ${BOUNCE_DIR}/BounceMessage.cpp) + target_include_directories( + zmqControlChannelTest PRIVATE ${CMAKE_BINARY_DIR}/_deps/cppzmq-src + ${ZMQ_INCLUDE_DIRS}) + target_link_libraries(zmqControlChannelTest PRIVATE ${ZMQ_LIBRARIES}) +endif() + +# ---- Real-NIXL tests (need NIXL_ROOT). The bounce DATA PLANE is always real +# NIXL — there is no loopback fake — so every test that drives the reactor +# (transport, failure-path, e2e) lives here, sharing one source set + the +# bounceTestNixlNode.h helpers. ---- +if(NIXL_ROOT) + # The reactor tests also need zmq for the control plane (cppzmq is fetched + # when NIXL is on); all of them share this source set + include/link sets. + # Same gate as the zmq-only tier above. + if(BOUNCE_CPPZMQ_AVAILABLE) + set(BOUNCE_REACTOR_SRCS + ${BOUNCE_DIR}/BounceTransport.cpp + ${BOUNCE_DIR}/BounceTransferPlan.cpp + ${BOUNCE_DIR}/BuddyAllocator.cpp + ${BOUNCE_DIR}/CreditScheduler.cpp + ${BOUNCE_DIR}/BounceMessage.cpp + ${BOUNCE_DIR}/BounceArena.cpp + ${BOUNCE_DIR}/ExecPool.cpp + ${BOUNCE_DIR}/GatherScatterKernel.cu + ${BOUNCE_DIR}/ZmqControlChannel.cpp) + set(BOUNCE_TEST_INCS ${CMAKE_BINARY_DIR}/_deps/cppzmq-src + ${ZMQ_INCLUDE_DIRS}) + set(BOUNCE_TEST_LIBS tensorrt_llm_nixl_wrapper NIXL::nixl + ${NIXL_BUILD_LIBRARY} ${ZMQ_LIBRARIES}) + + # One reactor test target: the bounce sources compiled in, plus the shared + # include/link sets. + function(add_bounce_reactor_test name src) + add_gtest(${name} ${src}) + target_sources(${name} PRIVATE ${BOUNCE_REACTOR_SRCS}) + target_include_directories(${name} PRIVATE ${BOUNCE_TEST_INCS}) + target_link_libraries(${name} PRIVATE ${BOUNCE_TEST_LIBS}) + endfunction() + + # End-to-end transport over real RDMA (gather -> RDMA write -> scatter + + # credit recycling), byte-exact. + add_bounce_reactor_test(bounceTransportTest bounceTransportTest.cpp) + + # Boundary, failure-path, and concurrency tests. + add_bounce_reactor_test(bounceTransportFailureTest + bounceTransportFailureTest.cpp) + + # Optional ThreadSanitizer variant of the reactor concurrency test. The + # bounce sources + the test are instrumented; the shared libtensorrt_llm + # (linked for the logger) and CUDA runtime are left uninstrumented (TSan + # partial-instrumentation mode). Enable with -DBOUNCE_TSAN=ON and run with a + # suppressions file that ignores CUDA/driver/libtllm frames. + option(BOUNCE_TSAN + "Build a ThreadSanitizer variant of the bounce reactor test" OFF) + if(BOUNCE_TSAN) + add_bounce_reactor_test(bounceTransportTsanTest + bounceTransportFailureTest.cpp) + target_compile_options( + bounceTransportTsanTest + PRIVATE $<$:-fsanitize=thread -g> + $<$:-Xcompiler=-fsanitize=thread>) + target_link_options(bounceTransportTsanTest PRIVATE -fsanitize=thread) + endif() + + # Production-path e2e: bounce engaged via + # NixlTransferAgent::submitTransferRequests. Bounce lives inside the + # wrapper; this test only uses the public agent API (the zmq include is + # needed only because the shared bounceTestNixlNode.h helpers pull in + # ZmqControlChannel.h at compile time). + add_gtest(bounceAgentE2ETest bounceAgentE2ETest.cpp) + target_include_directories(bounceAgentE2ETest PRIVATE ${BOUNCE_TEST_INCS}) + target_link_libraries( + bounceAgentE2ETest PRIVATE tensorrt_llm_nixl_wrapper NIXL::nixl + ${NIXL_BUILD_LIBRARY}) + endif() +endif() diff --git a/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp new file mode 100644 index 000000000000..60dcf5eb025c --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp @@ -0,0 +1,425 @@ +/* + * 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. + */ + +// Production-path e2e: bounce engaged transparently through NixlTransferAgent::submitTransferRequests +// with BaseAgentConfig::agentBufferSizeMb > 0. Real agents (2..N), real RDMA, wired exactly as production disagg +// does (tensorrt_llm/_torch/disaggregation/native/transfer.py): one-directional AgentDesc exchange +// (senders load the receiver, never the reverse) and the receiver self-bootstraps each sender from +// its WANT — no manual addPeer, no connection-info path. NOTE: the KV src/dst buffers are +// intentionally NOT NIXL-registered, so the standard path's createXferReq would fail on them -- +// a SUCCESS here proves the bounce fast path was taken (only the bounce arena is registered). + +#include "bounceTestNixlNode.h" + +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace kvc = tensorrt_llm::executor::kv_cache; + +using bounce_test::freeXferBufs; +using bounce_test::hasCuda; +using bounce_test::makeXferBufs; +using bounce_test::verifyXferBufs; +using bounce_test::XferBufs; + +namespace +{ +kvc::TransferRequest makeReq(XferBufs const& x, char const* dstPeer) +{ + return kvc::TransferRequest{kvc::TransferOp::kWRITE, x.srcDescs, x.dstDescs, dstPeer, std::nullopt}; +} + +// Poll a transfer status to a terminal state (bounce futures resolve once all chunks are +// scattered+ACKed); returns the last observed state (kIN_PROGRESS on deadline). +template +kvc::TransferState waitTerminal(StatusPtr& status, int seconds) +{ + auto st = kvc::TransferState::kIN_PROGRESS; + auto const deadline = std::chrono::steady_clock::now() + std::chrono::seconds(seconds); + while (std::chrono::steady_clock::now() < deadline) + { + st = status->wait(100); + if (st != kvc::TransferState::kIN_PROGRESS) + { + break; + } + } + return st; +} + +// Fan out one transfer per (sender agent, dst peer) route, each on its own thread, and return how +// many completed SUCCESS. bufs[i] backs routes[i]. +int runConcurrentFlows( + std::vector> const& routes, std::vector const& bufs) +{ + std::atomic ok{0}; + std::vector threads; + threads.reserve(routes.size()); + for (std::size_t i = 0; i < routes.size(); ++i) + { + threads.emplace_back( + [&, i] + { + auto req = makeReq(bufs[i], routes[i].second); + auto status = routes[i].first->submitTransferRequests(req); + if (status == nullptr) + { + return; + } + if (waitTerminal(status, 60) == kvc::TransferState::kSUCCESS) + { + ok.fetch_add(1); + } + }); + } + for (auto& t : threads) + { + t.join(); + } + return ok.load(); +} + +// Bounce-enabled agent config: agentBufferSizeMb switches bounce on, and the expert knobs ride +// bounceParams (dict > env > default) with thresholds tuned so a modest transfer engages bounce +// (small regions -> recycling). sizeMb == 0 keeps bounce off (no knobs attached). +kvc::BaseAgentConfig makeBounceConfig( + std::string name, std::size_t arenaSizeMb = 2, char const* granularityBytes = "256") +{ + kvc::BaseAgentConfig cfg{std::move(name), true, false, true}; + cfg.agentBufferSizeMb = arenaSizeMb; + if (arenaSizeMb > 0) + { + cfg.bounceParams = { + {"min_descriptor_count", "4"}, + {"max_chunk_size", "4096"}, + {"arena_allocation_granularity", granularityBytes}, + {"max_inflight_chunks_per_request", "2"}, + }; + } + return cfg; +} + +// Construct a NixlTransferAgent, returning nullptr (with the failure reason in skipMsg) when the +// NIXL agent/backend is unavailable — callers GTEST_SKIP on nullptr, mirroring makeNode(). +std::unique_ptr tryMakeAgent(kvc::BaseAgentConfig cfg, std::string& skipMsg) +{ + try + { + return std::make_unique(std::move(cfg)); + } + catch (std::exception const& e) + { + skipMsg = e.what(); + return nullptr; + } +} +} // namespace + +// BaseAgentConfig::agentBufferSizeMb (derived from CacheTransceiverConfig's +// agent_bounce_buffer_enable + kv_cache_bounce_size_mb) is the +// ONLY on/off switch: 0 keeps bounce disabled, >0 enables it. The legacy TRTLLM_NIXL_BOUNCE_ENABLE +// environment variable must have no effect anymore. +TEST(BounceAgentE2E, AgentBufferSizeControlsBounce) +{ + if (!hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; + } + setenv("TRTLLM_NIXL_BOUNCE_ENABLE", "1", 1); + try + { + EXPECT_FALSE(std::make_unique(makeBounceConfig("cfgOffAgent", 0))->isBounceEnabled()); + EXPECT_TRUE(std::make_unique(makeBounceConfig("cfgOnAgent", 2))->isBounceEnabled()); + // size 0 + non-empty params: the params must be ignored (warned about, not honored) and + // bounce must stay off. (The llm_args validator rejects this combination upfront; this + // covers the direct BaseAgentConfig entry point.) + auto orphanParams = makeBounceConfig("cfgOrphanAgent", 0); + orphanParams.bounceParams = {{"copy_stream_count", "2"}}; + EXPECT_FALSE(std::make_unique(orphanParams)->isBounceEnabled()); + } + catch (std::exception const& e) + { + unsetenv("TRTLLM_NIXL_BOUNCE_ENABLE"); + GTEST_SKIP() << "NIXL agent/backend unavailable: " << e.what(); + } + unsetenv("TRTLLM_NIXL_BOUNCE_ENABLE"); +} + +TEST(BounceAgentE2E, SubmitTransferRequestsUsesBounce) +{ + if (!hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; + } + std::string skipMsg; + auto a = tryMakeAgent(makeBounceConfig("bAgentA", /*arenaSizeMb=*/1, /*granularityBytes=*/"4096"), skipMsg); + auto b = a ? tryMakeAgent(makeBounceConfig("bAgentB", /*arenaSizeMb=*/1, /*granularityBytes=*/"4096"), skipMsg) + : nullptr; + if (!a || !b) + { + GTEST_SKIP() << "NIXL agent/backend unavailable: " << skipMsg; + } + + // Connect exactly as production disagg does (tensorrt_llm/_torch/disaggregation/native/transfer.py): + // the metadata exchange is ONE-DIRECTIONAL — only the KV sender loads the receiver's AgentDesc + // (get_local_agent_desc / loadRemoteAgent(AgentDesc)); the receiver never loads the sender. The + // bounce control endpoint rides inside that AgentDesc. B's reverse control path (GRANT/ACK back + // to A) is self-bootstrapped from A's WANT (BounceReceiver::onWant) — exercising it + // here is the whole point. (We deliberately do NOT touch the connection-info path.) + a->loadRemoteAgent("bAgentB", b->getLocalAgentDesc()); + + auto bufs = makeXferBufs(/*nDescs=*/24, /*descBytes=*/600, /*seed=*/7); + auto req = makeReq(bufs, "bAgentB"); + auto status = a->submitTransferRequests(req); + ASSERT_NE(status, nullptr); + EXPECT_EQ(waitTerminal(status, 30), kvc::TransferState::kSUCCESS) + << "bounce transfer via submitTransferRequests did not succeed"; + EXPECT_TRUE(verifyXferBufs(bufs)); + + a->shutdown(); + b->shutdown(); + freeXferBufs(bufs); +} + +// A non-power-of-two arena can clamp the effective chunk cap below the configured cap. A descriptor +// between those limits must use ordinary NIXL instead of entering bounce and failing plan creation. +TEST(BounceAgentE2E, EffectiveChunkCapFallsBackToStandardNixl) +{ + if (!hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; + } + // 3 MiB arena (non-power-of-two) -> buddy capacity is 2 MiB, below the configured 3 MiB chunk + // cap. The descriptor below sits between the two limits. + auto makeCapConfig = [](char const* name) + { + kvc::BaseAgentConfig cfg{name, true, false, true}; + cfg.agentBufferSizeMb = 3; + cfg.bounceParams = { + {"min_descriptor_count", "1"}, + {"max_chunk_size", "3145728"}, // configured 3 MiB + {"arena_allocation_granularity", "256"}, + {"max_average_descriptor_size", "8MB"}, + }; + return cfg; + }; + + std::string skipMsg; + auto a = tryMakeAgent(makeCapConfig("capAgentA"), skipMsg); + auto b = a ? tryMakeAgent(makeCapConfig("capAgentB"), skipMsg) : nullptr; + if (!a || !b) + { + GTEST_SKIP() << "NIXL agent/backend unavailable: " << skipMsg; + } + + // 2.5 MiB: above the effective (buddy) 2 MiB cap, below the configured 3 MiB cap. + auto bufs = makeXferBufs(/*nDescs=*/1, /*descBytes=*/2560 * 1024, /*seed=*/8); + auto req = makeReq(bufs, "capAgentB"); + a->registerMemory(req.getSrcDescs()); + b->registerMemory(req.getDstDescs()); + a->loadRemoteAgent("capAgentB", b->getLocalAgentDesc()); + + auto status = a->submitTransferRequests(req); + ASSERT_NE(status, nullptr); + EXPECT_EQ(waitTerminal(status, 30), kvc::TransferState::kSUCCESS); + EXPECT_EQ(a->getBounceSubmitCount(), 0); + EXPECT_TRUE(verifyXferBufs(bufs)); + EXPECT_TRUE(status->release()); + status.reset(); + + a->deregisterMemory(req.getSrcDescs()); + b->deregisterMemory(req.getDstDescs()); + a->shutdown(); + b->shutdown(); + freeXferBufs(bufs); +} + +// Production-path CONCURRENCY: many threads call submitTransferRequests on the SAME sender agent at +// once (mirrors transfer.py's KV_TRANSFER_NUM_THREADS>1 worker pool fanning out to one receiver). +// Each gets its own seeded buffers; all must complete SUCCESS + land byte-exact with no cross-talk, +// no hang/deadlock — the production-API counterpart of the transport-level concurrency tests. +TEST(BounceAgentE2E, ConcurrentSubmitUsesBounce) +{ + if (!hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; + } + std::string skipMsg; + auto a = tryMakeAgent(makeBounceConfig("cAgentA"), skipMsg); + auto b = a ? tryMakeAgent(makeBounceConfig("cAgentB"), skipMsg) : nullptr; + if (!a || !b) + { + GTEST_SKIP() << "NIXL agent/backend unavailable: " << skipMsg; + } + + // One-directional bootstrap, exactly like transfer.py: only the sender (A) loads the receiver + // (B); B self-bootstraps A from the first WANT (BounceReceiver::onWant). No manual addPeer — the agent's own + // bounce transport is wired through loadRemoteAgent(AgentDesc). + a->loadRemoteAgent("cAgentB", b->getLocalAgentDesc()); + + constexpr int kThreads = 8; + std::vector bufs; + std::vector> routes; + bufs.reserve(kThreads); + routes.reserve(kThreads); + for (int i = 0; i < kThreads; ++i) + { + bufs.push_back(makeXferBufs(/*nDescs=*/16, /*descBytes=*/500, /*seed=*/static_cast(i + 1))); + routes.emplace_back(a.get(), "cAgentB"); + } + + EXPECT_EQ(runConcurrentFlows(routes, bufs), kThreads) << "not all concurrent submitTransferRequests completed"; + for (auto const& x : bufs) + { + EXPECT_TRUE(verifyXferBufs(x)) << "byte mismatch / cross-talk for seed=" << x.seed; + } + + a->shutdown(); + b->shutdown(); + for (auto& x : bufs) + { + freeXferBufs(x); + } +} + +// Production-path BIDIRECTIONAL concurrency: two bounce-enabled agents each submit to the OTHER at +// once, so both arenas serve sender (gather) AND receiver (scatter) roles simultaneously. Both are +// senders here, so — exactly like transfer.py when two ranks each write to each other — each loads +// the other's AgentDesc; there is still NO manual addPeer (loadRemoteAgent wires each agent's own +// bounce transport). Every transfer must land byte-exact with no cross-talk or hang over the +// production API. +TEST(BounceAgentE2E, ConcurrentBidirectionalUsesBounce) +{ + if (!hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; + } + std::string skipMsg; + auto a = tryMakeAgent(makeBounceConfig("biAgentA"), skipMsg); + auto b = a ? tryMakeAgent(makeBounceConfig("biAgentB"), skipMsg) : nullptr; + if (!a || !b) + { + GTEST_SKIP() << "NIXL agent/backend unavailable: " << skipMsg; + } + + // Both agents send, so both load the other's AgentDesc (each direction's WANT self-bootstraps + // the reverse control path anyway; the redundant load is harmless and mirrors a two-way flow). + a->loadRemoteAgent("biAgentB", b->getLocalAgentDesc()); + b->loadRemoteAgent("biAgentA", a->getLocalAgentDesc()); + + constexpr int kThreads = 8; // 4 flows A->B + 4 flows B->A, all concurrent + std::vector bufs; + std::vector> routes; + bufs.reserve(kThreads); + routes.reserve(kThreads); + for (int i = 0; i < kThreads; ++i) + { + bufs.push_back(makeXferBufs(/*nDescs=*/16, /*descBytes=*/500, /*seed=*/static_cast(i + 1))); + // Even threads send A->B; odd threads send B->A (both arenas act as both roles). + bool const a2b = (i % 2 == 0); + routes.emplace_back(a2b ? a.get() : b.get(), a2b ? "biAgentB" : "biAgentA"); + } + + EXPECT_EQ(runConcurrentFlows(routes, bufs), kThreads) << "not all bidirectional submitTransferRequests completed"; + for (auto const& x : bufs) + { + EXPECT_TRUE(verifyXferBufs(x)) << "byte mismatch / cross-talk for seed=" << x.seed; + } + + a->shutdown(); + b->shutdown(); + for (auto& x : bufs) + { + freeXferBufs(x); + } +} + +// Production-path MULTI-AGENT (N>2): 1 receiver + S sender agents, every sender writing to the one +// receiver concurrently (the disagg "many context workers -> one gen" shape). This is the REAL +// one-directional bootstrap that transfer.py uses: each sender loads the receiver's AgentDesc; the +// receiver loads NOBODY and self-bootstraps every sender from its WANT (BounceReceiver::onWant) — +// so the reverse-control self-bootstrap is exercised across N distinct peers at once. +// Seed-distinct patterns -> any cross-talk fails; all must complete SUCCESS + land byte-exact. +TEST(BounceAgentE2E, MultiAgentManySendersToOneReceiver) +{ + if (!hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; + } + constexpr int kSenders = 3; // total agents = 1 receiver + 3 senders + std::string const recvName = "mnAgentR"; + + std::string skipMsg; + auto recv = tryMakeAgent(makeBounceConfig(recvName), skipMsg); + std::vector> senders; + for (int i = 1; recv && i <= kSenders; ++i) + { + auto s = tryMakeAgent(makeBounceConfig("mnAgentS" + std::to_string(i)), skipMsg); + if (!s) + { + break; + } + senders.push_back(std::move(s)); + } + if (!recv || senders.size() != static_cast(kSenders)) + { + GTEST_SKIP() << "NIXL agent/backend unavailable: " << skipMsg; + } + + // One-directional wiring: each sender loads the receiver; the receiver loads NOBODY. The + // receiver only ever hears about a sender when that sender's WANT arrives (self-bootstrap). + for (auto& s : senders) + { + s->loadRemoteAgent(recvName, recv->getLocalAgentDesc()); + } + + std::vector bufs; + std::vector> routes; + bufs.reserve(kSenders); + routes.reserve(kSenders); + for (int i = 0; i < kSenders; ++i) + { + bufs.push_back(makeXferBufs(/*nDescs=*/16, /*descBytes=*/500, /*seed=*/static_cast(100 + i))); + routes.emplace_back(senders[i].get(), recvName.c_str()); + } + + EXPECT_EQ(runConcurrentFlows(routes, bufs), kSenders) << "not all senders completed to the shared receiver"; + for (auto const& x : bufs) + { + EXPECT_TRUE(verifyXferBufs(x)) << "byte mismatch / cross-talk for seed=" << x.seed; + } + + for (auto& s : senders) + { + s->shutdown(); + } + recv->shutdown(); + for (auto& x : bufs) + { + freeXferBufs(x); + } +} diff --git a/cpp/tests/unit_tests/executor/bounce/bounceArenaTest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceArenaTest.cpp new file mode 100644 index 000000000000..cec0242fc484 --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/bounceArenaTest.cpp @@ -0,0 +1,72 @@ +/* + * 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. + */ + +// GPU test for BounceArena: the one shared bounce data buffer. Verifies allocation, the +// offset->address mapping the scheduler/transport rely on, and that the buffer is real device +// memory (writable/readable). Fabric is force-disabled so this runs on any CUDA device (CI/x86). + +#include "bounceTestUtils.h" + +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceArena.h" + +#include + +#include + +#include + +namespace b = tensorrt_llm::executor::kv_cache::bounce; + +using bounce_test::hasCuda; + +TEST(BounceArena, BaseAndOffsetMapping) +{ + if (!hasCuda()) + GTEST_SKIP(); + constexpr std::size_t kBytes = 4ULL << 20; // 4 MiB + b::BounceArena arena(kBytes, 0, /*allowFabric=*/false); + EXPECT_NE(arena.base(), nullptr); + EXPECT_EQ(arena.bytes(), kBytes); + EXPECT_EQ(arena.baseAddr(), reinterpret_cast(arena.base())); + EXPECT_FALSE(arena.isFabric()); // forced off + // at(offset) is exactly base + offset (what Grant.addr == baseAddr + offset relies on). + EXPECT_EQ(arena.at(0), arena.base()); + EXPECT_EQ(arena.at(1024), static_cast(arena.base()) + 1024); + EXPECT_EQ(reinterpret_cast(arena.at(4096)), arena.baseAddr() + 4096); +} + +TEST(BounceArena, BufferIsUsableDeviceMemory) +{ + if (!hasCuda()) + GTEST_SKIP(); + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + constexpr std::size_t kBytes = 1ULL << 20; + b::BounceArena arena(kBytes, 0, /*allowFabric=*/false); + + // Write a pattern at a non-zero offset region and read it back through at(). + constexpr std::uint64_t kOffset = 64 * 1024; + constexpr std::size_t kLen = 4096; + std::vector in(kLen); + for (std::size_t i = 0; i < kLen; ++i) + { + in[i] = static_cast((i * 31 + 7) & 0xFF); + } + ASSERT_EQ(cudaMemcpy(arena.at(kOffset), in.data(), kLen, cudaMemcpyHostToDevice), cudaSuccess); + std::vector out(kLen, 0); + ASSERT_EQ(cudaMemcpy(out.data(), arena.at(kOffset), kLen, cudaMemcpyDeviceToHost), cudaSuccess); + EXPECT_EQ(in, out); +} diff --git a/cpp/tests/unit_tests/executor/bounce/bounceConfigTest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceConfigTest.cpp new file mode 100644 index 000000000000..3af88e8d4f96 --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/bounceConfigTest.cpp @@ -0,0 +1,294 @@ +/* + * 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/BounceConfig.h" + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace b = tensorrt_llm::executor::kv_cache::bounce; + +namespace +{ +// Sets a group of env vars for one test and restores the previous values on destruction, so +// tests don't leak state into each other (or into a developer's shell-provided environment). +class ScopedEnv +{ +public: + void set(char const* name, char const* value) + { + char const* old = std::getenv(name); + mSaved.emplace_back(name, old != nullptr ? std::optional{old} : std::nullopt); + ::setenv(name, value, /*overwrite=*/1); + } + + ~ScopedEnv() + { + for (auto const& [name, old] : mSaved) + { + if (old.has_value()) + { + ::setenv(name.c_str(), old->c_str(), 1); + } + else + { + ::unsetenv(name.c_str()); + } + } + } + +private: + std::vector>> mSaved; +}; + +// Resolve one byte-valued knob through the params (dict) path. +std::size_t maxChunkSizeBytesForParam(char const* value) +{ + std::unordered_map const params{{"max_chunk_size", value}}; + return b::BounceConfig::fromParams(params, b::BounceConfig{}).maxChunkSizeBytes; +} +} // namespace + +// The on/off switch and the arena size come ONLY from CacheTransceiverConfig +// (agent_bounce_buffer_enable + kv_cache_bounce_size_mb); +// the legacy TRTLLM_NIXL_BOUNCE_ENABLE / TRTLLM_NIXL_BOUNCE_ARENA_SIZE_BYTES env vars are dead. +TEST(BounceConfig, EnableAndArenaAreNotEnvBacked) +{ + ScopedEnv env; + env.set("TRTLLM_NIXL_BOUNCE_ENABLE", "1"); + env.set("TRTLLM_NIXL_BOUNCE_ARENA_SIZE_BYTES", "1GB"); + auto const cfg = b::BounceConfig::fromEnv(); + EXPECT_EQ(cfg.arenaSizeBytes, b::BounceConfig{}.arenaSizeBytes); +} + +TEST(BounceConfig, ByteSuffixesParse) +{ + EXPECT_EQ(maxChunkSizeBytesForParam("12345"), 12345u); + EXPECT_EQ(maxChunkSizeBytesForParam("12345B"), 12345u); + EXPECT_EQ(maxChunkSizeBytesForParam("512K"), 512ULL << 10); + EXPECT_EQ(maxChunkSizeBytesForParam("512KB"), 512ULL << 10); + EXPECT_EQ(maxChunkSizeBytesForParam("512kib"), 512ULL << 10); + EXPECT_EQ(maxChunkSizeBytesForParam("256M"), 256ULL << 20); + EXPECT_EQ(maxChunkSizeBytesForParam("256MB"), 256ULL << 20); + EXPECT_EQ(maxChunkSizeBytesForParam("256mb"), 256ULL << 20); + EXPECT_EQ(maxChunkSizeBytesForParam("256MiB"), 256ULL << 20); + EXPECT_EQ(maxChunkSizeBytesForParam("1G"), 1ULL << 30); + EXPECT_EQ(maxChunkSizeBytesForParam("1GB"), 1ULL << 30); + EXPECT_EQ(maxChunkSizeBytesForParam("1gb"), 1ULL << 30); + EXPECT_EQ(maxChunkSizeBytesForParam("2GiB"), 2ULL << 30); +} + +TEST(BounceConfig, ByteGarbageFallsBackToDefault) +{ + std::size_t const def = b::BounceConfig{}.maxChunkSizeBytes; + EXPECT_EQ(maxChunkSizeBytesForParam("abc"), def); // no digits + EXPECT_EQ(maxChunkSizeBytesForParam("256XB"), def); // unknown suffix + EXPECT_EQ(maxChunkSizeBytesForParam("256 MB"), def); // space before suffix + EXPECT_EQ(maxChunkSizeBytesForParam("256MBx"), def); // trailing junk + EXPECT_EQ(maxChunkSizeBytesForParam(""), def); // empty -> default + EXPECT_EQ(maxChunkSizeBytesForParam("999999999999G"), def); // multiply would overflow u64 + EXPECT_EQ(maxChunkSizeBytesForParam("0"), def); // zero chunk cap is not usable +} + +TEST(BounceConfig, AllByteVarsAcceptSuffixFromEnv) +{ + ScopedEnv env; + env.set("TRTLLM_NIXL_BOUNCE_ARENA_ALLOCATION_GRANULARITY_BYTES", "2mb"); + env.set("TRTLLM_NIXL_BOUNCE_MAX_CHUNK_SIZE_BYTES", "64MB"); + env.set("TRTLLM_NIXL_BOUNCE_MAX_AVERAGE_DESCRIPTOR_SIZE_BYTES", "32kb"); + auto const cfg = b::BounceConfig::fromEnv(); + EXPECT_EQ(cfg.arenaAllocationGranularityBytes, 2ULL << 20); + EXPECT_EQ(cfg.maxChunkSizeBytes, 64ULL << 20); + EXPECT_EQ(cfg.maxAverageDescriptorSizeBytes, 32ULL << 10); +} + +TEST(BounceConfig, PlainCountsRejectSuffix) +{ + // Non-byte knobs keep strict integer parsing: a suffix is garbage -> default. Same rule on + // both the env and the params path. + ScopedEnv env; + env.set("TRTLLM_NIXL_BOUNCE_MAX_INFLIGHT_CHUNKS_PER_REQUEST", "4MB"); + auto const cfg = b::BounceConfig::fromEnv(); + EXPECT_EQ(cfg.maxInflightChunksPerRequest, b::BounceConfig{}.maxInflightChunksPerRequest); + + std::unordered_map const params{{"max_inflight_chunks_per_request", "4MB"}}; + auto const cfgP = b::BounceConfig::fromParams(params, b::BounceConfig{}); + EXPECT_EQ(cfgP.maxInflightChunksPerRequest, b::BounceConfig{}.maxInflightChunksPerRequest); +} + +TEST(BounceConfig, DescriptiveNamesParse) +{ + ScopedEnv env; + env.set("TRTLLM_NIXL_BOUNCE_MAX_INFLIGHT_CHUNKS_PER_REQUEST", "5"); + env.set("TRTLLM_NIXL_BOUNCE_COPY_STREAM_COUNT", "6"); + env.set("TRTLLM_NIXL_BOUNCE_SCATTER_WORKER_COUNT", "7"); + env.set("TRTLLM_NIXL_BOUNCE_MIN_DESCRIPTOR_COUNT", "8"); + env.set("TRTLLM_NIXL_BOUNCE_REQUEST_TIMEOUT_MS", "900"); + env.set("TRTLLM_NIXL_BOUNCE_DISABLE_FABRIC_MEMORY", "yes"); + env.set("TRTLLM_NIXL_BOUNCE_ENABLE_EAGER_GATHER", "false"); + env.set("TRTLLM_NIXL_BOUNCE_USE_ZERO_COPY_ARGUMENTS", "false"); + + auto const cfg = b::BounceConfig::fromEnv(); + EXPECT_EQ(cfg.maxInflightChunksPerRequest, 5u); + EXPECT_EQ(cfg.copyStreamCount, 6u); + EXPECT_EQ(cfg.scatterWorkerCount, 7u); + EXPECT_EQ(cfg.minDescriptorCount, 8u); + EXPECT_EQ(cfg.requestTimeoutMs, 900); + // Lease/quarantine are derived from the request timeout, not independent knobs. + EXPECT_EQ(cfg.receiverFlowTimeoutMs, 1800); + EXPECT_EQ(cfg.quarantineMs, 900); + EXPECT_TRUE(cfg.disableFabricMemory); + EXPECT_FALSE(cfg.enableEagerGather); + EXPECT_FALSE(cfg.useZeroCopyArguments); +} + +TEST(BounceConfig, InvalidResourceCountsFallBackToDefaults) +{ + ScopedEnv env; + env.set("TRTLLM_NIXL_BOUNCE_COPY_STREAM_COUNT", "0"); + env.set("TRTLLM_NIXL_BOUNCE_MAX_INFLIGHT_CHUNKS_PER_REQUEST", "4294967296"); + env.set("TRTLLM_NIXL_BOUNCE_REQUEST_TIMEOUT_MS", "2147483648"); + + auto const defaults = b::BounceConfig{}; + auto const cfg = b::BounceConfig::fromEnv(); + EXPECT_EQ(cfg.copyStreamCount, defaults.copyStreamCount); + EXPECT_EQ(cfg.maxInflightChunksPerRequest, defaults.maxInflightChunksPerRequest); + EXPECT_EQ(cfg.requestTimeoutMs, defaults.requestTimeoutMs); +} + +// The timeout parser admits any value up to INT_MAX, so the 2x lease derivation must clamp instead +// of overflowing: a wrapped-negative receiverFlowTimeoutMs reads as "lease disabled" and silently +// turns off the dead-sender region reclaim. +TEST(BounceConfig, LeaseDerivationClampsInsteadOfOverflowing) +{ + ScopedEnv env; + env.set("TRTLLM_NIXL_BOUNCE_REQUEST_TIMEOUT_MS", "2000000000"); // valid, but > INT_MAX / 2 + + auto const cfg = b::BounceConfig::fromEnv(); + EXPECT_EQ(cfg.requestTimeoutMs, 2000000000); + EXPECT_EQ(cfg.receiverFlowTimeoutMs, std::numeric_limits::max()); + EXPECT_EQ(cfg.quarantineMs, 2000000000); +} + +// agent_bounce_params dict > env var > default: a key present in both takes the dict value; keys +// only in the env keep the env value; everything else keeps the default. +TEST(BounceConfig, ParamsOverrideEnv) +{ + ScopedEnv env; + env.set("TRTLLM_NIXL_BOUNCE_COPY_STREAM_COUNT", "6"); + env.set("TRTLLM_NIXL_BOUNCE_SCATTER_WORKER_COUNT", "7"); + + std::unordered_map const params{ + {"copy_stream_count", "3"}, {"max_chunk_size", "64MB"}, // suffix parsing works on the params path too + }; + auto const cfg = b::BounceConfig::fromParams(params, b::BounceConfig::fromEnv()); + EXPECT_EQ(cfg.copyStreamCount, 3u); // dict wins over env + EXPECT_EQ(cfg.scatterWorkerCount, 7u); // env-only knob keeps env value + EXPECT_EQ(cfg.maxChunkSizeBytes, 64ULL << 20); // dict-only knob applies + EXPECT_EQ(cfg.maxInflightChunksPerRequest, // untouched knob keeps the default + b::BounceConfig{}.maxInflightChunksPerRequest); // +} + +// A dict value that fails to parse keeps the env/base value instead of clobbering it. "0" counts +// as unparsable for allowZero=false knobs (arena_allocation_granularity) but is a valid value for +// allowZero=true ones (max_average_descriptor_size: a zero gate routes everything to standard NIXL). +TEST(BounceConfig, ParamsGarbageKeepsBaseValue) +{ + ScopedEnv env; + env.set("TRTLLM_NIXL_BOUNCE_COPY_STREAM_COUNT", "6"); + std::unordered_map const params{ + {"copy_stream_count", "banana"}, + {"arena_allocation_granularity", "0"}, + {"max_average_descriptor_size", "0"}, + }; + auto const cfg = b::BounceConfig::fromParams(params, b::BounceConfig::fromEnv()); + EXPECT_EQ(cfg.copyStreamCount, 6u); + EXPECT_EQ(cfg.arenaAllocationGranularityBytes, b::BounceConfig{}.arenaAllocationGranularityBytes); + EXPECT_EQ(cfg.maxAverageDescriptorSizeBytes, 0u); +} + +// Unknown keys are skipped with a warning (must not crash or affect known knobs). +TEST(BounceConfig, UnknownParamKeyIsIgnored) +{ + std::unordered_map const params{ + {"definitely_not_a_bounce_knob", "42"}, + {"copy_stream_count", "3"}, + }; + auto const cfg = b::BounceConfig::fromParams(params, b::BounceConfig{}); + EXPECT_EQ(cfg.copyStreamCount, 3u); + EXPECT_EQ(cfg.maxChunkSizeBytes, b::BounceConfig{}.maxChunkSizeBytes); +} + +// A dict-provided request_timeout_ms must re-derive the receiver lease and quarantine values +// (otherwise the lease would keep the env/default-based derivation and break its "must exceed the +// peers' request timeout" invariant). +TEST(BounceConfig, ParamsRequestTimeoutRederivesLease) +{ + ScopedEnv env; + env.set("TRTLLM_NIXL_BOUNCE_REQUEST_TIMEOUT_MS", "900"); + + std::unordered_map const params{{"request_timeout_ms", "1000"}}; + auto const cfg = b::BounceConfig::fromParams(params, b::BounceConfig::fromEnv()); + EXPECT_EQ(cfg.requestTimeoutMs, 1000); + EXPECT_EQ(cfg.receiverFlowTimeoutMs, 2000); + EXPECT_EQ(cfg.quarantineMs, 1000); + + // The 64-bit clamp holds on the params path too. + std::unordered_map const bigParams{{"request_timeout_ms", "2000000000"}}; + auto const bigCfg = b::BounceConfig::fromParams(bigParams, b::BounceConfig::fromEnv()); + EXPECT_EQ(bigCfg.requestTimeoutMs, 2000000000); + EXPECT_EQ(bigCfg.receiverFlowTimeoutMs, std::numeric_limits::max()); + EXPECT_EQ(bigCfg.quarantineMs, 2000000000); + + // Edge values: "0" is accepted (timeout disabled, derived values 0), "-1" is rejected (strict + // unsigned parsing) and keeps the base value untouched. + auto const zeroCfg = b::BounceConfig::fromParams({{"request_timeout_ms", "0"}}, b::BounceConfig{}); + EXPECT_EQ(zeroCfg.requestTimeoutMs, 0); + EXPECT_EQ(zeroCfg.receiverFlowTimeoutMs, 0); + EXPECT_EQ(zeroCfg.quarantineMs, 0); + auto const negCfg = b::BounceConfig::fromParams({{"request_timeout_ms", "-1"}}, b::BounceConfig{}); + EXPECT_EQ(negCfg.requestTimeoutMs, b::BounceConfig{}.requestTimeoutMs); + EXPECT_EQ(negCfg.receiverFlowTimeoutMs, b::BounceConfig{}.receiverFlowTimeoutMs); +} + +// A dict WITHOUT request_timeout_ms must not re-derive: an env-derived lease survives, and +// directly-set lease/quarantine values on the base (as failure-injection tests use) stay intact. +TEST(BounceConfig, ParamsWithoutRequestTimeoutLeaveLeaseAlone) +{ + ScopedEnv env; + env.set("TRTLLM_NIXL_BOUNCE_REQUEST_TIMEOUT_MS", "900"); + std::unordered_map const params{{"copy_stream_count", "3"}}; + auto const cfg = b::BounceConfig::fromParams(params, b::BounceConfig::fromEnv()); + EXPECT_EQ(cfg.copyStreamCount, 3u); + EXPECT_EQ(cfg.requestTimeoutMs, 900); + EXPECT_EQ(cfg.receiverFlowTimeoutMs, 1800); + EXPECT_EQ(cfg.quarantineMs, 900); + + b::BounceConfig base; + base.receiverFlowTimeoutMs = 123456; + base.quarantineMs = 654321; + auto const cfgDirect = b::BounceConfig::fromParams(params, base); + EXPECT_EQ(cfgDirect.receiverFlowTimeoutMs, 123456); + EXPECT_EQ(cfgDirect.quarantineMs, 654321); +} diff --git a/cpp/tests/unit_tests/executor/bounce/bounceMessageCodecTest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceMessageCodecTest.cpp new file mode 100644 index 000000000000..e50e3ef8f6d5 --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/bounceMessageCodecTest.cpp @@ -0,0 +1,268 @@ +/* + * 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/BounceMessage.h" + +#include + +#include + +namespace b = tensorrt_llm::executor::kv_cache::bounce; + +namespace +{ +b::BounceMsgHeader decodeOk(std::string const& blob, b::BounceMsgType expectType) +{ + b::BounceMsgHeader h{}; + EXPECT_TRUE(b::decodeHeader(blob, h)); + EXPECT_EQ(h.msgType, static_cast(expectType)); + return h; +} +} // namespace + +TEST(BounceMessageCodec, WantCarriesChunkSizesAndEndpoint) +{ + std::vector sizes{4096, 8192, 1024}; + std::string const ep = "tcp://10.0.0.3:5170"; + auto blob = b::encodeWant(/*rid=*/42, sizes, ep); + auto h = decodeOk(blob, b::BounceMsgType::kWANT); + EXPECT_EQ(h.requestId, 42u); + EXPECT_EQ(h.count, 3u); + std::vector out; + std::string outEp; + ASSERT_TRUE(b::decodeWant(blob, h, out, outEp)); + EXPECT_EQ(out, sizes); + EXPECT_EQ(outEp, ep); +} + +TEST(BounceMessageCodec, CancelIsEmptyWantThatStillCarriesEndpoint) +{ + std::string const ep = "tcp://127.0.0.1:9999"; + auto blob = b::encodeCancel(42, ep); // same wire form as an empty-chunk WANT + auto h = decodeOk(blob, b::BounceMsgType::kWANT); + EXPECT_EQ(h.count, 0u); // no chunks -> cancel + std::vector out; + std::string outEp; + ASSERT_TRUE(b::decodeWant(blob, h, out, outEp)); + EXPECT_TRUE(out.empty()); + EXPECT_TRUE(b::isCancelWant(out)); // decoded WANT is recognized as a cancel + EXPECT_EQ(outEp, ep); // endpoint still travels so the receiver can bootstrap even on a bare cancel + + // A real (non-empty) WANT is NOT a cancel. + std::vector real; + std::string ep2; + auto wblob = b::encodeWant(43, std::vector{4096}, ep); + auto wh = decodeOk(wblob, b::BounceMsgType::kWANT); + ASSERT_TRUE(b::decodeWant(wblob, wh, real, ep2)); + EXPECT_FALSE(b::isCancelWant(real)); +} + +TEST(BounceMessageCodec, WantEmptyEndpointRoundTrips) +{ + auto blob = b::encodeWant(7, std::vector{16}, ""); + auto h = decodeOk(blob, b::BounceMsgType::kWANT); + std::vector out; + std::string outEp; + ASSERT_TRUE(b::decodeWant(blob, h, out, outEp)); + EXPECT_EQ(out, (std::vector{16})); + EXPECT_TRUE(outEp.empty()); +} + +TEST(BounceMessageCodec, GrantRoundTrip) +{ + std::vector credits{{0x3000, 256, 1, 3}, {0x7000, 512, 1, 7}}; + auto blob = b::encodeGrant(/*rid=*/99, credits); + auto h = decodeOk(blob, b::BounceMsgType::kGRANT); + EXPECT_EQ(h.requestId, 99u); + std::vector out; + ASSERT_TRUE(b::decodeCredits(blob, h, out)); + ASSERT_EQ(out.size(), 2u); + EXPECT_EQ(out[0].regionHandle, 3u); + EXPECT_EQ(out[0].devId, 1u); + EXPECT_EQ(out[1].addr, 0x7000u); + EXPECT_EQ(out[1].len, 512u); +} + +TEST(BounceMessageCodec, DataRoundTrip) +{ + // {bounceOffset, dstAddr, dstStride, bounceStride, pieceSize, count}: one plain extent + one + // 3-piece strided run. + std::vector entries{{0, 0xD000, 0, 0, 128, 1}, {128, 0xE000, 4096, 128, 64, 3}}; + auto blob = b::encodeData(/*rid=*/7, /*chunk=*/2, /*numChunks=*/5, /*regionHandle=*/4, entries); + auto h = decodeOk(blob, b::BounceMsgType::kDATA); + EXPECT_EQ(h.requestId, 7u); + EXPECT_EQ(h.chunkIdx, 2u); + EXPECT_EQ(h.numChunks, 5u); + EXPECT_EQ(h.regionHandle, 4u); + std::vector out; + ASSERT_TRUE(b::decodeScatter(blob, h, out)); + ASSERT_EQ(out.size(), 2u); + EXPECT_EQ(out[0].dstAddr, 0xD000u); + EXPECT_EQ(out[0].count, 1u); + EXPECT_EQ(out[1].bounceOffset, 128u); + EXPECT_EQ(out[1].dstStride, 4096u); + EXPECT_EQ(out[1].bounceStride, 128u); + EXPECT_EQ(out[1].pieceSize, 64u); + EXPECT_EQ(out[1].count, 3u); +} + +TEST(BounceMessageCodec, AckRoundTrip) +{ + auto blob = b::encodeAck(/*rid=*/7, /*chunk=*/3, /*regionHandle=*/9); + auto h = decodeOk(blob, b::BounceMsgType::kACK); + EXPECT_EQ(h.requestId, 7u); + EXPECT_EQ(h.chunkIdx, 3u); + EXPECT_EQ(h.regionHandle, 9u); +} + +TEST(BounceMessageCodec, EmptyEntriesSerialize) +{ + auto blob = b::encodeGrant(1, {}); + auto h = decodeOk(blob, b::BounceMsgType::kGRANT); + EXPECT_EQ(h.count, 0u); + EXPECT_EQ(h.payloadBytes, 0u); + std::vector out; + ASSERT_TRUE(b::decodeCredits(blob, h, out)); + EXPECT_TRUE(out.empty()); +} + +TEST(BounceMessageCodec, LargeScatterCountRoundTrips) +{ + // A single DATA chunk can carry many small descriptors (the exact case bounce optimizes). + // `count`/`payloadBytes` are 32-bit, so a count well past 65535 must round-trip intact — + // guards against any accidental 16-bit truncation of the entry count. + constexpr std::uint32_t kN = 100000; // 100k * 36B = 3.6 MB payload + std::vector entries(kN); + for (std::uint32_t i = 0; i < kN; ++i) + { + entries[i] = {static_cast(i) * 64, 0x10000000ULL + i, 0, 0, 64, 1}; + } + auto blob = b::encodeData(/*rid=*/1, /*chunk=*/0, /*numChunks=*/1, /*regionHandle=*/0, entries); + auto h = decodeOk(blob, b::BounceMsgType::kDATA); + EXPECT_EQ(h.count, kN); + EXPECT_EQ(h.payloadBytes, kN * sizeof(b::BounceScatterRun)); + std::vector out; + ASSERT_TRUE(b::decodeScatter(blob, h, out)); + ASSERT_EQ(out.size(), kN); + // Spot-check the boundaries and the >16-bit index region. + EXPECT_EQ(out[0].dstAddr, 0x10000000ULL); + EXPECT_EQ(out[65535].bounceOffset, static_cast(65535) * 64); + EXPECT_EQ(out[kN - 1].dstAddr, 0x10000000ULL + (kN - 1)); + EXPECT_EQ(out[kN - 1].count, 1u); +} + +TEST(BounceMessageCodec, LargeCreditCountRoundTrips) +{ + constexpr std::uint32_t kN = 70000; // > 65535 credits in one GRANT + std::vector credits(kN); + for (std::uint32_t i = 0; i < kN; ++i) + { + credits[i] = {0x2000ULL + i, 4096, 0, i}; // {addr, len, devId, regionHandle} + } + auto blob = b::encodeGrant(/*rid=*/5, credits); + auto h = decodeOk(blob, b::BounceMsgType::kGRANT); + EXPECT_EQ(h.count, kN); + EXPECT_EQ(h.payloadBytes, kN * sizeof(b::BounceCreditEntry)); + std::vector out; + ASSERT_TRUE(b::decodeCredits(blob, h, out)); + ASSERT_EQ(out.size(), kN); + EXPECT_EQ(out[kN - 1].regionHandle, kN - 1); + EXPECT_EQ(out[kN - 1].addr, 0x2000ULL + (kN - 1)); +} + +TEST(BounceMessageCodec, ShortBlobRejected) +{ + std::string tiny(8, '\0'); + b::BounceMsgHeader h{}; + EXPECT_FALSE(b::decodeHeader(tiny, h)); +} + +TEST(BounceMessageCodec, BadMagicRejected) +{ + auto blob = b::encodeAck(1, 0, 0); + blob[0] = 'X'; // corrupt magic + b::BounceMsgHeader h{}; + EXPECT_FALSE(b::decodeHeader(blob, h)); +} + +TEST(BounceMessageCodec, InflatedPayloadBytesRejected) +{ + auto blob = b::encodeAck(1, 0, 0); + // Corrupt payloadBytes to a value the blob can't satisfy. + std::uint32_t inflated = 4096; + std::memcpy(blob.data() + offsetof(b::BounceMsgHeader, payloadBytes), &inflated, sizeof(inflated)); + b::BounceMsgHeader h{}; + EXPECT_FALSE(b::decodeHeader(blob, h)); +} + +TEST(BounceMessageCodec, CrossTypeDecodeMismatchRejected) +{ + // A WANT blob (4-byte u32 chunk sizes) decoded as credits (24-byte) -> count*24 != payloadBytes. + auto blob = b::encodeWant(1, std::vector{4096, 8192}, "tcp://x:1"); + b::BounceMsgHeader h{}; + ASSERT_TRUE(b::decodeHeader(blob, h)); + std::vector wrong; + EXPECT_FALSE(b::decodeCredits(blob, h, wrong)); +} + +// ---- capability handshake (travels in AgentDesc, serialize_utils-based — not a control message) ---- + +TEST(BounceMessageCodec, HandshakeRoundTrip) +{ + b::BounceHandshake in; + in.wireVersion = b::kBounceVersion; + in.controlKind = b::BounceControlKind::kZMQ; + in.arenaUsableCapacityBytes = 256ULL << 20; + in.maxChunkSizeBytes = 32ULL << 20; + in.endpoint = "tcp://10.0.0.1:5555"; + b::BounceHandshake out; + ASSERT_TRUE(b::decodeHandshake(b::encodeHandshake(in), out)); + EXPECT_EQ(out.wireVersion, in.wireVersion); + EXPECT_EQ(out.controlKind, in.controlKind); + EXPECT_EQ(out.arenaUsableCapacityBytes, in.arenaUsableCapacityBytes); + EXPECT_EQ(out.maxChunkSizeBytes, in.maxChunkSizeBytes); + EXPECT_EQ(out.endpoint, in.endpoint); +} + +TEST(BounceMessageCodec, HandshakeEmptyEndpointRoundTrip) +{ + b::BounceHandshake in; // defaults; empty endpoint stays empty + b::BounceHandshake out; + ASSERT_TRUE(b::decodeHandshake(b::encodeHandshake(in), out)); + EXPECT_TRUE(out.endpoint.empty()); +} + +TEST(BounceMessageCodec, HandshakeGarbageRejected) +{ + b::BounceHandshake out; + EXPECT_FALSE(b::decodeHandshake("", out)); + EXPECT_FALSE(b::decodeHandshake("tcp://10.0.0.1:5555", out)); // legacy raw endpoint, no magic + EXPECT_FALSE(b::decodeHandshake(std::string(3, '\0'), out)); // shorter than the magic + auto blob = b::encodeHandshake(b::BounceHandshake{}); + blob[0] ^= 0x5A; // corrupt magic + EXPECT_FALSE(b::decodeHandshake(blob, out)); +} + +TEST(BounceMessageCodec, HandshakeTruncatedRejected) +{ + b::BounceHandshake in; + in.endpoint = "tcp://10.0.0.1:5555"; + auto blob = b::encodeHandshake(in); + b::BounceHandshake out; + EXPECT_FALSE(b::decodeHandshake(blob.substr(0, blob.size() - 4), out)); // endpoint cut short + EXPECT_FALSE(b::decodeHandshake(blob.substr(0, 10), out)); // fixed part cut short +} diff --git a/cpp/tests/unit_tests/executor/bounce/bounceTestNixlNode.h b/cpp/tests/unit_tests/executor/bounce/bounceTestNixlNode.h new file mode 100644 index 000000000000..2f4dc8c74a53 --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/bounceTestNixlNode.h @@ -0,0 +1,215 @@ +/* + * 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 + +// Shared helpers for the bounce v2 tests that drive the FULL pipeline over REAL NIXL RDMA: a "node" +// is a complete agent stack (NixlTransferAgent as the data plane + arena + exec + zmq control +// channel + BounceTransport), plus seeded device buffers and a byte-exact verifier. Used by +// bounceTransportTest, bounceTransportFailureTest, and bounceAgentE2ETest so they share one NIXL +// setup (no per-file copy, and no LocalCopy loopback fake — the data plane is always real NIXL). + +#include "bounceTestUtils.h" + +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceArena.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ExecPool.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ZmqControlChannel.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h" + +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace bounce_test +{ + +namespace b = tensorrt_llm::executor::kv_cache::bounce; +namespace kvc = tensorrt_llm::executor::kv_cache; + +// A `seed`-distinct byte pattern so concurrent transfers can't masquerade as each other. +inline unsigned char patSeed(std::uint32_t seed, std::size_t d, std::size_t i) +{ + return static_cast((seed * 131 + d * 191 + i * 23 + 5) & 0xFF); +} + +// One transfer's src (gather source) + dst (scatter target) device buffers + descriptor lists. +// These KV buffers are NOT NIXL-registered — only the bounce arena traverses RDMA. +struct XferBufs +{ + void* src{nullptr}; + void* dst{nullptr}; + std::uint32_t nDescs{}; + std::uint32_t descBytes{}; + std::uint32_t seed{}; + std::vector off; + std::uint64_t total{}; + kvc::TransferDescs srcDescs{kvc::MemoryType::kVRAM, {}}; + kvc::TransferDescs dstDescs{kvc::MemoryType::kVRAM, {}}; +}; + +inline XferBufs makeXferBufs(std::uint32_t nDescs, std::uint32_t descBytes, std::uint32_t seed) +{ + XferBufs x; + x.nDescs = nDescs; + x.descBytes = descBytes; + x.seed = seed; + x.off.resize(nDescs); + std::uint64_t cur = 0; + for (std::uint32_t i = 0; i < nDescs; ++i) + { + x.off[i] = cur; + cur = alignUp(cur + descBytes, 256); + } + x.total = cur; + EXPECT_EQ(cudaMalloc(&x.src, x.total), cudaSuccess); + EXPECT_EQ(cudaMalloc(&x.dst, x.total), cudaSuccess); + std::vector h(x.total, 0); + for (std::uint32_t i = 0; i < nDescs; ++i) + { + for (std::uint32_t j = 0; j < descBytes; ++j) + { + h[x.off[i] + j] = patSeed(seed, i, j); + } + } + EXPECT_EQ(cudaMemcpy(x.src, h.data(), x.total, cudaMemcpyHostToDevice), cudaSuccess); + EXPECT_EQ(cudaMemset(x.dst, 0, x.total), cudaSuccess); + auto a2 = [](void* p, std::uint64_t o) { return reinterpret_cast(static_cast(p) + o); }; + std::vector sd; + std::vector dd; + for (std::uint32_t i = 0; i < nDescs; ++i) + { + sd.emplace_back(a2(x.src, x.off[i]), descBytes, 0); + dd.emplace_back(a2(x.dst, x.off[i]), descBytes, 0); + } + x.srcDescs = kvc::TransferDescs{kvc::MemoryType::kVRAM, std::move(sd)}; + x.dstDescs = kvc::TransferDescs{kvc::MemoryType::kVRAM, std::move(dd)}; + return x; +} + +inline bool verifyXferBufs(XferBufs const& x) +{ + std::vector got(x.total, 0xEE); + if (cudaMemcpy(got.data(), x.dst, x.total, cudaMemcpyDeviceToHost) != cudaSuccess) + { + return false; + } + for (std::uint32_t i = 0; i < x.nDescs; ++i) + { + for (std::uint32_t j = 0; j < x.descBytes; ++j) + { + if (got[x.off[i] + j] != patSeed(x.seed, i, j)) + { + return false; + } + } + } + return true; +} + +inline void freeXferBufs(XferBufs& x) +{ + if (x.src) + { + cudaFree(x.src); + x.src = nullptr; + } + if (x.dst) + { + cudaFree(x.dst); + x.dst = nullptr; + } +} + +// One bounce node = a full agent stack (agent + arena + exec + control channel + transport). The +// data plane is the agent itself (postXferRequest / registerRegionImpl); `agent` may be a test +// subclass of NixlTransferAgent that overrides postXferRequest to inject transfer faults. +struct Node +{ + std::string name; + std::unique_ptr agent; + std::unique_ptr arena; + std::unique_ptr exec; + std::unique_ptr ch; + std::unique_ptr tx; + + ~Node() + { + tx.reset(); // join the transport threads before the arena/agent go away + if (agent && arena) + { + agent->deregisterRegionImpl(arena->base(), arena->bytes(), /*deviceId=*/0); + } + } +}; + +// Build one full bounce node (agent+arena+exec+channel+transport). Returns nullptr if the NIXL +// agent/backend can't init or the arena can't be registered (caller GTEST_SKIPs). `cfg` supplies +// arenaSizeBytes/arenaAllocationGranularityBytes/maxInflightChunksPerRequest etc., so callers +// control the scheduler/arena sizing. `makeAgent` lets failure tests substitute a fault-injecting +// NixlTransferAgent subclass. +inline std::unique_ptr makeNode(std::string const& name, b::BounceConfig const& cfg, std::size_t maxDescs, + std::function(kvc::BaseAgentConfig const&)> const& makeAgent = nullptr) +{ + auto n = std::make_unique(); + n->name = name; + try + { + kvc::BaseAgentConfig c{name, /*useProgThread=*/true, /*multiThread=*/false, /*useListenThread=*/true}; + n->agent = makeAgent ? makeAgent(c) : std::make_unique(c); + } + catch (std::exception const&) + { + return nullptr; + } + n->arena = std::make_unique(cfg.arenaSizeBytes, 0, /*allowFabric=*/false); + n->exec = std::make_unique(cfg.maxInflightChunksPerRequest + 4, maxDescs, 0, cfg.useZeroCopyArguments); + if (!n->agent->registerRegionImpl(n->arena->base(), n->arena->bytes(), /*deviceId=*/0)) + { + return nullptr; + } + n->ch = std::make_unique(name); + n->tx + = std::make_unique(n->name, cfg, 0, n->ch.get(), *n->agent, n->arena.get(), n->exec.get()); + return n; +} + +// Bidirectional connect for the white-box harness. Two SEPARATE wirings are needed here because the +// transport under test is hand-built (b::BounceTransport with its own b::ZmqControlChannel) and lives +// OUTSIDE the agent, so loadRemoteAgent cannot wire this standalone transport: +// - loadRemoteAgent(AgentDesc) exchanges only the NIXL metadata layer (so createXferReq can resolve +// the remote arena). The connection-info overload would not carry the structured metadata. +// - addPeer() wires the control-channel layer (the DEALER to the peer's ROUTER) on the standalone +// transport. NixlTransferAgent normally folds this into handshake registration plus WANT +// self-bootstrap; here it is manual. +inline void wirePair(Node& a, Node& b) +{ + a.agent->loadRemoteAgent(b.name, b.agent->getLocalAgentDesc()); + b.agent->loadRemoteAgent(a.name, a.agent->getLocalAgentDesc()); + a.tx->addPeer(b.name, b.ch->localEndpoint()); + b.tx->addPeer(a.name, a.ch->localEndpoint()); +} + +} // namespace bounce_test diff --git a/cpp/tests/unit_tests/executor/bounce/bounceTestUtils.h b/cpp/tests/unit_tests/executor/bounce/bounceTestUtils.h new file mode 100644 index 000000000000..1995a3d3d57f --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/bounceTestUtils.h @@ -0,0 +1,43 @@ +/* + * 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 + +// NIXL-free helpers shared by ALL bounce test tiers (pure-GPU tests include this directly; +// the real-NIXL tier gets it through bounceTestNixlNode.h): the CUDA availability probe and +// alignment math. The byte-pattern generators intentionally stay with their users — the +// gather/scatter kernel test and the NIXL node harness use DIFFERENT formulas. + +#include + +#include + +namespace bounce_test +{ + +inline bool hasCuda() +{ + int n = 0; + return cudaGetDeviceCount(&n) == cudaSuccess && n > 0; +} + +inline std::uint64_t alignUp(std::uint64_t v, std::uint64_t a) +{ + return (v + a - 1) / a * a; +} + +} // namespace bounce_test diff --git a/cpp/tests/unit_tests/executor/bounce/bounceTransferPlanTest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceTransferPlanTest.cpp new file mode 100644 index 000000000000..de678465ee3a --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/bounceTransferPlanTest.cpp @@ -0,0 +1,238 @@ +/* + * 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/BounceTransferPlan.h" + +#include + +#include +#include + +namespace b = tensorrt_llm::executor::kv_cache::bounce; +namespace kvc = tensorrt_llm::executor::kv_cache; + +namespace +{ +// Build a TransferDescs from (addr,len,dev) tuples. Addresses are synthetic; the planner never +// dereferences them, it only bins by length / device. +kvc::TransferDescs makeDescs(std::vector> const& t) +{ + std::vector v; + v.reserve(t.size()); + for (auto const& [a, l, d] : t) + { + v.emplace_back(a, l, d); + } + return kvc::TransferDescs{kvc::MemoryType::kVRAM, std::move(v)}; +} +} // namespace + +TEST(BounceTransferPlan, EmptyYieldsNoChunks) +{ + auto plan = b::BounceTransferPlan::build(makeDescs({}), makeDescs({}), /*maxChunkSizeBytes=*/1024, /*maxDescs=*/64); + EXPECT_EQ(plan.numChunks(), 0u); + EXPECT_EQ(plan.totalDescs(), 0u); + EXPECT_EQ(plan.totalBytes(), 0u); +} + +TEST(BounceTransferPlan, SingleDescOneChunk) +{ + auto plan = b::BounceTransferPlan::build(makeDescs({{0x1000, 100, 0}}), makeDescs({{0x9000, 100, 0}}), 1024, 64); + ASSERT_EQ(plan.numChunks(), 1u); + auto const& c = plan.chunks()[0]; + EXPECT_EQ(c.srcPtrs.size(), 1u); + EXPECT_EQ(c.bounceOffsets[0], 0u); + EXPECT_EQ(c.sizes[0], 100u); + EXPECT_EQ(c.totalBytes, 100u); + EXPECT_EQ(c.dstPtrs[0], 0x9000u); +} + +TEST(BounceTransferPlan, TwoDescsPackOneChunkWith32ByteAlignedOffsets) +{ + // len 100 -> next offset aligns up to 128 (multiple of 32). + auto plan = b::BounceTransferPlan::build( + makeDescs({{0x1000, 100, 0}, {0x2000, 50, 0}}), makeDescs({{0x9000, 100, 0}, {0xA000, 50, 0}}), 1024, 64); + ASSERT_EQ(plan.numChunks(), 1u); + auto const& c = plan.chunks()[0]; + EXPECT_EQ(c.bounceOffsets[0], 0u); + EXPECT_EQ(c.bounceOffsets[1], 128u); // alignUp(100,32)=128 + EXPECT_EQ(c.sizes, (std::vector{100, 50})); + EXPECT_EQ(c.totalBytes, 150u); +} + +TEST(BounceTransferPlan, OverflowSplitsIntoTwoChunks) +{ + // With a 256-byte chunk cap, two 200-byte descriptors cannot share a chunk. + auto plan = b::BounceTransferPlan::build( + makeDescs({{0x1000, 200, 0}, {0x2000, 200, 0}}), makeDescs({{0x9000, 200, 0}, {0xA000, 200, 0}}), 256, 64); + EXPECT_EQ(plan.numChunks(), 2u); + EXPECT_EQ(plan.chunks()[0].srcPtrs.size(), 1u); + EXPECT_EQ(plan.chunks()[1].srcPtrs.size(), 1u); +} + +TEST(BounceTransferPlan, DescExactlyChunkSizeIsOneChunk) +{ + auto plan = b::BounceTransferPlan::build(makeDescs({{0x1000, 256, 0}}), makeDescs({{0x9000, 256, 0}}), 256, 64); + ASSERT_EQ(plan.numChunks(), 1u); + EXPECT_EQ(plan.chunks()[0].totalBytes, 256u); +} + +TEST(BounceTransferPlan, DescLargerThanChunkThrows) +{ + EXPECT_ANY_THROW( + (void) b::BounceTransferPlan::build(makeDescs({{0x1000, 257, 0}}), makeDescs({{0x9000, 257, 0}}), 256, 64)); +} + +TEST(BounceTransferPlan, MaxChunkSizeBytesAboveU32Throws) +{ + // A chunk's packed size travels in 32-bit wire fields, so maxChunkSizeBytes must fit in 32 bits + // even though arena offsets are 64-bit. Building with a >4 GiB cap must be rejected. + EXPECT_ANY_THROW((void) b::BounceTransferPlan::build(makeDescs({{0x1000, 8, 0}}), makeDescs({{0x9000, 8, 0}}), + /*maxChunkSizeBytes=*/(std::size_t{1} << 32), 64)); + // Exactly 4 GiB - 1 is allowed. + EXPECT_NO_THROW((void) b::BounceTransferPlan::build(makeDescs({{0x1000, 8, 0}}), makeDescs({{0x9000, 8, 0}}), + /*maxChunkSizeBytes=*/(std::size_t{1} << 32) - 1, 64)); +} + +TEST(BounceTransferPlan, MaxDescsPerChunkBoundary) +{ + // 3 tiny descs, maxDescs=2 -> first chunk holds 2, second holds 1. + auto plan = b::BounceTransferPlan::build(makeDescs({{0x1000, 8, 0}, {0x2000, 8, 0}, {0x3000, 8, 0}}), + makeDescs({{0x9000, 8, 0}, {0xA000, 8, 0}, {0xB000, 8, 0}}), 4096, /*maxDescs=*/2); + ASSERT_EQ(plan.numChunks(), 2u); + EXPECT_EQ(plan.chunks()[0].srcPtrs.size(), 2u); + EXPECT_EQ(plan.chunks()[1].srcPtrs.size(), 1u); +} + +TEST(BounceTransferPlan, MixedSourceDeviceIdsThrow) +{ + EXPECT_ANY_THROW((void) b::BounceTransferPlan::build(makeDescs({{0x1000, 8, /*dev=*/0}, {0x2000, 8, /*dev=*/1}}), + makeDescs({{0x9000, 8, /*dev=*/0}, {0xA000, 8, /*dev=*/0}}), 4096, 64)); +} + +TEST(BounceTransferPlan, MixedDestinationDeviceIdsThrow) +{ + EXPECT_ANY_THROW((void) b::BounceTransferPlan::build(makeDescs({{0x1000, 8, /*dev=*/0}, {0x2000, 8, /*dev=*/0}}), + makeDescs({{0x9000, 8, /*dev=*/0}, {0xA000, 8, /*dev=*/1}}), 4096, 64)); +} + +TEST(BounceTransferPlan, ZeroLengthDescSkippedButCounted) +{ + auto plan = b::BounceTransferPlan::build( + makeDescs({{0x1000, 0, 0}, {0x2000, 16, 0}}), makeDescs({{0x9000, 0, 0}, {0xA000, 16, 0}}), 1024, 64); + ASSERT_EQ(plan.numChunks(), 1u); + EXPECT_EQ(plan.chunks()[0].srcPtrs.size(), 1u); // zero-len skipped from packing + EXPECT_EQ(plan.totalDescs(), 2u); // but still counted as seen + EXPECT_EQ(plan.totalBytes(), 16u); +} + +TEST(BounceTransferPlan, CountMismatchThrows) +{ + EXPECT_ANY_THROW((void) b::BounceTransferPlan::build(makeDescs({{0x1000, 8, 0}}), makeDescs({}), 1024, 64)); +} + +TEST(BounceTransferPlan, ContiguousSrcAndDstDescsMergeInPlace) +{ + // Both src and dst advance contiguously (and 32 divides 32, so the bounce cursor has no align + // gap) -> the two descs collapse into ONE plan desc covering 64 bytes. + auto plan = b::BounceTransferPlan::build( + makeDescs({{0x1000, 32, 0}, {0x1020, 32, 0}}), makeDescs({{0x9000, 32, 0}, {0x9020, 32, 0}}), 1024, 64); + ASSERT_EQ(plan.numChunks(), 1u); + auto const& c = plan.chunks()[0]; + ASSERT_EQ(c.srcPtrs.size(), 1u); + EXPECT_EQ(c.sizes[0], 64u); + EXPECT_EQ(c.totalBytes, 64u); + EXPECT_EQ(c.packedBytes, 64u); + EXPECT_EQ(plan.totalDescs(), 2u); // both input descs still counted as seen + EXPECT_EQ(plan.totalBytes(), 64u); +} + +TEST(BounceTransferPlan, ContiguousSrcOnlyDoesNotMergeDescs) +{ + // src contiguous but dst jumps -> per-desc arrays must stay separate (the gather is strided). + auto plan = b::BounceTransferPlan::build( + makeDescs({{0x1000, 32, 0}, {0x1020, 32, 0}}), makeDescs({{0x9000, 32, 0}, {0xA000, 32, 0}}), 1024, 64); + ASSERT_EQ(plan.numChunks(), 1u); + EXPECT_EQ(plan.chunks()[0].srcPtrs.size(), 2u); +} + +TEST(BounceTransferPlan, ScatterRunsCoalesceContiguousDst) +{ + // dst contiguous, src strided (e.g. ctx tp1 -> gen tp4: dst is the gen rank's dense head-slice + // pool): per-desc arrays keep 3 entries for the gather, but the scatter view collapses to ONE + // count==1 run whose pieceSize grew over the whole extent. + auto plan = b::BounceTransferPlan::build(makeDescs({{0x1000, 32, 0}, {0x3000, 32, 0}, {0x5000, 32, 0}}), + makeDescs({{0x9000, 32, 0}, {0x9020, 32, 0}, {0x9040, 32, 0}}), 1024, 64); + ASSERT_EQ(plan.numChunks(), 1u); + auto const& c = plan.chunks()[0]; + EXPECT_EQ(c.srcPtrs.size(), 3u); + ASSERT_EQ(c.scatterRuns.size(), 1u); + EXPECT_EQ(c.scatterRuns[0].dstAddr, 0x9000u); + EXPECT_EQ(c.scatterRuns[0].bounceOffset, 0u); + EXPECT_EQ(c.scatterRuns[0].pieceSize, 96u); + EXPECT_EQ(c.scatterRuns[0].count, 1u); +} + +TEST(BounceTransferPlan, ScatterRunsCoalesceUniformlyStridedDst) +{ + // dst uniformly strided (e.g. ctx tp-slice -> gen DP full-head pool: each 32B piece lands every + // 128B in the peer pool): ONE strided run of count 3. Bounce packing steps by exactly 32 + // (aligned), so bounceStride == pieceSize. + auto plan = b::BounceTransferPlan::build(makeDescs({{0x1000, 32, 0}, {0x3000, 32, 0}, {0x5000, 32, 0}}), + makeDescs({{0x9000, 32, 0}, {0x9080, 32, 0}, {0x9100, 32, 0}}), 1024, 64); + ASSERT_EQ(plan.numChunks(), 1u); + auto const& c = plan.chunks()[0]; + ASSERT_EQ(c.scatterRuns.size(), 1u); + EXPECT_EQ(c.scatterRuns[0].dstAddr, 0x9000u); + EXPECT_EQ(c.scatterRuns[0].dstStride, 0x80u); + EXPECT_EQ(c.scatterRuns[0].bounceStride, 32u); + EXPECT_EQ(c.scatterRuns[0].pieceSize, 32u); + EXPECT_EQ(c.scatterRuns[0].count, 3u); +} + +TEST(BounceTransferPlan, ScatterRunsBreakOnDstHoleOrAlignGap) +{ + // First pair: dst steps forward but the second desc's SIZE differs -> no stride latch -> two + // runs. Second pair: dst contiguous but the 100-byte desc aligns the cursor up to 128, leaving a + // bounce gap -> contiguous growth fails; the stride latch still absorbs it ONLY if sizes match — + // they don't (100 vs 32) -> two runs. + auto planHole = b::BounceTransferPlan::build( + makeDescs({{0x1000, 32, 0}, {0x3000, 16, 0}}), makeDescs({{0x9000, 32, 0}, {0xA000, 16, 0}}), 1024, 64); + ASSERT_EQ(planHole.numChunks(), 1u); + EXPECT_EQ(planHole.chunks()[0].scatterRuns.size(), 2u); + + auto planGap = b::BounceTransferPlan::build( + makeDescs({{0x1000, 100, 0}, {0x3000, 32, 0}}), makeDescs({{0x9000, 100, 0}, {0x9064, 32, 0}}), 1024, 64); + ASSERT_EQ(planGap.numChunks(), 1u); + auto const& c = planGap.chunks()[0]; + ASSERT_EQ(c.scatterRuns.size(), 2u); + EXPECT_EQ(c.scatterRuns[1].bounceOffset, 128u); // alignUp(100,32) +} + +TEST(BounceTransferPlan, ScatterRunsIrregularStrideBreaks) +{ + // Same sizes but NON-uniform dst steps (+0x80 then +0x40): the latch fixes stride 0x80 from the + // first pair; the third desc doesn't land on it -> it opens a new run (2 runs total, 3 pieces). + auto plan = b::BounceTransferPlan::build(makeDescs({{0x1000, 32, 0}, {0x3000, 32, 0}, {0x5000, 32, 0}}), + makeDescs({{0x9000, 32, 0}, {0x9080, 32, 0}, {0x90C0, 32, 0}}), 1024, 64); + ASSERT_EQ(plan.numChunks(), 1u); + auto const& c = plan.chunks()[0]; + ASSERT_EQ(c.scatterRuns.size(), 2u); + EXPECT_EQ(c.scatterRuns[0].count, 2u); + EXPECT_EQ(c.scatterRuns[1].count, 1u); + EXPECT_EQ(c.scatterRuns[1].dstAddr, 0x90C0u); +} diff --git a/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp new file mode 100644 index 000000000000..2e23efa73df2 --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/bounceTransportFailureTest.cpp @@ -0,0 +1,818 @@ +/* + * 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. + */ + +// Boundary / failure-path tests for BounceTransport: the reactor must always resolve a request +// (SUCCESS or FAILURE) and never hang — including when the peer never grants, the transfer engine +// fails, forgetPeer() drops a peer mid-transfer, or shutdown races in-flight requests. Plus a +// multi-threaded concurrent-submit test. +// The data plane is real NIXL (via bounceTestNixlNode helpers); the one exception is the +// transfer-failure path, which uses a FakeXferAgent (a NixlTransferAgent subclass overriding the +// postXferRequest primitive) because real NIXL cannot be made to deterministically fail a write. +// +// The happy-path / concurrency / bidirectional / multi-agent coverage lives in bounceAgentE2ETest, +// which drives the SAME pipeline through the production entry point (NixlTransferAgent:: +// submitTransferRequests) with one-directional AgentDesc bootstrap, so it is not duplicated here. + +#include "bounceTestNixlNode.h" + +#include + +#include +#include +#include +#include +#include + +namespace b = tensorrt_llm::executor::kv_cache::bounce; +namespace kvc = tensorrt_llm::executor::kv_cache; + +namespace +{ +// Test-controlled behavior for every write a FakeXferAgent posts. Shared by the agent and all the +// statuses it hands out. +struct XferControls +{ + std::atomic failWrites{false}; // every posted write polls as FAILURE + std::atomic allowTerminal{true}; // false -> posted writes stay IN_PROGRESS + std::atomic releaseSucceeds{true}; // false -> release() fails (handle retained) + std::atomic postCount{0}; + std::atomic releaseCount{0}; +}; + +// A TransferStatus driven by XferControls instead of a real backend handle. Unlike +// NixlTransferStatus it does NOT release in its destructor, keeping releaseCount deterministic. +class FakeXferStatus final : public kvc::TransferStatus +{ +public: + explicit FakeXferStatus(std::shared_ptr ctl) + : mCtl(std::move(ctl)) + { + } + + [[nodiscard]] bool isCompleted() const override + { + return wait(0) != kvc::TransferState::kIN_PROGRESS; + } + + [[nodiscard]] kvc::TransferState wait(int64_t) const override + { + if (mCtl->failWrites.load(std::memory_order_acquire)) + { + return kvc::TransferState::kFAILURE; + } + return mCtl->allowTerminal.load(std::memory_order_acquire) ? kvc::TransferState::kSUCCESS + : kvc::TransferState::kIN_PROGRESS; + } + + [[nodiscard]] bool release() override + { + if (mReleased) + { + return true; + } + mCtl->releaseCount.fetch_add(1, std::memory_order_relaxed); + if (!mCtl->releaseSucceeds.load(std::memory_order_acquire)) + { + return false; + } + mReleased = true; + return true; + } + +private: + std::shared_ptr mCtl; + bool mReleased{false}; +}; + +// The fault-injection seam: a REAL NixlTransferAgent (control plane / metadata untouched) whose +// low-level write primitive is overridden per XferControls — real NIXL cannot be coerced into a +// deterministic write failure. +class FakeXferAgent final : public kvc::NixlTransferAgent +{ +public: + FakeXferAgent(kvc::BaseAgentConfig const& config, std::shared_ptr ctl) + : kvc::NixlTransferAgent(config) + , mCtl(std::move(ctl)) + { + } + + [[nodiscard]] std::unique_ptr postXferRequest(kvc::TransferOp, kvc::TransferDescs const&, + kvc::TransferDescs const&, std::string const&, std::optional const&) override + { + mCtl->postCount.fetch_add(1, std::memory_order_relaxed); + return std::make_unique(mCtl); + } + +private: + std::shared_ptr mCtl; +}; + +b::BounceConfig cfg(int timeoutMs) +{ + b::BounceConfig c; + c.maxChunkSizeBytes = 4096; + c.maxInflightChunksPerRequest = 2; + c.scatterWorkerCount = 2; + c.arenaAllocationGranularityBytes = 256; + c.arenaSizeBytes = 1ULL << 20; + c.requestTimeoutMs = timeoutMs; + return c; +} + +// Shrink the arena to exactly `regionCap` max-size regions. Matching the allocation granularity to +// the maximum chunk size makes every region one buddy block. +void capRegions(b::BounceConfig& c, std::uint32_t regionCap) +{ + c.arenaAllocationGranularityBytes = c.maxChunkSizeBytes; + std::size_t arenaSizeBytes = c.arenaAllocationGranularityBytes; + while (arenaSizeBytes < static_cast(regionCap) * c.maxChunkSizeBytes) + { + arenaSizeBytes <<= 1; + } + c.arenaSizeBytes = arenaSizeBytes; +} + +// makeNode with a FakeXferAgent wired to `ctl`. +std::unique_ptr makeFakeNode( + std::string const& name, b::BounceConfig const& c, std::shared_ptr ctl) +{ + return bounce_test::makeNode(name, c, 1024, + [&ctl](kvc::BaseAgentConfig const& agentConfig) { return std::make_unique(agentConfig, ctl); }); +} + +// Poll `ch` until `budget` elapses, invoking `onMsg(header, blob)` for each decoded message; +// stops early when onMsg returns true (returns whether it did). +template +bool pumpChannel(b::ZmqControlChannel& ch, std::chrono::milliseconds budget, Fn&& onMsg) +{ + auto const deadline = std::chrono::steady_clock::now() + budget; + while (std::chrono::steady_clock::now() < deadline) + { + std::string peer; + std::string blob; + b::BounceMsgHeader h{}; + if (ch.recv(peer, blob, 20) && b::decodeHeader(blob, h) && onMsg(h, blob)) + { + return true; + } + } + return false; +} + +// Wait up to `budget` for a GRANT for `rid`, decoding its credits into `out`. +bool waitGrant(b::ZmqControlChannel& ch, std::uint64_t rid, std::chrono::milliseconds budget, + std::vector& out) +{ + return pumpChannel(ch, budget, + [&](b::BounceMsgHeader const& h, std::string const& blob) + { + if (static_cast(h.msgType) != b::BounceMsgType::kGRANT || h.requestId != rid) + { + return false; + } + EXPECT_TRUE(b::decodeCredits(blob, h, out)); + return !out.empty(); + }); +} + +// Count ACKs for `rid` over the FULL `budget` window (a negative-assertion helper: always runs to +// the deadline so extra/duplicate ACKs have time to show up). +int countAcks(b::ZmqControlChannel& ch, std::chrono::milliseconds budget, std::uint64_t rid) +{ + int n = 0; + pumpChannel(ch, budget, + [&](b::BounceMsgHeader const& h, std::string const&) + { + if (static_cast(h.msgType) == b::BounceMsgType::kACK && h.requestId == rid) + { + ++n; + } + return false; + }); + return n; +} +} // namespace + +// A peer that never GRANTs must fail the request on requestTimeoutMs, not hang. The WANT is +// DELIVERED to a live ROUTER that no transport ever recv()s on (a "ghost"), so this covers the +// stronger case (WANT accepted, nobody grants) as well as the weaker unknown-peer one (WANT +// dropped at send) — checkTimeouts must resolve the future kFAILURE either way. +TEST(BounceTransportFailure, NoGrantTimesOutNotHang) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/500); + auto t = bounce_test::makeNode("ngSolo", c, 1024); + if (!t) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + // A bound ROUTER that no transport ever recv()s on -> the WANT is delivered but never granted. + b::ZmqControlChannel ghost("ghostPeer"); + t->tx->addPeer("ghostPeer", ghost.localEndpoint()); + + auto bufs = bounce_test::makeXferBufs(8, 256, /*seed=*/1); + auto fut = t->tx->submit(bufs.srcDescs, bufs.dstDescs, "ghostPeer"); + ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "request hung with no grant"; + EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); + EXPECT_EQ(fut.get().reason, b::BounceFailReason::kNoProgressTimeout); + + t->tx->shutdown(); + bounce_test::freeXferBufs(bufs); +} + +// The posted write reports failure -> the request must FAIL (not hang, not falsely succeed). +// The sender is a FakeXferAgent whose writes always poll as FAILURE. +TEST(BounceTransportFailure, WriteFailureFailsRequest) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/5000); + capRegions(c, c.maxInflightChunksPerRequest); + auto ctl = std::make_shared(); + ctl->failWrites.store(true, std::memory_order_release); + auto A = makeFakeNode("feA", c, ctl); + auto B = bounce_test::makeNode("feB", c, 1024); + if (!A || !B) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + A->tx->addPeer("feB", B->ch->localEndpoint()); + B->tx->addPeer("feA", A->ch->localEndpoint()); + + auto bufs = bounce_test::makeXferBufs(8, 256, /*seed=*/1); + auto fut = A->tx->submit(bufs.srcDescs, bufs.dstDescs, "feB"); + ASSERT_EQ(fut.wait_for(std::chrono::seconds(10)), std::future_status::ready) << "write failure hung"; + EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); + EXPECT_EQ(fut.get().reason, b::BounceFailReason::kWriteFailed); + EXPECT_GE(ctl->postCount.load(std::memory_order_acquire), 1u); + + A->tx->shutdown(); + B->tx->shutdown(); + bounce_test::freeXferBufs(bufs); +} + +// DATA has no retransmission path in this protocol. Receiving it twice for one granted region must +// therefore produce one scatter and one ACK, rather than adding replay state for an impossible +// normal-flow event. +TEST(BounceTransportFailure, DuplicateDataProducesOneScatterAndAck) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/5000); + capRegions(c, c.maxInflightChunksPerRequest); + b::ZmqControlChannel sender("dupDataSender"); + auto receiver = bounce_test::makeNode("dupDataReceiver", c, 1024); // receiver posts no writes + if (!receiver) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + ASSERT_TRUE(sender.addPeer("dupDataReceiver", receiver->ch->localEndpoint())); + + // Hold every context so both DATA messages reach the reactor before the first scatter can finish. + std::vector heldExecContexts; + while (auto* ctx = receiver->exec->tryAcquire()) + { + heldExecContexts.push_back(ctx); + } + ASSERT_EQ(heldExecContexts.size(), receiver->exec->size()); + + constexpr std::uint64_t rid = 17; + sender.sendTo("dupDataReceiver", b::encodeWant(rid, {256}, sender.localEndpoint())); + + std::vector credits; + ASSERT_TRUE(waitGrant(sender, rid, std::chrono::seconds(5), credits)); + ASSERT_EQ(credits.size(), 1); + + void* dst = nullptr; + ASSERT_EQ(cudaMalloc(&dst, 256), cudaSuccess); + b::BounceScatterRun run{0, reinterpret_cast(dst), 0, 0, 256, 1}; + auto const data = b::encodeData(rid, /*chunkIdx=*/0, /*numChunks=*/1, credits.front().regionHandle, {run}); + sender.sendTo("dupDataReceiver", data); + sender.sendTo("dupDataReceiver", data); + // FIFO sentinel instead of a sleep: the reactor is a single FIFO consumer, so a GRANT for a + // fresh rid sent after both DATA messages proves it already consumed them (the second region is + // free under capRegions(2), so the sentinel WANT is grantable while rid=17 holds its region). + constexpr std::uint64_t sentinelRid = 18; + sender.sendTo("dupDataReceiver", b::encodeWant(sentinelRid, {256}, sender.localEndpoint())); + std::vector sentinelCredits; + ASSERT_TRUE(waitGrant(sender, sentinelRid, std::chrono::seconds(5), sentinelCredits)) + << "sentinel WANT never granted; DATA messages may not have reached the reactor"; + for (auto* ctx : heldExecContexts) + { + receiver->exec->release(ctx); + } + + EXPECT_EQ(countAcks(sender, std::chrono::seconds(2), rid), 1); + + receiver->tx->shutdown(); + EXPECT_EQ(cudaFree(dst), cudaSuccess); +} + +// A non-empty WANT has no retransmission path either (submit() sends it exactly once per fresh +// rid), so a replayed one for a live flow must be dropped: re-queueing would grant fresh regions +// over the still-held ones — the sender never writes the extras, so they leak — and the lease +// refresh would keep the flow off staleFlows() forever, defeating the reclaim that exists for +// exactly this state. +TEST(BounceTransportFailure, DuplicateWantIsDroppedWithoutRegrant) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/5000); + capRegions(c, c.maxInflightChunksPerRequest); + b::ZmqControlChannel sender("dupWantSender"); + auto receiver = bounce_test::makeNode("dupWantReceiver", c, 1024); // receiver posts no writes + if (!receiver) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + ASSERT_TRUE(sender.addPeer("dupWantReceiver", receiver->ch->localEndpoint())); + + constexpr std::uint64_t rid = 23; + sender.sendTo("dupWantReceiver", b::encodeWant(rid, {256}, sender.localEndpoint())); + std::vector credits; + ASSERT_TRUE(waitGrant(sender, rid, std::chrono::seconds(5), credits)); + ASSERT_EQ(credits.size(), 1u); + + // Replay the same WANT. The arena and the per-request cap both have room for a second region, + // so a re-queue WOULD be granted — the negative assertion below therefore fails pre-fix. + sender.sendTo("dupWantReceiver", b::encodeWant(rid, {256}, sender.localEndpoint())); + std::vector extra; + EXPECT_FALSE(waitGrant(sender, rid, std::chrono::milliseconds(500), extra)) << "duplicate WANT was re-granted"; + + receiver->tx->shutdown(); +} + +// WANT chunk sizes are untrusted peer input, and each of the two out-of-contract shapes is its own +// DoS pre-fix: a chunk of 0 can never be allocated, so the flow blocks forever, and once +// maybeActivateDrain() latches it as the drain flow, schedule() short-circuits before sweeping any +// other flow — every grant to every peer stops with no reclaim path (a pending-only flow holds no +// regions, so the lease sweep skips it). A chunk above the handshake-pinned maxChunkSizeBytes +// instead buddy-rounds up to the WHOLE usable arena (test config: 4097 -> 8192 = capacity), gets +// granted, and squats on it, starving all peers until the lease sweep. The receiver must reject +// both up front. The test drives each mode separately: park the zero-size flow, push the grant +// sequence past the drain bypass threshold with legitimate traffic, require fresh demand to still +// be granted; then send the oversized WANT and require the same. +TEST(BounceTransportFailure, MalformedWantChunkSizeIsRejectedNotWedged) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/5000); + capRegions(c, c.maxInflightChunksPerRequest); + b::ZmqControlChannel sender("badWantSender"); + auto receiver = bounce_test::makeNode("badWantReceiver", c, 1024); // receiver posts no writes + if (!receiver) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + ASSERT_TRUE(sender.addPeer("badWantReceiver", receiver->ch->localEndpoint())); + + sender.sendTo("badWantReceiver", b::encodeWant(/*rid=*/1, {0}, sender.localEndpoint())); + + // Drive enough legitimate grant/cancel cycles to push mGrantSequence past the drain bypass + // threshold (max(kMinimumBypassGrants, kBypassRounds * ringSize) = 8 here); pre-fix this + // latches rid=1 as the drain flow and the cycle after the eighth grant starves. + for (std::uint64_t rid = 100; rid < 112; ++rid) + { + sender.sendTo("badWantReceiver", b::encodeWant(rid, {256}, sender.localEndpoint())); + std::vector credits; + ASSERT_TRUE(waitGrant(sender, rid, std::chrono::seconds(5), credits)) << "legit WANT starved, rid=" << rid; + // Cancel to release the granted region and drop the flow from the ring. + sender.sendTo("badWantReceiver", b::encodeWant(rid, {}, sender.localEndpoint())); + } + + // The receiver must still grant fresh legitimate demand — pre-fix the drain latch wedges this. + sender.sendTo("badWantReceiver", b::encodeWant(/*rid=*/200, {256}, sender.localEndpoint())); + std::vector credits; + EXPECT_TRUE(waitGrant(sender, 200, std::chrono::seconds(5), credits)) << "receiver wedged after zero-size WANT"; + sender.sendTo("badWantReceiver", b::encodeWant(200, {}, sender.localEndpoint())); + + // Oversized mode: pre-fix this grant monopolizes the whole arena, so the follow-up legitimate + // WANT starves. + sender.sendTo("badWantReceiver", + b::encodeWant(/*rid=*/2, {static_cast(c.maxChunkSizeBytes) + 1}, sender.localEndpoint())); + sender.sendTo("badWantReceiver", b::encodeWant(/*rid=*/201, {256}, sender.localEndpoint())); + credits.clear(); + EXPECT_TRUE(waitGrant(sender, 201, std::chrono::seconds(5), credits)) << "arena monopolized by oversized WANT"; + + // The malformed flows themselves must never have been granted. The reactor is a single FIFO + // consumer, so the completed rid=200/201 grants above prove both malformed WANTs were already + // processed — these are not timing-window assertions. + std::vector badCredits; + EXPECT_FALSE(waitGrant(sender, 1, std::chrono::milliseconds(200), badCredits)) << "zero-size WANT was granted"; + EXPECT_FALSE(waitGrant(sender, 2, std::chrono::milliseconds(200), badCredits)) << "oversized WANT was granted"; + + receiver->tx->shutdown(); +} + +// The scatter worker sizes its plan by iterating the DATA run list; a hostile/corrupt run whose +// count is near 2^32 (bounceStride 0 keeps the span check happy) must be rejected up front, not +// counted piece by piece — pre-fix the worker (and the flow's region) is pinned for the whole +// count, which the rid=2 re-grant deadline below catches. +TEST(BounceTransportFailure, OversizedScatterRunListIsRejectedNotCounted) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/30000); + capRegions(c, /*regionCap=*/1); // arena holds exactly ONE region + b::ZmqControlChannel sender("bigRunSender"); + auto receiver = bounce_test::makeNode("bigRunReceiver", c, 1024); // receiver posts no writes + if (!receiver) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + ASSERT_TRUE(sender.addPeer("bigRunReceiver", receiver->ch->localEndpoint())); + + sender.sendTo("bigRunReceiver", b::encodeWant(/*rid=*/1, {256}, sender.localEndpoint())); + std::vector credits; + ASSERT_TRUE(waitGrant(sender, 1, std::chrono::seconds(5), credits)); + ASSERT_EQ(credits.size(), 1u); + + void* dst = nullptr; + ASSERT_EQ(cudaMalloc(&dst, 256), cudaSuccess); + // count=2^32-1 with bounceStride=0: every piece reads the same 256 in-bounds bytes, so the + // span check alone admits it; only the piece-count guard can reject it cheaply. + b::BounceScatterRun run{/*bounceOffset=*/0, reinterpret_cast(dst), /*dstStride=*/0, + /*bounceStride=*/0, /*pieceSize=*/256, /*count=*/0xFFFFFFFFu}; + sender.sendTo("bigRunReceiver", + b::encodeData(/*rid=*/1, /*chunkIdx=*/0, /*numChunks=*/1, credits.front().regionHandle, {run})); + + // The rejected scatter must not ACK... + EXPECT_EQ(countAcks(sender, std::chrono::seconds(2), /*rid=*/1), 0); + // ...and must release the single region promptly: rid=2 can only be granted once rid=1's + // region is freed, which a worker stuck counting 2^32 pieces cannot do within the deadline. + sender.sendTo("bigRunReceiver", b::encodeWant(/*rid=*/2, {256}, sender.localEndpoint())); + std::vector credits2; + EXPECT_TRUE(waitGrant(sender, 2, std::chrono::seconds(5), credits2)) << "region pinned by oversized run list"; + + receiver->tx->shutdown(); + EXPECT_EQ(cudaFree(dst), cudaSuccess); +} + +// A sender that takes a GRANT and then dies emits neither DATA nor a cancel, so nothing +// event-driven can ever reclaim its region — the receiver's grant lease must. After +// receiverFlowTimeoutMs of silence the flow is reclaimed and its region quarantined for +// quarantineMs, after which it must serve the next waiting flow; a late DATA for the expired +// grant must be dropped (no scatter, no ACK). +TEST(BounceTransportFailure, GrantLeaseExpiryReclaimsSilentSendersRegion) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/30000); // sender-side request timeout: irrelevant here + c.receiverFlowTimeoutMs = 400; + c.quarantineMs = 200; + capRegions(c, /*regionCap=*/1); // arena holds exactly ONE region + b::ZmqControlChannel sender("glSender"); + auto receiver = bounce_test::makeNode("glReceiver", c, 1024); // receiver posts no writes + if (!receiver) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + ASSERT_TRUE(sender.addPeer("glReceiver", receiver->ch->localEndpoint())); + + // rid=1 takes the single region... and goes silent (models a dead/unreachable sender). + sender.sendTo("glReceiver", b::encodeWant(/*rid=*/1, {256}, sender.localEndpoint())); + std::vector credit1; + ASSERT_TRUE(waitGrant(sender, 1, std::chrono::seconds(5), credit1)); + ASSERT_EQ(credit1.size(), 1u); + + // rid=2 wants the region too. Before the lease (400ms) + quarantine (200ms) can possibly have + // elapsed, it must NOT be granted (premature reclaim would race rid=1's hypothetical write). + sender.sendTo("glReceiver", b::encodeWant(/*rid=*/2, {256}, sender.localEndpoint())); + std::vector credit2; + EXPECT_FALSE(waitGrant(sender, 2, std::chrono::milliseconds(200), credit2)) + << "region re-granted before lease expiry"; + // ...but once the silent flow's lease expires and the quarantine passes, it must be. + ASSERT_TRUE(waitGrant(sender, 2, std::chrono::seconds(10), credit2)) << "silent sender's region never reclaimed"; + ASSERT_EQ(credit2.size(), 1u); + EXPECT_EQ(credit2.front().regionHandle, credit1.front().regionHandle); // the one region, recycled + + // Late DATA for the EXPIRED grant must be dropped (flow reclaimed -> heldByFlow false): exactly + // one ACK total, and it belongs to rid=2's legitimate DATA. + void* dst = nullptr; + ASSERT_EQ(cudaMalloc(&dst, 256), cudaSuccess); + b::BounceScatterRun run{0, reinterpret_cast(dst), 0, 0, 256, 1}; + sender.sendTo("glReceiver", b::encodeData(/*rid=*/1, 0, 1, credit1.front().regionHandle, {run})); + sender.sendTo("glReceiver", b::encodeData(/*rid=*/2, 0, 1, credit2.front().regionHandle, {run})); + int ackRid1 = 0; + int ackRid2 = 0; + pumpChannel(sender, std::chrono::seconds(2), + [&](b::BounceMsgHeader const& h, std::string const&) + { + if (static_cast(h.msgType) == b::BounceMsgType::kACK) + { + ackRid1 += h.requestId == 1 ? 1 : 0; + ackRid2 += h.requestId == 2 ? 1 : 0; + } + return false; + }); + EXPECT_EQ(ackRid1, 0) << "late DATA for an expired grant was scattered/ACKed"; + EXPECT_EQ(ackRid2, 1); + + receiver->tx->shutdown(); + EXPECT_EQ(cudaFree(dst), cudaSuccess); +} + +// GRANT and ACK messages carry peer-owned capabilities. Only the intended peer may grant, and an +// ACK is valid only for the matching region after that chunk has reached Sent. +TEST(BounceTransportFailure, RejectsWrongPeerGrantAndInvalidAcks) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/5000); + capRegions(c, c.maxInflightChunksPerRequest); + b::ZmqControlChannel legitimate("validateLegitimate"); + b::ZmqControlChannel attacker("validateAttacker"); + auto ctl = std::make_shared(); // defaults: writes complete instantly + auto sender = makeFakeNode("validateSender", c, ctl); + if (!sender) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + ASSERT_TRUE(sender->tx->addPeer("validateLegitimate", legitimate.localEndpoint())); + ASSERT_TRUE(legitimate.addPeer("validateSender", sender->ch->localEndpoint())); + ASSERT_TRUE(attacker.addPeer("validateSender", sender->ch->localEndpoint())); + + auto bufs = bounce_test::makeXferBufs(/*nDescs=*/1, /*descBytes=*/256, /*seed=*/4); + auto fut = sender->tx->submit(bufs.srcDescs, bufs.dstDescs, "validateLegitimate"); + b::BounceCreditEntry credit{/*addr=*/0, /*len=*/256, /*devId=*/0, /*regionHandle=*/77}; + + attacker.sendTo("validateSender", b::encodeGrant(/*requestId=*/1, {credit})); + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + EXPECT_EQ(ctl->postCount.load(std::memory_order_acquire), 0u); + + legitimate.sendTo("validateSender", b::encodeAck(/*requestId=*/1, /*chunkIdx=*/0, credit.regionHandle)); + EXPECT_EQ(fut.wait_for(std::chrono::milliseconds(200)), std::future_status::timeout); + legitimate.sendTo("validateSender", b::encodeGrant(/*requestId=*/1, {credit})); + + // Seeing DATA guarantees the matching sender chunk has transitioned to Sent. + ASSERT_TRUE(pumpChannel(legitimate, std::chrono::seconds(5), + [](b::BounceMsgHeader const& h, std::string const&) + { return static_cast(h.msgType) == b::BounceMsgType::kDATA; })); + EXPECT_EQ(ctl->postCount.load(std::memory_order_acquire), 1u); + + attacker.sendTo("validateSender", b::encodeAck(/*requestId=*/1, /*chunkIdx=*/0, credit.regionHandle)); + EXPECT_EQ(fut.wait_for(std::chrono::milliseconds(200)), std::future_status::timeout); + legitimate.sendTo("validateSender", b::encodeAck(/*requestId=*/1, /*chunkIdx=*/0, credit.regionHandle + 1)); + EXPECT_EQ(fut.wait_for(std::chrono::milliseconds(200)), std::future_status::timeout); + legitimate.sendTo("validateSender", b::encodeAck(/*requestId=*/1, /*chunkIdx=*/0, credit.regionHandle)); + legitimate.sendTo("validateSender", b::encodeAck(/*requestId=*/1, /*chunkIdx=*/0, credit.regionHandle)); + ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready); + EXPECT_EQ(fut.get().state, kvc::TransferState::kSUCCESS); + + sender->tx->shutdown(); + bounce_test::freeXferBufs(bufs); +} + +// shutdown() with an in-flight (stuck) request must resolve its future FAILURE, never leave wait() hanging. +TEST(BounceTransportFailure, ShutdownFailsInflight) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/0); // timeout disabled -> only shutdown can resolve it + auto t = bounce_test::makeNode("sdSolo", c, 1024); + if (!t) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + + auto bufs = bounce_test::makeXferBufs(8, 256, /*seed=*/1); + auto fut = t->tx->submit(bufs.srcDescs, bufs.dstDescs, "nobody"); // stuck (no grant, no timeout) + EXPECT_EQ(fut.wait_for(std::chrono::milliseconds(200)), std::future_status::timeout); // still pending + t->tx->shutdown(); + ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "shutdown left request hanging"; + EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); + EXPECT_EQ(fut.get().reason, b::BounceFailReason::kShutdown); + + bounce_test::freeXferBufs(bufs); +} + +// Shutdown uses the status object's bounded release instead of polling forever. If the backend +// cannot abort, the future still fails promptly; the release is attempted once in failAll and once +// in its retry pass (the production NixlTransferStatus destructor makes a final attempt). +TEST(BounceTransportFailure, ShutdownReleaseFailureDoesNotHangOrLoseHandle) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/0); + capRegions(c, c.maxInflightChunksPerRequest); + b::ZmqControlChannel peer("shutdownReleasePeer"); + auto ctl = std::make_shared(); + ctl->allowTerminal.store(false, std::memory_order_release); + ctl->releaseSucceeds.store(false, std::memory_order_release); + auto sender = makeFakeNode("shutdownReleaseSender", c, ctl); + if (!sender) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + ASSERT_TRUE(sender->tx->addPeer("shutdownReleasePeer", peer.localEndpoint())); + ASSERT_TRUE(peer.addPeer("shutdownReleaseSender", sender->ch->localEndpoint())); + + auto bufs = bounce_test::makeXferBufs(/*nDescs=*/1, /*descBytes=*/256, /*seed=*/5); + auto fut = sender->tx->submit(bufs.srcDescs, bufs.dstDescs, "shutdownReleasePeer"); + b::BounceCreditEntry credit{/*addr=*/0, /*len=*/256, /*devId=*/0, /*regionHandle=*/91}; + peer.sendTo("shutdownReleaseSender", b::encodeGrant(/*requestId=*/1, {credit})); + + auto const postDeadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (std::chrono::steady_clock::now() < postDeadline && ctl->postCount.load(std::memory_order_acquire) == 0) + { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_EQ(ctl->postCount.load(std::memory_order_acquire), 1u); + + auto const start = std::chrono::steady_clock::now(); + sender->tx->shutdown(); + EXPECT_LT(std::chrono::steady_clock::now() - start, std::chrono::seconds(5)); + ASSERT_EQ(fut.wait_for(std::chrono::seconds(1)), std::future_status::ready); + EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); + EXPECT_EQ(fut.get().reason, b::BounceFailReason::kShutdown); + // failAll attempts the failed release exactly twice (initial + bounded retry), never spins. + EXPECT_EQ(ctl->releaseCount.load(std::memory_order_acquire), 2u); + bounce_test::freeXferBufs(bufs); +} + +// forgetPeer() (the invalidateRemoteAgent path) must fail any in-flight request to the gone peer, +// even with the request timeout disabled, modeling a peer that drops out mid-transfer. +TEST(BounceTransportFailure, ForgetPeerFailsInflightRequest) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/0); // timeout disabled -> only forgetPeer can resolve it + auto t = bounce_test::makeNode("fpSolo", c, 1024); + if (!t) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + + auto bufs = bounce_test::makeXferBufs(8, 256, /*seed=*/1); + auto fut = t->tx->submit(bufs.srcDescs, bufs.dstDescs, "gonePeer"); // stuck: never granted + EXPECT_EQ(fut.wait_for(std::chrono::milliseconds(200)), std::future_status::timeout); // still pending + t->tx->forgetPeer("gonePeer"); + ASSERT_EQ(fut.wait_for(std::chrono::seconds(5)), std::future_status::ready) << "forgetPeer left request hanging"; + EXPECT_EQ(fut.get().state, kvc::TransferState::kFAILURE); + EXPECT_EQ(fut.get().reason, b::BounceFailReason::kPeerDropped); + + // forgetPeer for an unrelated/unknown peer must be a harmless no-op (no crash). + t->tx->forgetPeer("someoneElse"); + t->tx->shutdown(); + bounce_test::freeXferBufs(bufs); +} + +// forgetPeer() while a transfer is in flight over REAL RDMA must not hang, leak, or corrupt. We +// submit then immediately forgetPeer the target; the request must RESOLVE (SUCCESS if it beat the +// queued reclaim, else FAILURE) — never hang. Then several FRESH transfers to the same peer must +// still complete byte-exact, proving forgetPeer's reclaim returned the regions to the arena and +// left the reactor healthy (a small arena would quickly expose a leaked allocation). +TEST(BounceTransportFailure, ForgetPeerInFlightRecovers) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/5000); + std::size_t const maxDescs = std::max(1024ULL, c.maxChunkSizeBytes / 256ULL); + auto A = bounce_test::makeNode("fpA", c, maxDescs); + auto B = bounce_test::makeNode("fpB", c, maxDescs); + if (!A || !B) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + bounce_test::wirePair(*A, *B); + + auto bufs = bounce_test::makeXferBufs(/*nDescs=*/24, /*descBytes=*/600, /*seed=*/1); + auto fut = A->tx->submit(bufs.srcDescs, bufs.dstDescs, "fpB"); + A->tx->forgetPeer("fpB"); // drop the peer (queued; applied on A's IO thread) + ASSERT_EQ(fut.wait_for(std::chrono::seconds(10)), std::future_status::ready) << "request hung after forgetPeer"; + auto const res = fut.get(); + EXPECT_TRUE(res.state == kvc::TransferState::kSUCCESS || res.state == kvc::TransferState::kFAILURE) + << "unexpected state " << static_cast(res.state); + if (res.state == kvc::TransferState::kFAILURE) + { + EXPECT_EQ(res.reason, b::BounceFailReason::kPeerDropped); + } + bounce_test::freeXferBufs(bufs); + + // Recovery + no-leak: forgetPeer is a one-shot event; the scheduler/request reclaim is drained by + // the time fut resolved, so these new flows aren't reclaimed. forgetPeer ALSO drops the + // control-channel DEALER to fpB synchronously (on this thread), so re-establish it with addPeer + // before recovering — deterministic because forgetPeer's removePeer happens-before this addPeer + // (no async removePeer can race/erase the freshly re-added dealer). NIXL metadata persists, so + // only the dealer is re-added (no loadRemoteAgent / full re-wire). + A->tx->addPeer("fpB", B->ch->localEndpoint()); + for (int k = 0; k < 5; ++k) + { + auto rb + = bounce_test::makeXferBufs(/*nDescs=*/20, /*descBytes=*/600, /*seed=*/static_cast(50 + k)); + auto rf = A->tx->submit(rb.srcDescs, rb.dstDescs, "fpB"); + ASSERT_EQ(rf.wait_for(std::chrono::seconds(30)), std::future_status::ready) + << "post-forget transfer hung k=" << k; + EXPECT_EQ(rf.get().state, kvc::TransferState::kSUCCESS) << "post-forget transfer failed k=" << k; + EXPECT_TRUE(bounce_test::verifyXferBufs(rb)) << "post-forget byte mismatch k=" << k; + bounce_test::freeXferBufs(rb); + } + + A->tx->shutdown(); + B->tx->shutdown(); +} + +// One sender -> TWO receivers sharing ONE small outgoing arena, deliberately oversubscribed: B and C +// can each grant up to their per-request limit, so aggregate grants exceed sender staging capacity. +// The IO thread must NOT block in acquireLocal (it would deadlock: posted chunks could never send +// DATA -> no ACK -> no region freed). With non-blocking acquire + parked credits, both transfers +// must still complete. (Models KV_TRANSFER_NUM_THREADS>1 fanning out to many peers.) +TEST(BounceTransportFailure, MultiPeerSharedOutgoingPoolNoDeadlock) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/15000); + c.maxChunkSizeBytes = 4096; + c.maxInflightChunksPerRequest = 2; + // Arena holds only 2 max-size regions on every node: the sender's gather arena (2) is + // oversubscribed by B+C's aggregate grants (up to 4). + c.arenaAllocationGranularityBytes = c.maxChunkSizeBytes; + c.arenaSizeBytes = 2ULL * c.maxChunkSizeBytes; // exactly 2 regions (power of two) + std::size_t const maxDescs = std::max(1024ULL, c.maxChunkSizeBytes / 256ULL); + + auto A = bounce_test::makeNode("mpA", c, maxDescs); + auto B = bounce_test::makeNode("mpB", c, maxDescs); + auto C = bounce_test::makeNode("mpC", c, maxDescs); + if (!A || !B || !C) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + bounce_test::wirePair(*A, *B); + bounce_test::wirePair(*A, *C); + + // 48 descs * 512B = 24KiB -> ~6 chunks of 4KiB each, so each transfer needs many credits. + auto toB = bounce_test::makeXferBufs(48, 512, /*seed=*/2); + auto toC = bounce_test::makeXferBufs(48, 512, /*seed=*/3); + std::atomic ok{0}; + std::vector threads; + threads.emplace_back( + [&] + { + auto fut = A->tx->submit(toB.srcDescs, toB.dstDescs, B->name); + if (fut.wait_for(std::chrono::seconds(40)) == std::future_status::ready + && fut.get().state == kvc::TransferState::kSUCCESS) + ok.fetch_add(1); + }); + threads.emplace_back( + [&] + { + auto fut = A->tx->submit(toC.srcDescs, toC.dstDescs, C->name); + if (fut.wait_for(std::chrono::seconds(40)) == std::future_status::ready + && fut.get().state == kvc::TransferState::kSUCCESS) + ok.fetch_add(1); + }); + for (auto& th : threads) + th.join(); + EXPECT_EQ(ok.load(), 2) << "multi-peer oversubscribed outgoing arena did not both complete (deadlock?)"; + EXPECT_TRUE(bounce_test::verifyXferBufs(toB)) << "byte mismatch to B"; + EXPECT_TRUE(bounce_test::verifyXferBufs(toC)) << "byte mismatch to C"; + + A->tx->shutdown(); + B->tx->shutdown(); + C->tx->shutdown(); + bounce_test::freeXferBufs(toB); + bounce_test::freeXferBufs(toC); +} + +// Many threads submit concurrently to the same transport pair -> all complete (thread-safety). +TEST(BounceTransportFailure, ConcurrentMultiThreadedSubmit) +{ + if (!bounce_test::hasCuda()) + GTEST_SKIP() << "no CUDA device"; + auto c = cfg(/*timeoutMs=*/10000); + c.maxInflightChunksPerRequest = 4; + std::size_t const maxDescs = std::max(1024ULL, c.maxChunkSizeBytes / 256ULL); + + auto A = bounce_test::makeNode("cmtA", c, maxDescs); + auto B = bounce_test::makeNode("cmtB", c, maxDescs); + if (!A || !B) + GTEST_SKIP() << "NIXL agent/backend unavailable"; + bounce_test::wirePair(*A, *B); + + constexpr int kThreads = 8; + std::vector bufs(kThreads); + for (int i = 0; i < kThreads; ++i) + { + bufs[i] = bounce_test::makeXferBufs(6, 300, /*seed=*/static_cast(10 + i)); + } + std::atomic ok{0}; + std::vector threads; + for (int i = 0; i < kThreads; ++i) + { + threads.emplace_back( + [&, i] + { + auto fut = A->tx->submit(bufs[i].srcDescs, bufs[i].dstDescs, B->name); + if (fut.wait_for(std::chrono::seconds(30)) == std::future_status::ready + && fut.get().state == kvc::TransferState::kSUCCESS) + { + ok.fetch_add(1); + } + }); + } + for (auto& th : threads) + { + th.join(); + } + EXPECT_EQ(ok.load(), kThreads) << "not all concurrent submits succeeded"; + for (auto& bb : bufs) + { + EXPECT_TRUE(bounce_test::verifyXferBufs(bb)); + } + + A->tx->shutdown(); + B->tx->shutdown(); + for (auto& bb : bufs) + { + bounce_test::freeXferBufs(bb); + } +} diff --git a/cpp/tests/unit_tests/executor/bounce/bounceTransportTest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceTransportTest.cpp new file mode 100644 index 000000000000..19099c9ae1bb --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/bounceTransportTest.cpp @@ -0,0 +1,226 @@ +/* + * 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. + */ + +// End-to-end transport tests over REAL NIXL RDMA (two in-process agents, real UCX backend): drive +// the full bounce pipeline (gather -> RDMA write -> scatter + credit recycling) and verify every +// byte arrives. Sizing is chosen so the chunk count exceeds the per-request in-flight limit, +// forcing credit recycling. Skips if no CUDA device or the NIXL backend cannot initialize. + +#include "bounceTestNixlNode.h" + +#include + +#include +#include + +namespace kvc = tensorrt_llm::executor::kv_cache; +namespace b = tensorrt_llm::executor::kv_cache::bounce; + +namespace +{ +// One end-to-end transfer of `nDescs` x `descBytes` through the bounce pipeline between two real +// NIXL nodes: a sender built from `senderCfg` and a receiver built from `receiverCfg` (asymmetric +// configs let a test pin clamping/backpressure to one side). `tag` gives the two agents unique names. +void runTransfer(std::string const& tag, std::uint32_t nDescs, std::uint32_t descBytes, + b::BounceConfig const& senderCfg, b::BounceConfig const& receiverCfg, std::uint32_t seed) +{ + if (!bounce_test::hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; + } + std::size_t const maxDescs + = std::max(1024ULL, std::max(senderCfg.maxChunkSizeBytes, receiverCfg.maxChunkSizeBytes) / 256ULL); + + auto A = bounce_test::makeNode(tag + "A", senderCfg, maxDescs); + auto B = bounce_test::makeNode(tag + "B", receiverCfg, maxDescs); + if (!A || !B) + { + GTEST_SKIP() << "NIXL agent/backend unavailable"; + } + bounce_test::wirePair(*A, *B); + + auto bufs = bounce_test::makeXferBufs(nDescs, descBytes, seed); + auto fut = A->tx->submit(bufs.srcDescs, bufs.dstDescs, B->name); + ASSERT_EQ(fut.wait_for(std::chrono::seconds(30)), std::future_status::ready) << "transfer hung"; + EXPECT_EQ(fut.get().state, kvc::TransferState::kSUCCESS); + EXPECT_TRUE(bounce_test::verifyXferBufs(bufs)) << "byte mismatch"; + + A->tx->shutdown(); + B->tx->shutdown(); + bounce_test::freeXferBufs(bufs); +} + +// Thin wrapper: the same config on both ends, with maxChunkSizeBytes/maxInflightChunksPerRequest +// chosen so the chunk count exceeds the in-flight limit (forcing credit recycling). +void runTransfer(std::string const& tag, std::uint32_t nDescs, std::uint32_t descBytes, std::size_t maxChunkSizeBytes, + std::uint32_t maxInflightChunksPerRequest) +{ + b::BounceConfig cfg; + cfg.maxChunkSizeBytes = maxChunkSizeBytes; + cfg.maxInflightChunksPerRequest = maxInflightChunksPerRequest; + cfg.scatterWorkerCount = 2; + cfg.arenaAllocationGranularityBytes = 256; // matches the 256-aligned desc layout + // Arena large enough for the configured in-flight chunk count in both roles, with headroom for + // buddy rounding. A power-of-two size lets BuddyAllocator use the whole allocation. + std::size_t arenaSizeBytes = 1ULL << 20; + while (arenaSizeBytes < static_cast(maxInflightChunksPerRequest) * maxChunkSizeBytes * 4ULL) + { + arenaSizeBytes <<= 1; + } + cfg.arenaSizeBytes = arenaSizeBytes; + runTransfer(tag, nDescs, descBytes, cfg, cfg, /*seed=*/1); +} +} // namespace + +TEST(BounceTransport, SmallTransferFitsInflightLimit) +{ + // Four descriptors fit within the configured in-flight chunk limit. + runTransfer("btSmall", /*nDescs=*/4, /*descBytes=*/512, /*maxChunkSizeBytes=*/8192, /*maxInflightChunks=*/4); +} + +TEST(BounceTransport, LargeTransferRecyclesCredits) +{ + // Forty descriptors produce more chunks than the limit of two, forcing credit recycling. + runTransfer("btLarge", /*nDescs=*/40, /*descBytes=*/700, /*maxChunkSizeBytes=*/4096, /*maxInflightChunks=*/2); +} + +TEST(BounceTransport, ManySmallDescs) +{ + // Closer to the real KV pattern: many tiny descs. + runTransfer("btMany", /*nDescs=*/500, /*descBytes=*/256, /*maxChunkSizeBytes=*/16384, /*maxInflightChunks=*/4); +} + +TEST(BounceTransport, ConcurrentRequestsToSameReceiver) +{ + // Two independent transfers (distinct rids) over ONE transport pair, both submitted before + // either future is waited on, so the flows are truly concurrent on the same A->B connection. + if (!bounce_test::hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; + } + b::BounceConfig cfg; + cfg.maxChunkSizeBytes = 8192; + cfg.maxInflightChunksPerRequest = 3; + cfg.scatterWorkerCount = 2; + cfg.arenaAllocationGranularityBytes = 256; + cfg.arenaSizeBytes = 1ULL << 20; + std::size_t const maxDescs = std::max(1024ULL, cfg.maxChunkSizeBytes / 256ULL); + + auto A = bounce_test::makeNode("btConcA", cfg, maxDescs); + auto B = bounce_test::makeNode("btConcB", cfg, maxDescs); + if (!A || !B) + { + GTEST_SKIP() << "NIXL agent/backend unavailable"; + } + bounce_test::wirePair(*A, *B); + + // Seed-distinct payloads so any cross-talk between the two in-flight flows fails verification. + auto bufs1 = bounce_test::makeXferBufs(/*nDescs=*/8, /*descBytes=*/1024, /*seed=*/1); + auto bufs2 = bounce_test::makeXferBufs(/*nDescs=*/8, /*descBytes=*/1024, /*seed=*/2); + auto fut1 = A->tx->submit(bufs1.srcDescs, bufs1.dstDescs, B->name); + auto fut2 = A->tx->submit(bufs2.srcDescs, bufs2.dstDescs, B->name); + ASSERT_EQ(fut1.wait_for(std::chrono::seconds(30)), std::future_status::ready) << "first transfer hung"; + ASSERT_EQ(fut2.wait_for(std::chrono::seconds(30)), std::future_status::ready) << "second transfer hung"; + EXPECT_EQ(fut1.get().state, kvc::TransferState::kSUCCESS); + EXPECT_EQ(fut2.get().state, kvc::TransferState::kSUCCESS); + EXPECT_TRUE(bounce_test::verifyXferBufs(bufs1)) << "byte mismatch (first flow)"; + EXPECT_TRUE(bounce_test::verifyXferBufs(bufs2)) << "byte mismatch (second flow)"; + + A->tx->shutdown(); + B->tx->shutdown(); + bounce_test::freeXferBufs(bufs1); + bounce_test::freeXferBufs(bufs2); +} + +TEST(BounceTransport, ReplacementHandshakeClearsStaleCompatibility) +{ + if (!bounce_test::hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; + } + b::BounceConfig cfg; + cfg.arenaSizeBytes = 1ULL << 20; + cfg.maxChunkSizeBytes = 4096; + cfg.arenaAllocationGranularityBytes = 256; + auto node = bounce_test::makeNode("btHandshake", cfg, /*maxDescs=*/1024); + if (!node) + { + GTEST_SKIP() << "NIXL agent/backend unavailable"; + } + + auto compatible = node->tx->localHandshakeBlob(); + ASSERT_TRUE(node->tx->registerPeerHandshake("peer", compatible)); + EXPECT_TRUE(node->tx->hasPeerHandshake("peer")); + + b::BounceHandshake incompatible; + ASSERT_TRUE(b::decodeHandshake(compatible, incompatible)); + incompatible.maxChunkSizeBytes += 1; + EXPECT_FALSE(node->tx->registerPeerHandshake("peer", b::encodeHandshake(incompatible))); + EXPECT_FALSE(node->tx->hasPeerHandshake("peer")); + + ASSERT_TRUE(node->tx->registerPeerHandshake("peer", compatible)); + EXPECT_FALSE(node->tx->registerPeerHandshake("peer", {})); + EXPECT_FALSE(node->tx->hasPeerHandshake("peer")); + + b::BounceHandshake invalidEndpoint; + ASSERT_TRUE(b::decodeHandshake(compatible, invalidEndpoint)); + invalidEndpoint.endpoint.clear(); + EXPECT_ANY_THROW(node->tx->registerPeerHandshake("peer", b::encodeHandshake(invalidEndpoint))); + EXPECT_FALSE(node->tx->hasPeerHandshake("peer")); + + invalidEndpoint.endpoint = "not-a-zmq-endpoint"; + EXPECT_ANY_THROW(node->tx->registerPeerHandshake("peer", b::encodeHandshake(invalidEndpoint))); + EXPECT_FALSE(node->tx->hasPeerHandshake("peer")); + node->tx->shutdown(); +} + +// Regression: maxChunkSizeBytes can be no larger than the buddy allocator's usable capacity, which +// may be smaller than arenaSizeBytes. A 96 KiB arena with 256-byte granularity has only one 64 KiB +// top-level buddy block. The effective chunk cap must therefore become 64 KiB. +TEST(BounceTransport, MaxChunkSizeBytesClampedToUsableArena) +{ + b::BounceConfig cfg; + cfg.arenaSizeBytes = 96 * 1024; // buddy usable rounds DOWN to 64KiB (256<<8) + cfg.maxChunkSizeBytes = 96 * 1024; // exceeds the 64KiB usable -> must be clamped to 64KiB + cfg.arenaAllocationGranularityBytes = 256; + cfg.maxInflightChunksPerRequest = 2; + cfg.scatterWorkerCount = 2; + // 4 x 20KiB = 80KiB total > 64KiB usable. Unclamped, the planner packs all 80KiB into ONE chunk + // (<= 96KiB cap) that can never be allocated (rounds to 128KiB > 64KiB usable) -> hang. Clamped to + // 64KiB, it splits into chunks that each fit a drained arena and recycle through. + runTransfer("btClamp", /*nDescs=*/4, /*descBytes=*/20480, cfg, cfg, /*seed=*/9); +} + +// Sender-side arena backpressure: the receiver's arena and in-flight limit are generous +// credit up front) but the SENDER's arena only fits a few concurrent gather regions, so most +// credits get parked in pendingCredits and drain via drainPendingPosts as ACKs free regions. The +// transfer must still complete byte-exact (parked != dropped). This is also the path the +// `arenaStarved` NVTX span instruments. +TEST(BounceTransport, SenderArenaBackpressureParksCredits) +{ + b::BounceConfig small; // sender: 64KiB usable -> at most 4 in-flight 16KiB gather regions + small.maxChunkSizeBytes = 16 * 1024; + small.arenaAllocationGranularityBytes = 256; + small.maxInflightChunksPerRequest = 8; + small.scatterWorkerCount = 2; + small.arenaSizeBytes = 64 * 1024; + b::BounceConfig big = small; // receiver: room to grant all eight allowed credits at once + big.arenaSizeBytes = 1ULL << 20; + // 32 x 4KiB = 128KiB in ~8 chunks of 16KiB: double the sender's usable arena, so at least half + // the granted credits must park and retry. + runTransfer("btPark", /*nDescs=*/32, /*descBytes=*/4096, small, big, /*seed=*/11); +} diff --git a/cpp/tests/unit_tests/executor/bounce/buddyAllocatorTest.cpp b/cpp/tests/unit_tests/executor/bounce/buddyAllocatorTest.cpp new file mode 100644 index 000000000000..d9b66097a308 --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/buddyAllocatorTest.cpp @@ -0,0 +1,323 @@ +/* + * 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. + */ + +// Pure-logic tests for the bounce v2 buddy allocator: right-sizing, no overlap, coalescing on free, +// exhaustion/backpressure, and the "many small + one larger-than-buffer (recycled)" regimes. + +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BuddyAllocator.h" + +#include + +#include +#include +#include +#include +#include + +namespace b = tensorrt_llm::executor::kv_cache::bounce; + +namespace +{ +// Assert no two live [offset, offset+size) ranges overlap. +void checkNoOverlap(std::vector> const& live) +{ + std::vector> s(live); + std::sort(s.begin(), s.end()); + for (std::size_t i = 1; i < s.size(); ++i) + { + EXPECT_LE(s[i - 1].first + s[i - 1].second, s[i].first) + << "overlap: [" << s[i - 1].first << ",+" << s[i - 1].second << ") vs [" << s[i].first << ",...)"; + } +} +} // namespace + +TEST(BuddyAllocator, RoundsUpToPow2MultipleOfMinBlock) +{ + b::BuddyAllocator a(/*capacity=*/1024, /*minBlock=*/64); + EXPECT_EQ(a.capacity(), 1024u); + EXPECT_EQ(a.minBlock(), 64u); + EXPECT_EQ(a.maxAllocBytes(), 1024u); // whole arena free + // 1 byte -> a minBlock (64). 65 bytes -> 128. 64 -> 64. + auto o1 = a.alloc(1); + auto o2 = a.alloc(65); + auto o3 = a.alloc(64); + ASSERT_TRUE(o1 && o2 && o3); + EXPECT_EQ(*o1 % 64, 0u); + // freeBytes accounts for the rounded-up block sizes (64 + 128 + 64 used). + EXPECT_EQ(a.freeBytes(), 1024u - (64u + 128u + 64u)); + EXPECT_EQ(a.liveBlocks(), 3u); +} + +TEST(BuddyAllocator, ManySmallRegionsDistinctNoOverlap) +{ + // capacity 16 * minBlock -> up to 16 concurrent minBlock allocations (high concurrency of small + // requests, each right-sized to one block). + b::BuddyAllocator a(/*capacity=*/16 * 256, /*minBlock=*/256); + std::vector> live; + std::set seen; + for (int i = 0; i < 16; ++i) + { + auto o = a.alloc(200); // < 256 -> one minBlock + ASSERT_TRUE(o.has_value()) << "small alloc " << i << " failed"; + EXPECT_EQ(seen.count(*o), 0u) << "offset reused while live"; + seen.insert(*o); + live.emplace_back(*o, 256); + } + EXPECT_EQ(a.liveBlocks(), 16u); + EXPECT_EQ(a.freeBytes(), 0u); + EXPECT_FALSE(a.alloc(1).has_value()); // full -> backpressure (nullopt, no overlap/overcommit) + checkNoOverlap(live); + + for (auto& [off, sz] : live) + { + a.free(off); + } + EXPECT_EQ(a.freeBytes(), a.capacity()); // all coalesced back + EXPECT_EQ(a.maxAllocBytes(), a.capacity()); + EXPECT_EQ(a.liveBlocks(), 0u); +} + +TEST(BuddyAllocator, FreeCoalescesBuddies) +{ + b::BuddyAllocator a(/*capacity=*/1024, /*minBlock=*/256); // 4 minBlocks + auto o0 = a.alloc(256); + auto o1 = a.alloc(256); + auto o2 = a.alloc(256); + auto o3 = a.alloc(256); + ASSERT_TRUE(o0 && o1 && o2 && o3); + EXPECT_EQ(a.maxAllocBytes(), 0u); // fully split into minBlocks + // Free all -> must coalesce all the way back to one 1024 block (maxAlloc == capacity). + a.free(*o0); + a.free(*o1); + a.free(*o2); + a.free(*o3); + EXPECT_EQ(a.maxAllocBytes(), 1024u) << "buddies did not coalesce to the top order"; + // And a full-size alloc now succeeds. + auto big = a.alloc(1024); + ASSERT_TRUE(big.has_value()); + EXPECT_EQ(*big, 0u); +} + +TEST(BuddyAllocator, TooLargeRejected) +{ + b::BuddyAllocator a(/*capacity=*/1024, /*minBlock=*/256); + EXPECT_FALSE(a.alloc(2048).has_value()); // larger than the whole arena + EXPECT_TRUE(a.alloc(1024).has_value()); +} + +TEST(BuddyAllocator, ZeroAndOverflowSizeRejectedNoHang) +{ + b::BuddyAllocator a(/*capacity=*/1024, /*minBlock=*/256); + EXPECT_FALSE(a.alloc(0).has_value()); + // A near-SIZE_MAX request must be rejected WITHOUT hanging — roundUpPow2's doubling would + // otherwise overflow to 0 and spin forever. The early `bytes > usable` bound prevents it. + EXPECT_FALSE(a.alloc(std::numeric_limits::max()).has_value()); + EXPECT_FALSE(a.alloc(std::numeric_limits::max() - 1).has_value()); + EXPECT_EQ(a.freeBytes(), a.capacity()); // nothing consumed by the rejected requests + EXPECT_TRUE(a.alloc(256).has_value()); // still usable afterwards +} + +TEST(BuddyAllocator, LargeRecycledStreamLargerThanArena) +{ + // Models a single transfer whose total bytes exceed the arena: bounded chunks stream through + // a small arena with recycling (alloc chunk -> "send" -> free -> alloc next). The arena only + // ever holds a few chunks; an unbounded total streams through variable-size regions. + std::size_t const minBlock = 1024; + std::size_t const chunk = 4 * 1024; // maxChunkSizeBytes + b::BuddyAllocator a(/*capacity=*/4 * chunk, minBlock); // arena holds only 4 chunks at once + std::size_t streamed = 0; + std::vector inflight; + for (int i = 0; i < 1000; ++i) // 1000 chunks * 4KiB = 4MiB through a 16KiB arena + { + // Keep up to 4 chunks in flight. + if (inflight.size() == 4) + { + a.free(inflight.front()); + inflight.erase(inflight.begin()); + } + auto o = a.alloc(chunk); + ASSERT_TRUE(o.has_value()) << "chunk " << i << " did not fit despite recycling"; + inflight.push_back(*o); + streamed += chunk; + } + EXPECT_GT(streamed, a.capacity() * 10u); // streamed far more than the arena holds + for (auto o : inflight) + { + a.free(o); + } + EXPECT_EQ(a.freeBytes(), a.capacity()); + EXPECT_EQ(a.liveBlocks(), 0u); +} + +TEST(BuddyAllocator, MixedSmallAndLargeShareArena) +{ + // Small regions coexist with near-max chunks; after smalls free + coalesce a large alloc fits. + b::BuddyAllocator a(/*capacity=*/8 * 1024, /*minBlock=*/1024); // 8 blocks + auto big = a.alloc(4 * 1024); // 4 blocks + ASSERT_TRUE(big.has_value()); + std::vector smalls; + for (int i = 0; i < 4; ++i) + { + auto o = a.alloc(1024); + ASSERT_TRUE(o.has_value()); + smalls.push_back(*o); + } + EXPECT_EQ(a.freeBytes(), 0u); + EXPECT_FALSE(a.alloc(1024).has_value()); // full + // Free the smalls -> they coalesce -> a second 4*1024 chunk now fits. + for (auto o : smalls) + { + a.free(o); + } + auto big2 = a.alloc(4 * 1024); + ASSERT_TRUE(big2.has_value()); + a.free(*big); + a.free(*big2); + EXPECT_EQ(a.maxAllocBytes(), a.capacity()); +} + +TEST(BuddyAllocator, DoubleFreeIgnored) +{ + b::BuddyAllocator a(1024, 256); + auto o = a.alloc(256); + ASSERT_TRUE(o.has_value()); + a.free(*o); + a.free(*o); // ignored, no corruption + EXPECT_EQ(a.freeBytes(), a.capacity()); + EXPECT_EQ(a.liveBlocks(), 0u); +} + +// ---- boundary cases ---- + +TEST(BuddyAllocator, CapacityRoundedDownToMinBlockMultiple) +{ + // capacity is NOT a power-of-two multiple of minBlock -> usable rounds DOWN to 512 (256<<1), + // the trailing 1000-512 bytes are unusable. capacity() reports the usable size, not the arg. + b::BuddyAllocator a(/*capacity=*/1000, /*minBlock=*/256); + EXPECT_EQ(a.capacity(), 512u); + EXPECT_EQ(a.minBlock(), 256u); + EXPECT_EQ(a.maxAllocBytes(), 512u); + auto whole = a.alloc(512); + ASSERT_TRUE(whole.has_value()); + EXPECT_EQ(*whole, 0u); + EXPECT_FALSE(a.alloc(1).has_value()); // nothing left (the rounded-off tail is not allocatable) +} + +TEST(BuddyAllocator, MinBlockRoundedUpToPow2) +{ + // minBlock 100 -> rounded up to 128; capacity 512 -> orders 0..2 of 128/256/512. + b::BuddyAllocator a(/*capacity=*/512, /*minBlock=*/100); + EXPECT_EQ(a.minBlock(), 128u); + EXPECT_EQ(a.capacity(), 512u); + auto o1 = a.alloc(1); // -> 128 (one minBlock) + auto o2 = a.alloc(130); // -> 256 (next power of two) + ASSERT_TRUE(o1 && o2); + EXPECT_EQ(*o1 % 128, 0u); + EXPECT_EQ(*o2 % 256, 0u); // 256-block is 256-aligned + EXPECT_EQ(a.freeBytes(), 512u - (128u + 256u)); // 128 used + 256 used +} + +TEST(BuddyAllocator, SingleBlockArena) +{ + // capacity == minBlock -> exactly one order-0 block (maxOrder 0). + b::BuddyAllocator a(/*capacity=*/256, /*minBlock=*/256); + EXPECT_EQ(a.capacity(), 256u); + EXPECT_EQ(a.maxAllocBytes(), 256u); + auto o = a.alloc(10); + ASSERT_TRUE(o.has_value()); + EXPECT_EQ(*o, 0u); + EXPECT_FALSE(a.alloc(1).has_value()); // only one block -> full + EXPECT_EQ(a.maxAllocBytes(), 0u); + a.free(*o); + EXPECT_EQ(a.maxAllocBytes(), 256u); + EXPECT_TRUE(a.alloc(256).has_value()); // reusable +} + +TEST(BuddyAllocator, FreeUnknownOffsetIgnored) +{ + // Freeing an offset that was never allocated (out of range OR a valid-but-unallocated offset) + // must be a no-op, not corrupt the free lists or fabricate free space. + b::BuddyAllocator a(/*capacity=*/1024, /*minBlock=*/256); + auto o = a.alloc(256); + ASSERT_TRUE(o.has_value()); + auto const freeBefore = a.freeBytes(); + a.free(999999); // wildly out of range + a.free(*o + 64); // in range but not a block start / not allocated + EXPECT_EQ(a.freeBytes(), freeBefore); + EXPECT_EQ(a.liveBlocks(), 1u); // the one real allocation is untouched + a.free(*o); // the genuine free still works exactly once + EXPECT_EQ(a.freeBytes(), a.capacity()); + EXPECT_EQ(a.liveBlocks(), 0u); +} + +TEST(BuddyAllocator, BlockBytesReportsRoundedUpSize) +{ + // alloc returns only the offset; blockBytes() exposes the actual (rounded-up) block size for + // metrics. 0 for any non-live offset. + b::BuddyAllocator a(/*capacity=*/1024, /*minBlock=*/64); + auto o = a.alloc(65); // 65 -> rounded up to 128 + ASSERT_TRUE(o.has_value()); + EXPECT_EQ(a.blockBytes(*o), 128u); + EXPECT_EQ(a.blockBytes(*o + 999), 0u); // not a live block start + a.free(*o); + EXPECT_EQ(a.blockBytes(*o), 0u); // freed -> no longer live +} + +TEST(BuddyAllocator, FragmentationBlocksLargeAllocDespiteFreeBytes) +{ + // The defining buddy property: freeing every OTHER minBlock leaves half the arena free in bytes, + // but as scattered order-0 blocks with no free buddy pair -> a 2-block alloc must FAIL (no + // premature coalescing), and freeing the rest must then coalesce so it succeeds. + constexpr std::size_t kMin = 256; + b::BuddyAllocator a(/*capacity=*/8 * kMin, kMin); // 8 order-0 blocks + std::vector all; + for (int i = 0; i < 8; ++i) + { + auto o = a.alloc(kMin); + ASSERT_TRUE(o.has_value()); + all.push_back(*o); + } + EXPECT_EQ(a.freeBytes(), 0u); + // Free exactly the lower buddy of each pair (offset % 512 == 0) so no two freed blocks are + // buddies -> they cannot coalesce; their live buddies block every order-1 merge. + std::vector held; + for (auto o : all) + { + if (o % (2 * kMin) == 0) + { + a.free(o); + } + else + { + held.push_back(o); + } + } + EXPECT_EQ(a.freeBytes(), 4 * kMin); // half the arena is free... + EXPECT_EQ(a.maxAllocBytes(), kMin); // ...but only as isolated order-0 blocks + EXPECT_FALSE(a.alloc(2 * kMin).has_value()); // a 2-block request can't be satisfied (fragmented) + auto extra = a.alloc(kMin); // a 1-block request still works + ASSERT_TRUE(extra.has_value()); + // Return everything (the extra 1-block + the held buddies) -> all 8 blocks coalesce back so a + // full-capacity alloc fits again. + a.free(*extra); + for (auto o : held) + { + a.free(o); + } + EXPECT_EQ(a.liveBlocks(), 0u); + EXPECT_EQ(a.maxAllocBytes(), a.capacity()); +} diff --git a/cpp/tests/unit_tests/executor/bounce/creditSchedulerTest.cpp b/cpp/tests/unit_tests/executor/bounce/creditSchedulerTest.cpp new file mode 100644 index 000000000000..f94c3333dcaa --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/creditSchedulerTest.cpp @@ -0,0 +1,772 @@ +/* + * 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/CreditScheduler.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace b = tensorrt_llm::executor::kv_cache::bounce; + +namespace +{ + +// The variable-region scheduler carves byte offsets out of one arena. Most tests model an arena of +// N EQUAL regions by making the buddy min block == one region size, so each "want a region" is one +// minimum block and the arena holds exactly N of them — i.e. the old fixed-slot semantics, now +// expressed in bytes. (VariableSizeRegions exercises the genuinely variable case.) +constexpr std::uint64_t kBase = 0x100000ULL; // arena base device address (Grant.addr = kBase + offset) +constexpr std::uint32_t kRegion = 0x1000ULL; // 4096B: one "slot"-sized region (== buddy min block) + +// Scheduler over an arena holding exactly `nRegions` regions of kRegion bytes. +b::CreditScheduler makeSched(std::uint32_t nRegions, std::uint32_t maxInflightChunksPerRequest) +{ + return b::CreditScheduler( + kBase, static_cast(nRegions) * kRegion, kRegion, maxInflightChunksPerRequest); +} + +// "want n one-region chunks" -> n equal chunks of kRegion bytes (FIFO order). +std::vector want(std::uint32_t n) +{ + return std::vector(n, kRegion); +} + +// Free regions currently available (all equal-sized here, so bytes / region size). +std::size_t freeRegions(b::CreditScheduler& s) +{ + return s.freeBytes() / kRegion; +} + +// Model the Python transceiver's bounded worker queues: each sender has at most `threadsPerSender` +// synchronous submit+wait requests active, and a worker submits its next one-chunk request as soon as +// the previous request completes. The resulting small-request load is sustained without ever +// exceeding TRTLLM_KV_TRANSFER_NUM_THREADS active requests per sender. +void expectLargeChunkProgressWithBoundedSenderConcurrency(std::uint32_t numSenders, std::uint32_t threadsPerSender) +{ + struct WorkerSlot + { + std::uint32_t sender{}; + std::uint64_t nextRid{}; + std::string flow; + std::optional heldOffset; + }; + + std::uint32_t const numSmallFlows = numSenders * threadsPerSender; + ASSERT_GT(numSmallFlows, 1U); + // Both test instantiations use a power-of-two flow count, matching the buddy arena's full usable + // capacity. One one-region request per worker slot fills it exactly. + b::CreditScheduler s(kBase, static_cast(numSmallFlows) * kRegion, kRegion, + /*maxInflightChunksPerRequest=*/1); + + auto makeFlow = [](std::uint32_t sender, std::uint64_t rid) + { return std::string("smallPeer") + std::to_string(sender) + '\x1f' + std::to_string(rid); }; + + std::vector slots; + slots.reserve(numSmallFlows); + std::unordered_map flowToSlot; + for (std::uint32_t sender = 0; sender < numSenders; ++sender) + { + for (std::uint32_t thread = 0; thread < threadsPerSender; ++thread) + { + WorkerSlot slot; + slot.sender = sender; + slot.nextRid = thread; + slot.flow = makeFlow(sender, slot.nextRid); + slot.nextRid += threadsPerSender; + auto grants = s.onWant(slot.flow, want(1)); + ASSERT_EQ(grants.size(), 1U); + slot.heldOffset = grants.front().offset; + flowToSlot.emplace(slot.flow, slots.size()); + slots.push_back(std::move(slot)); + } + } + EXPECT_EQ(s.freeBytes(), 0U); + + std::string const largeFlow = std::string("largePeer\x1f") + "0"; + std::uint32_t const largeBytes = (numSmallFlows / 2) * kRegion; + EXPECT_TRUE(s.onWant(largeFlow, {largeBytes}).empty()); + + bool largeGranted = false; + std::size_t nextSlot = 0; + std::size_t smallCompletions = 0; + std::size_t const kMaxCompletions = static_cast(numSmallFlows) * 32; + auto recordGrants = [&](std::vector const& grants) + { + for (auto const& grant : grants) + { + if (grant.flow == largeFlow) + { + EXPECT_EQ(grant.len, largeBytes); + largeGranted = true; + continue; + } + auto const it = flowToSlot.find(grant.flow); + ASSERT_NE(it, flowToSlot.end()); + slots[it->second].heldOffset = grant.offset; + } + }; + + while (!largeGranted && smallCompletions < kMaxCompletions) + { + std::size_t checked = 0; + while (checked < slots.size() && !slots[nextSlot].heldOffset) + { + nextSlot = (nextSlot + 1) % slots.size(); + ++checked; + } + ASSERT_LT(checked, slots.size()) << "drain stopped every in-flight small request before the large grant"; + + auto& slot = slots[nextSlot]; + std::string const completedFlow = slot.flow; + std::uint64_t const completedOffset = *slot.heldOffset; + slot.heldOffset.reset(); + flowToSlot.erase(completedFlow); + recordGrants(s.onScatterDone(completedFlow, completedOffset)); + ++smallCompletions; + if (largeGranted) + { + break; + } + + // The synchronous Python worker can submit exactly one replacement after its old request + // completes. If receiver drain is active this new flow remains pending instead of refilling + // the just-freed region, allowing a large buddy block to coalesce. + slot.flow = makeFlow(slot.sender, slot.nextRid); + slot.nextRid += threadsPerSender; + flowToSlot.emplace(slot.flow, nextSlot); + recordGrants(s.onWant(slot.flow, want(1))); + nextSlot = (nextSlot + 1) % slots.size(); + } + + EXPECT_TRUE(largeGranted) << "sustained bounded-concurrency small requests starved the large chunk after " + << smallCompletions << " completions"; +} + +// Conservation: every region is either free, held by some flow, or locally held. With N equal +// regions this is the byte-budget invariant that, in v1, broke as a double-free on peer loss. +void checkConservation(b::CreditScheduler& s, std::vector const& flows, std::uint32_t nRegions) +{ + std::size_t held = 0; + for (auto const& f : flows) + { + held += s.heldCount(f); + } + EXPECT_EQ(freeRegions(s) + held + s.localHeldCount(), nRegions) << "region conservation violated"; +} + +// Track who-holds-what from the grant/return stream so we can assert no region is ever double- +// granted (held by two flows at once), and that grants carry a consistent addr/len. +struct Mirror +{ + std::map owner; // offset -> flow + + void grant(std::vector const& gs) + { + for (auto const& g : gs) + { + EXPECT_EQ(owner.count(g.offset), 0u) << "region " << g.offset << " granted while still held"; + owner[g.offset] = g.flow; + EXPECT_EQ(g.addr, kBase + g.offset) << "grant carried wrong addr"; + EXPECT_EQ(g.len, kRegion) << "grant carried wrong region length"; + } + } + + void free(std::uint64_t offset) + { + owner.erase(offset); + } + + // Free every region whose owning flow satisfies `pred` (mirrors a whole-flow/whole-peer reclaim). + template + void freeOwnedBy(Pred&& pred) + { + for (auto it = owner.begin(); it != owner.end();) + { + it = pred(it->second) ? owner.erase(it) : std::next(it); + } + } +}; + +} // namespace + +// One WANT against an (arena, per-request cap) pair grants exactly min(cap, nRegions): the +// per-request limit binds even with arena room to spare, a large limit fills the arena exactly +// once, and a want far beyond the arena neither over-grants nor loops. +TEST(CreditScheduler, SingleWantGrantsMinOfCapAndArena) +{ + struct Case + { + std::uint32_t nRegions; + std::uint32_t cap; + std::uint32_t wantChunks; + std::size_t expectGrants; + }; + + for (auto const& tc : { + Case{8, 4, 100, 4}, // capped by the per-request limit + Case{8, 16, 100, 8}, // large limit -> fills the arena + Case{8, 16, 2000, 8}, // huge want -> exactly N, no more, no hang + Case{8, 2, 100, 2}, // limit binds even with arena room + }) + { + auto s = makeSched(tc.nRegions, tc.cap); + auto g = s.onWant("A", want(tc.wantChunks)); + std::string const ctx = " nRegions=" + std::to_string(tc.nRegions) + " cap=" + std::to_string(tc.cap) + + " want=" + std::to_string(tc.wantChunks); + EXPECT_EQ(g.size(), tc.expectGrants) << ctx; + EXPECT_EQ(s.heldCount("A"), tc.expectGrants) << ctx; + EXPECT_EQ(freeRegions(s), tc.nRegions - tc.expectGrants) << ctx; + checkConservation(s, {"A"}, tc.nRegions); + } +} + +TEST(CreditScheduler, RecyclingOnScatterDone) +{ + auto s = makeSched(/*nRegions=*/4, /*maxInflightChunksPerRequest=*/16); + Mirror m; + m.grant(s.onWant("A", want(10))); // K=10 > N=4 -> gets 4 + EXPECT_EQ(s.heldCount("A"), 4u); + // Complete one chunk: its region frees and is immediately re-granted (remaining still > 0). + auto firstOff = m.owner.begin()->first; + m.free(firstOff); + auto re = s.onScatterDone("A", firstOff); + m.grant(re); + EXPECT_EQ(re.size(), 1u); + EXPECT_EQ(s.heldCount("A"), 4u); // in-flight allocation count stays at its cap + checkConservation(s, {"A"}, 4); +} + +TEST(CreditScheduler, BoundedInflightLimitGivesFairSplit) +{ + // With a per-request limit of four on an eight-region arena, two senders split it evenly. + auto s = makeSched(/*nRegions=*/8, /*maxInflightChunksPerRequest=*/4); + Mirror m; + m.grant(s.onWant("A", want(100))); + m.grant(s.onWant("B", want(100))); + EXPECT_EQ(s.heldCount("A"), 4u); + EXPECT_EQ(s.heldCount("B"), 4u); + EXPECT_EQ(freeRegions(s), 0u); + checkConservation(s, {"A", "B"}, 8); +} + +TEST(CreditScheduler, MoreSendersThanRegionsNoStarvation) +{ + auto s = makeSched(/*nRegions=*/2, /*maxInflightChunksPerRequest=*/16); + Mirror m; + m.grant(s.onWant("A", want(5))); + m.grant(s.onWant("B", want(5))); + m.grant(s.onWant("C", want(5))); + EXPECT_EQ(s.heldCount("A") + s.heldCount("B") + s.heldCount("C"), 2u); + EXPECT_EQ(freeRegions(s), 0u); + + // Recycling regions must serve the starved flows in turn — no flow is permanently starved. + std::set served; + for (auto const& p : {"A", "B", "C"}) + { + if (s.heldCount(p) > 0) + { + served.insert(p); + } + } + int guard = 0; + while (served.size() < 3 && guard++ < 50) + { + ASSERT_FALSE(m.owner.empty()) << "no region is held; cannot drive recycling"; + auto off = m.owner.begin()->first; + auto flow = m.owner.begin()->second; + m.free(off); + m.grant(s.onScatterDone(flow, off)); + for (auto const& p : {"A", "B", "C"}) + { + if (s.heldCount(p) > 0) + { + served.insert(p); + } + } + checkConservation(s, {"A", "B", "C"}, 2); + } + EXPECT_EQ(served.size(), 3u) << "some sender was permanently starved"; +} + +TEST(CreditScheduler, PeerGoneReclaimsNoDoubleFree) +{ + auto s = makeSched(/*nRegions=*/4, /*maxInflightChunksPerRequest=*/16); + (void) s.onWant("A", want(2)); + (void) s.onWant("B", want(2)); + checkConservation(s, {"A", "B"}, 4); + std::vector deferred; + (void) s.reclaimFlow("A", {}, deferred); // A's held regions return to the arena (none busy) + checkConservation(s, {"A", "B"}, 4); + EXPECT_EQ(s.heldCount("A"), 0u); + EXPECT_EQ(freeRegions(s) + s.heldCount("B"), 4u); +} + +TEST(CreditScheduler, PeerGoneMidRecycleHandsRegionsToWaiter) +{ + // A holds the whole arena; B is waiting. A disappears mid-flight -> its regions must be reclaimed + // AND immediately re-granted to the starved waiter B, with no double-free and no region lost. + auto s = makeSched(/*nRegions=*/4, /*maxInflightChunksPerRequest=*/16); + Mirror m; + m.grant(s.onWant("A", want(100))); // A grabs all 4 + EXPECT_EQ(s.heldCount("A"), 4u); + m.grant(s.onWant("B", want(100))); // B wants 4 but nothing free yet + EXPECT_EQ(s.heldCount("B"), 0u); + + m.freeOwnedBy([](std::string const& owner) { return owner == "A"; }); // mirror: A no longer owns these + std::vector deferred; + auto re = s.reclaimFlow("A", {}, deferred); + m.grant(re); + EXPECT_EQ(s.heldCount("A"), 0u); + EXPECT_EQ(s.heldCount("B"), 4u) << "reclaimed regions did not flow to the waiting sender"; + EXPECT_EQ(freeRegions(s), 0u); + checkConservation(s, {"A", "B"}, 4); +} + +TEST(CreditScheduler, ReclaimByPrefixDropsAllFlowsOfPeer) +{ + // Transport keys flows as "peer\x1f rid". reclaimByPrefix("p1\x1f") must drop EVERY p1 flow + // (multiple concurrent requests from one peer) and hand the freed regions to an unrelated peer. + constexpr char sep = '\x1f'; + std::string const p1a = std::string("p1") + sep + "1"; + std::string const p1b = std::string("p1") + sep + "2"; + std::string const p2 = std::string("p2") + sep + "1"; + auto s = makeSched(/*nRegions=*/4, /*maxInflightChunksPerRequest=*/16); + Mirror m; + m.grant(s.onWant(p1a, want(2))); + m.grant(s.onWant(p1b, want(2))); // p1 now holds all 4 across two flows + m.grant(s.onWant(p2, want(4))); // p2 waits (nothing free) + EXPECT_EQ(s.heldCount(p1a) + s.heldCount(p1b), 4u); + EXPECT_EQ(s.heldCount(p2), 0u); + + m.freeOwnedBy([&](std::string const& owner) { return owner == p1a || owner == p1b; }); + std::vector deferred; + auto re = s.reclaimByPrefix(std::string("p1") + sep, {}, deferred); + m.grant(re); + EXPECT_EQ(s.heldCount(p1a), 0u); + EXPECT_EQ(s.heldCount(p1b), 0u); + EXPECT_EQ(s.heldCount(p2), 4u) << "freed regions did not go to the surviving peer"; + checkConservation(s, {p1a, p1b, p2}, 4); + + // A non-matching prefix is a no-op. + auto none = s.reclaimByPrefix(std::string("nomatch") + sep, {}, deferred); + EXPECT_TRUE(none.empty()); + EXPECT_EQ(s.heldCount(p2), 4u); + + // An EMPTY prefix must NOT reclaim everything (compare(0,0,"")==0 matches all keys) — guarded. + auto empty = s.reclaimByPrefix(std::string(), {}, deferred); + EXPECT_TRUE(empty.empty()); + EXPECT_EQ(s.heldCount(p2), 4u) << "empty prefix wrongly reclaimed a live flow"; +} + +TEST(CreditScheduler, CompletedFlowsReclaimedNoTombstoneLeak) +{ + // The transport keys each request as "peer\x1f rid" with a monotonic, never-reused rid. + // A long-running server runs many flows; completed flows MUST be reclaimed, else mFlows/mOrder + // grow without bound and schedule() degrades to O(historical requests). + auto s = makeSched(/*nRegions=*/4, /*maxInflightChunksPerRequest=*/16); + for (int rid = 0; rid < 1000; ++rid) + { + std::string const key = std::string("peerA\x1f") + std::to_string(rid); + auto g = s.onWant(key, want(2)); + ASSERT_EQ(g.size(), 2u); + (void) s.onScatterDone(key, g[0].offset); // one still held -> flow kept + EXPECT_EQ(s.trackedFlows(), 1u); + (void) s.onScatterDone(key, g[1].offset); // last held region drains -> flow reclaimed + EXPECT_EQ(s.trackedFlows(), 0u) << "completed flow left a tombstone at rid=" << rid; + } + EXPECT_EQ(s.trackedFlows(), 0u); // bounded by in-flight flows, NOT lifetime request count + EXPECT_EQ(freeRegions(s), 4u); +} + +TEST(CreditScheduler, LastFlowLeavingRingResetsCursorAndRingRefills) +{ + // Regression: dropFromRing used to run `mCursor %= mRing.size()` AFTER erasing the LAST ring + // element -> modulo by zero (UB; a deterministic SIGFPE in -O0 builds). Any normal completion + // of the only active flow hits it. Drain a single flow to empty the ring, then refill and + // verify scheduling still rotates fairly from a sane cursor. + auto s = makeSched(/*nRegions=*/4, /*maxInflightChunksPerRequest=*/16); + std::string const k1 = std::string("p\x1f") + "1"; + auto g = s.onWant(k1, want(1)); + ASSERT_EQ(g.size(), 1u); + (void) s.onScatterDone(k1, g[0].offset); // last held drains -> flow erased -> ring now EMPTY + EXPECT_EQ(s.trackedFlows(), 0u); + EXPECT_EQ(freeRegions(s), 4u); + + // Ring refills after going empty; grants keep alternating (cursor was reset, not left dangling). + std::string const k2 = std::string("p\x1f") + "2"; + std::string const k3 = std::string("p\x1f") + "3"; + auto g2 = s.onWant(k2, want(2)); + auto g3 = s.onWant(k3, want(2)); + ASSERT_EQ(g2.size() + g3.size(), 4u); + Mirror m; + m.grant(g2); + m.grant(g3); // no double-grant across the empty->refill transition + for (auto const& x : g2) + { + (void) s.onScatterDone(k2, x.offset); + } + for (auto const& x : g3) + { + (void) s.onScatterDone(k3, x.offset); + } + EXPECT_EQ(s.trackedFlows(), 0u); // BOTH flows drain -> ring empties again (erase-behind-cursor path) + EXPECT_EQ(freeRegions(s), 4u); +} + +TEST(CreditScheduler, CancelledFlowReclaimed) +{ + auto s = makeSched(/*nRegions=*/4, /*maxInflightChunksPerRequest=*/16); + std::string const k1 = std::string("p\x1f") + "1"; + std::string const k2 = std::string("p\x1f") + "2"; + // Empty WANT with nothing in flight -> no tombstone created. + (void) s.onWant(k1, want(0)); + EXPECT_EQ(s.trackedFlows(), 0u); + // Grant, then cancel while regions are still held -> kept until they drain, then reclaimed. + auto g = s.onWant(k2, want(2)); + ASSERT_EQ(g.size(), 2u); + (void) s.onWant(k2, want(0)); // cancel; 2 regions still in flight + EXPECT_EQ(s.trackedFlows(), 1u); // must stay until held drains (regions still leased) + (void) s.onScatterDone(k2, g[0].offset); + (void) s.onScatterDone(k2, g[1].offset); + EXPECT_EQ(s.trackedFlows(), 0u); // reclaimed after the last held region returns + EXPECT_EQ(freeRegions(s), 4u); +} + +TEST(CreditScheduler, ReclaimDefersBusyRegionThenFreeOrphan) +{ + // A region whose scatter is still running on the receiver must NOT be freed/re-granted when the + // peer is reclaimed (forgetPeer) — else another sender's RDMA write races the worker's read. + // reclaimByPrefix defers such "busy" regions; freeOrphanRegion releases them once scatter is done. + auto s = makeSched(/*nRegions=*/2, /*maxInflightChunksPerRequest=*/16); + std::string const keyA = std::string("A\x1f") + "1"; + std::string const prefixA = std::string("A\x1f"); + std::string const keyB = std::string("B\x1f") + "1"; + + auto g = s.onWant(keyA, want(2)); + ASSERT_EQ(g.size(), 2u); // A holds both regions + std::uint64_t const busyOff = g[0].offset; + std::uint64_t const idleOff = g[1].offset; + + std::unordered_set busy{busyOff}; // a scatter is still reading busyOff + std::vector deferred; + auto re = s.reclaimByPrefix(prefixA, busy, deferred); + EXPECT_TRUE(re.empty()); // no waiter yet + ASSERT_EQ(deferred.size(), 1u); + EXPECT_EQ(deferred[0], busyOff); + EXPECT_EQ(s.trackedFlows(), 0u); // flow A erased + EXPECT_EQ(freeRegions(s), 1u); // only the idle region is free; busy one is in limbo + + // A new sender B wanting 2 may take only the idle region — NEVER the busy (in-limbo) one. + auto gb = s.onWant(keyB, want(2)); + ASSERT_EQ(gb.size(), 1u); + EXPECT_EQ(gb[0].offset, idleOff); + EXPECT_EQ(freeRegions(s), 0u); + + // Scatter finishes -> free the orphan -> now B can finally get it. + auto re2 = s.freeOrphanRegion(busyOff); + ASSERT_EQ(re2.size(), 1u); + EXPECT_EQ(re2[0].offset, busyOff); + EXPECT_EQ(s.heldCount(keyB), 2u); + checkConservation(s, {keyA, keyB}, 2); + + // A second freeOrphanRegion for the same (now non-orphan) offset is a no-op — must NOT free + // busyOff again, which is now live under keyB. + auto re3 = s.freeOrphanRegion(busyOff); + EXPECT_TRUE(re3.empty()); + EXPECT_EQ(s.heldCount(keyB), 2u); + checkConservation(s, {keyA, keyB}, 2); +} + +TEST(CreditScheduler, FreeOrphanRegionIgnoresNonOrphan) +{ + // freeOrphanRegion must only free regions actually deferred as orphans; a call for a LIVE, + // never-deferred region is a no-op (defense in depth against a stray caller). + auto s = makeSched(/*nRegions=*/2, /*maxInflightChunksPerRequest=*/16); + std::string const flow = std::string("A\x1f") + "1"; + auto g = s.onWant(flow, want(1)); + ASSERT_EQ(g.size(), 1u); + auto const live = g[0].offset; // a live region, never reclaim-deferred + EXPECT_EQ(s.heldCount(flow), 1u); + + auto re = s.freeOrphanRegion(live); + EXPECT_TRUE(re.empty()); + EXPECT_EQ(s.heldCount(flow), 1u); // still held — NOT freed + EXPECT_EQ(freeRegions(s), 1u); // unchanged + + (void) s.onScatterDone(flow, live); // normal path still frees it exactly once + EXPECT_EQ(freeRegions(s), 2u); + EXPECT_EQ(s.heldCount(flow), 0u); +} + +TEST(CreditScheduler, ReclaimFlowFreesHeldDefersBusyAndHeldByFlow) +{ + // Explicit cancel of one flow (empty WANT path): free its granted-but-unwritten regions now, + // defer any whose scatter is still running; heldByFlow lets the transport drop late DATA. + auto s = makeSched(/*nRegions=*/4, /*maxInflightChunksPerRequest=*/16); + std::string const flow = std::string("p\x1f") + "1"; + auto g = s.onWant(flow, want(3)); + ASSERT_EQ(g.size(), 3u); // 3 held, 1 free + + EXPECT_TRUE(s.heldByFlow(flow, g[0].offset)); + EXPECT_FALSE(s.heldByFlow(flow, 0x999999)); // not a held offset + EXPECT_FALSE( + s.heldByFlow("p\x1f" + "2", + g[0].offset)); // held by a different flow key + + std::unordered_set busy{g[1].offset}; // a scatter is still reading g[1] + std::vector deferred; + auto re = s.reclaimFlow(flow, busy, deferred); + EXPECT_TRUE(re.empty()); // no other flow waiting + EXPECT_EQ(s.trackedFlows(), 0u); + ASSERT_EQ(deferred.size(), 1u); + EXPECT_EQ(deferred[0], g[1].offset); + EXPECT_FALSE(s.heldByFlow(flow, g[0].offset)); // flow erased + EXPECT_EQ(freeRegions(s), 3u); // g0,g2 freed + the originally-free one; g1 in limbo + + auto re2 = s.freeOrphanRegion(g[1].offset); // scatter finished -> free the deferred one + EXPECT_EQ(freeRegions(s), 4u); + + // Cancelling an unknown flow is a harmless no-op. + std::vector none; + EXPECT_TRUE(s.reclaimFlow("nope", {}, none).empty()); + EXPECT_TRUE(none.empty()); +} + +// Scheduler over `nRegions` regions with a controllable fake clock: tests advance *now instead of +// sleeping, keeping the lease/quarantine tests deterministic and instant. +b::CreditScheduler makeTimedSched( + std::uint32_t nRegions, std::uint32_t maxInflight, std::shared_ptr now) +{ + return b::CreditScheduler( + kBase, static_cast(nRegions) * kRegion, kRegion, maxInflight, [now] { return *now; }); +} + +TEST(CreditScheduler, ReceiverReclaimQuarantinesUnwrittenRegionsUntilReaped) +{ + // Receiver-initiated reclaim (peer loss / lease expiry): granted-but-unwritten regions may STILL + // be RDMA-written by the peer, so quarantineFor > 0 must keep them out of the arena — not free + // them — until reapQuarantine() passes the time barrier. + auto now = std::make_shared(std::chrono::steady_clock::now()); + auto s = makeTimedSched(/*nRegions=*/2, /*maxInflight=*/16, now); + std::string const flowA = std::string("A\x1f") + "1"; + std::string const flowB = std::string("B\x1f") + "1"; + auto g = s.onWant(flowA, want(2)); + ASSERT_EQ(g.size(), 2u); + + std::vector deferred; + auto re = s.reclaimFlow(flowA, /*busy=*/{}, deferred, std::chrono::seconds(30)); + EXPECT_TRUE(re.empty()); // no waiter yet + EXPECT_TRUE(deferred.empty()); + EXPECT_EQ(s.trackedFlows(), 0u); // flow gone... + EXPECT_FALSE(s.heldByFlow(flowA, g[0].offset)); // ...so late DATA gets dropped + EXPECT_EQ(freeRegions(s), 0u); // ...but NOTHING went back to the arena + + // A new flow must not be granted a quarantined region, and reaping early frees nothing. + EXPECT_TRUE(s.onWant(flowB, want(1)).empty()); + EXPECT_TRUE(s.reapQuarantine().empty()); + EXPECT_EQ(freeRegions(s), 0u); + + // Quarantine over -> the regions re-enter circulation; one serves the waiter. + *now += std::chrono::seconds(31); + auto re2 = s.reapQuarantine(); + ASSERT_EQ(re2.size(), 1u); + EXPECT_EQ(re2[0].flow, flowB); + EXPECT_EQ(s.heldCount(flowB), 1u); + EXPECT_EQ(freeRegions(s), 1u); + + // Reaping again is a harmless no-op (the reaped offsets are live/free now, not re-freed). + EXPECT_TRUE(s.reapQuarantine().empty()); + EXPECT_EQ(freeRegions(s), 1u); +} + +TEST(CreditScheduler, ReceiverReclaimBusyRegionsDeferAsOrphansNotQuarantine) +{ + // Busy takes precedence: a region a scatter worker still READS follows the orphan path + // (freeOrphanRegion on scatter completion), never the quarantine path — reapQuarantine must not + // free it even after the quarantine window. + auto now = std::make_shared(std::chrono::steady_clock::now()); + auto s = makeTimedSched(/*nRegions=*/2, /*maxInflight=*/16, now); + std::string const flow = std::string("A\x1f") + "1"; + auto g = s.onWant(flow, want(2)); + ASSERT_EQ(g.size(), 2u); + + std::unordered_set busy{g[0].offset}; + std::vector deferred; + (void) s.reclaimFlow(flow, busy, deferred, std::chrono::seconds(30)); + ASSERT_EQ(deferred.size(), 1u); + EXPECT_EQ(deferred[0], g[0].offset); + + *now += std::chrono::seconds(31); + EXPECT_TRUE(s.reapQuarantine().empty()); + EXPECT_EQ(freeRegions(s), 1u); // only the quarantined region came back; the orphan is still busy + + (void) s.freeOrphanRegion(g[0].offset); + EXPECT_EQ(freeRegions(s), 2u); +} + +TEST(CreditScheduler, StaleFlowsReportsIdleFlowsOnly) +{ + // The lease that makes a dead sender observable: a flow with no progress (WANT / grant issued / + // scatter completed) beyond the idle limit is reported stale; any progress renews the lease. + auto now = std::make_shared(std::chrono::steady_clock::now()); + auto s = makeTimedSched(/*nRegions=*/2, /*maxInflight=*/1, now); + std::string const dead = std::string("A\x1f") + "1"; + std::string const live = std::string("B\x1f") + "1"; + auto ga = s.onWant(dead, want(1)); + ASSERT_EQ(ga.size(), 1u); + auto gb = s.onWant(live, want(2)); // one granted, one pending (in-flight cap 1) + ASSERT_EQ(gb.size(), 1u); + + // t+40s: `live` completes its first chunk -> frees the region + gets its next grant = progress. + *now += std::chrono::seconds(40); + ASSERT_EQ(s.onScatterDone(live, gb[0].offset).size(), 1u); + + // t+70s: `dead` has been silent for 70s -> stale; `live` progressed 30s ago -> not stale. + *now += std::chrono::seconds(30); + auto stale = s.staleFlows(std::chrono::seconds(60)); + ASSERT_EQ(stale.size(), 1u); + EXPECT_EQ(stale[0], dead); + EXPECT_TRUE(s.staleFlows(std::chrono::seconds(90)).empty()); // longer lease: nobody is stale yet + + // Reclaim + quarantine + reap: the arena fully recovers the dead flow's region. + std::vector deferred; + (void) s.reclaimFlow(dead, /*busy=*/{}, deferred, std::chrono::seconds(30)); + EXPECT_EQ(freeRegions(s), 0u); // quarantined, not yet back + *now += std::chrono::seconds(31); + (void) s.reapQuarantine(); + EXPECT_EQ(freeRegions(s), 1u); // recovered (live still holds its in-flight second chunk) +} + +TEST(CreditScheduler, SharedArenaLocalAndRemoteShareOneAllocator) +{ + // Shared single arena: the local sender (gather staging, acquireLocal) and remote flows (grants) + // draw from the SAME allocator. Conservation holds across {free, remote-held, local-held}; a + // released local region flows to a waiting remote flow. + auto s = makeSched(/*nRegions=*/4, /*maxInflightChunksPerRequest=*/16); + std::string const flow = std::string("p\x1f") + "1"; + + auto a0 = s.acquireLocal(kRegion); + auto a1 = s.acquireLocal(kRegion); + ASSERT_TRUE(a0.has_value() && a1.has_value()); + EXPECT_EQ(s.localHeldCount(), 2u); + EXPECT_EQ(freeRegions(s), 2u); + + // Remote flow wants 4 but only 2 are free now (local holds 2) -> gets 2. + auto g = s.onWant(flow, want(4)); + EXPECT_EQ(g.size(), 2u); + EXPECT_EQ(s.heldCount(flow), 2u); + EXPECT_EQ(freeRegions(s), 0u); + checkConservation(s, {flow}, 4); + + // No region free -> acquireLocal returns nullopt (caller would park and retry). + EXPECT_FALSE(s.acquireLocal(kRegion).has_value()); + + // Release one local region -> the remote flow (still wants 2 more) immediately grabs it. + auto re = s.releaseLocal(*a0); + EXPECT_EQ(re.size(), 1u); + EXPECT_EQ(s.heldCount(flow), 3u); + EXPECT_EQ(s.localHeldCount(), 1u); + checkConservation(s, {flow}, 4); + + // Release the last local region -> remote flow reaches its wanted 4. + (void) s.releaseLocal(*a1); + EXPECT_EQ(s.heldCount(flow), 4u); + EXPECT_EQ(s.localHeldCount(), 0u); + EXPECT_EQ(freeRegions(s), 0u); + checkConservation(s, {flow}, 4); + + // releaseLocal on a region that isn't locally held is an idempotent no-op. + auto none = s.releaseLocal(*a0); + EXPECT_TRUE(none.empty()); + checkConservation(s, {flow}, 4); +} + +TEST(CreditScheduler, ScatterDoneIdempotentForUnknownRegion) +{ + auto s = makeSched(/*nRegions=*/4, /*maxInflightChunksPerRequest=*/16); + (void) s.onWant("A", want(1)); // A holds 1, 3 free + auto before = freeRegions(s); + auto re = s.onScatterDone("A", 0x999000); // an offset A never held + EXPECT_TRUE(re.empty()); + EXPECT_EQ(freeRegions(s), before); // no spurious free + checkConservation(s, {"A"}, 4); +} + +TEST(CreditScheduler, WantEmptyStopsGranting) +{ + auto s = makeSched(/*nRegions=*/8, /*maxInflightChunksPerRequest=*/16); + (void) s.onWant("A", want(4)); + EXPECT_EQ(s.heldCount("A"), 4u); + auto g = s.onWant("A", want(0)); // cancel further grants + EXPECT_TRUE(g.empty()); + EXPECT_EQ(s.activeFlows(), 0u); +} + +TEST(CreditScheduler, VariableSizeRegionsPackAndBackpressure) +{ + // The genuinely variable case: a flow's chunks have different byte sizes, each granted a region + // of exactly that size, packed densely. A chunk that cannot fit right now is parked (no grant) + // while smaller following chunks are NOT reordered ahead of it (FIFO per flow). + constexpr std::size_t kArena = 8 * kRegion; // 32 KiB, min block kRegion + b::CreditScheduler s(kBase, kArena, kRegion, /*maxInflightChunksPerRequest=*/16); + + // Chunk sizes: 1, 2, 4, 1 regions (total 8 -> exactly fills the arena). + std::vector sizes{kRegion, 2 * kRegion, 4 * kRegion, kRegion}; + auto g = s.onWant("A", sizes); + ASSERT_EQ(g.size(), 4u); + EXPECT_EQ(g[0].len, kRegion); + EXPECT_EQ(g[1].len, 2 * kRegion); + EXPECT_EQ(g[2].len, 4 * kRegion); + EXPECT_EQ(g[3].len, kRegion); + // Distinct, non-overlapping regions. + for (std::size_t i = 0; i < g.size(); ++i) + { + for (std::size_t j = i + 1; j < g.size(); ++j) + { + bool const disjoint = g[i].offset + g[i].len <= g[j].offset || g[j].offset + g[j].len <= g[i].offset; + EXPECT_TRUE(disjoint) << "regions " << i << " and " << j << " overlap"; + } + } + EXPECT_EQ(s.freeBytes(), 0u); // arena exactly full + + // Free the 4-region chunk; a fresh flow wanting a 4-region chunk can now be granted. + (void) s.onScatterDone("A", g[2].offset); + auto g2 = s.onWant("B", std::vector{4 * kRegion}); + ASSERT_EQ(g2.size(), 1u); + EXPECT_EQ(g2[0].len, 4 * kRegion); +} + +TEST(CreditScheduler, ManySingleThreadSendersCannotStarveLargeChunk) +{ + expectLargeChunkProgressWithBoundedSenderConcurrency(/*numSenders=*/8, /*threadsPerSender=*/1); +} + +TEST(CreditScheduler, FourThreadsPerSenderCannotStarveLargeChunk) +{ + expectLargeChunkProgressWithBoundedSenderConcurrency(/*numSenders=*/4, /*threadsPerSender=*/4); +} diff --git a/cpp/tests/unit_tests/executor/bounce/execPoolTest.cpp b/cpp/tests/unit_tests/executor/bounce/execPoolTest.cpp new file mode 100644 index 000000000000..d90687c5df80 --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/execPoolTest.cpp @@ -0,0 +1,93 @@ +/* + * 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 "bounceTestUtils.h" + +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/ExecPool.h" + +#include + +#include + +#include +#include +#include + +namespace b = tensorrt_llm::executor::kv_cache::bounce; + +using bounce_test::hasCuda; + +TEST(ExecPool, AcquireReleaseExhaustion) +{ + if (!hasCuda()) + GTEST_SKIP(); + b::ExecPool pool(/*count=*/3, /*maxDescsPerChunk=*/64, /*dev=*/0); + EXPECT_EQ(pool.size(), 3u); + EXPECT_EQ(pool.freeCount(), 3u); + + std::vector held; + std::set ids; + std::set scratch; + for (int i = 0; i < 3; ++i) + { + auto* c = pool.tryAcquire(); + ASSERT_NE(c, nullptr); + EXPECT_NE(c->stream, nullptr); + EXPECT_NE(c->event, nullptr); + EXPECT_NE(c->scratch, nullptr); + EXPECT_NE(c->hostPinned, nullptr); + ids.insert(c->id); + scratch.insert(c->scratch); + held.push_back(c); + } + EXPECT_EQ(ids.size(), 3u); // distinct contexts + EXPECT_EQ(scratch.size(), 3u); // distinct scratch buffers + EXPECT_EQ(pool.freeCount(), 0u); + EXPECT_EQ(pool.tryAcquire(), nullptr); // exhausted -> non-blocking nullptr + + pool.release(held.back()); + held.pop_back(); + EXPECT_EQ(pool.freeCount(), 1u); + auto* again = pool.tryAcquire(); // the freed one is reusable + ASSERT_NE(again, nullptr); + pool.release(again); + for (auto* c : held) + { + pool.release(c); + } + EXPECT_EQ(pool.freeCount(), 3u); +} + +TEST(ExecPool, StreamScratchAreUsable) +{ + if (!hasCuda()) + GTEST_SKIP(); + b::ExecPool pool(1, 64, 0); + auto* c = pool.tryAcquire(); + ASSERT_NE(c, nullptr); + // Use the ctx's hostPinned -> scratch H2D on its stream, then record/poll its event. + std::vector in(128, 0xCD); + ASSERT_LE(128u, c->scratchBytes); + std::memcpy(c->hostPinned, in.data(), 128); + ASSERT_EQ(cudaMemcpyAsync(c->scratch, c->hostPinned, 128, cudaMemcpyHostToDevice, c->stream), cudaSuccess); + ASSERT_EQ(cudaEventRecord(c->event, c->stream), cudaSuccess); + ASSERT_EQ(cudaEventSynchronize(c->event), cudaSuccess); + std::vector out(128, 0); + ASSERT_EQ(cudaMemcpy(out.data(), c->scratch, 128, cudaMemcpyDeviceToHost), cudaSuccess); + EXPECT_EQ(in, out); + pool.release(c); +} diff --git a/cpp/tests/unit_tests/executor/bounce/gatherScatterKernelTest.cpp b/cpp/tests/unit_tests/executor/bounce/gatherScatterKernelTest.cpp new file mode 100644 index 000000000000..8eb4369a0958 --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/gatherScatterKernelTest.cpp @@ -0,0 +1,165 @@ +/* + * 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 "bounceTestUtils.h" + +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/GatherScatterKernel.h" + +#include + +#include + +#include +#include +#include +#include + +namespace b = tensorrt_llm::executor::kv_cache::bounce; + +#define CUDA_OK(call) ASSERT_EQ((call), cudaSuccess) << "CUDA call failed: " #call + +using bounce_test::alignUp; +using bounce_test::hasCuda; + +namespace +{ +// Distinct per-(buffer,byte) pattern so any mis-routing/overlap is caught. +unsigned char pattern(std::size_t buf, std::size_t idx) +{ + return static_cast((buf * 131 + idx * 7 + 13) & 0xFF); +} +} // namespace + +TEST(GatherScatterKernel, ZeroBuffersIsNoop) +{ + if (!hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; + } + cudaStream_t stream{}; + CUDA_OK(cudaStreamCreate(&stream)); + EXPECT_EQ(b::launchBatchedCopy(nullptr, nullptr, nullptr, 0, stream), cudaSuccess); + CUDA_OK(cudaStreamSynchronize(stream)); + cudaStreamDestroy(stream); +} + +namespace +{ +// Round-trip driver shared by both plan-placement tests: gather scattered source buffers into one +// packed arena region, then scatter that region into fresh dst buffers, asserting byte-exact. +// Mixes 16B-aligned sizes (uint4 path) and unaligned sizes (byte path). `mappedPlan` puts the +// [srcs|dsts|sizes] plan arrays in MAPPED host memory (the zero-copy-args path) instead of device +// memory. +void runRoundTrip(bool mappedPlan) +{ + if (!hasCuda()) + { + GTEST_SKIP() << "no CUDA device"; + } + std::vector sizes{64, 100, 256, 16, 4096, 17, 1, 32, 3, 512}; + auto const n = static_cast(sizes.size()); + std::vector off(n); + std::uint64_t cur = 0; + for (std::uint32_t i = 0; i < n; ++i) + { + off[i] = cur; + cur = alignUp(cur + sizes[i], 256); + } + std::uint64_t const total = cur; + std::vector srcHost(total, 0); + for (std::uint32_t i = 0; i < n; ++i) + for (std::uint32_t j = 0; j < sizes[i]; ++j) + srcHost[off[i] + j] = pattern(i, j); + + void *dSrc = nullptr, *dPackedRegion = nullptr, *dDst = nullptr; + CUDA_OK(cudaMalloc(&dSrc, total)); + CUDA_OK(cudaMalloc(&dPackedRegion, total)); + CUDA_OK(cudaMalloc(&dDst, total)); + CUDA_OK(cudaMemcpy(dSrc, srcHost.data(), total, cudaMemcpyHostToDevice)); + CUDA_OK(cudaMemset(dDst, 0, total)); + auto base = [](void* p, std::uint64_t o) { return reinterpret_cast(static_cast(p) + o); }; + std::vector gSrc(n), gDst(n), sSrc(n), sDst(n); + for (std::uint32_t i = 0; i < n; ++i) + { + gSrc[i] = base(dSrc, off[i]); + gDst[i] = base(dPackedRegion, off[i]); + sSrc[i] = base(dPackedRegion, off[i]); + sDst[i] = base(dDst, off[i]); + } + + // Lay a plan array on device (cudaMalloc + H2D) or in mapped host (zero-copy alias). + std::vector devBufs; // device-path allocations, freed below + std::vector mappedHosts; // mapped-path host allocations, freed below + auto place = [&](void const* data, std::size_t bytes) -> void* + { + if (mappedPlan) + { + void* hp = nullptr; + EXPECT_EQ(cudaHostAlloc(&hp, bytes, cudaHostAllocMapped), cudaSuccess); + std::memcpy(hp, data, bytes); + mappedHosts.push_back(hp); + void* dp = nullptr; + EXPECT_EQ(cudaHostGetDevicePointer(&dp, hp, 0), cudaSuccess); // device-accessible alias + return dp; + } + void* d = nullptr; + EXPECT_EQ(cudaMalloc(&d, bytes), cudaSuccess); + EXPECT_EQ(cudaMemcpy(d, data, bytes, cudaMemcpyHostToDevice), cudaSuccess); + devBufs.push_back(d); + return d; + }; + auto* dgSrc = static_cast(place(gSrc.data(), n * sizeof(std::uint64_t))); + auto* dgDst = static_cast(place(gDst.data(), n * sizeof(std::uint64_t))); + auto* dsSrc = static_cast(place(sSrc.data(), n * sizeof(std::uint64_t))); + auto* dsDst = static_cast(place(sDst.data(), n * sizeof(std::uint64_t))); + auto* dSizes = static_cast(place(sizes.data(), n * sizeof(std::uint32_t))); + + cudaStream_t stream{}; + CUDA_OK(cudaStreamCreate(&stream)); + CUDA_OK(b::launchBatchedCopy(dgSrc, dgDst, dSizes, n, stream)); // gather + CUDA_OK(b::launchBatchedCopy(dsSrc, dsDst, dSizes, n, stream)); // scatter + CUDA_OK(cudaStreamSynchronize(stream)); + + std::vector dstHost(total, 0xEE); + CUDA_OK(cudaMemcpy(dstHost.data(), dDst, total, cudaMemcpyDeviceToHost)); + for (std::uint32_t i = 0; i < n; ++i) + for (std::uint32_t j = 0; j < sizes[i]; ++j) + ASSERT_EQ(dstHost[off[i] + j], pattern(i, j)) << "mismatch buf=" << i << " byte=" << j; + + cudaFree(dSrc); + cudaFree(dPackedRegion); + cudaFree(dDst); + for (void* p : devBufs) + cudaFree(p); + for (void* p : mappedHosts) + cudaFreeHost(p); + cudaStreamDestroy(stream); +} +} // namespace + +// Plan arrays staged in device memory (TRTLLM_NIXL_BOUNCE_USE_ZERO_COPY_ARGUMENTS=0). +TEST(GatherScatterKernel, GatherThenScatterRoundTrip) +{ + runRoundTrip(/*mappedPlan=*/false); +} + +// Zero-copy plan args: kernel reads [srcs|dsts|sizes] from mapped host +// (TRTLLM_NIXL_BOUNCE_USE_ZERO_COPY_ARGUMENTS=1, the default). +TEST(GatherScatterKernel, ZeroCopyPlanArgsRoundTrip) +{ + runRoundTrip(/*mappedPlan=*/true); +} diff --git a/cpp/tests/unit_tests/executor/bounce/zmqControlChannelTest.cpp b/cpp/tests/unit_tests/executor/bounce/zmqControlChannelTest.cpp new file mode 100644 index 000000000000..953b7db21c20 --- /dev/null +++ b/cpp/tests/unit_tests/executor/bounce/zmqControlChannelTest.cpp @@ -0,0 +1,190 @@ +/* + * 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/ZmqControlChannel.h" +#include "tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceMessage.h" + +#include + +#include +#include +#include +#include +#include + +namespace b = tensorrt_llm::executor::kv_cache::bounce; + +namespace +{ +// recv with a few retries to absorb zmq connection-establishment latency. +bool recvRetry(b::ZmqControlChannel& ch, std::string& peer, std::string& blob, int totalMs) +{ + int waited = 0; + while (waited < totalMs) + { + if (ch.recv(peer, blob, 100)) + { + return true; + } + waited += 100; + } + return false; +} +} // namespace + +TEST(ZmqControlChannel, BidirectionalMessagesOverTcp) +{ + b::ZmqControlChannel a("agentA"); + b::ZmqControlChannel bch("agentB"); + a.addPeer("agentB", bch.localEndpoint()); + bch.addPeer("agentA", a.localEndpoint()); + + // A -> B : WANT(rid=42, chunk sizes [4096, 8192, 1024]) + a.sendTo("agentB", b::encodeWant(42, std::vector{4096, 8192, 1024}, a.localEndpoint())); + std::string from, blob; + ASSERT_TRUE(recvRetry(bch, from, blob, 4000)) << "B never received A's WANT"; + EXPECT_EQ(from, "agentA"); + b::BounceMsgHeader h{}; + ASSERT_TRUE(b::decodeHeader(blob, h)); + EXPECT_EQ(h.msgType, static_cast(b::BounceMsgType::kWANT)); + EXPECT_EQ(h.requestId, 42u); + std::vector sizes; + std::string ep; + ASSERT_TRUE(b::decodeWant(blob, h, sizes, ep)); + EXPECT_EQ(sizes, (std::vector{4096, 8192, 1024})); + EXPECT_EQ(ep, a.localEndpoint()); + + // B -> A : GRANT(rid=42, {region handle 3 @ 0x3000, len 4096}) -- {addr, len, devId, regionHandle} + bch.sendTo("agentA", b::encodeGrant(42, {{0x3000, 4096, 0, 3}})); + ASSERT_TRUE(recvRetry(a, from, blob, 4000)) << "A never received B's GRANT"; + EXPECT_EQ(from, "agentB"); + ASSERT_TRUE(b::decodeHeader(blob, h)); + EXPECT_EQ(h.msgType, static_cast(b::BounceMsgType::kGRANT)); + std::vector credits; + ASSERT_TRUE(b::decodeCredits(blob, h, credits)); + ASSERT_EQ(credits.size(), 1u); + EXPECT_EQ(credits[0].regionHandle, 3u); + EXPECT_EQ(credits[0].addr, 0x3000u); +} + +TEST(ZmqControlChannel, RecvTimesOutWhenIdle) +{ + b::ZmqControlChannel a("idleA"); + std::string peer, blob; + EXPECT_FALSE(a.recv(peer, blob, 50)); // nothing sent -> timeout, no spurious message +} + +TEST(ZmqControlChannel, InvalidEndpointsThrow) +{ + b::ZmqControlChannel a("invalidEndpointA"); + EXPECT_ANY_THROW(a.addPeer("emptyPeer", "")); + EXPECT_ANY_THROW(a.addPeer("malformedPeer", "not-a-zmq-endpoint")); + + b::ZmqControlChannel valid("validEndpointB"); + EXPECT_TRUE(a.addPeer("validPeer", valid.localEndpoint())); +} + +TEST(ZmqControlChannel, ManyMessagesPreserveOrderAndContent) +{ + b::ZmqControlChannel a("ordA"); + b::ZmqControlChannel bch("ordB"); + a.addPeer("ordB", bch.localEndpoint()); + + constexpr int kN = 50; + for (int i = 0; i < kN; ++i) + { + a.sendTo("ordB", b::encodeAck(/*rid=*/static_cast(i), /*chunk=*/i, /*regionHandle=*/i)); + } + // Per-peer FIFO: messages from a single DEALER arrive in order. + for (int i = 0; i < kN; ++i) + { + std::string from, blob; + ASSERT_TRUE(recvRetry(bch, from, blob, 4000)) << "missing msg " << i; + b::BounceMsgHeader h{}; + ASSERT_TRUE(b::decodeHeader(blob, h)); + EXPECT_EQ(h.requestId, static_cast(i)); + EXPECT_EQ(h.chunkIdx, static_cast(i)); + } +} + +TEST(ZmqControlChannel, SendToDoesNotBlockWhenPeerQueueFull) +{ + // sendTo() runs on the reactor IO thread (and under the dealer mutex that also gates submit()), + // so it must NEVER block. Blast far more messages than any HWM/TCP buffer can hold to a peer that + // never reads: a blocking send would wedge the caller forever (the reactor stall the design + // forbids); the non-blocking send must drop and let the loop finish promptly. + auto a = std::make_shared("floodA"); + b::ZmqControlChannel bch("floodB"); // bound ROUTER, but we deliberately never recv() on it + a->addPeer("floodB", bch.localEndpoint()); + + std::promise done; + auto fut = done.get_future(); + // Capture `a` by value (shared_ptr) so that if the bug is present and this thread stays wedged in + // sendTo when we give up and detach below, the channel outlives this scope (no use-after-free). + std::thread flood( + [a, &done] + { + for (int i = 0; i < 200000; ++i) + { + a->sendTo("floodB", b::encodeAck(static_cast(i), i, i)); + } + done.set_value(); + }); + + bool const finished = fut.wait_for(std::chrono::seconds(10)) == std::future_status::ready; + EXPECT_TRUE(finished) << "sendTo blocked when the peer's queue filled (it must drop, not block)"; + if (finished) + { + flood.join(); + } + else + { + flood.detach(); // bug present: leave the wedged thread; the shared_ptr keeps `a` alive + } +} + +TEST(ZmqControlChannel, BidirectionalMessagesOverIPv6) +{ + // Binding an IPv6 endpoint requires ZMQ_IPV6 on the ROUTER, and connecting to one requires it on + // the DEALER (both off by default). Bind the IPv6 loopback; skip cleanly where IPv6 is unavailable. + std::unique_ptr a; + std::unique_ptr bch; + try + { + a = std::make_unique("v6A", "tcp://[::1]:*"); + bch = std::make_unique("v6B", "tcp://[::1]:*"); + } + catch (zmq::error_t const& e) + { + GTEST_SKIP() << "IPv6 unavailable: " << e.what(); + } + // The bound endpoint must be IPv6 (bracketed) — proves the IPv6 bind actually took. + ASSERT_NE(a->localEndpoint().find('['), std::string::npos) << "expected IPv6 endpoint, got " << a->localEndpoint(); + a->addPeer("v6B", bch->localEndpoint()); + + a->sendTo("v6B", b::encodeWant(7, std::vector{2048}, a->localEndpoint())); + std::string from, blob; + ASSERT_TRUE(recvRetry(*bch, from, blob, 4000)) << "B never received A's WANT over IPv6"; + EXPECT_EQ(from, "v6A"); + b::BounceMsgHeader h{}; + ASSERT_TRUE(b::decodeHeader(blob, h)); + std::vector sizes; + std::string ep; + ASSERT_TRUE(b::decodeWant(blob, h, sizes, ep)); + EXPECT_EQ(sizes, (std::vector{2048})); + EXPECT_EQ(ep, a->localEndpoint()); +} diff --git a/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp b/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp index e3c0071756d7..8da2e3cd3c8d 100644 --- a/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp +++ b/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp @@ -1081,7 +1081,9 @@ TEST(SerializeUtilsTest, MethodReturnType) TEST(SerializeUtilsTest, CacheTransceiverConfig) { texec::CacheTransceiverConfig cacheTransceiverConfig( - tensorrt_llm::executor::CacheTransceiverConfig::BackendType::UCX, 1024, 100, 1000); + tensorrt_llm::executor::CacheTransceiverConfig::BackendType::UCX, 1024, 100, 1000, + texec::CacheTransceiverConfig::kDefaultKvTransferPollIntervalMs, 512, + std::map{{"max_chunk_size", "32MB"}, {"copy_stream_count", "4"}}); auto cacheTransceiverConfig2 = serializeDeserialize(cacheTransceiverConfig); EXPECT_EQ(cacheTransceiverConfig.getBackendType(), cacheTransceiverConfig2.getBackendType()); EXPECT_EQ(cacheTransceiverConfig.getMaxTokensInBuffer(), cacheTransceiverConfig2.getMaxTokensInBuffer()); @@ -1090,6 +1092,8 @@ TEST(SerializeUtilsTest, CacheTransceiverConfig) cacheTransceiverConfig2.getKvTransferSenderFutureTimeoutMs()); EXPECT_EQ( cacheTransceiverConfig.getKvTransferPollIntervalMs(), cacheTransceiverConfig2.getKvTransferPollIntervalMs()); + EXPECT_EQ(cacheTransceiverConfig.getAgentBufferSizeMb(), cacheTransceiverConfig2.getAgentBufferSizeMb()); + EXPECT_EQ(cacheTransceiverConfig.getAgentBounceParams(), cacheTransceiverConfig2.getAgentBounceParams()); } TEST(SerializeUtilsTest, BlockKeyBasic) diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 21632f9358e2..68eb74c9890c 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -23,7 +23,7 @@ import weakref from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, List, Optional, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Union import msgpack import numpy as np @@ -80,8 +80,10 @@ AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType LlmRequestType = tensorrt_llm.bindings.internal.batch_manager.LlmRequestType -# Number of worker threads for KV transfer queues (default: 1) -KV_TRANSFER_NUM_THREADS = int(os.environ.get("TRTLLM_KV_TRANSFER_NUM_THREADS", "1")) +# Number of worker threads for KV transfer queues (default: 4). Each worker owns one FIFO and +# submits transfers synchronously (blocking wait per slice), so this value is the per-rank cap on +# concurrent in-flight KV transfer slices. +KV_TRANSFER_NUM_THREADS = int(os.environ.get("TRTLLM_KV_TRANSFER_NUM_THREADS", "4")) # Keep standalone TxSession waits responsive to cancellation even when callers # do not configure a sender-future wait slice. @@ -2425,7 +2427,13 @@ def __exit__(self, _exc_type, _exc_val, _exc_tb): self.shutdown() -def _create_nixl_agent(name: str, rank: int, world_size: int) -> NixlTransferAgent: +def _create_nixl_agent( + name: str, + rank: int, + world_size: int, + agent_buffer_size_mb: int = 0, + agent_bounce_params: Optional[Dict[str, str]] = None, +) -> NixlTransferAgent: num_threads = int(os.environ.get("TRTLLM_NIXL_NUM_THREADS", "8")) kwargs = {} if "TRTLLM_NIXL_SPLIT_BATCH_SIZE" in os.environ: @@ -2436,6 +2444,8 @@ def _create_nixl_agent(name: str, rank: int, world_size: int) -> NixlTransferAge num_threads=num_threads, rank=rank, world_size=world_size, + agent_buffer_size_mb=agent_buffer_size_mb, + agent_bounce_params=agent_bounce_params, **kwargs, ) @@ -2466,6 +2476,15 @@ class TransferWorkerConfig: rx_timeout_s: Optional[float] = None bounce: Optional["Config"] = None tx_overall_timeout_s: Optional[float] = None + # Transfer-agent staging-buffer (bounce v2) arena size in MiB; 0 disables the + # fast path. Derived upstream from CacheTransceiverConfig: it equals + # kv_cache_bounce_size_mb when agent_bounce_buffer_enable is set (in which + # case `bounce` above is off) and 0 otherwise. + agent_buffer_size_mb: int = 0 + # Expert bounce knobs (TRTLLM_NIXL_BOUNCE_* names without the prefix and the + # trailing _BYTES, lowercased); dict > env > default. Ignored when + # agent_buffer_size_mb == 0. + agent_bounce_params: Optional[Dict[str, str]] = None class TransferWorker: @@ -2536,6 +2555,8 @@ def _setup_transfer_engine(self): self._rank_info.instance_name + str(self._rank_info.instance_rank), rank=mapping.rank, world_size=mapping.world_size, + agent_buffer_size_mb=self._config.agent_buffer_size_mb, + agent_bounce_params=self._config.agent_bounce_params, ) self._registered_mem: list = [] try: diff --git a/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py b/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py index d2290d311c10..65a71a63dbd2 100644 --- a/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py +++ b/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py @@ -84,6 +84,8 @@ def __init__( enable_telemetry: bool = False, rank: int | None = None, world_size: int | None = None, + agent_buffer_size_mb: int = 0, + agent_bounce_params: dict[str, str] | None = None, **kwargs, ): backend_params = kwargs @@ -107,6 +109,8 @@ def __init__( "(expected a positive integer)" ) + # Bounce config stays OUT of backend_params: backend_params is forwarded verbatim to the + # NIXL backend plugin, while the bounce knobs ride their own BaseAgentConfig fields. config = BaseAgentConfig( name, use_prog_thread, @@ -116,6 +120,8 @@ def __init__( backend_params=backend_params, rank=rank, world_size=world_size, + agent_buffer_size_mb=agent_buffer_size_mb, + bounce_params=agent_bounce_params or {}, ) self._cpp_agent = CppNixlTransferAgent(config) self.name = name @@ -123,6 +129,8 @@ def __init__( f"BindingsNixlTransferAgent init: agent={name} " f"use_prog_thread={use_prog_thread} num_threads={num_threads} " f"enable_telemetry={enable_telemetry} backend_params={backend_params} " + f"agent_buffer_size_mb={agent_buffer_size_mb} " + f"agent_bounce_params={agent_bounce_params or {}} " f"UCX_TLS={os.environ.get('UCX_TLS', '')} " f"UCX_LOG_LEVEL={os.environ.get('UCX_LOG_LEVEL', '')}" ) @@ -160,6 +168,16 @@ def load_remote_agent_by_connection(self, name: str, connection_info: str): """Load a remote agent by connection info.""" self._cpp_agent.load_remote_agent_by_connection(name, connection_info) + @property + def bounce_enabled(self) -> bool: + """Whether the NIXL bounce v2 fast path is active on this agent.""" + return self._cpp_agent.bounce_enabled + + @property + def bounce_submit_count(self) -> int: + """Number of transfer requests routed to the bounce fast path so far.""" + return self._cpp_agent.bounce_submit_count + def get_local_agent_desc(self) -> bytes: """Get the local agent descriptor as bytes.""" agent_desc = self._cpp_agent.get_local_agent_desc() diff --git a/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py b/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py index dee500b8815c..fa48f5f7ba44 100644 --- a/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py +++ b/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py @@ -91,6 +91,8 @@ def __init__( num_threads: int = 1, rank: int | None = None, world_size: int | None = None, + agent_buffer_size_mb: int = 0, + agent_bounce_params: dict[str, str] | None = None, **kwargs, ): """ @@ -100,7 +102,18 @@ def __init__( :param num_workers: Specify number of threads for the supported multi-threaded backends. :param rank: Process rank, used only to keep the shared agent interface consistent. :param world_size: Process count, used only to keep the shared agent interface consistent. + :param agent_buffer_size_mb: Accepted for interface consistency; the Python nixl-library + agent has no staging-buffer (bounce) support, so a positive size is ignored with a + warning. + :param agent_bounce_params: Accepted for interface consistency; ignored alongside + agent_buffer_size_mb. """ + if agent_buffer_size_mb > 0: + logger.warning( + "The transfer-agent bounce buffer (agent_bounce_buffer_enable) is not " + "supported by the Python nixl-library agent; ignoring the arena size " + "(and any agent_bounce_params)" + ) if (rank is None) != (world_size is None): raise ValueError("rank and world_size must be specified together") if rank is not None: diff --git a/tensorrt_llm/_torch/disaggregation/nixl/bounce_knobs.py b/tensorrt_llm/_torch/disaggregation/nixl/bounce_knobs.py new file mode 100644 index 000000000000..105b93dd3cf2 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/nixl/bounce_knobs.py @@ -0,0 +1,34 @@ +# 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. +"""Valid agent_bounce_params keys. Must stay dependency-free: imported from the +llm_args validator, where the transfer-agent bindings may not be loadable.""" + +# TRTLLM_NIXL_BOUNCE_* env-var names without the prefix and trailing _BYTES, lowercased. +# KEEP IN SYNC with kEnvKnobs in cpp/.../nixl_utils/bounce/BounceConfig.h (sync-tested). +AGENT_BOUNCE_PARAM_KEYS: frozenset = frozenset( + { + "max_chunk_size", + "arena_allocation_granularity", + "max_inflight_chunks_per_request", + "copy_stream_count", + "scatter_worker_count", + "min_descriptor_count", + "max_average_descriptor_size", + "request_timeout_ms", + "disable_fabric_memory", + "enable_eager_gather", + "use_zero_copy_arguments", + } +) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index be3acc832558..817f8e2cbedc 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -126,10 +126,23 @@ def __init__( tx_timeout_s=sender_wait_slice_s, tx_overall_timeout_s=transfer_timeout_s, rx_timeout_s=transfer_timeout_s, - # Size 0 turns bounce off; the per-transfer size gates are internal (tuned via - # env: TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS for plain-KV payloads, - # TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES for recurrent-state payloads). - bounce=bounce_config_from_size(cache_transceiver_config.kv_cache_bounce_size_mb), + # kv_cache_bounce_size_mb is the shared bounce capacity; agent_bounce_buffer_enable + # routes it to exactly one implementation: the Python bounce below (per-region, + # size 0 = off; the per-transfer size gates are internal, tuned via env: + # TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS for plain-KV payloads, + # TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES for recurrent-state payloads) or the C++ + # transfer-agent staging arena (bounce v2, one shared arena). + bounce=bounce_config_from_size( + 0 + if cache_transceiver_config.agent_bounce_buffer_enable + else cache_transceiver_config.kv_cache_bounce_size_mb + ), + agent_buffer_size_mb=( + cache_transceiver_config.kv_cache_bounce_size_mb + if cache_transceiver_config.agent_bounce_buffer_enable + else 0 + ), + agent_bounce_params=cache_transceiver_config.agent_bounce_params, ) ) logger.info( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index b54068d8e33e..ce4d8ce1e045 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3283,11 +3283,14 @@ def maybe_to_pybind(ins): return ins @staticmethod - def mirror_pybind_fields(pybind_class): + def mirror_pybind_fields(pybind_class, excluded_fields=None): """Class decorator that ensures Python class fields mirror those of a C++ class. Args: pybind_class: The C++ class whose fields should be mirrored + excluded_fields: Optional set of C++ field names exempt from the + mirror check, for internal pipeline values that the Python + class derives from other fields instead of exposing directly Returns: A decorator function that validates field mirroring @@ -3301,6 +3304,8 @@ def decorator(cls): # Check if all C++ fields exist in the Python class for field in cpp_fields: + if excluded_fields and field in excluded_fields: + continue if field not in python_fields: raise ValueError( f"Field {field} is not mirrored in Python class {cls.__name__} from C++ class {pybind_class.__name__}. Please update the class." @@ -4294,7 +4299,11 @@ def _to_pybind(self): ) -@PybindMirror.mirror_pybind_fields(_CacheTransceiverConfig) +# agent_buffer_size_mb is exempt from the mirror check: the C++ config keeps it +# as the internal pipeline value (agent arena MiB), which _to_pybind derives +# from kv_cache_bounce_size_mb + agent_bounce_buffer_enable. +@PybindMirror.mirror_pybind_fields(_CacheTransceiverConfig, + excluded_fields={"agent_buffer_size_mb"}) class CacheTransceiverConfig(StrictBaseModel, PybindMirror): """Configuration for the cache transceiver.""" @@ -4345,8 +4354,69 @@ class CacheTransceiverConfig(StrictBaseModel, PybindMirror): default=0, ge=0, description= - "Per-region size in MiB of the native-disagg KV-cache bounce buffer (one for send, one for recv). Bounce coalesces a request's scattered per-block KV into one contiguous fabric-VMM buffer and issues a single multi-rail NIXL write. The size doubles as the on/off switch: 0 (default) keeps the per-block path, >0 enables bounce at that capacity. Only used by the Python (v2) transceiver." - ) + "Capacity in MiB of the native-disagg KV-cache bounce buffer, which " + "coalesces a request's scattered per-block KV for a single multi-rail " + "NIXL write. The size doubles as the on/off switch: 0 (default) keeps " + "the per-block path, >0 enables bounce at that capacity. With the " + "default Python implementation it is per-region (one buffer for send, " + "one for recv); with agent_bounce_buffer_enable it is one shared " + "arena. The Python implementation requires the Python (v2) " + "transceiver runtime; on the C++ transceiver path the size is ignored " + "unless agent_bounce_buffer_enable is set.") + + agent_bounce_buffer_enable: bool = Field( + default=False, + description= + "Use the C++ transfer-agent bounce implementation instead of the " + "Python one. kv_cache_bounce_size_mb is then a single arena shared by " + "send and recv (about half the memory footprint of the Python " + "per-region pair for the same number). Configure the context and " + "generation instances consistently.") + + agent_bounce_params: Optional[Dict[str, str]] = Field( + default=None, + description= + "Expert tuning knobs for the C++ transfer-agent bounce pipeline; see " + "tensorrt_llm/_torch/disaggregation/nixl/bounce_knobs.py for the valid " + "keys. Byte-valued knobs accept a KB/MB/GB suffix (e.g. '32MB'). " + "Precedence: this dict > environment variable > built-in default. " + "Requires agent_bounce_buffer_enable.") + + @field_validator('agent_bounce_params', mode='before') + @classmethod + def coerce_agent_bounce_params_to_str(cls, v): + """Coerce agent_bounce_params values to strings (YAML often yields ints/bools).""" + if v is None: + return v + if not isinstance(v, dict): + raise ValueError( + f"agent_bounce_params must be a dict of str to str, got " + f"{type(v).__name__}") + return {str(k): str(val) for k, val in v.items()} + + @model_validator(mode='after') + def validate_bounce_config(self) -> 'CacheTransceiverConfig': + if self.agent_bounce_buffer_enable and self.kv_cache_bounce_size_mb == 0: + raise ValueError( + "agent_bounce_buffer_enable selects the C++ transfer-agent " + "bounce implementation, but kv_cache_bounce_size_mb is 0 " + "(bounce disabled); set a positive capacity.") + if self.agent_bounce_params: + if not self.agent_bounce_buffer_enable: + raise ValueError( + "agent_bounce_params only applies to the C++ transfer-agent " + "bounce implementation; set agent_bounce_buffer_enable=True " + "or drop the params.") + # Lazy: importing tensorrt_llm._torch at module scope is circular. + from tensorrt_llm._torch.disaggregation.nixl.bounce_knobs import \ + AGENT_BOUNCE_PARAM_KEYS + unknown = sorted( + set(self.agent_bounce_params) - AGENT_BOUNCE_PARAM_KEYS) + if unknown: + raise ValueError( + f"Unknown agent_bounce_params key(s) {unknown}; valid keys " + f"are {sorted(AGENT_BOUNCE_PARAM_KEYS)}.") + return self def _resolve_default_backend(self) -> Tuple[Optional[str], Optional[str]]: """Effective backend after resolving "DEFAULT" against legacy env vars. @@ -4369,7 +4439,11 @@ def _to_pybind(self): kv_transfer_timeout_ms=self.kv_transfer_timeout_ms, kv_transfer_sender_future_timeout_ms=self. kv_transfer_sender_future_timeout_ms, - kv_transfer_poll_interval_ms=self.kv_transfer_poll_interval_ms) + kv_transfer_poll_interval_ms=self.kv_transfer_poll_interval_ms, + # The C++ side keeps a single knob: the agent arena MiB (0 = off). + agent_buffer_size_mb=(self.kv_cache_bounce_size_mb + if self.agent_bounce_buffer_enable else 0), + agent_bounce_params=self.agent_bounce_params or {}) @dataclass diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 1dfc185d4a6e..8410570b459a 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -146,6 +146,13 @@ "kind": "value", "path": "batch_wait_timeout_ms" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "cache_transceiver_config.agent_bounce_buffer_enable" + }, { "allowed_values": [ "DEFAULT", diff --git a/tests/integration/defs/cpp/test_unit_tests.py b/tests/integration/defs/cpp/test_unit_tests.py index 730ebbf389ed..d169374e0493 100644 --- a/tests/integration/defs/cpp/test_unit_tests.py +++ b/tests/integration/defs/cpp/test_unit_tests.py @@ -3,23 +3,38 @@ import defs.cpp.cpp_common as _cpp import pytest +_TEST_GROUP_DIRS = { + "executor_bounce": "executor/bounce", +} + @pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"], indirect=True) @pytest.mark.parametrize("test_group", [ - "batch_manager", "common", "executor", "kernels", "layers", "runtime", - "thop" + "batch_manager", "common", "executor", "executor_bounce", "kernels", + "layers", "runtime", "thop" ]) def test_unit_tests(build_google_tests, test_group, build_dir, lora_setup): xml_name = f"results-unit-tests-{test_group}.xml" + test_group_dir = _TEST_GROUP_DIRS.get(test_group, test_group) + + if test_group == "executor_bounce": + # Real-NIXL bounce tests are conditionally built when NIXL and ZMQ are + # available. Require one here so pre-merge cannot pass with only the + # dependency-free subset after silently losing its E2E coverage. + required_test = (build_dir / "tests/unit_tests/executor/bounce" / + "bounceAgentE2ETest") + if not required_test.is_file(): + pytest.fail( + f"Required NIXL bounce E2E test was not built: {required_test}") # Discover and run the actual gtests ctest_command = [ "ctest", "--output-on-failure", "--test-dir", - f"{build_dir}/tests/unit_tests/{test_group}", + f"{build_dir}/tests/unit_tests/{test_group_dir}", "--output-junit", f"{build_dir}/{xml_name}", ] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 79d77502baf4..e23fcc9285c6 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -109,6 +109,8 @@ l0_h100: # DSA indexer-K side cache (V1, REPLICATED). - unittest/disaggregated/test_cache_transceiver_single_process.py::test_cache_transceiver_v1_dsa_indexer - unittest/disaggregated/test_cache_transceiver_single_process.py::test_cache_transceiver_v1_masked_dsa_indexer_across_asymmetric_pp + # Python transceiver with the C++ NIXL agent's bounce v2 enabled. + - unittest/disaggregated/test_cache_transceiver_single_process.py::test_python_nixl_cache_transceiver_uses_cpp_bounce - unittest/disaggregated/test_cache_transceiver_harness_report.py - unittest/disaggregated/test_cache_transceiver_harness.py - unittest/disaggregated/test_cache_transceiver_precheck_e2e.py @@ -287,6 +289,7 @@ l0_h100: tests: # ------------- CPP tests --------------- - cpp/test_unit_tests.py::test_unit_tests[common-90] + - cpp/test_unit_tests.py::test_unit_tests[executor_bounce-90] - cpp/test_unit_tests.py::test_unit_tests[kernels-90] - cpp/test_unit_tests.py::test_unit_tests[layers-90] - cpp/test_unit_tests.py::test_unit_tests[thop-90] diff --git a/tests/unittest/bindings/test_executor_bindings.py b/tests/unittest/bindings/test_executor_bindings.py index 70bc6f53e251..7bd057de1abd 100644 --- a/tests/unittest/bindings/test_executor_bindings.py +++ b/tests/unittest/bindings/test_executor_bindings.py @@ -1388,11 +1388,36 @@ def test_cache_transceiver_config_pickle(): config = trtllm.CacheTransceiverConfig( backend=trtllm.CacheTransceiverBackendType.UCX, max_tokens_in_buffer=1024, - kv_transfer_poll_interval_ms=5000) + kv_transfer_poll_interval_ms=5000, + agent_buffer_size_mb=512, + agent_bounce_params={ + "max_chunk_size": "32MB", + "copy_stream_count": "4" + }) config_copy = pickle.loads(pickle.dumps(config)) assert config_copy.backend == config.backend assert config_copy.max_tokens_in_buffer == config.max_tokens_in_buffer assert config_copy.kv_transfer_poll_interval_ms == config.kv_transfer_poll_interval_ms + assert config_copy.agent_buffer_size_mb == 512 + assert config_copy.agent_bounce_params == { + "max_chunk_size": "32MB", + "copy_stream_count": "4" + } + # Defaults roundtrip too (size 0 = bounce off, empty params). + default_copy = pickle.loads(pickle.dumps(trtllm.CacheTransceiverConfig())) + assert default_copy.agent_buffer_size_mb == 0 + assert default_copy.agent_bounce_params == {} + + # Legacy 6-tuple states carried the retired Optional[bool] agent_buffer_enable at index 5; + # any such value (True/False/None) must deserialize as "bounce off" (size 0) — in particular + # a legacy True must NOT be number-converted into a 1 MiB arena. + for legacy_agent_buffer_enable in (True, False, None): + legacy = trtllm.CacheTransceiverConfig.__new__( + trtllm.CacheTransceiverConfig) + legacy.__setstate__((trtllm.CacheTransceiverBackendType.UCX, 1024, 100, + 1000, 5000, legacy_agent_buffer_enable)) + assert legacy.agent_buffer_size_mb == 0, legacy_agent_buffer_enable + assert legacy.agent_bounce_params == {} @pytest.mark.cpu_only diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index 418ea4352634..b13490a69aec 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py @@ -21,6 +21,8 @@ """ import os +import subprocess +import sys import threading import uuid from types import SimpleNamespace @@ -981,6 +983,7 @@ def run_transfer_test( request_lengths: Optional[List[int]] = None, enable_indexer_k_cache: bool = False, indexer_k_cache_layer_mask: list[bool] | None = None, + expect_cpp_bounce: bool = False, ) -> None: """Run a full KV transfer test using KvCacheTransceiverV2.""" if request_lengths is None: @@ -1024,6 +1027,18 @@ def run_transfer_test( backend="NIXL", transceiver_runtime="PYTHON", max_tokens_in_buffer=512, + # Keep the Python-native bounce layer disabled by default (size 0). + # Dedicated C++ bounce coverage enables the NIXL agent's bounce v2 via + # agent_bounce_buffer_enable + kv_cache_bounce_size_mb plus + # agent_bounce_params, which exercises the config plumbing end-to-end. + kv_cache_bounce_size_mb=64 if expect_cpp_bounce else 0, + agent_bounce_buffer_enable=expect_cpp_bounce, + agent_bounce_params={ + "min_descriptor_count": "1", + "max_average_descriptor_size": "1MB", + } + if expect_cpp_bounce + else None, ) ctx_tcs = create_instance_transceivers( ctx_tp, ctx_pp, ctx_enable_dp, ctx_managers, config, is_mla @@ -1157,6 +1172,22 @@ def run_transfer_test( num_layers=num_layers, ) + # Programmatic bounce check (no log parsing): every agent must have the C++ + # bounce v2 transport active, and at least one transfer must have been routed + # through it (only KV senders submit WRITEs, so we assert on the total). + if expect_cpp_bounce: + # The shared capacity must be routed to exactly one implementation: with + # agent_bounce_buffer_enable the Python-native bounce stays off. + assert all(tc._transfer_worker._config.bounce is None for tc in ctx_tcs + gen_tcs), ( + "Python-native bounce must stay disabled when the C++ agent bounce is selected" + ) + agents = [tc._transfer_worker._agent for tc in ctx_tcs + gen_tcs] + assert all(getattr(agent, "bounce_enabled", False) for agent in agents), ( + "C++ bounce v2 transport is not active on every NIXL agent" + ) + total_bounce_submits = sum(agent.bounce_submit_count for agent in agents) + assert total_bounce_submits > 0, "no transfer was routed through the bounce fast path" + # 9. Cleanup if use_v2: torch.cuda.current_stream().synchronize() @@ -1547,6 +1578,104 @@ def test_cache_transceiver_v1_masked_dsa_indexer_across_asymmetric_pp() -> None: ) +_NIXL_BOUNCE_SUBPROCESS_ENV = "TRTLLM_TEST_NIXL_BOUNCE_SUBPROCESS" + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize( + ( + "ctx_tp", + "ctx_pp", + "gen_tp", + "gen_pp", + "is_mla", + "use_v2", + "enable_indexer_k_cache", + ), + [ + pytest.param(1, 1, 1, 1, False, True, False, id="v2_mha"), + pytest.param(2, 1, 2, 1, True, True, False, id="v2_mla_tp2"), + pytest.param(1, 1, 1, 1, True, False, True, id="v1_dsa_indexer"), + pytest.param(2, 1, 1, 2, False, True, False, id="v2_ctx_tp2_gen_pp2"), + ], +) +def test_python_nixl_cache_transceiver_uses_cpp_bounce( + ctx_tp: int, + ctx_pp: int, + gen_tp: int, + gen_pp: int, + is_mla: bool, + use_v2: bool, + enable_indexer_k_cache: bool, + request: pytest.FixtureRequest, +) -> None: + """Exercise C++ bounce v2 through representative Python NIXL transceiver paths.""" + # BindingsNixlTransferStatus.last_status_str() resolves this exact attribute on the C++ + # status; if the binding name drifts, failure details silently degrade to "". + from tensorrt_llm.tensorrt_llm_transfer_agent_binding import TransferStatus + + assert hasattr(TransferStatus, "get_last_status_str") + + if os.environ.get(_NIXL_BOUNCE_SUBPROCESS_ENV) == "1": + run_transfer_test( + ctx_tp=ctx_tp, + ctx_pp=ctx_pp, + gen_tp=gen_tp, + gen_pp=gen_pp, + ctx_enable_dp=False, + gen_enable_dp=False, + is_mla=is_mla, + use_v2=use_v2, + request_lengths=[30, 60], + enable_indexer_k_cache=enable_indexer_k_cache, + expect_cpp_bounce=True, + ) + return + + env = os.environ.copy() + for name in tuple(env): + if name.startswith("TRTLLM_NIXL_BOUNCE_"): + del env[name] + env.update( + { + _NIXL_BOUNCE_SUBPROCESS_ENV: "1", + "TRTLLM_USE_PY_NIXL_KVCACHE": "0", + # Bounce is enabled (agent_bounce_buffer_enable + kv_cache_bounce_size_mb) + # and tuned (expert knobs) via CacheTransceiverConfig / agent_bounce_params inside + # run_transfer_test — no TRTLLM_NIXL_BOUNCE_* environment variables — + # covering the config path end-to-end. + # Pin the transport deterministically: the child would otherwise inherit + # developer/CI UCX settings (an injected UCX_NET_DEVICES or differing + # UCX_TLS can fail the transfer or oversubscribe CPUs). + "UCX_TLS": "^ib,gdr_copy", + "TRTLLM_NIXL_NUM_THREADS": "1", + } + ) + # Deliberately NOT setting TLLM_LOG_LEVEL_BY_MODULE: bounce engagement is verified + # programmatically inside the child (agent.bounce_enabled / bounce_submit_count), + # and a non-empty per-module level map can hang the child at exit (static-destruction + # order: CudaMemPool's deleter logs through an already-destroyed Logger module map). + env.pop("TLLM_LOG_LEVEL_BY_MODULE", None) + env.pop("UCX_NET_DEVICES", None) + # Do not override TRTLLM_NIXL_BOUNCE_DISABLE_FABRIC_MEMORY: the test must + # exercise the production default, including automatic cudaMalloc fallback + # on devices without fabric-memory support. + env.pop("TRTLLM_NIXL_BOUNCE_DISABLE_FABRIC_MEMORY", None) + + node_id = f"{__file__}::{request.node.name}" + result = subprocess.run( + [sys.executable, "-m", "pytest", "-s", "-q", node_id], + env=env, + capture_output=True, + text=True, + timeout=180, + ) + output = result.stdout + result.stderr + assert result.returncode == 0, output + # A skipped child exits successfully too; require the selected bounce case to run. + assert "1 passed" in output, f"Child pytest did not run the bounce case:\n{output}" + + if __name__ == "__main__": # Quick smoke test run_transfer_test(1, 1, 1, 1, False, False, False, False) diff --git a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py index 9185d9ef2283..ae2eadd3fc82 100644 --- a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py +++ b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py @@ -696,20 +696,28 @@ def test_tx_session_first_send_anchors_deadline_once(monkeypatch) -> None: session.close() -@pytest.mark.parametrize( - ("transfer_timeout_ms", "sender_wait_ms", "expected_timeout_s", "expected_slice_s"), - [ - (60_000, 1_000, 60.0, 1.0), - (60_000, None, 60.0, None), - ], -) -def test_transceiver_wires_separate_sender_slice_and_overall_timeout( - monkeypatch, - transfer_timeout_ms: Optional[int], - sender_wait_ms: Optional[int], - expected_timeout_s: Optional[float], - expected_slice_s: Optional[float], -) -> None: +def _make_cache_config(**overrides) -> SimpleNamespace: + """SimpleNamespace standing in for CacheTransceiverConfig. + + Every attribute the KvCacheTransceiverV2 constructor reads must be present here. + """ + fields = dict( + kv_transfer_timeout_ms=60_000, + kv_transfer_poll_interval_ms=5_000, + kv_transfer_sender_future_timeout_ms=1_000, + kv_cache_bounce_size_mb=0, + agent_bounce_buffer_enable=False, + agent_bounce_params=None, + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _construct_worker_config(monkeypatch, cache_config) -> TransferWorkerConfig: + """Build a KvCacheTransceiverV2 with all collectives/native setup mocked out. + + Returns the TransferWorkerConfig it wired. + """ worker = SimpleNamespace(page_table=None) worker_constructor = Mock(return_value=worker) monkeypatch.setattr( @@ -720,9 +728,11 @@ def test_transceiver_wires_separate_sender_slice_and_overall_timeout( "tensorrt_llm._torch.disaggregation.transceiver.create_cache_reuse_adapter", Mock(return_value=Mock()), ) + # Echo the routed size: None for 0 (off, mirroring the real helper), a sentinel + # carrying the size otherwise, so tests can assert which implementation got it. monkeypatch.setattr( "tensorrt_llm._torch.disaggregation.transceiver.bounce_config_from_size", - Mock(return_value=None), + Mock(side_effect=lambda size_mb: ("bounce", size_mb) if size_mb > 0 else None), ) monkeypatch.setattr( "tensorrt_llm._torch.disaggregation.transceiver.torch.cuda.current_device", @@ -746,12 +756,6 @@ def test_transceiver_wires_separate_sender_slice_and_overall_timeout( tp_size=1, enable_attention_dp=False, ) - cache_config = SimpleNamespace( - kv_transfer_timeout_ms=transfer_timeout_ms, - kv_transfer_poll_interval_ms=5_000, - kv_transfer_sender_future_timeout_ms=sender_wait_ms, - kv_cache_bounce_size_mb=0, - ) KvCacheTransceiverV2( mapping=mapping, @@ -762,11 +766,64 @@ def test_transceiver_wires_separate_sender_slice_and_overall_timeout( worker_config = worker_constructor.call_args.args[0] assert isinstance(worker_config, TransferWorkerConfig) + return worker_config + + +@pytest.mark.parametrize( + ("transfer_timeout_ms", "sender_wait_ms", "expected_timeout_s", "expected_slice_s"), + [ + (60_000, 1_000, 60.0, 1.0), + (60_000, None, 60.0, None), + ], +) +def test_transceiver_wires_separate_sender_slice_and_overall_timeout( + monkeypatch, + transfer_timeout_ms: Optional[int], + sender_wait_ms: Optional[int], + expected_timeout_s: Optional[float], + expected_slice_s: Optional[float], +) -> None: + worker_config = _construct_worker_config( + monkeypatch, + _make_cache_config( + kv_transfer_timeout_ms=transfer_timeout_ms, + kv_transfer_sender_future_timeout_ms=sender_wait_ms, + ), + ) assert worker_config.tx_timeout_s == expected_slice_s assert worker_config.tx_overall_timeout_s == expected_timeout_s assert worker_config.rx_timeout_s == expected_timeout_s +@pytest.mark.parametrize( + ("bounce_size_mb", "agent_enable", "expected_python_bounce", "expected_agent_size_mb"), + [ + # Shared capacity, Python implementation (default): per-region bounce on, agent off. + (384, False, ("bounce", 384), 0), + # Shared capacity, C++ agent implementation: agent arena gets the size, Python off. + (384, True, None, 384), + # Size 0 keeps both implementations off. + (0, False, None, 0), + ], +) +def test_transceiver_routes_bounce_capacity_to_one_implementation( + monkeypatch, + bounce_size_mb: int, + agent_enable: bool, + expected_python_bounce, + expected_agent_size_mb: int, +) -> None: + worker_config = _construct_worker_config( + monkeypatch, + _make_cache_config( + kv_cache_bounce_size_mb=bounce_size_mb, + agent_bounce_buffer_enable=agent_enable, + ), + ) + assert worker_config.bounce == expected_python_bounce + assert worker_config.agent_buffer_size_mb == expected_agent_size_mb + + def test_transceiver_rejects_unset_transfer_timeout() -> None: cache_config = SimpleNamespace( kv_transfer_timeout_ms=None, diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 529851947843..2f2ed87495aa 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import math +import re import tempfile from collections import defaultdict from dataclasses import is_dataclass @@ -2600,21 +2601,114 @@ def test_cache_transceiver_config_arbitrary_args(self): assert config.max_tokens_in_buffer == 1024 assert config.kv_transfer_poll_interval_ms == 5000 - # The bounce on/off switch defaults to off (0), accepts a positive size, and rejects - # negatives at the Pydantic boundary (ge=0). It is a Python-only field consumed directly by - # the v2 transceiver, so it is intentionally not part of _to_pybind(). + # The shared bounce capacity defaults to off (0), accepts a positive size, and rejects + # negatives at the Pydantic boundary (ge=0). assert config.kv_cache_bounce_size_mb == 0 assert CacheTransceiverConfig( kv_cache_bounce_size_mb=384).kv_cache_bounce_size_mb == 384 with pytest.raises(pydantic_core._pydantic_core.ValidationError): CacheTransceiverConfig(kv_cache_bounce_size_mb=-1) + # agent_bounce_buffer_enable defaults to the Python implementation (False); enabling the + # C++ transfer-agent implementation with a zero capacity is a contradiction. + assert config.agent_bounce_buffer_enable is False + assert config.agent_bounce_params is None + assert CacheTransceiverConfig( + kv_cache_bounce_size_mb=512, + agent_bounce_buffer_enable=True).agent_bounce_buffer_enable is True + with pytest.raises(pydantic_core._pydantic_core.ValidationError, + match="kv_cache_bounce_size_mb is 0"): + CacheTransceiverConfig(agent_bounce_buffer_enable=True) + + # _to_pybind routes the shared capacity to the C++ agent arena only when the agent + # implementation is selected. backend must be set: from_string(None) is a pre-existing + # _to_pybind limit. + assert CacheTransceiverConfig(backend="NIXL", + kv_cache_bounce_size_mb=384, + agent_bounce_buffer_enable=True + )._to_pybind().agent_buffer_size_mb == 384 + assert CacheTransceiverConfig( + backend="NIXL", + kv_cache_bounce_size_mb=384)._to_pybind().agent_buffer_size_mb == 0 + # Arbitrary arguments should be rejected with pytest.raises( pydantic_core._pydantic_core.ValidationError) as exc_info: CacheTransceiverConfig(backend="UCX", invalid_config="should_fail") assert "invalid_config" in str(exc_info.value) + def test_cache_transceiver_config_agent_bounce_params(self): + """agent_bounce_params coercion, key/consistency validation, pybind passthrough.""" + # Values are coerced to strings (YAML often yields ints/bools). + config = CacheTransceiverConfig(kv_cache_bounce_size_mb=512, + agent_bounce_buffer_enable=True, + agent_bounce_params={ + "max_chunk_size": 4096, + "enable_eager_gather": False + }) + assert config.agent_bounce_params == { + "max_chunk_size": "4096", + "enable_eager_gather": "False" + } + + # Params without the C++ agent implementation are a contradiction: they + # would be silently ignored, so validation rejects them outright. + with pytest.raises(pydantic_core._pydantic_core.ValidationError, + match="agent_bounce_buffer_enable"): + CacheTransceiverConfig( + kv_cache_bounce_size_mb=512, + agent_bounce_params={"copy_stream_count": "2"}) + + # Unknown keys (here the env-var-style typo WITH the trailing _bytes) are + # rejected, and the message lists every valid key (spot-check one). + with pytest.raises(pydantic_core._pydantic_core.ValidationError, + match="max_chunk_size_bytes.*min_descriptor_count"): + CacheTransceiverConfig( + kv_cache_bounce_size_mb=512, + agent_bounce_buffer_enable=True, + agent_bounce_params={"max_chunk_size_bytes": "4096"}) + + # Non-dict input fails cleanly in the coercion validator, not with a bare + # AttributeError. + with pytest.raises(pydantic_core._pydantic_core.ValidationError, + match="must be a dict"): + CacheTransceiverConfig(kv_cache_bounce_size_mb=512, + agent_bounce_buffer_enable=True, + agent_bounce_params="max_chunk_size=4096") + + # _to_pybind converts the shared capacity into the C++ agent arena size and + # passes the params through (defaulting to an empty dict). + pybind_config = CacheTransceiverConfig(backend="NIXL", + kv_cache_bounce_size_mb=512, + agent_bounce_buffer_enable=True, + agent_bounce_params={ + "max_chunk_size": "4096" + })._to_pybind() + assert pybind_config.agent_buffer_size_mb == 512 + assert pybind_config.agent_bounce_params == {"max_chunk_size": "4096"} + # backend must be set: from_string(None) is a pre-existing _to_pybind limit. + assert CacheTransceiverConfig( + backend="NIXL")._to_pybind().agent_bounce_params == {} + + def test_agent_bounce_param_keys_match_cpp_env_knobs(self): + """AGENT_BOUNCE_PARAM_KEYS must mirror kEnvKnobs in BounceConfig.h (parsed here).""" + from tensorrt_llm._torch.disaggregation.nixl.bounce_knobs import \ + AGENT_BOUNCE_PARAM_KEYS + bounce_config_h = (Path(__file__).resolve().parents[3] / + "cpp/tensorrt_llm/executor/cache_transmission/" + "nixl_utils/bounce/BounceConfig.h") + if not bounce_config_h.is_file(): + pytest.skip("C++ sources not present (wheel-only checkout)") + cpp_keys = set( + re.findall(r'\{"([a-z0-9_]+)",\s*"TRTLLM_NIXL_BOUNCE_', + bounce_config_h.read_text())) + assert cpp_keys, "failed to parse kEnvKnobs from BounceConfig.h" + python_keys = set(AGENT_BOUNCE_PARAM_KEYS) + assert python_keys == cpp_keys, ( + f"agent_bounce_params allowlist drifted from kEnvKnobs: " + f"only in Python: {sorted(python_keys - cpp_keys)}, " + f"only in C++: {sorted(cpp_keys - python_keys)}") + def test_torch_compile_config_arbitrary_args(self): """Test that TorchCompileConfig rejects arbitrary arguments.""" # Valid arguments should work