[TRTLLM-15344][feat] nixl bounce v2 (1/6): foundation components

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

@@ -0,0 +1,63 @@
+/*
+ * 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.
+    mIsFabric = allowFabric && tensorrt_llm::batch_manager::kv_cache_manager::FabricMemory::supportFabricMemory();
+    if (mIsFabric)
+    {
+        mFabric = std::make_unique<tensorrt_llm::batch_manager::kv_cache_manager::FabricMemory>(bytes);
+        mBase = mFabric->getPtr();
+        TLLM_LOG_DEBUG("BounceArena: %zuB backed by fabric memory", bytes);
+    }
+    else
+    {
+        TLLM_CUDA_CHECK(cudaMalloc(&mBase, bytes));
+    }
+}
+
+BounceArena::~BounceArena()
+{
+    if (!mIsFabric && 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

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

@@ -0,0 +1,104 @@
+/*
+ * 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 <cstddef>
+#include <cstdint>
+#include <memory>
+
+// 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<std::uint64_t>(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<char*>(mBase) + offset;
+    }
+
+    [[nodiscard]] bool isFabric() const noexcept
+    {
+        return mIsFabric;
+    }
+
+private:
+    int mDeviceId{0};
+    std::size_t mBytes{0};
+    void* mBase{nullptr}; // the registered RDMA src/dst buffer (offset 0)
+    // Non-null (and mIsFabric=true) on MNNVL when fabric-backed; else null and mBase is cudaMalloc'd.
+    std::unique_ptr<tensorrt_llm::batch_manager::kv_cache_manager::FabricMemory> mFabric;
+    bool mIsFabric{false};
+};
+
+} // namespace tensorrt_llm::executor::kv_cache::bounce

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

@@ -0,0 +1,161 @@
+/*
+ * 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 <cstdarg>
+#include <cstdint>
+#include <cstdio>
+
+#ifndef NVTX_DISABLE
+#include <sys/syscall.h>
+#include <unistd.h>
+#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 push/pop, 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<BounceNvtxDomain>(msg, ::nvtx3::color{argb});
+    }
+
+    ~BounceNvtxScope()
+    {
+        ::nvtx3::end_range_in<BounceNvtxDomain>(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<BounceNvtxDomain>(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 — PushPop ranges (like the scatter span)
+/// appear under the thread that pushed them, which is otherwise an anonymous worker.
+inline void bounceNameThread(char const* name)
+{
+#ifndef NVTX_DISABLE
+    nvtxNameOsThreadA(static_cast<std::uint32_t>(::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<BounceNvtxDomain>(::nvtx3::range_handle{handle});
+    }
+#endif
+    handle = 0;
+}
+
+} // namespace tensorrt_llm::executor::kv_cache::bounce