[None][fix] page-lock large host KV cache pools in chunks

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.cpp

@@ -310,9 +310,7 @@ void HostMem::parallelPrefault(int numThreads)
 void HostMem::registerToCuda()
 {
     TLLM_CHECK_DEBUG(mNumRegisteredChunks == 0);
-    static bool chunked = shouldUseChunkedRegistration();
-
-    size_t chunkSize = (chunked && mSize > kChunkSize) ? kChunkSize : mSize;
+    size_t chunkSize = std::min(kChunkSize, mSize);
     for (size_t offset = 0; offset < mSize; offset += chunkSize)
     {
         size_t sz = std::min(chunkSize, mSize - offset);
@@ -325,8 +323,7 @@ void HostMem::registerToCuda()
 
 void HostMem::unregisterFromCuda()
 {
-    static bool chunked = shouldUseChunkedRegistration();
-    size_t chunkSize = (chunked && mSize > kChunkSize) ? kChunkSize : mSize;
+    size_t chunkSize = std::min(kChunkSize, mSize);
     for (size_t offset = 0; offset < mSize && mNumRegisteredChunks > 0; offset += chunkSize)
     {
         cuMemHostUnregister(reinterpret_cast<void*>(mAddr + offset));

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.h

@@ -34,8 +34,10 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager_v2
 //   - Optionally prefaulted in parallel before CUDA registration
 //   - Registered to CUDA as page-locked (CU_MEMHOSTREGISTER_DEVICEMAP)
 //
-// On kernels 6.11/6.12/6.13, pinning is chunked in 2GB pieces to work around
-// a kernel bug that prevents pinning more than 2GB in one call.
+// Pinning is always chunked in 2GB pieces. It works around a kernel bug on
+// 6.11/6.12/6.13 that prevents pinning more than 2GB in one call, and it bounds
+// how long a single cuMemHostRegister holds the driver-global locks that every
+// other process on the node contends on for its own CUDA/NVML calls.
 // ---------------------------------------------------------------------------
 class HostMem
 {

cpp/tensorrt_llm/runtime/tllmBuffers.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2024, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,10 +15,120 @@
  */
 
 #include "tensorrt_llm/runtime/tllmBuffers.h"
+#include "tensorrt_llm/common/envUtils.h"
 #include "tensorrt_llm/common/tllmDataType.h"
 
+#include <cstdint>
+#include <new>
+
 namespace tensorrt_llm::runtime
 {
+namespace
+{
+//! \brief Rounds \p size up to a whole number of host pages, as required by ::cudaHostRegister.
+std::size_t pageAlignedSize(std::size_t size)
+{
+    return common::ceilDiv(size, PinnedAllocator::kHostPageSize) * PinnedAllocator::kHostPageSize;
+}
+
+//! \brief Whether an allocation of \p n bytes is backed by host memory that is page-locked in chunks.
+bool useChunkedPinning(std::size_t n)
+{
+    auto const chunkSize = PinnedAllocator::getPinChunkSize();
+    return chunkSize != 0 && n > chunkSize;
+}
+} // namespace
+
+void hostRegisterChunked(void* ptr, std::size_t size, std::size_t chunkSize)
+{
+    TLLM_CHECK_WITH_INFO(chunkSize > 0, "Page-locking chunk size must be positive");
+    auto* const base = static_cast<std::uint8_t*>(ptr);
+    std::size_t offset{0};
+    try
+    {
+        for (; offset < size; offset += chunkSize)
+        {
+            TLLM_CUDA_CHECK(
+                ::cudaHostRegister(base + offset, std::min(chunkSize, size - offset), cudaHostRegisterDefault));
+        }
+    }
+    catch (...)
+    {
+        // Leave no partially registered range behind. Warn instead of throwing on the unwind so that the original
+        // registration failure is the one that reaches the caller.
+        for (std::size_t undone{0}; undone < offset; undone += chunkSize)
+        {
+            TLLM_CUDA_CHECK_WARN(::cudaHostUnregister(base + undone));
+        }
+        throw;
+    }
+}
+
+void hostUnregisterChunked(void* ptr, std::size_t size, std::size_t chunkSize)
+{
+    TLLM_CHECK_WITH_INFO(chunkSize > 0, "Page-locking chunk size must be positive");
+    auto* const base = static_cast<std::uint8_t*>(ptr);
+    for (std::size_t offset{0}; offset < size; offset += chunkSize)
+    {
+        TLLM_CUDA_CHECK_FREE_RESOURCE(::cudaHostUnregister(base + offset));
+    }
+}
+
+std::size_t PinnedAllocator::getPinChunkSize()
+{
+    static std::size_t const chunkSize
+        = common::getUInt64Env("TRTLLM_HOST_PIN_CHUNK_BYTES").value_or(kDefaultPinChunkSize);
+    return chunkSize;
+}
+
+PinnedAllocator::PointerType PinnedAllocator::allocateChunkPinned(std::size_t n, std::size_t chunkSize)
+{
+    auto const lockedBytes = pageAlignedSize(n);
+    TLLM_LOG_DEBUG("PinnedAllocator: page-locking %zu B in chunks of %zu B", lockedBytes, chunkSize);
+
+    auto* const base = std::aligned_alloc(kHostPageSize, lockedBytes);
+    if (base == nullptr)
+    {
+        throw std::bad_alloc();
+    }
+    try
+    {
+        hostRegisterChunked(base, lockedBytes, chunkSize);
+    }
+    catch (...)
+    {
+        std::free(base);
+        throw;
+    }
+    return base;
+}
+
+void PinnedAllocator::deallocateChunkPinned(PointerType ptr, std::size_t n, std::size_t chunkSize)
+{
+    hostUnregisterChunked(ptr, pageAlignedSize(n), chunkSize);
+    std::free(ptr);
+}
+
+void PinnedAllocator::allocateImpl(PointerType* ptr, std::size_t n)
+{
+    if (!useChunkedPinning(n))
+    {
+        TLLM_CUDA_CHECK(::cudaHostAlloc(ptr, n, cudaHostAllocDefault));
+        return;
+    }
+    *ptr = allocateChunkPinned(n, getPinChunkSize());
+}
+
+void PinnedAllocator::deallocateImpl(PointerType ptr, std::size_t n)
+{
+    if (!useChunkedPinning(n))
+    {
+        TLLM_CUDA_CHECK_FREE_RESOURCE(::cudaFreeHost(ptr));
+        return;
+    }
+    deallocateChunkPinned(ptr, n, getPinChunkSize());
+}
+
 template <typename TAllocator>
 typename PoolAllocator<TAllocator>::PoolType& PoolAllocator<TAllocator>::getPool()
 {

cpp/tensorrt_llm/runtime/tllmBuffers.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2024, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -31,6 +31,7 @@
 #include <cuda_runtime_api.h>
 
 #include <algorithm>
+#include <cstddef>
 #include <cstdlib>
 #include <list>
 #include <memory>
@@ -42,6 +43,20 @@
 namespace tensorrt_llm::runtime
 {
 
+//! \brief Page-locks [ptr, ptr + size) with CUDA, issuing one ::cudaHostRegister call per \p chunkSize bytes.
+//!
+//! Page-locking host memory serializes on driver-global locks. A single multi-hundred-GB request keeps those locks
+//! held for many minutes, which blocks CUDA and NVML calls issued by every other process sharing the node. Splitting
+//! the request bounds how long any one call holds the locks, so unrelated processes keep making progress.
+//!
+//! \p ptr and \p chunkSize must be multiples of the host page size. Already registered chunks are unregistered again
+//! before rethrowing if a chunk fails.
+void hostRegisterChunked(void* ptr, std::size_t size, std::size_t chunkSize);
+
+//! \brief Reverses hostRegisterChunked() using the same chunk boundaries. Must be called with the \p size and
+//! \p chunkSize that were passed to hostRegisterChunked().
+void hostUnregisterChunked(void* ptr, std::size_t size, std::size_t chunkSize);
+
 // CRTP base class
 template <typename TDerived, MemoryType memoryType, bool count = true>
 class BaseAllocator
@@ -159,6 +174,12 @@ class UVMAllocator : public BaseAllocator<UVMAllocator, MemoryType::kUVM>
     }
 };
 
+//! \brief Allocator for page-locked host memory.
+//!
+//! Allocations larger than getPinChunkSize() are backed by ordinary host memory that is page-locked in chunks rather
+//! than by a single ::cudaHostAlloc. ::cudaHostAlloc page-locks the whole range in one go, so a large allocation
+//! (e.g. a several-hundred-GB KV cache host offload pool) holds driver-global locks for the entire allocation and
+//! stalls the CUDA and NVML calls of every other process on the node for just as long. See hostRegisterChunked().
 class PinnedAllocator : public BaseAllocator<PinnedAllocator, MemoryType::kPINNED>
 {
     friend class BaseAllocator<PinnedAllocator, MemoryType::kPINNED>;
@@ -167,17 +188,31 @@ class PinnedAllocator : public BaseAllocator<PinnedAllocator, MemoryType::kPINNE
     using Base = BaseAllocator<PinnedAllocator, MemoryType::kPINNED>;
     PinnedAllocator() noexcept = default;
 
+    //! \brief Host page size that chunked allocations and their chunk boundaries are aligned to, as required by
+    //! ::cudaHostRegister.
+    static std::size_t constexpr kHostPageSize{4096};
+
+    //! \brief Default value of getPinChunkSize().
+    static std::size_t constexpr kDefaultPinChunkSize{std::size_t{1} << 30}; // 1 GiB
+
+    //! \brief Number of bytes page-locked per ::cudaHostRegister call, and the size up to which allocations keep
+    //! using a single ::cudaHostAlloc.
+    //!
+    //! Read once from TRTLLM_HOST_PIN_CHUNK_BYTES, defaulting to kDefaultPinChunkSize. Setting it to 0 disables
+    //! chunking and restores plain ::cudaHostAlloc for every allocation size.
+    [[nodiscard]] static std::size_t getPinChunkSize();
+
+    //! \brief Allocates \p n bytes of host memory and page-locks it in \p chunkSize pieces.
+    [[nodiscard]] static PointerType allocateChunkPinned(std::size_t n, std::size_t chunkSize);
+
+    //! \brief Frees memory obtained from allocateChunkPinned(). \p n and \p chunkSize must match the allocation.
+    static void deallocateChunkPinned(PointerType ptr, std::size_t n, std::size_t chunkSize);
+
 protected:
-    void allocateImpl(PointerType* ptr, std::size_t n) // NOLINT(readability-convert-member-functions-to-static)
-    {
-        TLLM_CUDA_CHECK(::cudaHostAlloc(ptr, n, cudaHostAllocDefault));
-    }
+    void allocateImpl(PointerType* ptr, std::size_t n); // NOLINT(readability-convert-member-functions-to-static)
 
-    void deallocateImpl( // NOLINT(readability-convert-member-functions-to-static)
-        PointerType ptr, [[maybe_unused]] std::size_t n)
-    {
-        TLLM_CUDA_CHECK_FREE_RESOURCE(::cudaFreeHost(ptr));
-    }
+    void deallocateImpl(                                // NOLINT(readability-convert-member-functions-to-static)
+        PointerType ptr, std::size_t n);
 };
 
 class HostAllocator : public BaseAllocator<HostAllocator, MemoryType::kCPU>

cpp/tests/unit_tests/runtime/tllmBuffersTest.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2024, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -27,6 +27,8 @@
 #include "tensorrt_llm/runtime/tllmBuffers.h"
 
 #include <algorithm>
+#include <cstdint>
+#include <cstdlib>
 #include <limits>
 #include <memory>
 #include <random>
@@ -113,6 +115,74 @@ TEST_F(TllmBuffersTest, PinnedAllocator)
     EXPECT_THROW(allocator.deallocate(ptr, size), std::runtime_error);
 }
 
+TEST_F(TllmBuffersTest, HostRegisterChunked)
+{
+    if (mDeviceCount == 0)
+    {
+        GTEST_SKIP() << noDeviceSkipReason;
+    }
+
+    // A range page-locked chunk by chunk must behave exactly like one page-locked in a single call: every chunk,
+    // including a partial trailing one, is reported as host memory and the whole range is usable as the source of an
+    // async copy across chunk boundaries.
+    auto constexpr chunkSize = std::size_t{1} << 20; // 1 MB
+    auto constexpr size = 4 * chunkSize + PinnedAllocator::kHostPageSize;
+
+    auto* const host = static_cast<std::uint8_t*>(std::aligned_alloc(PinnedAllocator::kHostPageSize, size));
+    ASSERT_NE(host, nullptr);
+    for (std::size_t i = 0; i < size; ++i)
+    {
+        host[i] = static_cast<std::uint8_t>(i);
+    }
+
+    hostRegisterChunked(host, size, chunkSize);
+
+    for (std::size_t offset = 0; offset < size; offset += chunkSize)
+    {
+        cudaPointerAttributes attributes{};
+        TLLM_CUDA_CHECK(::cudaPointerGetAttributes(&attributes, host + offset));
+        EXPECT_EQ(attributes.type, cudaMemoryTypeHost) << "chunk at offset " << offset << " was not page-locked";
+    }
+
+    auto device = BufferManager::gpuSync(size, tensorrt_llm::DataType::kUINT8);
+    std::vector<std::uint8_t> roundTrip(size, 0);
+    TLLM_CUDA_CHECK(::cudaMemcpyAsync(device->data(), host, size, cudaMemcpyHostToDevice, mStream->get()));
+    TLLM_CUDA_CHECK(::cudaMemcpyAsync(roundTrip.data(), device->data(), size, cudaMemcpyDeviceToHost, mStream->get()));
+    mStream->synchronize();
+    EXPECT_TRUE(std::equal(roundTrip.begin(), roundTrip.end(), host));
+
+    hostUnregisterChunked(host, size, chunkSize);
+
+    cudaPointerAttributes attributes{};
+    TLLM_CUDA_CHECK(::cudaPointerGetAttributes(&attributes, host));
+    EXPECT_EQ(attributes.type, cudaMemoryTypeUnregistered);
+
+    std::free(host);
+}
+
+TEST_F(TllmBuffersTest, PinnedAllocatorChunkPinnedAllocation)
+{
+    if (mDeviceCount == 0)
+    {
+        GTEST_SKIP() << noDeviceSkipReason;
+    }
+
+    // The chunk-pinned path that PinnedAllocator takes for allocations larger than the chunk size. The requested
+    // size is deliberately neither chunk- nor page-aligned, so the tail is rounded up to a whole page.
+    auto constexpr chunkSize = std::size_t{1} << 20; // 1 MB
+    auto constexpr size = 2 * chunkSize + 1234;
+
+    auto* const ptr = PinnedAllocator::allocateChunkPinned(size, chunkSize);
+    ASSERT_NE(ptr, nullptr);
+    EXPECT_EQ(reinterpret_cast<std::uintptr_t>(ptr) % PinnedAllocator::kHostPageSize, std::uintptr_t{0});
+
+    cudaPointerAttributes attributes{};
+    TLLM_CUDA_CHECK(::cudaPointerGetAttributes(&attributes, static_cast<std::uint8_t*>(ptr) + size - 1));
+    EXPECT_EQ(attributes.type, cudaMemoryTypeHost);
+
+    EXPECT_NO_THROW(PinnedAllocator::deallocateChunkPinned(ptr, size, chunkSize));
+}
+
 TEST_F(TllmBuffersTest, HostAllocator)
 {
     auto constexpr size = 1024;

tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py

@@ -21,7 +21,6 @@
 import itertools
 import operator
 import os
-import platform
 import sys
 import traceback
 import warnings
@@ -502,11 +501,11 @@ class HostMem:
     __slots__ = ("_address", "_size", "_num_registered_chunks")
     _address: int
     _size: int
-    # If True and _size > 2GB, use multiple chunks to register pinned memory due to a Linux kernel
-    # 6.11/6.12/6.13 bug preventing pinning more than 2GB of host memory in one operation.
-    _CHUNKED_REGISTRATION: ClassVar[bool] = platform.system() == "Linux" and platform.release()[
-        :4
-    ] in ["6.11", "6.12", "6.13"]
+    # Pinning is always chunked. It works around a Linux kernel 6.11/6.12/6.13 bug preventing
+    # pinning more than 2GB of host memory in one operation, and it bounds how long a single
+    # cuMemHostRegister holds the driver-global locks that every other process on the node
+    # contends on for its own CUDA/NVML calls -- one unbroken multi-hundred-GB registration
+    # stalls colocated workers for the whole duration.
     _CHUNK_SIZE: ClassVar[int] = 2 << 30
     _num_registered_chunks: int
 
@@ -624,9 +623,11 @@ def _unregister_from_cuda(self) -> None:
         assert self._num_registered_chunks == 0
 
     def _iterate_chunks(self) -> Iterator[tuple[int, int]]:
+        if self._size == 0:
+            return
         start = self._address
         end = start + self._size
-        chunk_size = self._CHUNK_SIZE if self._CHUNKED_REGISTRATION else self._size
+        chunk_size = min(self._CHUNK_SIZE, self._size)
         for addr in range(start, end, chunk_size):
             yield addr, min(end - addr, chunk_size)