diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.cpp index 7faa97c7d9aa..680fd9f2091d 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.cpp +++ b/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(mAddr + offset)); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.h index fa1f88e65364..6a4895ae175f 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.h +++ b/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 { diff --git a/cpp/tensorrt_llm/runtime/tllmBuffers.cpp b/cpp/tensorrt_llm/runtime/tllmBuffers.cpp index 4876d5b87bf6..38615462b560 100644 --- a/cpp/tensorrt_llm/runtime/tllmBuffers.cpp +++ b/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 +#include + 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(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(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 PoolAllocator::PoolType& PoolAllocator::getPool() { diff --git a/cpp/tensorrt_llm/runtime/tllmBuffers.h b/cpp/tensorrt_llm/runtime/tllmBuffers.h index d023823de5b2..4391012ff6fc 100644 --- a/cpp/tensorrt_llm/runtime/tllmBuffers.h +++ b/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 #include +#include #include #include #include @@ -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 class BaseAllocator @@ -159,6 +174,12 @@ class UVMAllocator : public BaseAllocator } }; +//! \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 { friend class BaseAllocator; @@ -167,17 +188,31 @@ class PinnedAllocator : public BaseAllocator; 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 diff --git a/cpp/tests/unit_tests/runtime/tllmBuffersTest.cpp b/cpp/tests/unit_tests/runtime/tllmBuffersTest.cpp index 4080a0f29e9e..dc840a94371b 100644 --- a/cpp/tests/unit_tests/runtime/tllmBuffersTest.cpp +++ b/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 +#include +#include #include #include #include @@ -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::aligned_alloc(PinnedAllocator::kHostPageSize, size)); + ASSERT_NE(host, nullptr); + for (std::size_t i = 0; i < size; ++i) + { + host[i] = static_cast(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 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(ptr) % PinnedAllocator::kHostPageSize, std::uintptr_t{0}); + + cudaPointerAttributes attributes{}; + TLLM_CUDA_CHECK(::cudaPointerGetAttributes(&attributes, static_cast(ptr) + size - 1)); + EXPECT_EQ(attributes.type, cudaMemoryTypeHost); + + EXPECT_NO_THROW(PinnedAllocator::deallocateChunkPinned(ptr, size, chunkSize)); +} + TEST_F(TllmBuffersTest, HostAllocator) { auto constexpr size = 1024; diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py index 2f970ca347b7..3ecfb6bc38bb 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py +++ b/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) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_host_mem_registration.py b/tests/unittest/kv_cache_manager_v2_tests/test_host_mem_registration.py new file mode 100644 index 000000000000..9189a97d16c9 --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_host_mem_registration.py @@ -0,0 +1,77 @@ +# 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 unit tests for how HostMem splits CUDA page-locking into chunks.""" + +import unittest +from importlib.util import find_spec +from typing import TYPE_CHECKING + +import pytest + +pytestmark = pytest.mark.cpu_only + + +if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None: + from kv_cache_manager_v2._utils import HostMem +else: + from tensorrt_llm.runtime.kv_cache_manager_v2._utils import HostMem + + +class _FakeHostMem: + """Stand-in exposing only what _iterate_chunks reads, so no memory is mapped or pinned.""" + + _CHUNK_SIZE = HostMem._CHUNK_SIZE + + def __init__(self, address: int, size: int) -> None: + self._address = address + self._size = size + + +def _chunks(size: int, address: int = 0x1000) -> list[tuple[int, int]]: + """Enumerate the ranges HostMem would page-lock, without allocating anything.""" + return list(HostMem._iterate_chunks(_FakeHostMem(address, size))) + + +class TestHostMemChunking(unittest.TestCase): + def test_pinning_is_always_chunked(self) -> None: + # A pool much larger than the chunk size must never be page-locked in a single + # cuMemHostRegister call: that holds driver-global locks for the whole operation and + # stalls the CUDA/NVML calls of every other process on the node. + size = 5 * HostMem._CHUNK_SIZE + chunks = _chunks(size) + self.assertEqual(len(chunks), 5) + self.assertTrue(all(length <= HostMem._CHUNK_SIZE for _, length in chunks)) + + def test_chunks_tile_the_range_exactly(self) -> None: + # Registration and unregistration walk the same boundaries, so the chunks must be + # contiguous, non-overlapping and cover the range exactly -- including a partial tail. + address = 0x1000 + size = 2 * HostMem._CHUNK_SIZE + 4096 + chunks = _chunks(size, address) + self.assertEqual(chunks[0][0], address) + self.assertEqual(sum(length for _, length in chunks), size) + for (addr, length), (next_addr, _) in zip(chunks, chunks[1:]): + self.assertEqual(addr + length, next_addr) + self.assertEqual(chunks[-1][1], 4096) + + def test_range_smaller_than_a_chunk_is_a_single_chunk(self) -> None: + self.assertEqual(_chunks(4096, 0x1000), [(0x1000, 4096)]) + + def test_empty_range_yields_no_chunks(self) -> None: + self.assertEqual(_chunks(0), []) + + +if __name__ == "__main__": + unittest.main()