diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 691c49e545aa..348c60fa787f 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & +# SPDX-FileCopyrightText: Copyright (c) 2022-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 diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index f61e58e16b28..f04647f9d059 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -80,6 +80,10 @@ endif() include(${CMAKE_CURRENT_SOURCE_DIR}/kv_cache_manager_v2/CMakeLists.txt) list(APPEND SRCS ${KV_CACHE_MANAGER_V2_SRCS}) +# Include the KVCM cold-page compression codec bridge sources. +include(${CMAKE_CURRENT_SOURCE_DIR}/kv_cache_compression/CMakeLists.txt) +list(APPEND SRCS ${KV_CACHE_COMPRESSION_SRCS}) + add_library(${BATCH_MANAGER_STATIC_TARGET} STATIC ${SRCS}) target_include_directories( ${BATCH_MANAGER_STATIC_TARGET} diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_compression/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/kv_cache_compression/CMakeLists.txt new file mode 100644 index 000000000000..12018b632383 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_compression/CMakeLists.txt @@ -0,0 +1,18 @@ +# 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. + +# Sources for the KVCM cold-page compression codec bridge. These are added to +# the tensorrt_llm_batch_manager_static target by the parent CMakeLists.txt. +set(KV_CACHE_COMPRESSION_SRCS kv_cache_compression/nativeColdPageCodec.cpp) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_compression/nativeColdPageCodec.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_compression/nativeColdPageCodec.cpp new file mode 100644 index 000000000000..5beed9e998bb --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_compression/nativeColdPageCodec.cpp @@ -0,0 +1,298 @@ +/* + * 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/batch_manager/kv_cache_compression/nativeColdPageCodec.h" + +#include "kv_cache_manager_v2/utils/hostMem.h" +#include "tensorrt_llm/common/logger.h" + +#include +#include +#include +#include +#include + +namespace tensorrt_llm::kv_cache_compression +{ +namespace +{ + +ResolvedHotLifecycle resolveLifecycle(kv::PoolGroupDesc const& gpuDesc, kv::SlotDescVariant const& variant) +{ + ResolvedHotLifecycle result{variant.lifeCycleId, {}}; + for (kv::PoolIndex poolIndex{0}; poolIndex < variant.coalescedBuffers.size(); ++poolIndex) + { + auto const& coalesced = variant.coalescedBuffers.at(poolIndex); + auto const& pool = gpuDesc.pools.at(poolIndex); + std::size_t offset = 0; + for (auto const& bufferId : coalesced.bufferIds) + { + auto& layer = result.layers[bufferId.layerId]; + if (!layer + .emplace(bufferId.role, + ResolvedHotBuffer{pool.baseAddress + offset, pool.slotBytes, coalesced.singleBufferSize}) + .second) + { + throw std::invalid_argument("GPU lifecycle contains a duplicate buffer role"); + } + offset += coalesced.singleBufferSize; + } + } + return result; +} + +void drainAfterProviderFailure(cudaStream_t stream) noexcept +{ + auto const status = cudaStreamSynchronize(stream); + if (status != cudaSuccess) + { + TLLM_LOG_ERROR("Cold-page provider rollback drain failed: %s", cudaGetErrorString(status)); + std::terminate(); + } +} + +} // namespace + +NativeColdPageCodec::NativeColdPageCodec(std::set layerIds) + : mLayerIds(std::move(layerIds)) +{ +} + +bool NativeColdPageCodec::configure(kv::PoolGroupDesc const* gpuDescs, kv::PoolGroupIndex numGpuDescs) noexcept +{ + try + { + auto losslessCodec = kv::createDefaultKvCacheColdPageCodec(); + if (!losslessCodec->configure(gpuDescs, numGpuDescs)) + { + throw std::invalid_argument("Default lossless codec rejected GPU layouts"); + } + + std::map pendingGroups; + std::vector providerLifecycles; + std::set consumedLayers; + + for (kv::PoolGroupIndex poolGroupIndex{0}; poolGroupIndex < numGpuDescs; ++poolGroupIndex) + { + auto const& gpuDesc = gpuDescs[kv::toSizeT(poolGroupIndex)]; + for (auto const& variant : gpuDesc.slotDesc.variants) + { + auto resolved = resolveLifecycle(gpuDesc, variant); + auto const providerLayerCount = static_cast(std::count_if(resolved.layers.begin(), + resolved.layers.end(), [this](auto const& layer) { return mLayerIds.count(layer.first) != 0U; })); + + LayerGroupState state; + if (providerLayerCount == 0U) + { + state.coldPageBytes = losslessCodec->queryColdPageBytes(variant.lifeCycleId); + state.pageIndexLocation = losslessCodec->queryPageIndexLocation(variant.lifeCycleId); + } + else + { + if (providerLayerCount != resolved.layers.size()) + { + throw std::invalid_argument("A lifecycle cannot mix provider-owned and fallback layers"); + } + for (auto const& [layerId, buffers] : resolved.layers) + { + static_cast(buffers); + if (!consumedLayers.emplace(layerId).second) + { + throw std::invalid_argument("A provider layer appears in multiple lifecycles"); + } + } + state.lifecycleIndex = providerLifecycles.size(); + providerLifecycles.push_back(std::move(resolved)); + } + + if (!pendingGroups.emplace(variant.lifeCycleId, std::move(state)).second) + { + throw std::invalid_argument("GPU lifecycle ID appears in multiple pool groups"); + } + } + } + if (consumedLayers != mLayerIds) + { + throw std::invalid_argument("A provider layer is absent from all GPU descriptors"); + } + + // Fail closed until KVCM replaces the batched cuMemcpyBatchAsync copies with kernels: on host + // kernels that need chunked pinned-memory registration (Linux 6.11-6.13), the embedded lossless + // codec cannot split its copies at registration boundaries when wrapped by this codec. + bool hasFallbackLifecycle = false; + for (auto const& [lifeCycleId, state] : pendingGroups) + { + static_cast(lifeCycleId); + if (!state.lifecycleIndex) + { + hasFallbackLifecycle = true; + break; + } + } + if (hasFallbackLifecycle && kv::HostMem::shouldUseChunkedRegistration()) + { + throw std::invalid_argument( + "Cold-page compression is not supported for models with lossless-fallback lifecycles (SSM/GDN) on " + "this host kernel: chunked pinned-memory registration (Linux 6.11-6.13) breaks the fallback codec's " + "batched copies. Disable KV cache compression for this model or use a different host kernel."); + } + + auto const properties = configureProvider(providerLifecycles); + if (properties.size() != providerLifecycles.size()) + { + throw std::invalid_argument("Cold-page provider returned an unexpected lifecycle count"); + } + for (std::size_t index = 0; index < properties.size(); ++index) + { + auto const& lifecycle = properties[index]; + if (lifecycle.coldPageBytes == 0U || lifecycle.pageIndexLocation == kv::PageIndexLocation::kBadLocation) + { + throw std::invalid_argument("Cold-page provider returned invalid storage properties"); + } + auto& state = pendingGroups.at(providerLifecycles[index].lifeCycleId); + state.coldPageBytes = lifecycle.coldPageBytes; + state.pageIndexLocation = lifecycle.pageIndexLocation; + } + + mLayerGroups = std::move(pendingGroups); + mLosslessCodec = std::move(losslessCodec); + return true; + } + catch (std::exception const& error) + { + TLLM_LOG_ERROR("NativeColdPageCodec::configure rejected GPU layouts: %s", error.what()); + return false; + } + catch (...) + { + TLLM_LOG_ERROR("NativeColdPageCodec::configure rejected GPU layouts: unknown error"); + return false; + } +} + +NativeColdPageCodec::LayerGroupState const* NativeColdPageCodec::findLayerGroup( + kv::LayerGroupId layerGroupId) const noexcept +{ + auto const found = mLayerGroups.find(layerGroupId); + return found == mLayerGroups.end() ? nullptr : &found->second; +} + +std::size_t NativeColdPageCodec::queryColdPageBytes(kv::LayerGroupId layerGroupId) const noexcept +{ + auto const* state = findLayerGroup(layerGroupId); + return state == nullptr ? 0U : state->coldPageBytes; +} + +kv::LayerGroupId NativeColdPageCodec::getBatchingLayerGroupId(kv::LayerGroupId layerGroupId) const noexcept +{ + return findLayerGroup(layerGroupId) == nullptr ? kv::LayerGroupId{-1} : layerGroupId; +} + +kv::PageIndexLocation NativeColdPageCodec::queryPageIndexLocation(kv::LayerGroupId layerGroupId) const noexcept +{ + auto const* state = findLayerGroup(layerGroupId); + return state == nullptr ? kv::PageIndexLocation::kBadLocation : state->pageIndexLocation; +} + +bool NativeColdPageCodec::encode(kv::LayerGroupId layerGroupId, void* dstBasePtr, kv::PageIndexPair const* pageIndices, + std::size_t numBasePages, cudaStream_t stream) noexcept +{ + bool providerStarted = false; + try + { + auto const* state = findLayerGroup(layerGroupId); + if (state == nullptr || (numBasePages != 0U && (dstBasePtr == nullptr || pageIndices == nullptr))) + { + throw std::invalid_argument("encode received an invalid lifecycle or Page batch"); + } + if (numBasePages == 0U) + { + return true; + } + if (!state->lifecycleIndex) + { + return mLosslessCodec->encode(layerGroupId, dstBasePtr, pageIndices, numBasePages, stream); + } + providerStarted = true; + encodeProvider(*state->lifecycleIndex, dstBasePtr, pageIndices, numBasePages, stream); + return true; + } + catch (std::exception const& error) + { + if (providerStarted) + { + drainAfterProviderFailure(stream); + } + TLLM_LOG_ERROR("NativeColdPageCodec::encode failed before completion fencing: %s", error.what()); + return false; + } + catch (...) + { + if (providerStarted) + { + drainAfterProviderFailure(stream); + } + TLLM_LOG_ERROR("NativeColdPageCodec::encode failed before completion fencing: unknown error"); + return false; + } +} + +bool NativeColdPageCodec::decode(kv::LayerGroupId layerGroupId, void const* srcBasePtr, + kv::PageIndexPair const* pageIndices, std::size_t numBasePages, cudaStream_t stream) noexcept +{ + bool providerStarted = false; + try + { + auto const* state = findLayerGroup(layerGroupId); + if (state == nullptr || (numBasePages != 0U && (srcBasePtr == nullptr || pageIndices == nullptr))) + { + throw std::invalid_argument("decode received an invalid lifecycle or Page batch"); + } + if (numBasePages == 0U) + { + return true; + } + if (!state->lifecycleIndex) + { + return mLosslessCodec->decode(layerGroupId, srcBasePtr, pageIndices, numBasePages, stream); + } + providerStarted = true; + decodeProvider(*state->lifecycleIndex, srcBasePtr, pageIndices, numBasePages, stream); + return true; + } + catch (std::exception const& error) + { + if (providerStarted) + { + drainAfterProviderFailure(stream); + } + TLLM_LOG_ERROR("NativeColdPageCodec::decode failed before completion fencing: %s", error.what()); + return false; + } + catch (...) + { + if (providerStarted) + { + drainAfterProviderFailure(stream); + } + TLLM_LOG_ERROR("NativeColdPageCodec::decode failed before completion fencing: unknown error"); + return false; + } +} + +} // namespace tensorrt_llm::kv_cache_compression diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_compression/nativeColdPageCodec.h b/cpp/tensorrt_llm/batch_manager/kv_cache_compression/nativeColdPageCodec.h new file mode 100644 index 000000000000..34e6515cd295 --- /dev/null +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_compression/nativeColdPageCodec.h @@ -0,0 +1,108 @@ +/* + * 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 "kv_cache_manager_v2/coldPageCodec.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::kv_cache_compression +{ + +namespace kv = batch_manager::kv_cache_manager_v2; + +//! One hot buffer resolved from KVCM's authoritative pool descriptors. +struct ResolvedHotBuffer +{ + std::uintptr_t rawBase = 0; + std::size_t rawSlotBytes = 0; + std::size_t rawBytes = 0; +}; + +using ResolvedHotLayer = std::map; + +//! One KVCM lifecycle resolved into its hot buffers. +struct ResolvedHotLifecycle +{ + kv::LifeCycleId lifeCycleId{-1}; + std::map layers; +}; + +//! Storage properties produced while an algorithm prepares one lifecycle. +struct ColdPageLifecycleProperties +{ + std::size_t coldPageBytes = 0; + kv::PageIndexLocation pageIndexLocation = kv::PageIndexLocation::kBadLocation; +}; + +//! Resolves KVCM layouts and routes lifecycles for one native compression method. +class NativeColdPageCodec : public kv::IKvCacheColdPageCodec +{ +public: + explicit NativeColdPageCodec(std::set layerIds); + + bool configure(kv::PoolGroupDesc const* gpuDescs, kv::PoolGroupIndex numGpuDescs) noexcept final; + + [[nodiscard]] std::size_t queryColdPageBytes(kv::LayerGroupId layerGroupId) const noexcept final; + + [[nodiscard]] kv::LayerGroupId getBatchingLayerGroupId(kv::LayerGroupId layerGroupId) const noexcept final; + + [[nodiscard]] kv::PageIndexLocation queryPageIndexLocation(kv::LayerGroupId layerGroupId) const noexcept final; + + bool encode(kv::LayerGroupId layerGroupId, void* dstBasePtr, kv::PageIndexPair const* pageIndices, + std::size_t numBasePages, cudaStream_t stream) noexcept final; + + bool decode(kv::LayerGroupId layerGroupId, void const* srcBasePtr, kv::PageIndexPair const* pageIndices, + std::size_t numBasePages, cudaStream_t stream) noexcept final; + +private: + virtual std::vector configureProvider( + std::vector const& lifecycles) + = 0; + + //! Enqueue only on stream; this codec drains partial submissions after a throw. + virtual void encodeProvider(std::size_t lifecycleIndex, void* coldBase, kv::PageIndexPair const* pageIndices, + std::size_t numPages, cudaStream_t stream) + = 0; + + virtual void decodeProvider(std::size_t lifecycleIndex, void const* coldBase, kv::PageIndexPair const* pageIndices, + std::size_t numPages, cudaStream_t stream) + = 0; + + struct LayerGroupState + { + std::optional lifecycleIndex; + std::size_t coldPageBytes = 0; + kv::PageIndexLocation pageIndexLocation = kv::PageIndexLocation::kBadLocation; + }; + + [[nodiscard]] LayerGroupState const* findLayerGroup(kv::LayerGroupId layerGroupId) const noexcept; + + std::set mLayerIds; + std::map mLayerGroups; + std::unique_ptr mLosslessCodec; +}; + +} // namespace tensorrt_llm::kv_cache_compression diff --git a/cpp/tensorrt_llm/kernels/nvfp4ColdPageKernels.cu b/cpp/tensorrt_llm/kernels/nvfp4ColdPageKernels.cu new file mode 100644 index 000000000000..57d2d3b9433f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/nvfp4ColdPageKernels.cu @@ -0,0 +1,873 @@ +/* + * 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/kernels/nvfp4ColdPageKernels.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/envUtils.h" +#include "tensorrt_llm/kernels/cudaAsyncOps.cuh" +#include "tensorrt_llm/kernels/quantization.cuh" + +#include +#include +#include +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace +{ + +// Match KVCM V2's mapped-Host copy CTA and one-split policy. +constexpr std::uint32_t kThreadsPerBlock = 128; +constexpr std::uint32_t kAsyncStages = 4; +// Mapped-Host reads use eight cp.async stages; GPU-resident input uses four. +constexpr std::uint32_t kHostLoadAsyncStages = 8; +constexpr std::uint32_t kMappedHostGridSplits = 1; +constexpr std::uint32_t kMaxTasksPerLaunch = 256; +constexpr std::uint32_t kElementsPerHalfGroup = 8; +constexpr std::uint32_t kElementsPerScaleGroup = 16; +constexpr std::uint32_t kHalfGroupsPerScaleGroup = kElementsPerScaleGroup / kElementsPerHalfGroup; +// Bound per-tile shared scale staging to 1 KiB. +constexpr std::uint32_t kMaxScaleBytesPerTile = 1024; +constexpr std::uint32_t kMaxHalfGroupsPerTile = kHalfGroupsPerScaleGroup * kMaxScaleBytesPerTile; +constexpr std::size_t kKernelParameterLimitBytes = 32764; + +enum WideField : std::uint32_t +{ + kRawBase, + kRawSlotBytes, + kRawBytes, + kColdDataOffset, + kColdScaleOffset, + kColdPaddingOffset, +}; + +enum IntegerField : std::uint32_t +{ + kColdPaddingBytes, + kTransform, + kNumKvHeads, + kTokensPerPage, + kHeadDim, +}; + +enum ScaleField : std::uint32_t +{ + kNvfp4ScaleOrigQuant, + kNvfp4ScaleQuantOrig, + kFp8ScaleOrigQuant, + kFp8ScaleQuantOrig, +}; + +struct Nvfp4ColdPageKernelParams +{ + std::int32_t numKvHeads; + std::int32_t tokensPerPage; + std::int32_t headDim; + float nvfp4ScaleOrigQuant; + float nvfp4ScaleQuantOrig; + float fp8ScaleOrigQuant; + float fp8ScaleQuantOrig; +}; + +enum class Nvfp4ColdPageTransform : std::int32_t +{ + kNvfp4 = 0, + kLosslessCopy = 1, +}; + +struct Nvfp4ColdPageBuffer +{ + std::uintptr_t rawBase; + std::size_t rawSlotBytes; + std::size_t rawBytes; + std::size_t coldDataOffset; + std::size_t coldScaleOffset; + std::size_t coldPaddingOffset; + std::uint32_t coldPaddingBytes; + Nvfp4ColdPageTransform transform; + Nvfp4ColdPageKernelParams params; +}; + +// Private KVCM Page-index view; its matching layout is asserted here and at the callback boundary. +struct alignas(8) PageIndexPairView +{ + std::int32_t dst; + std::int32_t src; +}; + +static_assert(sizeof(PageIndexPairView) == 8); +static_assert(alignof(PageIndexPairView) == 8); +static_assert(offsetof(PageIndexPairView, dst) == 0); +static_assert(offsetof(PageIndexPairView, src) == 4); +static_assert(std::is_trivially_copyable_v); + +// Keep every CTA iteration and shared-memory tile on complete 16-value scale groups. +static_assert(kElementsPerScaleGroup % kElementsPerHalfGroup == 0, + "An NVFP4 scale group must contain a whole number of half-groups"); +static_assert(kHalfGroupsPerScaleGroup == 2U, "NVFP4 stores one scale for two eight-value half-groups"); +static_assert(kThreadsPerBlock % kHalfGroupsPerScaleGroup == 0, "A CTA iteration must not split an NVFP4 scale group"); +static_assert( + kMaxScaleBytesPerTile % sizeof(uint4) == 0, "A full scale tile must preserve the 16-byte transfer fast path"); + +// Keep both kernel argument packs within CUDA's modern 32,764-byte limit. +static_assert(sizeof(std::array) + sizeof(Nvfp4ColdPageWideTable) + + sizeof(Nvfp4ColdPageIntegerTable) + sizeof(Nvfp4ColdPageScaleTable) + 2U * sizeof(std::uintptr_t) + <= kKernelParameterLimitBytes, + "Cold-page kernel arguments exceed CUDA's parameter limit"); + +// Device data path. + +// Issue one predicated 16-byte cp.async load. +template +__device__ __forceinline__ void copyAsyncGlobalToShared(T* shared, T const* global, bool valid) +{ + static_assert(sizeof(T) == 16, "Cold-page transfer grains must match batchedCopy's 16-byte width"); + if (valid) + { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" + : + : "l"(__cvta_generic_to_shared(shared)), "l"(global) + : "memory"); + } +} + +struct OffloadBufferTask +{ + std::uint8_t const* raw; + std::uint8_t* coldData; + std::uint8_t* coldScale; + std::uint8_t* coldPadding; +}; + +struct OnboardBufferTask +{ + std::uint8_t const* coldData; + std::uint8_t const* coldScale; + std::uint8_t* raw; +}; + +__device__ Nvfp4ColdPageBuffer loadBuffer(std::uint32_t index, Nvfp4ColdPageWideTable const& wide, + Nvfp4ColdPageIntegerTable const& integers, Nvfp4ColdPageScaleTable const& scales) +{ + auto const& w = wide[index]; + auto const& i = integers[index]; + auto const& s = scales[index]; + return {static_cast(w[kRawBase]), static_cast(w[kRawSlotBytes]), + static_cast(w[kRawBytes]), static_cast(w[kColdDataOffset]), + static_cast(w[kColdScaleOffset]), static_cast(w[kColdPaddingOffset]), + static_cast(i[kColdPaddingBytes]), static_cast(i[kTransform]), + {i[kNumKvHeads], i[kTokensPerPage], i[kHeadDim], s[kNvfp4ScaleOrigQuant], s[kNvfp4ScaleQuantOrig], + s[kFp8ScaleOrigQuant], s[kFp8ScaleQuantOrig]}}; +} + +__device__ OffloadBufferTask resolveOffloadTask( + PageIndexPairView const& page, Nvfp4ColdPageBuffer const& buffer, std::uint8_t* coldBase, std::size_t coldPageBytes) +{ + std::size_t const gpuPage = static_cast(page.src); + auto* coldPage = coldBase + static_cast(page.dst) * coldPageBytes; + return {reinterpret_cast(buffer.rawBase + gpuPage * buffer.rawSlotBytes), + coldPage + buffer.coldDataOffset, coldPage + buffer.coldScaleOffset, coldPage + buffer.coldPaddingOffset}; +} + +__device__ OnboardBufferTask resolveOnboardTask(PageIndexPairView const& page, Nvfp4ColdPageBuffer const& buffer, + std::uint8_t const* coldBase, std::size_t coldPageBytes) +{ + std::size_t const gpuPage = static_cast(page.dst); + auto const* coldPage = coldBase + static_cast(page.src) * coldPageBytes; + return {coldPage + buffer.coldDataOffset, coldPage + buffer.coldScaleOffset, + reinterpret_cast(buffer.rawBase + gpuPage * buffer.rawSlotBytes)}; +} + +// One eight-element half-group produces one uint32_t of packed E2M1 data. +__host__ __device__ constexpr std::uint32_t packedBytesForHalfGroups(std::uint32_t halfGroupCount) +{ + return halfGroupCount * sizeof(std::uint32_t); +} + +// Store one E4M3 scale byte per pair of half-groups. +__host__ __device__ constexpr std::uint32_t scaleBytesForHalfGroups(std::uint32_t halfGroupCount) +{ + return halfGroupCount / kHalfGroupsPerScaleGroup; +} + +// Align packed shared staging to uint4; this padding is not stored in the cold record. +__host__ __device__ constexpr std::uint32_t packedStageBytesForHalfGroups(std::uint32_t halfGroupCount) +{ + return (packedBytesForHalfGroups(halfGroupCount) + sizeof(uint4) - 1U) / sizeof(uint4) * sizeof(uint4); +} + +// Return shared-memory bytes for aligned packed data followed by scale bytes. +__host__ __device__ constexpr std::uint32_t compactStageBytesForHalfGroups(std::uint32_t halfGroupCount) +{ + return packedStageBytesForHalfGroups(halfGroupCount) + scaleBytesForHalfGroups(halfGroupCount); +} + +// Flatten one buffer's [head, token, dim] geometry into eight-element half-groups. +__host__ __device__ constexpr std::uint32_t halfGroupCount(Nvfp4ColdPageKernelParams const& params) +{ + return static_cast(params.numKvHeads) * static_cast(params.tokensPerPage) + * (static_cast(params.headDim) / kElementsPerHalfGroup); +} + +// Cap a buffer tile at shared-memory capacity without splitting a scale group. +__host__ __device__ constexpr std::uint32_t tileHalfGroupCount(Nvfp4ColdPageKernelParams const& params) +{ + // Avoid std::min: its reference return ODR-uses this host/device constexpr. + auto const halfGroups = halfGroupCount(params); + return halfGroups < kMaxHalfGroupsPerTile ? halfGroups : kMaxHalfGroupsPerTile; +} + +// Flush vectorized packed values and scales, then copy remaining tails bytewise. +__device__ void flushCompactRangeToHost(std::uint8_t const* compactStages, OffloadBufferTask const& task, + std::uint32_t packedStageCapacityBytes, std::uint32_t packedDestinationOffset, std::uint32_t packedBytes, + std::uint32_t scaleDestinationOffset, std::uint32_t scaleBytes) +{ + auto const* packedSource = compactStages; + auto* packedDestination = task.coldData + packedDestinationOffset; + bool const alignedPacked = reinterpret_cast(packedSource) % sizeof(uint4) == 0 + && reinterpret_cast(packedDestination) % sizeof(uint4) == 0; + bool const alignedPackedPair = reinterpret_cast(packedSource) % sizeof(uint2) == 0 + && reinterpret_cast(packedDestination) % sizeof(uint2) == 0; + std::uint32_t const packedVectorBytes = alignedPacked ? packedBytes - packedBytes % sizeof(uint4) : 0; + std::uint32_t const packedPairBytes = alignedPackedPair ? packedBytes - packedBytes % sizeof(uint2) : 0; + for (std::uint32_t grain = threadIdx.x; grain < packedVectorBytes / sizeof(uint4); grain += blockDim.x) + { + reinterpret_cast(packedDestination)[grain] = reinterpret_cast(packedSource)[grain]; + } + for (std::uint32_t pair = packedVectorBytes / sizeof(uint2) + threadIdx.x; pair < packedPairBytes / sizeof(uint2); + pair += blockDim.x) + { + reinterpret_cast(packedDestination)[pair] = reinterpret_cast(packedSource)[pair]; + } + for (std::uint32_t byte = packedPairBytes + threadIdx.x; byte < packedBytes; byte += blockDim.x) + { + packedDestination[byte] = packedSource[byte]; + } + + auto const* scaleSource = compactStages + packedStageCapacityBytes; + auto* scaleDestination = task.coldScale + scaleDestinationOffset; + bool const alignedScale = reinterpret_cast(scaleSource) % sizeof(uint4) == 0 + && reinterpret_cast(scaleDestination) % sizeof(uint4) == 0; + std::uint32_t const scaleVectorBytes = alignedScale ? scaleBytes - scaleBytes % sizeof(uint4) : 0; + for (std::uint32_t grain = threadIdx.x; grain < scaleVectorBytes / sizeof(uint4); grain += blockDim.x) + { + reinterpret_cast(scaleDestination)[grain] = reinterpret_cast(scaleSource)[grain]; + } + for (std::uint32_t byte = scaleVectorBytes + threadIdx.x; byte < scaleBytes; byte += blockDim.x) + { + scaleDestination[byte] = scaleSource[byte]; + } +} + +// Zero codec-specified record padding so persisted cold Slots are deterministic. +__device__ void clearColdPadding(OffloadBufferTask const& task, Nvfp4ColdPageBuffer const& buffer) +{ + if (blockIdx.x != 0U) + { + return; + } + for (std::uint32_t byte = threadIdx.x; byte < buffer.coldPaddingBytes; byte += blockDim.x) + { + task.coldPadding[byte] = 0U; + } +} + +// CTA-uniform byte-exact copy used by lossless side buffers. Vectorize only when both +// endpoints permit it; arbitrary descriptor offsets and byte tails remain supported. +__device__ void copyLosslessBytes(std::uint8_t const* source, std::uint8_t* destination, std::size_t bytes) +{ + bool const aligned = reinterpret_cast(source) % sizeof(uint4) == 0 + && reinterpret_cast(destination) % sizeof(uint4) == 0; + std::size_t const vectorBytes = aligned ? bytes - bytes % sizeof(uint4) : 0U; + for (std::size_t grain = threadIdx.x; grain < vectorBytes / sizeof(uint4); grain += blockDim.x) + { + reinterpret_cast(destination)[grain] = reinterpret_cast(source)[grain]; + } + for (std::size_t byte = vectorBytes + threadIdx.x; byte < bytes; byte += blockDim.x) + { + destination[byte] = source[byte]; + } +} + +__device__ __forceinline__ uint4 collectTwoFp8Words(std::uint64_t first, std::uint64_t second) +{ + return make_uint4(static_cast(first), static_cast(first >> 32U), + static_cast(second), static_cast(second >> 32U)); +} + +// E2M1-to-FP16x2 PTX is adapted from arcquantFP4.cu and fusedMoeCommKernels.cu. +__device__ void unpackE2m1ToFloat(std::uint32_t packed, float2 (&values)[4]) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + std::uint32_t fp16Pairs[4]; + asm volatile( + "{\n" + ".reg .b8 b0, b1, b2, b3;\n" + "mov.b32 {b0, b1, b2, b3}, %4;\n" + "cvt.rn.f16x2.e2m1x2 %0, b0;\n" + "cvt.rn.f16x2.e2m1x2 %1, b1;\n" + "cvt.rn.f16x2.e2m1x2 %2, b2;\n" + "cvt.rn.f16x2.e2m1x2 %3, b3;\n" + "}\n" + : "=r"(fp16Pairs[0]), "=r"(fp16Pairs[1]), "=r"(fp16Pairs[2]), "=r"(fp16Pairs[3]) + : "r"(packed)); + +#pragma unroll + for (std::uint32_t i = 0; i < 4; ++i) + { + values[i] = __half22float2(reinterpret_cast<__half2&>(fp16Pairs[i])); + } +#endif +} + +template +__device__ void store16BitValues(T* output, std::uint32_t elementOffset, float2 const (&values)[4], float scale) +{ + // Store through uint4 to preserve STG.128; nvcc may scalarize PackedVec. + std::uint32_t outputWords[4]; +#pragma unroll + for (std::uint32_t i = 0; i < 4; ++i) + { + float2 const scaled = make_float2(values[i].x * scale, values[i].y * scale); + if constexpr (std::is_same_v) + { + half2 const pair = __float22half2_rn(scaled); + outputWords[i] = reinterpret_cast(pair); + } + else + { + __nv_bfloat162 const pair = __float22bfloat162_rn(scaled); + outputWords[i] = reinterpret_cast(pair); + } + } + uint4 const outputGrain = make_uint4(outputWords[0], outputWords[1], outputWords[2], outputWords[3]); + reinterpret_cast(output + elementOffset)[0] = outputGrain; +} + +// Preserve independent source-FP8 and destination-NVFP4 global scales. +// Restore in FP32, round pairs to FP16, then reduce each 16-value grain. +__device__ uint2 quantizeFp8GrainToNvfp4(PackedVec<__nv_fp8_e4m3> const& grain, float fp8ScaleQuantOrig, + float nvfp4ScaleOrigQuant, std::uint8_t* scaleOutput) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + // Use the production packed FP8x2 conversion surface. + PackedVec restored[2]; +#pragma unroll + for (std::uint32_t pair = 0; pair < 8; ++pair) + { + float2 values = static_cast(grain.elts[pair]); + values.x *= fp8ScaleQuantOrig; + values.y *= fp8ScaleQuantOrig; + restored[pair / 4U].elts[pair % 4U] = __float22half2_rn(values); + } + + auto firstHalfMax = cuda_abs(restored[0].elts[0]); + auto secondHalfMax = cuda_abs(restored[1].elts[0]); +#pragma unroll + for (std::uint32_t i = 1; i < 4; ++i) + { + firstHalfMax = cuda_max(firstHalfMax, cuda_abs(restored[0].elts[i])); + secondHalfMax = cuda_max(secondHalfMax, cuda_abs(restored[1].elts[i])); + } + auto const localMax = cuda_max(firstHalfMax, secondHalfMax); + float const vecMax = static_cast(cuda_max(localMax.x, localMax.y)); + + float scaleValue = nvfp4ScaleOrigQuant * (vecMax * reciprocal_approximate_ftz(6.0F)); + __nv_fp8_e4m3 const roundedScale(scaleValue); + *scaleOutput = roundedScale.__x; + scaleValue = static_cast(roundedScale); + float const outputScale = vecMax != 0.0F + ? reciprocal_approximate_ftz(scaleValue * reciprocal_approximate_ftz(nvfp4ScaleOrigQuant)) + : 0.0F; + + std::uint32_t packed[2]; +#pragma unroll + for (std::uint32_t halfGroup = 0; halfGroup < 2; ++halfGroup) + { + float2 values[4]; +#pragma unroll + for (std::uint32_t i = 0; i < 4; ++i) + { + values[i] = __half22float2(restored[halfGroup].elts[i]); + values[i].x *= outputScale; + values[i].y *= outputScale; + } + packed[halfGroup] = fp32_vec_to_e2m1(values); + } + return make_uint2(packed[0], packed[1]); +#else + static_cast(grain); + static_cast(fp8ScaleQuantOrig); + static_cast(nvfp4ScaleOrigQuant); + static_cast(scaleOutput); + return make_uint2(0U, 0U); +#endif +} + +// Restore one natural 16-value NVFP4 scale group. +template +__device__ float onboardDequantScale(std::uint8_t encodedScale, Nvfp4ColdPageKernelParams const& params) +{ + __nv_fp8_e4m3 blockScale; + blockScale.__x = encodedScale; + float scale = static_cast(blockScale) * params.nvfp4ScaleQuantOrig; + if constexpr (std::is_same_v) + { + scale *= params.fp8ScaleOrigQuant; + } + return scale; +} + +template +__device__ void restoreNvfp4Pair(uint2 packedPair, T* output, std::uint32_t firstHalfGroup, float dequantScale) +{ + std::uint32_t const packedWords[2] = {packedPair.x, packedPair.y}; + if constexpr (!std::is_same_v) + { +#pragma unroll + for (std::uint32_t laneInScale = 0; laneInScale < 2; ++laneInScale) + { + float2 values[4]; + unpackE2m1ToFloat(packedWords[laneInScale], values); + store16BitValues(output, (firstHalfGroup + laneInScale) * kElementsPerHalfGroup, values, dequantScale); + } + } + else + { + std::uint64_t packedFp8[2]; +#pragma unroll + for (std::uint32_t laneInScale = 0; laneInScale < 2; ++laneInScale) + { + float2 values[4]; + unpackE2m1ToFloat(packedWords[laneInScale], values); +#pragma unroll + for (std::uint32_t i = 0; i < 4; ++i) + { + values[i].x *= dequantScale; + values[i].y *= dequantScale; + } + packedFp8[laneInScale] = fp32_vec_to_e4m3(values); + } + reinterpret_cast(output)[firstHalfGroup / 2U] = collectTwoFp8Words(packedFp8[0], packedFp8[1]); + } +} + +// FP16/BF16 GPU Page -> mapped-Host NVFP4 in bounded tiles. +template +__global__ void offloadFrom16BitTiledKernel( + std::array const __grid_constant__ pages, + Nvfp4ColdPageWideTable const __grid_constant__ wide, Nvfp4ColdPageIntegerTable const __grid_constant__ integers, + Nvfp4ColdPageScaleTable const __grid_constant__ scales, std::uint8_t* coldBase, std::size_t coldPageBytes) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm volatile("griddepcontrol.launch_dependents;\n"); + + std::uint32_t const bufferIndex = blockIdx.y; + auto const buffer = loadBuffer(bufferIndex, wide, integers, scales); + auto const task = resolveOffloadTask(pages[blockIdx.z], buffer, coldBase, coldPageBytes); + + asm volatile("griddepcontrol.wait;\n" : : : "memory"); + if (buffer.transform == Nvfp4ColdPageTransform::kLosslessCopy) + { + copyLosslessBytes(task.raw, task.coldData, buffer.rawBytes); + clearColdPadding(task, buffer); + return; + } + + auto const params = buffer.params; + std::uint32_t const halfGroupsPerBuffer = halfGroupCount(params); + std::uint32_t const tileHalfGroups = tileHalfGroupCount(params); + std::uint32_t const packedStageCapacityBytes = packedStageBytesForHalfGroups(tileHalfGroups); + + // Use a four-stage 16-byte cp.async ring and tile-bounded compact staging. + __shared__ __align__(16) PackedVec rawStages[kAsyncStages][kThreadsPerBlock]; + extern __shared__ __align__(16) std::uint8_t compactStages[]; + auto* packedStages = reinterpret_cast(compactStages); + auto* scaleStages = compactStages + packedStageCapacityBytes; + + for (std::uint32_t firstHalfGroup = blockIdx.x * tileHalfGroups; firstHalfGroup < halfGroupsPerBuffer; + firstHalfGroup += gridDim.x * tileHalfGroups) + { + std::uint32_t const halfGroups = std::min(tileHalfGroups, halfGroupsPerBuffer - firstHalfGroup); + std::uint32_t const iterations = (halfGroups + kThreadsPerBlock - 1U) / kThreadsPerBlock; + + for (std::uint32_t iteration = 0; iteration < iterations + kAsyncStages; ++iteration) + { + std::uint32_t const stage = iteration % kAsyncStages; + if (iteration >= kAsyncStages) + { + std::uint32_t const transformIteration = iteration - kAsyncStages; + std::uint32_t const localHalfGroup = kThreadsPerBlock * transformIteration + threadIdx.x; + cp_async_wait_group(); + if (localHalfGroup < halfGroups) + { + std::uint32_t const globalHalfGroup = firstHalfGroup + localHalfGroup; + std::uint32_t const laneInScale = globalHalfGroup & 1U; + + // Even tile boundaries preserve the 16-value scale groups. + std::uint32_t const localScaleOffset = localHalfGroup >> 1U; + std::uint8_t* scale = laneInScale == 0 ? scaleStages + localScaleOffset : nullptr; + PackedVec input = rawStages[stage][threadIdx.x]; + packedStages[localHalfGroup] = cvt_warp_fp16_to_fp4( + input, params.nvfp4ScaleOrigQuant, scale); + } + } + + std::uint32_t const localLoadHalfGroup = kThreadsPerBlock * iteration + threadIdx.x; + bool const valid = localLoadHalfGroup < halfGroups; + auto const* rawInput = reinterpret_cast const*>(task.raw); + auto const* source = valid ? rawInput + firstHalfGroup + localLoadHalfGroup : rawInput; + copyAsyncGlobalToShared(&rawStages[stage][threadIdx.x], source, valid); + cp_async_commit_group(); + } + + // Publish quant results and finish the flush before reusing staging. + cp_async_wait_group<0>(); + __syncthreads(); + flushCompactRangeToHost(compactStages, task, packedStageCapacityBytes, packedBytesForHalfGroups(firstHalfGroup), + packedBytesForHalfGroups(halfGroups), scaleBytesForHalfGroups(firstHalfGroup), + scaleBytesForHalfGroups(halfGroups)); + __syncthreads(); + } + clearColdPadding(task, buffer); +#endif +} + +// FP8 E4M3 GPU Page -> mapped-Host NVFP4 in bounded tiles. +__global__ void offloadFromFp8TiledKernel( + std::array const __grid_constant__ pages, + Nvfp4ColdPageWideTable const __grid_constant__ wide, Nvfp4ColdPageIntegerTable const __grid_constant__ integers, + Nvfp4ColdPageScaleTable const __grid_constant__ scales, std::uint8_t* coldBase, std::size_t coldPageBytes) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm volatile("griddepcontrol.launch_dependents;\n"); + + std::uint32_t const bufferIndex = blockIdx.y; + auto const buffer = loadBuffer(bufferIndex, wide, integers, scales); + auto const task = resolveOffloadTask(pages[blockIdx.z], buffer, coldBase, coldPageBytes); + + asm volatile("griddepcontrol.wait;\n" : : : "memory"); + if (buffer.transform == Nvfp4ColdPageTransform::kLosslessCopy) + { + copyLosslessBytes(task.raw, task.coldData, buffer.rawBytes); + clearColdPadding(task, buffer); + return; + } + + auto const params = buffer.params; + std::uint32_t const halfGroupsPerBuffer = halfGroupCount(params); + std::uint32_t const tileHalfGroups = tileHalfGroupCount(params); + std::uint32_t const packedStageCapacityBytes = packedStageBytesForHalfGroups(tileHalfGroups); + + // Each cp.async moves one 16-byte grain in the production PackedVec layout. + __shared__ __align__(16) PackedVec<__nv_fp8_e4m3> rawStages[kAsyncStages][kThreadsPerBlock]; + extern __shared__ __align__(16) std::uint8_t compactStages[]; + auto* packedStages = reinterpret_cast(compactStages); + auto* scaleStages = compactStages + packedStageCapacityBytes; + + for (std::uint32_t firstHalfGroup = blockIdx.x * tileHalfGroups; firstHalfGroup < halfGroupsPerBuffer; + firstHalfGroup += gridDim.x * tileHalfGroups) + { + std::uint32_t const halfGroups = std::min(tileHalfGroups, halfGroupsPerBuffer - firstHalfGroup); + std::uint32_t const firstGrain = firstHalfGroup / 2U; + std::uint32_t const grains = halfGroups / 2U; + std::uint32_t const iterations = (grains + kThreadsPerBlock - 1U) / kThreadsPerBlock; + + for (std::uint32_t iteration = 0; iteration < iterations + kAsyncStages; ++iteration) + { + std::uint32_t const stage = iteration % kAsyncStages; + if (iteration >= kAsyncStages) + { + std::uint32_t const transformIteration = iteration - kAsyncStages; + cp_async_wait_group(); + // One lane owns the complete 16-value scale group. + std::uint32_t const localGrain = kThreadsPerBlock * transformIteration + threadIdx.x; + if (localGrain < grains) + { + uint2 const packed = quantizeFp8GrainToNvfp4(rawStages[stage][threadIdx.x], + params.fp8ScaleQuantOrig, params.nvfp4ScaleOrigQuant, scaleStages + localGrain); + reinterpret_cast(packedStages)[localGrain] = packed; + } + } + + std::uint32_t const localLoadGrain = kThreadsPerBlock * iteration + threadIdx.x; + bool const valid = localLoadGrain < grains; + auto const* rawInput = reinterpret_cast const*>(task.raw); + auto const* source = valid ? rawInput + firstGrain + localLoadGrain : rawInput; + copyAsyncGlobalToShared(&rawStages[stage][threadIdx.x], source, valid); + cp_async_commit_group(); + } + + cp_async_wait_group<0>(); + __syncthreads(); + flushCompactRangeToHost(compactStages, task, packedStageCapacityBytes, packedBytesForHalfGroups(firstHalfGroup), + packedBytesForHalfGroups(halfGroups), scaleBytesForHalfGroups(firstHalfGroup), + scaleBytesForHalfGroups(halfGroups)); + __syncthreads(); + } + clearColdPadding(task, buffer); +#endif +} + +// Load packed values and one-byte scales from mapped Host memory into shared. +// Use uint4/uint2 when aligned and byte tails otherwise. +__device__ void loadCompactRangeFromHost(std::uint8_t* compactStages, OnboardBufferTask const& task, + std::uint32_t packedStageCapacityBytes, std::uint32_t packedSourceOffset, std::uint32_t packedBytes, + std::uint32_t scaleSourceOffset, std::uint32_t scaleBytes) +{ + auto const* packedSource = task.coldData + packedSourceOffset; + auto* packedDestination = compactStages; + auto const* scaleSource = task.coldScale + scaleSourceOffset; + auto* scaleDestination = compactStages + packedStageCapacityBytes; + bool const alignedPacked = reinterpret_cast(packedSource) % sizeof(uint4) == 0 + && reinterpret_cast(packedDestination) % sizeof(uint4) == 0; + bool const alignedPackedPair = reinterpret_cast(packedSource) % sizeof(uint2) == 0 + && reinterpret_cast(packedDestination) % sizeof(uint2) == 0; + std::uint32_t const packedVectorBytes = alignedPacked ? packedBytes - packedBytes % sizeof(uint4) : 0; + std::uint32_t const packedPairBytes = alignedPackedPair ? packedBytes - packedBytes % sizeof(uint2) : 0; + bool const alignedScale = reinterpret_cast(scaleSource) % sizeof(uint4) == 0 + && reinterpret_cast(scaleDestination) % sizeof(uint4) == 0; + std::uint32_t const scaleVectorBytes = alignedScale ? scaleBytes - scaleBytes % sizeof(uint4) : 0; + std::uint32_t const packedGrains = packedVectorBytes / sizeof(uint4); + std::uint32_t const scaleGrains = scaleVectorBytes / sizeof(uint4); + std::uint32_t const totalGrains = packedGrains + scaleGrains; + std::uint32_t const iterations = (totalGrains + blockDim.x - 1U) / blockDim.x; + + for (std::uint32_t iteration = 0; iteration < iterations; ++iteration) + { + if (iteration >= kHostLoadAsyncStages) + { + cp_async_wait_group(); + } + std::uint32_t const grain = blockDim.x * iteration + threadIdx.x; + bool const valid = grain < totalGrains; + auto const* source = reinterpret_cast(packedSource); + auto* destination = reinterpret_cast(packedDestination); + if (valid && grain < packedGrains) + { + source += grain; + destination += grain; + } + else if (valid) + { + source = reinterpret_cast(scaleSource) + grain - packedGrains; + destination = reinterpret_cast(scaleDestination) + grain - packedGrains; + } + copyAsyncGlobalToShared(destination, source, valid); + cp_async_commit_group(); + } + cp_async_wait_group<0>(); + + // headDim % 16 makes packed intervals exact uint2 scale groups. + for (std::uint32_t pair = packedVectorBytes / sizeof(uint2) + threadIdx.x; pair < packedPairBytes / sizeof(uint2); + pair += blockDim.x) + { + reinterpret_cast(packedDestination)[pair] = reinterpret_cast(packedSource)[pair]; + } + for (std::uint32_t byte = packedPairBytes + threadIdx.x; byte < packedBytes; byte += blockDim.x) + { + packedDestination[byte] = packedSource[byte]; + } + for (std::uint32_t byte = scaleVectorBytes + threadIdx.x; byte < scaleBytes; byte += blockDim.x) + { + scaleDestination[byte] = scaleSource[byte]; + } + __syncthreads(); +} + +// Mapped-Host NVFP4 -> runtime GPU Page in bounded tiles. +template +__global__ void onboardTiledKernel(std::array const __grid_constant__ pages, + Nvfp4ColdPageWideTable const __grid_constant__ wide, Nvfp4ColdPageIntegerTable const __grid_constant__ integers, + Nvfp4ColdPageScaleTable const __grid_constant__ scales, std::uint8_t const* coldBase, std::size_t coldPageBytes) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm volatile("griddepcontrol.launch_dependents;\n"); + + std::uint32_t const bufferIndex = blockIdx.y; + auto const buffer = loadBuffer(bufferIndex, wide, integers, scales); + auto const task = resolveOnboardTask(pages[blockIdx.z], buffer, coldBase, coldPageBytes); + + asm volatile("griddepcontrol.wait;\n" : : : "memory"); + if (buffer.transform == Nvfp4ColdPageTransform::kLosslessCopy) + { + copyLosslessBytes(task.coldData, task.raw, buffer.rawBytes); + return; + } + + auto const params = buffer.params; + std::uint32_t const halfGroupsPerBuffer = halfGroupCount(params); + std::uint32_t const tileHalfGroups = tileHalfGroupCount(params); + std::uint32_t const packedStageCapacityBytes = packedStageBytesForHalfGroups(tileHalfGroups); + extern __shared__ __align__(16) std::uint8_t compactStages[]; + auto* rawOutput = reinterpret_cast(task.raw); + + for (std::uint32_t firstHalfGroup = blockIdx.x * tileHalfGroups; firstHalfGroup < halfGroupsPerBuffer; + firstHalfGroup += gridDim.x * tileHalfGroups) + { + std::uint32_t const halfGroups = std::min(tileHalfGroups, halfGroupsPerBuffer - firstHalfGroup); + std::uint32_t const packedBytes = packedBytesForHalfGroups(halfGroups); + std::uint32_t const scaleBytes = scaleBytesForHalfGroups(halfGroups); + + // Stage packed data and scales before dequantization. + loadCompactRangeFromHost(compactStages, task, packedStageCapacityBytes, + packedBytesForHalfGroups(firstHalfGroup), packedBytes, scaleBytesForHalfGroups(firstHalfGroup), scaleBytes); + + auto const* packedStages = reinterpret_cast(compactStages); + auto const* scaleStages = compactStages + packedStageCapacityBytes; + std::uint32_t const packedGrains = packedBytes / sizeof(uint4); + for (std::uint32_t localGrain = threadIdx.x; localGrain < packedGrains; localGrain += blockDim.x) + { + uint4 const packedGrain = packedStages[localGrain]; + std::uint32_t const packedWords[4] = {packedGrain.x, packedGrain.y, packedGrain.z, packedGrain.w}; +#pragma unroll + for (std::uint32_t pair = 0; pair < 2; ++pair) + { + std::uint32_t const localScaleGroup = localGrain * 2U + pair; + std::uint32_t const firstPairHalfGroup = firstHalfGroup + localScaleGroup * 2U; + + restoreNvfp4Pair(make_uint2(packedWords[pair * 2U], packedWords[pair * 2U + 1U]), rawOutput, + firstPairHalfGroup, onboardDequantScale(scaleStages[localScaleGroup], params)); + } + } + if (packedBytes % sizeof(uint4) != 0U && threadIdx.x == 0) + { + std::uint32_t const localScaleGroup = packedGrains * 2U; + std::uint32_t const firstPairHalfGroup = firstHalfGroup + localScaleGroup * 2U; + restoreNvfp4Pair(reinterpret_cast(compactStages)[localScaleGroup], rawOutput, + firstPairHalfGroup, onboardDequantScale(scaleStages[localScaleGroup], params)); + } + + // Finish consumers before reusing shared memory. + __syncthreads(); + } +#endif +} + +// Submit one whole KVCM Page batch through the fixed 256-descriptor kernel ABI. +template +void launchPageChunks(Kernel kernel, void const* pages, std::size_t numPages, std::int64_t const* wide, + std::int32_t const* integers, float const* scales, std::uint32_t numBuffers, std::uint32_t maxHalfGroupsPerTile, + std::size_t coldPageBytes, ColdPointer coldBase, cudaStream_t stream) +{ + dim3 const block(kThreadsPerBlock); + auto const* pageBytes = static_cast(pages); + std::size_t offset = 0; + while (offset < numPages) + { + std::uint32_t const numChunkPages + = static_cast(std::min(numPages - offset, kMaxTasksPerLaunch)); + auto const* chunkPages = pageBytes + offset * sizeof(PageIndexPairView); + + // CUDA copies the full by-value array, so pad only the final partial chunk. + std::array paddedPages{}; + if (numChunkPages < kMaxTasksPerLaunch) + { + std::memcpy(paddedPages.data(), chunkPages, numChunkPages * sizeof(PageIndexPairView)); + chunkPages = reinterpret_cast(paddedPages.data()); + } + + cudaLaunchAttribute attribute{}; + attribute.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute.val.programmaticStreamSerializationAllowed = common::getEnvEnablePDL() ? 1 : 0; + + cudaLaunchConfig_t config{}; + config.gridDim = dim3(kMappedHostGridSplits, numBuffers, numChunkPages); + config.blockDim = block; + config.dynamicSmemBytes = compactStageBytesForHalfGroups(maxHalfGroupsPerTile); + config.stream = stream; + config.attrs = &attribute; + config.numAttrs = 1; + + void* arguments[] = {const_cast(chunkPages), const_cast(wide), + const_cast(integers), const_cast(scales), &coldBase, &coldPageBytes}; + TLLM_CUDA_CHECK(cudaLaunchKernelExC(&config, reinterpret_cast(kernel), arguments)); + offset += numChunkPages; + } +} + +} // namespace + +void invokeNvfp4ColdPageEncode(void const* pages, std::size_t numPages, std::int64_t const* wide, + std::int32_t const* integers, float const* scales, std::uint32_t numBuffers, std::uint32_t maxHalfGroupsPerTile, + std::size_t coldPageBytes, Nvfp4ColdPageRuntimeType runtimeType, void* coldBase, cudaStream_t stream) +{ + if (numPages == 0) + { + return; + } + TLLM_CHECK_WITH_INFO(pages != nullptr, "pages must not be null"); + TLLM_CHECK_WITH_INFO(coldBase != nullptr, "coldBase must not be null"); + switch (runtimeType) + { + case Nvfp4ColdPageRuntimeType::kFloat16: + launchPageChunks(offloadFrom16BitTiledKernel, pages, numPages, wide, integers, scales, numBuffers, + maxHalfGroupsPerTile, coldPageBytes, static_cast(coldBase), stream); + break; + case Nvfp4ColdPageRuntimeType::kBfloat16: + launchPageChunks(offloadFrom16BitTiledKernel<__nv_bfloat16>, pages, numPages, wide, integers, scales, + numBuffers, maxHalfGroupsPerTile, coldPageBytes, static_cast(coldBase), stream); + break; + case Nvfp4ColdPageRuntimeType::kFp8E4m3: + launchPageChunks(offloadFromFp8TiledKernel, pages, numPages, wide, integers, scales, numBuffers, + maxHalfGroupsPerTile, coldPageBytes, static_cast(coldBase), stream); + break; + default: TLLM_THROW("Unsupported NVFP4 cold-page runtime type"); + } +} + +void invokeNvfp4ColdPageDecode(void const* pages, std::size_t numPages, std::int64_t const* wide, + std::int32_t const* integers, float const* scales, std::uint32_t numBuffers, std::uint32_t maxHalfGroupsPerTile, + std::size_t coldPageBytes, Nvfp4ColdPageRuntimeType runtimeType, void const* coldBase, cudaStream_t stream) +{ + if (numPages == 0) + { + return; + } + TLLM_CHECK_WITH_INFO(pages != nullptr, "pages must not be null"); + TLLM_CHECK_WITH_INFO(coldBase != nullptr, "coldBase must not be null"); + switch (runtimeType) + { + case Nvfp4ColdPageRuntimeType::kFloat16: + launchPageChunks(onboardTiledKernel, pages, numPages, wide, integers, scales, numBuffers, + maxHalfGroupsPerTile, coldPageBytes, static_cast(coldBase), stream); + break; + case Nvfp4ColdPageRuntimeType::kBfloat16: + launchPageChunks(onboardTiledKernel<__nv_bfloat16>, pages, numPages, wide, integers, scales, numBuffers, + maxHalfGroupsPerTile, coldPageBytes, static_cast(coldBase), stream); + break; + case Nvfp4ColdPageRuntimeType::kFp8E4m3: + launchPageChunks(onboardTiledKernel<__nv_fp8_e4m3>, pages, numPages, wide, integers, scales, numBuffers, + maxHalfGroupsPerTile, coldPageBytes, static_cast(coldBase), stream); + break; + default: TLLM_THROW("Unsupported NVFP4 cold-page runtime type"); + } +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/nvfp4ColdPageKernels.h b/cpp/tensorrt_llm/kernels/nvfp4ColdPageKernels.h new file mode 100644 index 000000000000..6a249b246cd9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/nvfp4ColdPageKernels.h @@ -0,0 +1,65 @@ +/* + * 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/config.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +//! Active GPU representation encoded into cold Pages. +enum class Nvfp4ColdPageRuntimeType : std::uint8_t +{ + kFloat16 = 0, + kBfloat16 = 1, + kFp8E4m3 = 2, +}; + +inline constexpr std::uint32_t kNvfp4ColdPageMaxBuffersPerLaunch = 256; +inline constexpr std::uint32_t kNvfp4ColdPageWideFields = 6; +inline constexpr std::uint32_t kNvfp4ColdPageIntegerFields = 5; +inline constexpr std::uint32_t kNvfp4ColdPageScaleFields = 4; + +using Nvfp4ColdPageWideTable + = std::array, kNvfp4ColdPageMaxBuffersPerLaunch>; +using Nvfp4ColdPageIntegerTable + = std::array, kNvfp4ColdPageMaxBuffersPerLaunch>; +using Nvfp4ColdPageScaleTable + = std::array, kNvfp4ColdPageMaxBuffersPerLaunch>; + +//! Compress one whole KVCM Page-index batch; the launcher performs 256-Page chunking internally. +void invokeNvfp4ColdPageEncode(void const* pages, std::size_t numPages, std::int64_t const* wide, + std::int32_t const* integers, float const* scales, std::uint32_t numBuffers, std::uint32_t maxHalfGroupsPerTile, + std::size_t coldPageBytes, Nvfp4ColdPageRuntimeType runtimeType, void* coldBase, cudaStream_t stream); + +//! Restore one whole KVCM Page-index batch; the launcher performs 256-Page chunking internally. +void invokeNvfp4ColdPageDecode(void const* pages, std::size_t numPages, std::int64_t const* wide, + std::int32_t const* integers, float const* scales, std::uint32_t numBuffers, std::uint32_t maxHalfGroupsPerTile, + std::size_t coldPageBytes, Nvfp4ColdPageRuntimeType runtimeType, void const* coldBase, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/nanobind/CMakeLists.txt b/cpp/tensorrt_llm/nanobind/CMakeLists.txt index 5dc1de84a308..982fe0369431 100755 --- a/cpp/tensorrt_llm/nanobind/CMakeLists.txt +++ b/cpp/tensorrt_llm/nanobind/CMakeLists.txt @@ -17,6 +17,7 @@ set(SRCS executor/bindings.cpp executor/executorConfig.cpp executor/request.cpp + kvCacheCompression/bindings.cpp process_group/bindings.cpp runtime/bindings.cpp runtime/hostfunc.cpp diff --git a/cpp/tensorrt_llm/nanobind/bindings.cpp b/cpp/tensorrt_llm/nanobind/bindings.cpp index a2054dbd7217..0371c48d5954 100644 --- a/cpp/tensorrt_llm/nanobind/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/bindings.cpp @@ -44,6 +44,7 @@ #include "tensorrt_llm/nanobind/batch_manager/llmRequest.h" #include "tensorrt_llm/nanobind/common/tllmExceptions.h" #include "tensorrt_llm/nanobind/executor/bindings.h" +#include "tensorrt_llm/nanobind/kvCacheCompression/bindings.h" #include "tensorrt_llm/nanobind/process_group/bindings.h" #include "tensorrt_llm/nanobind/runtime/bindings.h" #include "tensorrt_llm/nanobind/suffixAutomaton/bindings.h" @@ -138,6 +139,8 @@ NB_MODULE(TRTLLM_NB_MODULE, m) = mInternalBatchManager.def_submodule("kv_cache_manager_v2_utils", "KV Cache Manager V2 Utils bindings"); auto mInternalBatchManagerKvCacheV2 = mInternalBatchManager.def_submodule("kv_cache_manager_v2", "KV Cache Manager V2 bindings"); + auto mInternalKvCacheCompression + = mInternal.def_submodule("kv_cache_compression", "KV cache compression internal bindings"); tensorrt_llm::nanobind::batch_manager::KvCacheManagerV2Bindings::initBindings(mInternalBatchManagerKvCacheV2); auto mInternalThop = mInternal.def_submodule("thop", "Torch op internal bindings"); auto mExceptions = m.def_submodule("exceptions", "Exceptions internal bindings"); @@ -145,6 +148,7 @@ NB_MODULE(TRTLLM_NB_MODULE, m) tensorrt_llm::nanobind::executor::initBindings(mExecutor); tensorrt_llm::nanobind::runtime::initBindingsEarly(mInternalRuntime); tensorrt_llm::nanobind::common::initExceptionsBindings(mExceptions); + tensorrt_llm::nanobind::kv_cache_compression::initBindings(mInternalKvCacheCompression); tensorrt_llm::nanobind::thop::initBindings(mInternalThop); auto buildInfo = m.def_submodule("BuildInfo"); diff --git a/cpp/tensorrt_llm/nanobind/kvCacheCompression/bindings.cpp b/cpp/tensorrt_llm/nanobind/kvCacheCompression/bindings.cpp new file mode 100644 index 000000000000..352721a7d1ca --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/kvCacheCompression/bindings.cpp @@ -0,0 +1,206 @@ +/* + * 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 "bindings.h" +#include "tensorrt_llm/batch_manager/kv_cache_compression/nativeColdPageCodec.h" +#include "tensorrt_llm/kernels/nvfp4ColdPageKernels.h" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace nb = nanobind; +namespace compression = tensorrt_llm::kv_cache_compression; +namespace kv = tensorrt_llm::batch_manager::kv_cache_manager_v2; + +namespace tensorrt_llm::nanobind::kv_cache_compression +{ +namespace +{ + +static_assert(sizeof(kv::PageIndexPair) == 8); +static_assert(alignof(kv::PageIndexPair) == 8); +static_assert(offsetof(kv::PageIndexPair, dst) == 0); +static_assert(offsetof(kv::PageIndexPair, src) == 4); +static_assert(std::is_trivially_copyable_v); + +//! Algorithm-neutral adapter from KVCM migration calls to a Python provider. +class PythonColdPageCodec final : public compression::NativeColdPageCodec +{ +public: + PythonColdPageCodec(nb::handle provider, nb::handle codecState) + : NativeColdPageCodec(readLayerIds(codecState)) + , mProvider(provider.ptr()) + , mCodecState(codecState.ptr()) + { + Py_INCREF(mProvider); + Py_INCREF(mCodecState); + } + + ~PythonColdPageCodec() override + { + if (Py_IsInitialized()) + { + nb::gil_scoped_acquire acquire; + Py_DECREF(mCodecState); + Py_DECREF(mProvider); + } + } + +private: + static std::set readLayerIds(nb::handle codecState) + { + if (codecState.is_none()) + { + throw std::invalid_argument("Cold-page codec state must not be None"); + } + auto const layerIds = nb::cast>(codecState.attr("layer_ids")); + std::set result(layerIds.begin(), layerIds.end()); + if (result.size() != layerIds.size()) + { + throw std::invalid_argument("Cold-page codec state layer IDs must be unique"); + } + return result; + } + + std::vector configureProvider( + std::vector const& lifecycles) override + { + nb::gil_scoped_acquire acquire; + try + { + return nb::cast>( + nb::borrow(mProvider).attr("configure")(nb::borrow(mCodecState), lifecycles)); + } + catch (nb::python_error const& error) + { + throw std::runtime_error(error.what()); + } + } + + void encodeProvider(std::size_t lifecycleIndex, void* coldBase, kv::PageIndexPair const* pageIndices, + std::size_t numPages, cudaStream_t stream) override + { + invoke("encode_cold_pages", lifecycleIndex, coldBase, pageIndices, numPages, stream); + } + + void decodeProvider(std::size_t lifecycleIndex, void const* coldBase, kv::PageIndexPair const* pageIndices, + std::size_t numPages, cudaStream_t stream) override + { + invoke("decode_cold_pages", lifecycleIndex, coldBase, pageIndices, numPages, stream); + } + + template + void invoke(char const* method, std::size_t lifecycleIndex, ColdPointer coldBase, + kv::PageIndexPair const* pageIndices, std::size_t numPages, cudaStream_t stream) + { + // Forward the complete KVCM batch once. The native launcher owns any chunking. + nb::gil_scoped_acquire acquire; + try + { + nb::borrow(mProvider).attr(method)(nb::borrow(mCodecState), lifecycleIndex, + reinterpret_cast(coldBase), reinterpret_cast(pageIndices), numPages, + reinterpret_cast(stream)); + } + catch (nb::python_error const& error) + { + throw std::runtime_error(error.what()); + } + } + + PyObject* mProvider; + PyObject* mCodecState; +}; + +} // namespace + +void initBindings(nb::module_& module) +{ + nb::enum_(module, "ColdPageIndexLocation") + .value("BAD_LOCATION", kv::PageIndexLocation::kBadLocation) + .value("HOST", kv::PageIndexLocation::kHost) + .value("DEVICE", kv::PageIndexLocation::kDevice); + + nb::class_(module, "ResolvedHotBuffer") + .def_ro("raw_base", &compression::ResolvedHotBuffer::rawBase) + .def_ro("raw_slot_bytes", &compression::ResolvedHotBuffer::rawSlotBytes) + .def_ro("raw_bytes", &compression::ResolvedHotBuffer::rawBytes); + + nb::class_(module, "ResolvedHotLifecycle") + .def_prop_ro("life_cycle_id", + [](compression::ResolvedHotLifecycle const& lifecycle) { return lifecycle.lifeCycleId.value(); }) + .def_ro("layers", &compression::ResolvedHotLifecycle::layers); + + nb::class_(module, "ColdPageLifecycleProperties") + .def(nb::init<>()) + .def_rw("cold_page_bytes", &compression::ColdPageLifecycleProperties::coldPageBytes) + .def_rw("page_index_location", &compression::ColdPageLifecycleProperties::pageIndexLocation); + + module.def( + "create_python_cold_page_codec", + [](nb::handle provider, nb::handle codecState) -> std::unique_ptr + { return std::make_unique(provider, codecState); }, + nb::arg("provider"), nb::arg("codec_state")); + + // The Python provider traffics in raw KVCM addresses, so the launcher trampolines take scalar integers. + module.def("nvfp4_cold_page_encode", + [](std::int64_t pageIndices, std::int64_t numPages, std::int64_t wide, std::int64_t integers, + std::int64_t scales, std::int64_t numBuffers, std::int64_t maxHalfGroupsPerTile, std::int64_t coldPageBytes, + std::int64_t runtimeType, std::int64_t coldBase, std::int64_t stream) + { + tensorrt_llm::kernels::invokeNvfp4ColdPageEncode( + reinterpret_cast(static_cast(pageIndices)), + static_cast(numPages), + reinterpret_cast(static_cast(wide)), + reinterpret_cast(static_cast(integers)), + reinterpret_cast(static_cast(scales)), + static_cast(numBuffers), static_cast(maxHalfGroupsPerTile), + static_cast(coldPageBytes), + static_cast(runtimeType), + reinterpret_cast(static_cast(coldBase)), + reinterpret_cast(static_cast(stream))); + }); + + module.def("nvfp4_cold_page_decode", + [](std::int64_t pageIndices, std::int64_t numPages, std::int64_t wide, std::int64_t integers, + std::int64_t scales, std::int64_t numBuffers, std::int64_t maxHalfGroupsPerTile, std::int64_t coldPageBytes, + std::int64_t runtimeType, std::int64_t coldBase, std::int64_t stream) + { + tensorrt_llm::kernels::invokeNvfp4ColdPageDecode( + reinterpret_cast(static_cast(pageIndices)), + static_cast(numPages), + reinterpret_cast(static_cast(wide)), + reinterpret_cast(static_cast(integers)), + reinterpret_cast(static_cast(scales)), + static_cast(numBuffers), static_cast(maxHalfGroupsPerTile), + static_cast(coldPageBytes), + static_cast(runtimeType), + reinterpret_cast(static_cast(coldBase)), + reinterpret_cast(static_cast(stream))); + }); +} + +} // namespace tensorrt_llm::nanobind::kv_cache_compression diff --git a/cpp/tensorrt_llm/nanobind/kvCacheCompression/bindings.h b/cpp/tensorrt_llm/nanobind/kvCacheCompression/bindings.h new file mode 100644 index 000000000000..f85d098097e8 --- /dev/null +++ b/cpp/tensorrt_llm/nanobind/kvCacheCompression/bindings.h @@ -0,0 +1,27 @@ +/* + * 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 + +namespace tensorrt_llm::nanobind::kv_cache_compression +{ + +void initBindings(::nanobind::module_& module); + +} // namespace tensorrt_llm::nanobind::kv_cache_compression diff --git a/cpp/tests/unit_tests/CMakeLists.txt b/cpp/tests/unit_tests/CMakeLists.txt index 9d22bd03b52a..554be48db1da 100644 --- a/cpp/tests/unit_tests/CMakeLists.txt +++ b/cpp/tests/unit_tests/CMakeLists.txt @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & +# SPDX-FileCopyrightText: Copyright (c) 2023-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 diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index 468542c1dc03..e74699fbf418 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -38,6 +38,9 @@ add_gtest(kvCacheManagerV2ColdPageCopyTest kvCacheManagerV2ColdPageCopyTest.cu) target_include_directories( kvCacheManagerV2ColdPageCopyTest PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager) +add_gtest(coldPageCodecTest coldPageCodecTest.cpp) +target_include_directories( + coldPageCodecTest PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager) add_gtest(kvCacheManagerV2TypedIndexTest kvCacheManagerV2TypedIndexTest.cpp) target_include_directories( kvCacheManagerV2TypedIndexTest diff --git a/cpp/tests/unit_tests/batch_manager/coldPageCodecTest.cpp b/cpp/tests/unit_tests/batch_manager/coldPageCodecTest.cpp new file mode 100644 index 000000000000..05ba3affa22a --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/coldPageCodecTest.cpp @@ -0,0 +1,270 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "tensorrt_llm/batch_manager/kv_cache_compression/nativeColdPageCodec.h" + +#include "kv_cache_manager_v2/utils/hostMem.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tensorrt_llm::kv_cache_compression +{ +namespace +{ + +namespace kv = batch_manager::kv_cache_manager_v2; + +static_assert(std::is_abstract_v); +static_assert(std::is_base_of_v); + +constexpr std::uintptr_t kGpuKBase = 0x100000; +constexpr std::uintptr_t kGpuVBase = 0x200000; +constexpr std::uintptr_t kColdBase = 0x300000; +constexpr std::uintptr_t kStreamValue = 0x7000; +constexpr std::size_t kRawBytes = 320; + +kv::PoolGroupDesc makeAttentionDesc(kv::PoolGroupIndex poolGroupIndex = kv::PoolGroupIndex{0}, + kv::LayerGroupId lifeCycle = kv::LayerGroupId{0}, std::size_t count = 1U, int firstLayer = 0, + std::uintptr_t keyBase = kGpuKBase, std::uintptr_t valueBase = kGpuVBase) +{ + kv::CoalescedBuffer keys{kRawBytes, {}}; + kv::CoalescedBuffer values{kRawBytes, {}}; + for (std::size_t index = 0; index < count; ++index) + { + auto const layerId = firstLayer + static_cast(index); + keys.bufferIds.push_back({layerId, "key"}); + values.bufferIds.push_back({layerId, "value"}); + } + auto const slotBytes = count * kRawBytes; + kv::SlotDescVariant variant{ + lifeCycle, kv::TypedVec{std::move(keys), std::move(values)}}; + return kv::PoolGroupDesc{poolGroupIndex, kv::SlotCount{512}, kv::SlotDesc{{std::move(variant)}}, + kv::TypedVec{ + {kv::PoolIndex{0}, keyBase, slotBytes}, {kv::PoolIndex{1}, valueBase, slotBytes}}}; +} + +kv::PoolGroupDesc makeLosslessDesc(kv::PoolGroupIndex poolGroupIndex, kv::LayerGroupId lifeCycle) +{ + kv::SlotDescVariant variant{lifeCycle, + kv::TypedVec{ + kv::CoalescedBuffer{64U, {{10, "ssm_state"}}}, kv::CoalescedBuffer{32U, {{10, "conv_state"}}}}}; + return {poolGroupIndex, kv::SlotCount{8}, kv::SlotDesc{{std::move(variant)}}, + kv::TypedVec{ + {kv::PoolIndex{0}, 0x400000, 64U}, {kv::PoolIndex{1}, 0x500000, 32U}}}; +} + +bool configureOne(kv::IKvCacheColdPageCodec& codec, kv::PoolGroupDesc const& desc) +{ + return codec.configure(&desc, kv::PoolGroupIndex{1}); +} + +class RecordingCodec final : public NativeColdPageCodec +{ +public: + explicit RecordingCodec(std::set layerIds) + : NativeColdPageCodec(std::move(layerIds)) + { + } + + std::vector configureProvider( + std::vector const& lifecycles) override + { + resolved = lifecycles; + if (failConfigure) + { + throw std::runtime_error("requested configure failure"); + } + return std::vector( + lifecycles.size(), ColdPageLifecycleProperties{777U, kv::PageIndexLocation::kHost}); + } + + void encodeProvider(std::size_t lifecycleIndex, void* coldBase, kv::PageIndexPair const* pageIndices, + std::size_t numPages, cudaStream_t stream) override + { + if (failBatches) + { + enqueueFailureMarker(stream); + throw std::runtime_error("requested batch failure"); + } + ++encodeCalls; + lastLifecycleIndex = lifecycleIndex; + lastColdBase = coldBase; + lastIndices.assign(pageIndices, pageIndices + numPages); + lastStream = stream; + } + + void decodeProvider(std::size_t lifecycleIndex, void const* coldBase, kv::PageIndexPair const* pageIndices, + std::size_t numPages, cudaStream_t stream) override + { + if (failBatches) + { + enqueueFailureMarker(stream); + throw std::runtime_error("requested batch failure"); + } + ++decodeCalls; + lastLifecycleIndex = lifecycleIndex; + lastColdBase = coldBase; + lastIndices.assign(pageIndices, pageIndices + numPages); + lastStream = stream; + } + + bool failConfigure = false; + bool failBatches = false; + std::atomic_bool* failureMarker = nullptr; + int encodeCalls = 0; + int decodeCalls = 0; + std::size_t lastLifecycleIndex = 0; + void const* lastColdBase = nullptr; + cudaStream_t lastStream{}; + std::vector lastIndices; + std::vector resolved; + +private: + void enqueueFailureMarker(cudaStream_t stream) + { + if (failureMarker != nullptr + && cudaLaunchHostFunc( + stream, [](void* marker) { static_cast(marker)->store(true); }, failureMarker) + != cudaSuccess) + { + throw std::runtime_error("failed to enqueue the requested batch failure marker"); + } + } +}; + +TEST(NativeColdPageCodecTest, ResolvesKvcManagerLayoutAndForwardsWholeBatchOnce) +{ + RecordingCodec codec{{0, 1}}; + + ASSERT_TRUE(configureOne(codec, makeAttentionDesc(kv::PoolGroupIndex{0}, kv::LayerGroupId{3}, 2U))); + ASSERT_EQ(codec.resolved.size(), 1U); + EXPECT_EQ(codec.resolved.front().lifeCycleId, kv::LifeCycleId{3}); + auto const& layers = codec.resolved.front().layers; + EXPECT_EQ(layers.at(0).at("key").rawBase, kGpuKBase); + EXPECT_EQ(layers.at(0).at("value").rawBase, kGpuVBase); + EXPECT_EQ(layers.at(1).at("key").rawBase, kGpuKBase + kRawBytes); + EXPECT_EQ(layers.at(1).at("key").rawSlotBytes, 2U * kRawBytes); + EXPECT_EQ(layers.at(1).at("key").rawBytes, kRawBytes); + EXPECT_EQ(codec.queryColdPageBytes(kv::LayerGroupId{3}), 777U); + + std::vector indices(4096); + for (std::size_t index = 0; index < indices.size(); ++index) + { + indices[index] = {static_cast(index + 1U), static_cast(index)}; + } + auto const stream = reinterpret_cast(kStreamValue); + ASSERT_TRUE( + codec.encode(kv::LayerGroupId{3}, reinterpret_cast(kColdBase), indices.data(), indices.size(), stream)); + EXPECT_EQ(codec.encodeCalls, 1); + EXPECT_EQ(codec.lastLifecycleIndex, 0U); + EXPECT_EQ(codec.lastIndices.size(), 4096U); + EXPECT_EQ(codec.lastIndices.back().src, 4095); + EXPECT_EQ(codec.lastColdBase, reinterpret_cast(kColdBase)); + EXPECT_EQ(codec.lastStream, stream); + + ASSERT_TRUE(codec.decode( + kv::LayerGroupId{3}, reinterpret_cast(kColdBase), indices.data(), indices.size(), stream)); + EXPECT_EQ(codec.decodeCalls, 1); +} + +TEST(NativeColdPageCodecTest, UnownedLifecycleUsesLosslessFallback) +{ + RecordingCodec codec{{0}}; + std::array descs{makeAttentionDesc(), makeLosslessDesc(kv::PoolGroupIndex{1}, kv::LayerGroupId{1})}; + + if (kv::HostMem::shouldUseChunkedRegistration()) + { + // Fallback lifecycles fail closed on host kernels with chunked pinned-memory + // registration until KVCM replaces the batched copies with kernels. + EXPECT_FALSE(codec.configure(descs.data(), kv::PoolGroupIndex{2})); + return; + } + + ASSERT_TRUE(codec.configure(descs.data(), kv::PoolGroupIndex{2})); + EXPECT_EQ(codec.resolved.size(), 1U); + EXPECT_EQ(codec.queryColdPageBytes(kv::LayerGroupId{0}), 777U); + EXPECT_EQ(codec.queryColdPageBytes(kv::LayerGroupId{1}), 96U); + EXPECT_EQ(codec.queryPageIndexLocation(kv::LayerGroupId{1}), kv::PageIndexLocation::kHost); +} + +TEST(NativeColdPageCodecTest, RejectsMixedMissingAndDuplicateLifecycleMappings) +{ + { + RecordingCodec codec{{0}}; + EXPECT_FALSE(configureOne(codec, makeAttentionDesc(kv::PoolGroupIndex{0}, kv::LayerGroupId{0}, 2U))); + } + { + RecordingCodec codec{{0, 1}}; + EXPECT_FALSE(configureOne(codec, makeAttentionDesc())); + } + { + RecordingCodec codec{{0, 1}}; + std::array descs{makeAttentionDesc(), + makeAttentionDesc(kv::PoolGroupIndex{1}, kv::LayerGroupId{0}, 1U, 1, 0x600000, 0x700000)}; + EXPECT_FALSE(codec.configure(descs.data(), kv::PoolGroupIndex{2})); + } +} + +TEST(NativeColdPageCodecTest, CatchesProviderConfigureFailuresAndInvalidBatches) +{ + RecordingCodec codec{{0}}; + codec.failConfigure = true; + EXPECT_FALSE(configureOne(codec, makeAttentionDesc())); + + RecordingCodec validCodec{{0}}; + ASSERT_TRUE(configureOne(validCodec, makeAttentionDesc())); + EXPECT_TRUE(validCodec.encode(kv::LayerGroupId{0}, nullptr, nullptr, 0U, nullptr)); + EXPECT_TRUE(validCodec.decode(kv::LayerGroupId{0}, nullptr, nullptr, 0U, nullptr)); + kv::PageIndexPair const indices[]{{0, 0}}; + EXPECT_FALSE(validCodec.encode(kv::LayerGroupId{0}, nullptr, indices, 1U, nullptr)); + EXPECT_FALSE(validCodec.decode(kv::LayerGroupId{0}, nullptr, indices, 1U, nullptr)); + EXPECT_EQ(validCodec.encodeCalls, 0); + EXPECT_EQ(validCodec.decodeCalls, 0); +} + +TEST(NativeColdPageCodecTest, ProviderFailureUsesTheSuppliedCudaStreamForRollback) +{ + int deviceCount = 0; + if (cudaGetDeviceCount(&deviceCount) != cudaSuccess || deviceCount == 0) + { + GTEST_SKIP() << "Failure draining requires a CUDA device"; + } + + RecordingCodec codec{{0}}; + ASSERT_TRUE(configureOne(codec, makeAttentionDesc())); + codec.failBatches = true; + std::atomic_bool completed = false; + codec.failureMarker = &completed; + cudaStream_t stream{}; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + kv::PageIndexPair const indices[]{{0, 0}}; + EXPECT_FALSE(codec.encode(kv::LayerGroupId{0}, reinterpret_cast(kColdBase), indices, 1U, stream)); + EXPECT_TRUE(completed.exchange(false)); + EXPECT_FALSE(codec.decode(kv::LayerGroupId{0}, reinterpret_cast(kColdBase), indices, 1U, stream)); + EXPECT_TRUE(completed.load()); + EXPECT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +TEST(NativeColdPageCodecTest, UnknownLifecycleUsesFailureSentinels) +{ + RecordingCodec codec{{0}}; + ASSERT_TRUE(configureOne(codec, makeAttentionDesc())); + EXPECT_EQ(codec.queryColdPageBytes(kv::LayerGroupId{99}), 0U); + EXPECT_EQ(codec.getBatchingLayerGroupId(kv::LayerGroupId{99}), kv::LayerGroupId{-1}); + EXPECT_EQ(codec.queryPageIndexLocation(kv::LayerGroupId{99}), kv::PageIndexLocation::kBadLocation); +} + +} // namespace +} // namespace tensorrt_llm::kv_cache_compression diff --git a/cpp/tests/unit_tests/kernels/CMakeLists.txt b/cpp/tests/unit_tests/kernels/CMakeLists.txt index 6b0e5a118211..9a4a4a471387 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/cpp/tests/unit_tests/kernels/CMakeLists.txt @@ -113,3 +113,20 @@ endif() add_gtest(eaglePackDataTest eaglePackDataTest.cpp) add_gtest(sparseKvCacheTest sparseKvCacheTest.cu) add_gtest(prepareCustomMaskTest prepareCustomMaskTest.cpp) + +set(NVFP4_COLD_PAGE_KERNEL_TEST_SRC + nvfp4ColdPageKernelsTest.cpp + ${PROJECT_SOURCE_DIR}/tensorrt_llm/kernels/nvfp4ColdPageKernels.cu + ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.cpp + ${PROJECT_SOURCE_DIR}/tensorrt_llm/common/assert.cpp + ${PROJECT_SOURCE_DIR}/tensorrt_llm/common/envUtils.cpp + ${PROJECT_SOURCE_DIR}/tensorrt_llm/common/logger.cpp + ${PROJECT_SOURCE_DIR}/tensorrt_llm/common/stringUtils.cpp + ${PROJECT_SOURCE_DIR}/tensorrt_llm/common/tllmException.cpp) +add_gtest(nvfp4ColdPageKernelsTest "${NVFP4_COLD_PAGE_KERNEL_TEST_SRC}" + NO_TLLM_LINKAGE) +target_include_directories( + nvfp4ColdPageKernelsTest + PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager) +target_link_libraries(nvfp4ColdPageKernelsTest PRIVATE CUDA::cudart + CUDA::cuda_driver) diff --git a/cpp/tests/unit_tests/kernels/nvfp4ColdPageKernelsTest.cpp b/cpp/tests/unit_tests/kernels/nvfp4ColdPageKernelsTest.cpp new file mode 100644 index 000000000000..9a552d1e83af --- /dev/null +++ b/cpp/tests/unit_tests/kernels/nvfp4ColdPageKernelsTest.cpp @@ -0,0 +1,1297 @@ +/* + * 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/kernels/nvfp4ColdPageKernels.h" +#include "kv_cache_manager_v2/coldPageCodec.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/hostMem.h" +#include "tensorrt_llm/common/cudaUtils.h" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace +{ + +using tensorrt_llm::batch_manager::kv_cache_manager_v2::HostMem; +using tensorrt_llm::batch_manager::kv_cache_manager_v2::MemAddress; +using tensorrt_llm::batch_manager::kv_cache_manager_v2::PageIndexPair; +using tensorrt_llm::kernels::Nvfp4ColdPageIntegerTable; +using tensorrt_llm::kernels::Nvfp4ColdPageRuntimeType; +using tensorrt_llm::kernels::Nvfp4ColdPageScaleTable; +using tensorrt_llm::kernels::Nvfp4ColdPageWideTable; + +struct Nvfp4ColdPageKernelParams +{ + std::int32_t numKvHeads; + std::int32_t tokensPerPage; + std::int32_t headDim; + float nvfp4ScaleOrigQuant; + float nvfp4ScaleQuantOrig; + float fp8ScaleOrigQuant; + float fp8ScaleQuantOrig; +}; + +enum class Nvfp4ColdPageTransform : std::int32_t +{ + kNvfp4 = 0, + kLosslessCopy = 1, +}; + +struct Nvfp4ColdPageTestBuffer +{ + std::uintptr_t rawBase; + std::size_t rawSlotBytes; + std::size_t rawBytes; + std::size_t coldDataOffset; + std::size_t coldScaleOffset; + std::size_t coldPaddingOffset; + std::uint32_t coldPaddingBytes; + Nvfp4ColdPageTransform transform; + Nvfp4ColdPageKernelParams params; +}; + +struct Nvfp4ColdPageTestMetadata +{ + Nvfp4ColdPageWideTable wide{}; + Nvfp4ColdPageIntegerTable integers{}; + Nvfp4ColdPageScaleTable scales{}; + std::uint32_t numBuffers{}; + std::uint32_t maxHalfGroupsPerTile{}; + std::size_t coldPageBytes{}; + Nvfp4ColdPageRuntimeType runtimeType{}; +}; + +Nvfp4ColdPageTestMetadata makeNvfp4ColdPageTestMetadata(std::vector const& buffers, + std::size_t coldPageBytes, Nvfp4ColdPageRuntimeType runtimeType) +{ + Nvfp4ColdPageTestMetadata metadata; + metadata.numBuffers = static_cast(buffers.size()); + metadata.coldPageBytes = coldPageBytes; + metadata.runtimeType = runtimeType; + for (std::size_t index = 0; index < buffers.size(); ++index) + { + auto const& buffer = buffers[index]; + metadata.wide[index] + = {static_cast(buffer.rawBase), static_cast(buffer.rawSlotBytes), + static_cast(buffer.rawBytes), static_cast(buffer.coldDataOffset), + static_cast(buffer.coldScaleOffset), static_cast(buffer.coldPaddingOffset)}; + metadata.integers[index] + = {static_cast(buffer.coldPaddingBytes), static_cast(buffer.transform), + buffer.params.numKvHeads, buffer.params.tokensPerPage, buffer.params.headDim}; + metadata.scales[index] = {buffer.params.nvfp4ScaleOrigQuant, buffer.params.nvfp4ScaleQuantOrig, + buffer.params.fp8ScaleOrigQuant, buffer.params.fp8ScaleQuantOrig}; + if (buffer.transform == Nvfp4ColdPageTransform::kNvfp4) + { + auto const halfGroups = static_cast(buffer.params.numKvHeads) + * static_cast(buffer.params.tokensPerPage) + * (static_cast(buffer.params.headDim) / 8U); + metadata.maxHalfGroupsPerTile = std::max(metadata.maxHalfGroupsPerTile, std::min(halfGroups, 2048U)); + } + } + return metadata; +} + +void invokeNvfp4ColdPageEncode(void const* pages, std::size_t numPages, Nvfp4ColdPageTestMetadata const& metadata, + void* coldBase, cudaStream_t stream) +{ + tensorrt_llm::kernels::invokeNvfp4ColdPageEncode(pages, numPages, metadata.wide.front().data(), + metadata.integers.front().data(), metadata.scales.front().data(), metadata.numBuffers, + metadata.maxHalfGroupsPerTile, metadata.coldPageBytes, metadata.runtimeType, coldBase, stream); +} + +void invokeNvfp4ColdPageDecode(void const* pages, std::size_t numPages, Nvfp4ColdPageTestMetadata const& metadata, + void const* coldBase, cudaStream_t stream) +{ + tensorrt_llm::kernels::invokeNvfp4ColdPageDecode(pages, numPages, metadata.wide.front().data(), + metadata.integers.front().data(), metadata.scales.front().data(), metadata.numBuffers, + metadata.maxHalfGroupsPerTile, metadata.coldPageBytes, metadata.runtimeType, coldBase, stream); +} + +constexpr std::size_t kGuardBytes = 64; +constexpr std::uint8_t kCanary = 0xA5; +constexpr std::size_t kDefaultNumPages = 3; +constexpr std::size_t kCrossLaunchNumPages = 257; + +struct PageGeometry +{ + std::int32_t numHeads; + std::int32_t tokensPerPage; + std::int32_t headDim; +}; + +constexpr PageGeometry kDefaultGeometry{2, 8, 32}; +constexpr PageGeometry kMinimumCompactGeometry{1, 1, 16}; +constexpr PageGeometry kPackedBodyAndTailGeometry{1, 3, 16}; +constexpr PageGeometry kSmallVectorGeometry{1, 4, 16}; +constexpr PageGeometry kLinearScaleTailGeometry{1, 5, 32}; +constexpr PageGeometry kTiledLinearScaleTailGeometry{1, 4097, 16}; +constexpr PageGeometry kCrossRowTileGeometry{1, 343, 48}; +constexpr PageGeometry kLargeHeadDimTailGeometry{1, 1, 65552}; +constexpr PageGeometry kModelLikeGeometry{8, 64, 128}; +constexpr std::array kValidTokenCounts{1, 16, 17, 63, 64}; + +enum class RawKind +{ + kFloat16, + kBfloat16, + kFp8, +}; + +enum class InputPattern +{ + kDense, + kAllZero, + kSparseOutlier, + kRoundingMargins, +}; + +std::size_t roundUp(std::size_t value, std::size_t alignment) +{ + return (value + alignment - 1) / alignment * alignment; +} + +class CudaStream +{ +public: + CudaStream() + { + TLLM_CUDA_CHECK(cudaStreamCreateWithFlags(&mStream, cudaStreamNonBlocking)); + } + + ~CudaStream() + { + if (mStream != nullptr) + { + cudaStreamDestroy(mStream); + } + } + + operator cudaStream_t() const + { + return mStream; + } + + CudaStream(CudaStream const&) = delete; + CudaStream& operator=(CudaStream const&) = delete; + +private: + cudaStream_t mStream{}; +}; + +//! Device allocation guarded by canaries to catch descriptor or vector-tail out-of-bounds writes. +class DeviceRegion +{ +public: + explicit DeviceRegion(std::size_t payloadBytes) + : mPayloadBytes(payloadBytes) + , mTotalBytes(payloadBytes + 2 * kGuardBytes) + { + TLLM_CUDA_CHECK(cudaMalloc(&mBase, mTotalBytes)); + TLLM_CUDA_CHECK(cudaMemset(mBase, kCanary, mTotalBytes)); + } + + ~DeviceRegion() + { + if (mBase != nullptr) + { + cudaFree(mBase); + } + } + + DeviceRegion(DeviceRegion const&) = delete; + DeviceRegion& operator=(DeviceRegion const&) = delete; + + void* data() const + { + return static_cast(mBase) + kGuardBytes; + } + + void copyFrom(std::vector const& bytes) + { + ASSERT_EQ(bytes.size(), mPayloadBytes); + ASSERT_EQ(cudaMemcpy(data(), bytes.data(), bytes.size(), cudaMemcpyHostToDevice), cudaSuccess); + } + + void copyFrom(std::size_t offset, std::vector const& bytes) + { + ASSERT_LE(offset + bytes.size(), mPayloadBytes); + ASSERT_EQ( + cudaMemcpy(static_cast(data()) + offset, bytes.data(), bytes.size(), cudaMemcpyHostToDevice), + cudaSuccess); + } + + std::vector copyToHost() const + { + std::vector bytes(mPayloadBytes); + EXPECT_EQ(cudaMemcpy(bytes.data(), data(), bytes.size(), cudaMemcpyDeviceToHost), cudaSuccess); + return bytes; + } + + std::vector copyToHost(std::size_t offset, std::size_t bytes) const + { + EXPECT_LE(offset + bytes, mPayloadBytes); + std::vector result(bytes); + EXPECT_EQ( + cudaMemcpy(result.data(), static_cast(data()) + offset, bytes, cudaMemcpyDeviceToHost), + cudaSuccess); + return result; + } + + void expectCanaries() const + { + std::vector bytes(mTotalBytes); + ASSERT_EQ(cudaMemcpy(bytes.data(), mBase, bytes.size(), cudaMemcpyDeviceToHost), cudaSuccess); + EXPECT_TRUE(std::all_of( + bytes.begin(), bytes.begin() + kGuardBytes, [](std::uint8_t value) { return value == kCanary; })); + EXPECT_TRUE( + std::all_of(bytes.end() - kGuardBytes, bytes.end(), [](std::uint8_t value) { return value == kCanary; })); + } + +private: + void* mBase{}; + std::size_t mPayloadBytes{}; + std::size_t mTotalBytes{}; +}; + +//! CUDA-mapped HostMem matching KVCM V2's Host carrier. +class MappedHostRegion +{ +public: + explicit MappedHostRegion(std::size_t payloadBytes) + : mMemory(roundUp(kGuardBytes + payloadBytes + kGuardBytes, HostMem::kAlignment)) + , mPayloadBytes(payloadBytes) + { + TLLM_CHECK_WITH_INFO(kGuardBytes + payloadBytes + kGuardBytes <= mMemory.size(), + "Mapped Host test allocation is too small for payload and canaries"); + std::memset(reinterpret_cast(mMemory.address()), kCanary, mMemory.size()); + } + + void* data() const + { + return reinterpret_cast(mMemory.address() + kGuardBytes); + } + + std::uint8_t* bytes() const + { + return static_cast(data()); + } + + std::vector payload() const + { + return {bytes(), bytes() + mPayloadBytes}; + } + + void expectCanaries() const + { + auto const* base = reinterpret_cast(mMemory.address()); + EXPECT_TRUE(std::all_of(base, base + kGuardBytes, [](std::uint8_t value) { return value == kCanary; })); + EXPECT_TRUE(std::all_of(base + kGuardBytes + mPayloadBytes, base + mMemory.size(), + [](std::uint8_t value) { return value == kCanary; })); + } + +private: + HostMem mMemory; + std::size_t mPayloadBytes{}; +}; + +struct LayerBuffers +{ + explicit LayerBuffers(std::size_t rawBytes) + : rawK(rawBytes) + , rawV(rawBytes) + { + } + + DeviceRegion rawK; + DeviceRegion rawV; +}; + +std::size_t numElements(PageGeometry const& geometry) +{ + return static_cast(geometry.numHeads) * geometry.tokensPerPage * geometry.headDim; +} + +std::size_t rawBytes(RawKind kind, PageGeometry const& geometry) +{ + return numElements(geometry) * (kind == RawKind::kFp8 ? 1 : 2); +} + +std::size_t rawElementBytes(RawKind kind) +{ + return kind == RawKind::kFp8 ? 1U : 2U; +} + +std::size_t packedBytes(PageGeometry const& geometry) +{ + return numElements(geometry) / 2; +} + +std::size_t scaleBytes(PageGeometry const& geometry) +{ + return numElements(geometry) / 16; +} + +Nvfp4ColdPageKernelParams makeParams(PageGeometry const& geometry = kDefaultGeometry, std::uint32_t role = 0U) +{ + Nvfp4ColdPageKernelParams params{}; + params.numKvHeads = geometry.numHeads; + params.tokensPerPage = geometry.tokensPerPage; + params.headDim = geometry.headDim; + params.nvfp4ScaleOrigQuant = role == 0U ? 1.0F : 2.0F; + params.nvfp4ScaleQuantOrig = role == 0U ? 1.0F : 0.5F; + params.fp8ScaleOrigQuant = role == 0U ? 2.0F : 4.0F; + params.fp8ScaleQuantOrig = role == 0U ? 0.5F : 0.25F; + return params; +} + +Nvfp4ColdPageRuntimeType runtimeType(RawKind kind) +{ + switch (kind) + { + case RawKind::kFloat16: return Nvfp4ColdPageRuntimeType::kFloat16; + case RawKind::kBfloat16: return Nvfp4ColdPageRuntimeType::kBfloat16; + case RawKind::kFp8: return Nvfp4ColdPageRuntimeType::kFp8E4m3; + } + return Nvfp4ColdPageRuntimeType::kFloat16; +} + +template +void storeScalar(std::vector& bytes, std::size_t index, T value) +{ + std::memcpy(bytes.data() + index * sizeof(T), &value, sizeof(T)); +} + +template +T loadScalar(std::vector const& bytes, std::size_t index) +{ + T value; + std::memcpy(&value, bytes.data() + index * sizeof(T), sizeof(T)); + return value; +} + +void storeRawValue(std::vector& bytes, RawKind kind, std::size_t index, float value, + Nvfp4ColdPageKernelParams const& params) +{ + switch (kind) + { + case RawKind::kFloat16: storeScalar(bytes, index, __float2half(value)); break; + case RawKind::kBfloat16: storeScalar(bytes, index, __float2bfloat16(value)); break; + case RawKind::kFp8: storeScalar(bytes, index, __nv_fp8_e4m3(value * params.fp8ScaleOrigQuant)); break; + } +} + +float loadRawValue( + std::vector const& bytes, RawKind kind, std::size_t index, Nvfp4ColdPageKernelParams const& params) +{ + switch (kind) + { + case RawKind::kFloat16: return __half2float(loadScalar(bytes, index)); + case RawKind::kBfloat16: return __bfloat162float(loadScalar<__nv_bfloat16>(bytes, index)); + case RawKind::kFp8: return static_cast(loadScalar<__nv_fp8_e4m3>(bytes, index)) * params.fp8ScaleQuantOrig; + } + return 0.0F; +} + +std::uint32_t linearScaleOffset(std::uint32_t row, std::uint32_t scaleInRow, PageGeometry const& geometry) +{ + std::uint32_t const scalesPerRow = static_cast(geometry.headDim) / 16; + return row * scalesPerRow + scaleInRow; +} + +constexpr std::array kE2m1Levels{0.0F, 0.5F, 1.0F, 1.5F, 2.0F, 3.0F, 4.0F, 6.0F}; + +float e2m1Value(std::uint8_t nibble) +{ + float const value = kE2m1Levels[nibble & 0x7U]; + return (nibble & 0x8U) != 0 ? -value : value; +} + +//! Independent nearest-level oracle; fixtures avoid ties instead of duplicating production tie rules. +std::uint8_t quantizeE2m1(float value) +{ + bool const negative = std::signbit(value); + float const magnitude = std::abs(value); + std::uint8_t best = 0; + float bestDistance = std::abs(magnitude - kE2m1Levels[0]); + for (std::uint8_t index = 1; index < kE2m1Levels.size(); ++index) + { + float const distance = std::abs(magnitude - kE2m1Levels[index]); + if (distance < bestDistance) + { + best = index; + bestDistance = distance; + } + } + return static_cast(best | (negative ? 0x8U : 0U)); +} + +//! Exactly representable E2M1 values and E4M3 scales keep byte comparisons deterministic. +std::vector makeRawPage(RawKind kind, std::size_t page, std::uint32_t role, + Nvfp4ColdPageKernelParams const& params, PageGeometry const& geometry, InputPattern inputPattern) +{ + constexpr std::array densePattern{ + 0.0F, 0.5F, -1.0F, 1.5F, -2.0F, 3.0F, -4.0F, 6.0F, -0.5F, 1.0F, -1.5F, 2.0F, -3.0F, 4.0F, -6.0F, 0.5F}; + constexpr std::array firstLaneOutlierPattern{ + 6.0F, -0.5F, 0.0F, 0.0F, 0.0F, 0.0F, 0.0F, 0.0F, 0.0F, 0.5F, 0.0F, 0.0F, 0.0F, 0.0F, 0.0F, 0.0F}; + constexpr std::array secondLaneOutlierPattern{ + 0.0F, -0.5F, 0.0F, 0.0F, 0.0F, 0.0F, 0.0F, 0.0F, 0.0F, 0.5F, 0.0F, 0.0F, 0.0F, 0.0F, 0.0F, -6.0F}; + constexpr std::array roundingMarginsPattern{ + 6.0F, 0.20F, 0.30F, 0.65F, 0.85F, 1.10F, 1.40F, 1.60F, 1.90F, 2.30F, 2.70F, 3.20F, 3.80F, 4.50F, 5.20F, -0.30F}; + constexpr std::array blockScales{0.25F, 0.5F, 1.0F, 2.0F}; + + std::vector bytes(rawBytes(kind, geometry)); + for (std::size_t index = 0; index < numElements(geometry); ++index) + { + std::size_t const scaleGroup = index / 16; + float const blockScale = blockScales[(scaleGroup + page + role) % blockScales.size()]; + float normalizedValue = 0.0F; + if (inputPattern == InputPattern::kDense) + { + normalizedValue = densePattern[index % densePattern.size()]; + } + else if (inputPattern == InputPattern::kSparseOutlier) + { + auto const& pattern = (scaleGroup & 1U) == 0 ? firstLaneOutlierPattern : secondLaneOutlierPattern; + normalizedValue = pattern[index % pattern.size()]; + } + else if (inputPattern == InputPattern::kRoundingMargins) + { + normalizedValue = roundingMarginsPattern[index % roundingMarginsPattern.size()]; + } + // kAllZero intentionally keeps the zero initializer. + if (((page / 32) & 1U) != 0) + { + normalizedValue = -normalizedValue; + } + float const value = normalizedValue * blockScale / params.nvfp4ScaleOrigQuant; + storeRawValue(bytes, kind, index, value, params); + } + return bytes; +} + +struct ReferenceNvfp4 +{ + std::vector packed; + std::vector scales; +}; + +ReferenceNvfp4 compressReference(std::vector const& raw, RawKind kind, + Nvfp4ColdPageKernelParams const& params, PageGeometry const& geometry) +{ + ReferenceNvfp4 result{{}, {}}; + result.packed.resize(packedBytes(geometry)); + result.scales.resize(scaleBytes(geometry)); + std::uint32_t const scalesPerRow = static_cast(geometry.headDim) / 16; + std::uint32_t const rows = static_cast(geometry.numHeads * geometry.tokensPerPage); + + for (std::uint32_t row = 0; row < rows; ++row) + { + for (std::uint32_t scaleInRow = 0; scaleInRow < scalesPerRow; ++scaleInRow) + { + std::size_t const blockStart = static_cast(row) * geometry.headDim + scaleInRow * 16; + float amax = 0.0F; + for (std::uint32_t i = 0; i < 16; ++i) + { + amax = std::max(amax, std::abs(loadRawValue(raw, kind, blockStart + i, params))); + } + + __nv_fp8_e4m3 blockScale(params.nvfp4ScaleOrigQuant * amax / 6.0F); + result.scales[linearScaleOffset(row, scaleInRow, geometry)] = blockScale.__x; + float const blockScaleFloat = static_cast(blockScale); + float const outputScale = blockScaleFloat == 0.0F ? 0.0F : params.nvfp4ScaleOrigQuant / blockScaleFloat; + for (std::uint32_t i = 0; i < 16; i += 2) + { + std::uint8_t const lo = quantizeE2m1(loadRawValue(raw, kind, blockStart + i, params) * outputScale); + std::uint8_t const hi = quantizeE2m1(loadRawValue(raw, kind, blockStart + i + 1, params) * outputScale); + result.packed[(blockStart + i) / 2] = static_cast(lo | (hi << 4)); + } + } + } + return result; +} + +std::vector decompressReference(ReferenceNvfp4 const& compressed, RawKind kind, + Nvfp4ColdPageKernelParams const& params, PageGeometry const& geometry) +{ + std::vector raw(rawBytes(kind, geometry)); + std::uint32_t const scalesPerRow = static_cast(geometry.headDim) / 16; + std::uint32_t const rows = static_cast(geometry.numHeads * geometry.tokensPerPage); + for (std::uint32_t row = 0; row < rows; ++row) + { + for (std::uint32_t scaleInRow = 0; scaleInRow < scalesPerRow; ++scaleInRow) + { + __nv_fp8_e4m3 blockScale; + blockScale.__x = compressed.scales[linearScaleOffset(row, scaleInRow, geometry)]; + float const dequantScale = static_cast(blockScale) * params.nvfp4ScaleQuantOrig; + std::size_t const blockStart = static_cast(row) * geometry.headDim + scaleInRow * 16; + for (std::uint32_t i = 0; i < 16; ++i) + { + std::uint8_t const byte = compressed.packed[(blockStart + i) / 2]; + std::uint8_t const nibble = (i & 1U) == 0 ? byte & 0xFU : byte >> 4; + storeRawValue(raw, kind, blockStart + i, e2m1Value(nibble) * dequantScale, params); + } + } + } + return raw; +} + +void runColdPageRoundTrip(RawKind kind, PageGeometry const& geometry = kDefaultGeometry, + std::size_t numPages = kDefaultNumPages, InputPattern inputPattern = InputPattern::kDense, + bool synchronizeBetweenDirections = true, bool repeatRoundTrip = false, std::size_t coldBaseOffset = 0) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + if (!tensorrt_llm::common::isSM100Family()) + { + GTEST_SKIP() << "NVFP4 cold-page kernels require an SM100-family GPU"; + } + std::array const params{makeParams(geometry, 0U), makeParams(geometry, 1U)}; + CudaStream stream; + std::size_t const rawSlotBytes = rawBytes(kind, geometry); + // Align compact Slot strides while independently testing an arbitrary staging-base offset. + std::size_t const compactSlotBytes = roundUp(2U * (packedBytes(geometry) + scaleBytes(geometry)), alignof(uint4)); + // Use alternate Slots to cover non-contiguous KVCM Page indices. + std::size_t const slotCapacity = 2U * numPages; + DeviceRegion rawInputK(slotCapacity * rawSlotBytes); + DeviceRegion rawInputV(slotCapacity * rawSlotBytes); + DeviceRegion rawOutputK(slotCapacity * rawSlotBytes); + DeviceRegion rawOutputV(slotCapacity * rawSlotBytes); + MappedHostRegion compactPages(coldBaseOffset + slotCapacity * compactSlotBytes); + auto* compactBase = compactPages.bytes() + coldBaseOffset; + std::vector, 2>> rawHost(numPages); + std::vector offloadTasks; + offloadTasks.reserve(numPages); + for (std::size_t page = 0; page < numPages; ++page) + { + std::size_t const rawSlot = 2U * page; + std::size_t const coldSlot = rawSlot + 1U; + rawHost[page][0] = makeRawPage(kind, page, 0, params[0], geometry, inputPattern); + rawHost[page][1] = makeRawPage(kind, page, 1, params[1], geometry, inputPattern); + rawInputK.copyFrom(rawSlot * rawSlotBytes, rawHost[page][0]); + rawInputV.copyFrom(rawSlot * rawSlotBytes, rawHost[page][1]); + offloadTasks.push_back({static_cast(coldSlot), static_cast(rawSlot)}); + } + + std::size_t const packed = packedBytes(geometry); + std::size_t const scale = scaleBytes(geometry); + std::size_t const payloadBytes = 2U * (packed + scale); + std::uint32_t const paddingBytes = static_cast(compactSlotBytes - payloadBytes); + std::vector const inputBuffers{ + {reinterpret_cast(rawInputK.data()), rawSlotBytes, rawSlotBytes, 0U, 2U * packed, 0U, 0U, + Nvfp4ColdPageTransform::kNvfp4, params[0]}, + {reinterpret_cast(rawInputV.data()), rawSlotBytes, rawSlotBytes, packed, 2U * packed + scale, + payloadBytes, paddingBytes, Nvfp4ColdPageTransform::kNvfp4, params[1]}}; + + auto const inputMetadata = makeNvfp4ColdPageTestMetadata(inputBuffers, compactSlotBytes, runtimeType(kind)); + std::vector> references(numPages); + for (std::size_t page = 0; page < numPages; ++page) + { + references[page][0] = compressReference(rawHost[page][0], kind, params[0], geometry); + references[page][1] = compressReference(rawHost[page][1], kind, params[1], geometry); + } + + auto const verifyCompressedPages = [&] + { + auto const payload = compactPages.payload(); + EXPECT_TRUE(std::all_of(payload.begin(), payload.begin() + static_cast(coldBaseOffset), + [](std::uint8_t value) { return value == kCanary; })); + for (std::size_t page = 0; page < numPages; ++page) + { + std::size_t const base = coldBaseOffset + (2U * page + 1U) * compactSlotBytes; + auto const region = [&](std::size_t offset, std::size_t bytes) + { + return std::vector(payload.begin() + static_cast(base + offset), + payload.begin() + static_cast(base + offset + bytes)); + }; + EXPECT_EQ(region(0, packed), references[page][0].packed); + EXPECT_EQ(region(packed, packed), references[page][1].packed); + EXPECT_EQ(region(2 * packed, scale), references[page][0].scales); + EXPECT_EQ(region(2 * packed + scale, scale), references[page][1].scales); + auto const padding = region(2U * (packed + scale), compactSlotBytes - 2U * (packed + scale)); + EXPECT_TRUE(std::all_of(padding.begin(), padding.end(), [](std::uint8_t value) { return value == 0U; })); + + std::size_t const unusedBase = coldBaseOffset + 2U * page * compactSlotBytes; + EXPECT_TRUE(std::all_of(payload.begin() + static_cast(unusedBase), + payload.begin() + static_cast(unusedBase + compactSlotBytes), + [](std::uint8_t value) { return value == kCanary; })); + } + }; + + invokeNvfp4ColdPageEncode(offloadTasks.data(), offloadTasks.size(), inputMetadata, compactBase, stream); + + if (synchronizeBetweenDirections) + { + // Read the Host Slot only after StorageManager-style event fencing. + cudaEvent_t offloadComplete{}; + ASSERT_EQ(cudaEventCreateWithFlags(&offloadComplete, cudaEventDisableTiming), cudaSuccess); + ASSERT_EQ(cudaEventRecord(offloadComplete, stream), cudaSuccess); + ASSERT_EQ(cudaEventSynchronize(offloadComplete), cudaSuccess); + ASSERT_EQ(cudaEventDestroy(offloadComplete), cudaSuccess); + verifyCompressedPages(); + + std::size_t const compactPayloadBytes = 2U * (packedBytes(geometry) + scaleBytes(geometry)); + if (compactPayloadBytes != compactSlotBytes) + { + // Re-encode poisoned recycled Slots to verify deterministic payload and padding bytes. + auto const firstSerialization = compactPages.payload(); + for (std::size_t page = 0; page < numPages; ++page) + { + std::memset(compactBase + (2U * page + 1U) * compactSlotBytes, 0x5A, compactSlotBytes); + } + invokeNvfp4ColdPageEncode(offloadTasks.data(), offloadTasks.size(), inputMetadata, compactBase, stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + EXPECT_EQ(compactPages.payload(), firstSerialization); + verifyCompressedPages(); + } + } + + std::vector onboardTasks; + onboardTasks.reserve(numPages); + for (std::size_t page = 0; page < numPages; ++page) + { + std::size_t const rawSlot = 2U * page; + std::size_t const coldSlot = rawSlot + 1U; + onboardTasks.push_back({static_cast(rawSlot), static_cast(coldSlot)}); + } + std::vector const outputBuffers{ + {reinterpret_cast(rawOutputK.data()), rawSlotBytes, rawSlotBytes, 0U, 2U * packed, 0U, 0U, + Nvfp4ColdPageTransform::kNvfp4, params[0]}, + {reinterpret_cast(rawOutputV.data()), rawSlotBytes, rawSlotBytes, packed, 2U * packed + scale, + payloadBytes, paddingBytes, Nvfp4ColdPageTransform::kNvfp4, params[1]}}; + auto const outputMetadata = makeNvfp4ColdPageTestMetadata(outputBuffers, compactSlotBytes, runtimeType(kind)); + invokeNvfp4ColdPageDecode(onboardTasks.data(), onboardTasks.size(), outputMetadata, compactBase, stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + if (!synchronizeBetweenDirections) + { + // Verify back-to-back offload/onboard without an intervening Host fence. + verifyCompressedPages(); + } + + if (repeatRoundTrip) + { + // A second lossy round trip catches stale descriptors and validates Q(D(Q(D(Q(x))))). + for (std::size_t page = 0; page < numPages; ++page) + { + for (std::uint32_t role = 0; role < 2; ++role) + { + auto const restored = decompressReference(references[page][role], kind, params[role], geometry); + references[page][role] = compressReference(restored, kind, params[role], geometry); + } + } + invokeNvfp4ColdPageEncode(offloadTasks.data(), offloadTasks.size(), outputMetadata, compactBase, stream); + invokeNvfp4ColdPageDecode(onboardTasks.data(), onboardTasks.size(), inputMetadata, compactBase, stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + verifyCompressedPages(); + } + + for (std::size_t page = 0; page < numPages; ++page) + { + std::size_t const slotOffset = 2U * page * rawSlotBytes; + auto const& finalK = repeatRoundTrip ? rawInputK : rawOutputK; + auto const& finalV = repeatRoundTrip ? rawInputV : rawOutputV; + EXPECT_EQ(finalK.copyToHost(slotOffset, rawSlotBytes), + decompressReference(references[page][0], kind, params[0], geometry)); + EXPECT_EQ(finalV.copyToHost(slotOffset, rawSlotBytes), + decompressReference(references[page][1], kind, params[1], geometry)); + } + rawInputK.expectCanaries(); + rawInputV.expectCanaries(); + rawOutputK.expectCanaries(); + rawOutputV.expectCanaries(); + compactPages.expectCanaries(); +} + +std::vector makePartialRawPage(RawKind kind, std::int32_t validTokens, bool zeroTail, std::uint32_t role, + Nvfp4ColdPageKernelParams const& params, PageGeometry const& geometry) +{ + std::vector bytes(rawBytes(kind, geometry)); + for (std::int32_t head = 0; head < geometry.numHeads; ++head) + { + for (std::int32_t token = 0; token < geometry.tokensPerPage; ++token) + { + for (std::int32_t dim = 0; dim < geometry.headDim; ++dim) + { + std::size_t const index + = (static_cast(head) * geometry.tokensPerPage + token) * geometry.headDim + dim; + float value = 0.0F; + if (token < validTokens) + { + // Zero-tail and stale-tail fixtures share valid rows with distinct K/V values and scales. + value = static_cast((dim % 13) - 6) * 0.125F + static_cast(head) * 0.03125F + + static_cast(token) * 0.0078125F + static_cast(role) * 0.0625F; + } + else if (!zeroTail) + { + // Poison inactive rows; 16-value groups stay within a token row and cannot affect the prefix. + if (dim % 31 == 0) + { + value = std::numeric_limits::quiet_NaN(); + } + else if (dim % 29 == 0) + { + value = std::numeric_limits::infinity(); + } + else + { + value = static_cast((dim % 9) - 4) * 0.25F + static_cast(token) * 0.015625F + + static_cast(head + 3 * role) * 0.046875F; + } + } + storeRawValue(bytes, kind, index, value, params); + } + } + } + return bytes; +} + +void expectSameValidPrefix(std::vector const& lhs, std::vector const& rhs, RawKind kind, + std::int32_t validTokens, PageGeometry const& geometry) +{ + std::size_t const rowBytes = static_cast(geometry.headDim) * rawElementBytes(kind); + for (std::int32_t head = 0; head < geometry.numHeads; ++head) + { + for (std::int32_t token = 0; token < validTokens; ++token) + { + std::size_t const offset = (static_cast(head) * geometry.tokensPerPage + token) * rowBytes; + EXPECT_EQ(std::memcmp(lhs.data() + offset, rhs.data() + offset, rowBytes), 0) + << "valid prefix differs at head=" << head << " token=" << token; + } + } +} + +void expectZeroTail( + std::vector const& bytes, RawKind kind, std::int32_t validTokens, PageGeometry const& geometry) +{ + std::size_t const rowBytes = static_cast(geometry.headDim) * rawElementBytes(kind); + std::vector const zero(rowBytes, 0); + for (std::int32_t head = 0; head < geometry.numHeads; ++head) + { + for (std::int32_t token = validTokens; token < geometry.tokensPerPage; ++token) + { + std::size_t const offset = (static_cast(head) * geometry.tokensPerPage + token) * rowBytes; + EXPECT_EQ(std::memcmp(bytes.data() + offset, zero.data(), rowBytes), 0) + << "zero tail changed at head=" << head << " token=" << token; + } + } +} + +void runPartialPageTailIsolation(RawKind kind) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + if (!tensorrt_llm::common::isSM100Family()) + { + GTEST_SKIP() << "NVFP4 cold-page kernels require an SM100-family GPU"; + } + + PageGeometry constexpr geometry = kModelLikeGeometry; + std::size_t constexpr pageVariants = 2U; + std::size_t const numPages = pageVariants * kValidTokenCounts.size(); + std::array const params{makeParams(geometry, 0U), makeParams(geometry, 1U)}; + std::size_t const rawSlotBytes = rawBytes(kind, geometry); + std::size_t const compactSlotBytes = roundUp(2U * (packedBytes(geometry) + scaleBytes(geometry)), alignof(uint4)); + + DeviceRegion rawInputK(numPages * rawSlotBytes); + DeviceRegion rawInputV(numPages * rawSlotBytes); + DeviceRegion rawOutputK(numPages * rawSlotBytes); + DeviceRegion rawOutputV(numPages * rawSlotBytes); + MappedHostRegion compactPages(numPages * compactSlotBytes); + std::vector offloadTasks; + std::vector onboardTasks; + offloadTasks.reserve(numPages); + onboardTasks.reserve(numPages); + for (std::size_t page = 0; page < numPages; ++page) + { + std::int32_t const validTokens = kValidTokenCounts[page / pageVariants]; + bool const zeroTail = page % pageVariants == 0U; + rawInputK.copyFrom( + page * rawSlotBytes, makePartialRawPage(kind, validTokens, zeroTail, 0, params[0], geometry)); + rawInputV.copyFrom( + page * rawSlotBytes, makePartialRawPage(kind, validTokens, zeroTail, 1, params[1], geometry)); + auto const pageIndex = static_cast(page); + offloadTasks.push_back({pageIndex, pageIndex}); + onboardTasks.push_back({pageIndex, pageIndex}); + } + + std::size_t const packed = packedBytes(geometry); + std::size_t const scale = scaleBytes(geometry); + std::size_t const payloadBytes = 2U * (packed + scale); + auto const makeBuffers = [&](DeviceRegion const& rawK, DeviceRegion const& rawV) + { + return std::vector{ + {reinterpret_cast(rawK.data()), rawSlotBytes, rawSlotBytes, 0U, 2U * packed, 0U, 0U, + Nvfp4ColdPageTransform::kNvfp4, params[0]}, + {reinterpret_cast(rawV.data()), rawSlotBytes, rawSlotBytes, packed, 2U * packed + scale, + payloadBytes, static_cast(compactSlotBytes - payloadBytes), + Nvfp4ColdPageTransform::kNvfp4, params[1]}}; + }; + auto const inputMetadata + = makeNvfp4ColdPageTestMetadata(makeBuffers(rawInputK, rawInputV), compactSlotBytes, runtimeType(kind)); + auto const outputMetadata + = makeNvfp4ColdPageTestMetadata(makeBuffers(rawOutputK, rawOutputV), compactSlotBytes, runtimeType(kind)); + CudaStream stream; + invokeNvfp4ColdPageEncode(offloadTasks.data(), offloadTasks.size(), inputMetadata, compactPages.data(), stream); + invokeNvfp4ColdPageDecode(onboardTasks.data(), onboardTasks.size(), outputMetadata, compactPages.data(), stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + for (std::size_t pair = 0; pair < kValidTokenCounts.size(); ++pair) + { + std::int32_t const validTokens = kValidTokenCounts[pair]; + for (std::uint32_t role = 0; role < 2U; ++role) + { + auto const& output = role == 0U ? rawOutputK : rawOutputV; + auto const zeroOutput = output.copyToHost(pageVariants * pair * rawSlotBytes, rawSlotBytes); + auto const staleOutput = output.copyToHost((pageVariants * pair + 1U) * rawSlotBytes, rawSlotBytes); + expectSameValidPrefix(zeroOutput, staleOutput, kind, validTokens, geometry); + expectZeroTail(zeroOutput, kind, validTokens, geometry); + } + } + + rawInputK.expectCanaries(); + rawInputV.expectCanaries(); + rawOutputK.expectCanaries(); + rawOutputV.expectCanaries(); + compactPages.expectCanaries(); +} + +struct RoundTripCase +{ + char const* name; + RawKind kind; + PageGeometry geometry{kDefaultGeometry}; + std::size_t numPages{kDefaultNumPages}; + InputPattern inputPattern{InputPattern::kDense}; + bool synchronizeBetweenDirections{true}; + bool repeatRoundTrip{false}; + std::size_t coldBaseOffset{0}; +}; + +RoundTripCase constexpr kRoundTripCases[]{ + {"DefaultFloat16", RawKind::kFloat16}, + {"DefaultBfloat16", RawKind::kBfloat16}, + {"DefaultFp8IndependentScales", RawKind::kFp8}, + {"SmallVectorFloat16", RawKind::kFloat16, kSmallVectorGeometry, 1}, + {"MinimumFloat16", RawKind::kFloat16, kMinimumCompactGeometry, 1}, + {"MinimumBfloat16", RawKind::kBfloat16, kMinimumCompactGeometry, 1}, + {"MinimumFp8", RawKind::kFp8, kMinimumCompactGeometry, 1}, + {"PackedScaleTailFloat16", RawKind::kFloat16, kPackedBodyAndTailGeometry, 1}, + {"PackedScaleTailFp8", RawKind::kFp8, kPackedBodyAndTailGeometry, 1}, + {"ByteAlignedColdBaseBfloat16", RawKind::kBfloat16, kPackedBodyAndTailGeometry, 1, InputPattern::kDense, true, + false, 1}, + {"ByteAlignedColdBaseFp8", RawKind::kFp8, kPackedBodyAndTailGeometry, 1, InputPattern::kDense, true, false, 1}, + {"OddTokenFloat16", RawKind::kFloat16, kLinearScaleTailGeometry, 1}, + {"OddTokenBfloat16", RawKind::kBfloat16, kLinearScaleTailGeometry, 1}, + {"OddTokenFp8", RawKind::kFp8, kLinearScaleTailGeometry, 1}, + {"TiledTailsFloat16", RawKind::kFloat16, kTiledLinearScaleTailGeometry, 1}, + {"TiledTailsFp8", RawKind::kFp8, kTiledLinearScaleTailGeometry, 1}, + {"CrossRowTileBfloat16", RawKind::kBfloat16, kCrossRowTileGeometry, 1, InputPattern::kDense, true, false, 1}, + {"LargeHeadDimFp8", RawKind::kFp8, kLargeHeadDimTailGeometry, 1}, + {"ModelLikeBfloat16", RawKind::kBfloat16, kModelLikeGeometry, 1}, + {"ModelLikeFp8", RawKind::kFp8, kModelLikeGeometry, 1}, + {"ZeroGroupsFloat16", RawKind::kFloat16, kDefaultGeometry, 1, InputPattern::kAllZero}, + {"ZeroGroupsBfloat16", RawKind::kBfloat16, kDefaultGeometry, 1, InputPattern::kAllZero}, + {"ZeroGroupsFp8", RawKind::kFp8, kDefaultGeometry, 1, InputPattern::kAllZero}, + {"WarpLaneAmaxFloat16", RawKind::kFloat16, kDefaultGeometry, 1, InputPattern::kSparseOutlier}, + {"WarpLaneAmaxBfloat16", RawKind::kBfloat16, kDefaultGeometry, 1, InputPattern::kSparseOutlier}, + {"WarpLaneAmaxFp8", RawKind::kFp8, kDefaultGeometry, 1, InputPattern::kSparseOutlier}, + {"ReuseDefaultFloat16", RawKind::kFloat16, kDefaultGeometry, 3, InputPattern::kDense, true, true}, + {"ReuseDefaultBfloat16", RawKind::kBfloat16, kDefaultGeometry, 3, InputPattern::kDense, true, true}, + {"ReuseDefaultFp8", RawKind::kFp8, kDefaultGeometry, 3, InputPattern::kDense, true, true}, + {"ReuseModelLikeFloat16", RawKind::kFloat16, kModelLikeGeometry, 2, InputPattern::kDense, true, true}, + {"ReuseModelLikeBfloat16", RawKind::kBfloat16, kModelLikeGeometry, 2, InputPattern::kDense, true, true}, + {"ReuseModelLikeFp8", RawKind::kFp8, kModelLikeGeometry, 2, InputPattern::kDense, true, true}, + {"RoundingMarginsFloat16", RawKind::kFloat16, kDefaultGeometry, 1, InputPattern::kRoundingMargins}, + {"CrossLaunchBfloat16", RawKind::kBfloat16, kSmallVectorGeometry, kCrossLaunchNumPages}, + {"CrossLaunchFp8", RawKind::kFp8, kSmallVectorGeometry, kCrossLaunchNumPages}, + {"PdlBfloat16", RawKind::kBfloat16, kSmallVectorGeometry, 65, InputPattern::kDense, false}, + {"PdlFp8", RawKind::kFp8, kSmallVectorGeometry, 65, InputPattern::kDense, false}, +}; + +class Nvfp4ColdPageRoundTripTest : public testing::TestWithParam +{ +}; + +TEST_P(Nvfp4ColdPageRoundTripTest, MatchesReference) +{ + auto const& test = GetParam(); + runColdPageRoundTrip(test.kind, test.geometry, test.numPages, test.inputPattern, test.synchronizeBetweenDirections, + test.repeatRoundTrip, test.coldBaseOffset); +} + +std::string roundTripCaseName(testing::TestParamInfo const& info) +{ + return info.param.name; +} + +INSTANTIATE_TEST_SUITE_P(Scenarios, Nvfp4ColdPageRoundTripTest, testing::ValuesIn(kRoundTripCases), roundTripCaseName); + +void runUnaryMlaWithLosslessSideRoundTrip(RawKind kind) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + if (!tensorrt_llm::common::isSM100Family()) + { + GTEST_SKIP() << "NVFP4 cold-page kernels require an SM100-family GPU"; + } + + PageGeometry constexpr geometry{1, 64, 576}; + std::size_t constexpr numPages = 2; + std::size_t constexpr sideRawBytes = 64U * (128U + 4U); + std::size_t constexpr sideSlotBytes = sideRawBytes + 13U; + std::size_t constexpr coldBaseOffset = 1; + auto const params = makeParams(geometry); + std::size_t const mlaRawBytes = rawBytes(kind, geometry); + std::size_t const mlaPackedBytes = packedBytes(geometry); + std::size_t const mlaScaleBytes = scaleBytes(geometry); + std::size_t const mlaPayloadBytes = mlaPackedBytes + mlaScaleBytes; + std::size_t constexpr gapBeforeSide = 3; + std::size_t const sideColdOffset = mlaPayloadBytes + gapBeforeSide; + std::size_t const sideColdEnd = sideColdOffset + sideRawBytes; + std::size_t const coldPageBytes = roundUp(sideColdEnd, alignof(uint4)); + + DeviceRegion mlaInput(numPages * mlaRawBytes); + DeviceRegion mlaOutput(numPages * mlaRawBytes); + DeviceRegion sideInput(numPages * sideSlotBytes); + DeviceRegion sideOutput(numPages * sideSlotBytes); + MappedHostRegion coldStorage(coldBaseOffset + numPages * coldPageBytes); + auto* coldBase = coldStorage.bytes() + coldBaseOffset; + + std::array, numPages> mlaHost; + std::array, numPages> sideHost; + std::array references; + std::vector offloadTasks; + std::vector onboardTasks; + for (std::size_t page = 0; page < numPages; ++page) + { + mlaHost[page] = makeRawPage(kind, page, 0U, params, geometry, InputPattern::kDense); + references[page] = compressReference(mlaHost[page], kind, params, geometry); + sideHost[page].resize(sideRawBytes); + for (std::size_t byte = 0; byte < sideRawBytes; ++byte) + { + sideHost[page][byte] = static_cast((17U * byte + 53U * page + 11U) & 0xFFU); + } + mlaInput.copyFrom(page * mlaRawBytes, mlaHost[page]); + sideInput.copyFrom(page * sideSlotBytes, sideHost[page]); + auto const pageIndex = static_cast(page); + offloadTasks.push_back({pageIndex, pageIndex}); + onboardTasks.push_back({pageIndex, pageIndex}); + } + + auto const makePlans = [&](DeviceRegion const& mla, DeviceRegion const& side) + { + return std::vector{ + {reinterpret_cast(mla.data()), mlaRawBytes, mlaRawBytes, 0U, mlaPackedBytes, + mlaPayloadBytes, static_cast(gapBeforeSide), Nvfp4ColdPageTransform::kNvfp4, params}, + {reinterpret_cast(side.data()), sideSlotBytes, sideRawBytes, sideColdOffset, 0U, + sideColdEnd, static_cast(coldPageBytes - sideColdEnd), + Nvfp4ColdPageTransform::kLosslessCopy, {}}}; + }; + auto const inputMetadata + = makeNvfp4ColdPageTestMetadata(makePlans(mlaInput, sideInput), coldPageBytes, runtimeType(kind)); + auto const outputMetadata + = makeNvfp4ColdPageTestMetadata(makePlans(mlaOutput, sideOutput), coldPageBytes, runtimeType(kind)); + + CudaStream stream; + invokeNvfp4ColdPageEncode(offloadTasks.data(), offloadTasks.size(), inputMetadata, coldBase, stream); + invokeNvfp4ColdPageDecode(onboardTasks.data(), onboardTasks.size(), outputMetadata, coldBase, stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + auto const cold = coldStorage.payload(); + EXPECT_EQ(cold.front(), kCanary); + for (std::size_t page = 0; page < numPages; ++page) + { + std::size_t const coldPage = coldBaseOffset + page * coldPageBytes; + auto const coldRegion = [&](std::size_t offset, std::size_t bytes) + { + return std::vector(cold.begin() + static_cast(coldPage + offset), + cold.begin() + static_cast(coldPage + offset + bytes)); + }; + EXPECT_EQ(coldRegion(0U, mlaPackedBytes), references[page].packed); + EXPECT_EQ(coldRegion(mlaPackedBytes, mlaScaleBytes), references[page].scales); + EXPECT_TRUE(std::all_of(cold.begin() + static_cast(coldPage + mlaPayloadBytes), + cold.begin() + static_cast(coldPage + sideColdOffset), + [](std::uint8_t value) { return value == 0U; })); + EXPECT_EQ(coldRegion(sideColdOffset, sideRawBytes), sideHost[page]); + EXPECT_TRUE(std::all_of(cold.begin() + static_cast(coldPage + sideColdEnd), + cold.begin() + static_cast(coldPage + coldPageBytes), + [](std::uint8_t value) { return value == 0U; })); + + EXPECT_EQ(mlaOutput.copyToHost(page * mlaRawBytes, mlaRawBytes), + decompressReference(references[page], kind, params, geometry)); + EXPECT_EQ(sideOutput.copyToHost(page * sideSlotBytes, sideRawBytes), sideHost[page]); + for (auto const* side : {&sideInput, &sideOutput}) + { + auto const slotTail = side->copyToHost(page * sideSlotBytes + sideRawBytes, sideSlotBytes - sideRawBytes); + EXPECT_TRUE( + std::all_of(slotTail.begin(), slotTail.end(), [](std::uint8_t value) { return value == kCanary; })); + } + } + mlaInput.expectCanaries(); + mlaOutput.expectCanaries(); + sideInput.expectCanaries(); + sideOutput.expectCanaries(); + coldStorage.expectCanaries(); +} + +class Nvfp4ColdPageMlaSideTest : public testing::TestWithParam +{ +}; + +TEST_P(Nvfp4ColdPageMlaSideTest, MlaPageAndDefaultDsaIndexKeyRoundTripExactly) +{ + runUnaryMlaWithLosslessSideRoundTrip(GetParam()); +} + +INSTANTIATE_TEST_SUITE_P( + AllRuntimeTypes, Nvfp4ColdPageMlaSideTest, testing::Values(RawKind::kFloat16, RawKind::kBfloat16, RawKind::kFp8)); + +TEST(Nvfp4ColdPageWholePageTest, DifferentLayerScalesRemainInOneCompletePageBatch) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + if (!tensorrt_llm::common::isSM100Family()) + { + GTEST_SKIP() << "NVFP4 cold-page kernels require an SM100-family GPU"; + } + + constexpr std::size_t numLayers = 2; + RawKind constexpr kind = RawKind::kBfloat16; + PageGeometry constexpr geometry = kMinimumCompactGeometry; + std::size_t const rawSlotBytes = rawBytes(kind, geometry); + std::size_t const layerRecordBytes = 2U * (packedBytes(geometry) + scaleBytes(geometry)); + std::size_t const layerRecordStride = roundUp(layerRecordBytes, alignof(uint4)); + std::size_t const coldPageBytes = numLayers * layerRecordStride; + + std::array, numLayers> rawInputK; + std::array, numLayers> rawInputV; + std::array, numLayers> rawOutputK; + std::array, numLayers> rawOutputV; + std::array, numLayers> params{ + std::array{makeParams(geometry, 0U), makeParams(geometry, 1U)}, + std::array{makeParams(geometry, 0U), makeParams(geometry, 1U)}}; + // Distinct per-layer K/V scales verify blockIdx.y selects immutable launch metadata. + params[1][0].nvfp4ScaleOrigQuant = 0.5F; + params[1][0].nvfp4ScaleQuantOrig = 2.0F; + params[1][1].nvfp4ScaleOrigQuant = 4.0F; + params[1][1].nvfp4ScaleQuantOrig = 0.25F; + + std::array, 2>, numLayers> rawHost; + std::array, numLayers> references; + std::vector inputMetadatas; + std::vector outputMetadatas; + inputMetadatas.reserve(2U * numLayers); + outputMetadatas.reserve(2U * numLayers); + for (std::size_t layer = 0; layer < numLayers; ++layer) + { + rawInputK[layer] = std::make_unique(rawSlotBytes); + rawInputV[layer] = std::make_unique(rawSlotBytes); + rawOutputK[layer] = std::make_unique(rawSlotBytes); + rawOutputV[layer] = std::make_unique(rawSlotBytes); + for (std::uint32_t role = 0; role < 2; ++role) + { + rawHost[layer][role] = makeRawPage(kind, layer, role, params[layer][role], geometry, InputPattern::kDense); + references[layer][role] = compressReference(rawHost[layer][role], kind, params[layer][role], geometry); + } + rawInputK[layer]->copyFrom(rawHost[layer][0]); + rawInputV[layer]->copyFrom(rawHost[layer][1]); + std::size_t const base = layer * layerRecordStride; + std::size_t const packed = packedBytes(geometry); + std::size_t const scale = scaleBytes(geometry); + auto const appendPlans = [&](auto& plans, DeviceRegion const& rawK, DeviceRegion const& rawV) + { + plans.push_back({reinterpret_cast(rawK.data()), rawSlotBytes, rawSlotBytes, base, + base + 2U * packed, 0U, 0U, Nvfp4ColdPageTransform::kNvfp4, params[layer][0]}); + plans.push_back({reinterpret_cast(rawV.data()), rawSlotBytes, rawSlotBytes, base + packed, + base + 2U * packed + scale, base + layerRecordBytes, + static_cast(layerRecordStride - layerRecordBytes), Nvfp4ColdPageTransform::kNvfp4, + params[layer][1]}); + }; + appendPlans(inputMetadatas, *rawInputK[layer], *rawInputV[layer]); + appendPlans(outputMetadatas, *rawOutputK[layer], *rawOutputV[layer]); + } + + MappedHostRegion compactPage(coldPageBytes); + CudaStream stream; + auto const inputMetadata + = makeNvfp4ColdPageTestMetadata(inputMetadatas, coldPageBytes, Nvfp4ColdPageRuntimeType::kBfloat16); + auto const outputMetadata + = makeNvfp4ColdPageTestMetadata(outputMetadatas, coldPageBytes, Nvfp4ColdPageRuntimeType::kBfloat16); + PageIndexPair const page{0, 0}; + invokeNvfp4ColdPageEncode(&page, 1U, inputMetadata, compactPage.data(), stream); + invokeNvfp4ColdPageDecode(&page, 1U, outputMetadata, compactPage.data(), stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + auto const compact = compactPage.payload(); + std::size_t const packed = packedBytes(geometry); + std::size_t const scale = scaleBytes(geometry); + auto const compactRegion = [&](std::size_t offset, std::size_t bytes) + { + return std::vector(compact.begin() + static_cast(offset), + compact.begin() + static_cast(offset + bytes)); + }; + for (std::size_t layer = 0; layer < numLayers; ++layer) + { + std::size_t const base = layer * layerRecordStride; + EXPECT_EQ(compactRegion(base, packed), references[layer][0].packed); + EXPECT_EQ(compactRegion(base + packed, packed), references[layer][1].packed); + EXPECT_EQ(compactRegion(base + 2U * packed, scale), references[layer][0].scales); + EXPECT_EQ(compactRegion(base + 2U * packed + scale, scale), references[layer][1].scales); + auto const padding = compactRegion(base + layerRecordBytes, layerRecordStride - layerRecordBytes); + EXPECT_TRUE(std::all_of(padding.begin(), padding.end(), [](std::uint8_t value) { return value == 0U; })); + EXPECT_EQ(rawOutputK[layer]->copyToHost(), + decompressReference(references[layer][0], kind, params[layer][0], geometry)); + EXPECT_EQ(rawOutputV[layer]->copyToHost(), + decompressReference(references[layer][1], kind, params[layer][1], geometry)); + } + compactPage.expectCanaries(); +} + +void expectWholePageLaunchTopology(std::size_t numPages, std::vector expectedGridZ) +{ + constexpr std::size_t numLayers = 2; + RawKind constexpr kind = RawKind::kBfloat16; + std::size_t const rawSlotBytes = rawBytes(kind, kSmallVectorGeometry); + std::size_t const recordBytes = 2U * (packedBytes(kSmallVectorGeometry) + scaleBytes(kSmallVectorGeometry)); + std::size_t const recordStride = roundUp(recordBytes, alignof(uint4)); + std::size_t const coldPageBytes = numLayers * recordStride; + + std::array, numLayers> rawK; + std::array, numLayers> rawV; + std::vector buffers; + buffers.reserve(2U * numLayers); + for (std::size_t layer = 0; layer < numLayers; ++layer) + { + rawK[layer] = std::make_unique(numPages * rawSlotBytes); + rawV[layer] = std::make_unique(numPages * rawSlotBytes); + auto kParams = makeParams(kSmallVectorGeometry, 0U); + auto const vParams = makeParams(kSmallVectorGeometry, 1U); + kParams.nvfp4ScaleOrigQuant *= static_cast(layer + 1U); + kParams.nvfp4ScaleQuantOrig /= static_cast(layer + 1U); + std::size_t const base = layer * recordStride; + std::size_t const packed = packedBytes(kSmallVectorGeometry); + std::size_t const scale = scaleBytes(kSmallVectorGeometry); + buffers.push_back({reinterpret_cast(rawK[layer]->data()), rawSlotBytes, rawSlotBytes, base, + base + 2U * packed, 0U, 0U, Nvfp4ColdPageTransform::kNvfp4, kParams}); + buffers.push_back({reinterpret_cast(rawV[layer]->data()), rawSlotBytes, rawSlotBytes, + base + packed, base + 2U * packed + scale, base + recordBytes, + static_cast(recordStride - recordBytes), Nvfp4ColdPageTransform::kNvfp4, vParams}); + } + + MappedHostRegion coldPages(numPages * coldPageBytes); + std::vector offloadPages; + std::vector onboardPages; + offloadPages.reserve(numPages); + onboardPages.reserve(numPages); + for (std::size_t page = 0; page < numPages; ++page) + { + auto const pageIndex = static_cast(page); + offloadPages.push_back({pageIndex, pageIndex}); + onboardPages.push_back({pageIndex, pageIndex}); + } + + CudaStream stream; + auto const plan = makeNvfp4ColdPageTestMetadata(buffers, coldPageBytes, Nvfp4ColdPageRuntimeType::kBfloat16); + auto const expectWholePageKernels = [&](auto const& enqueue) + { + cudaGraph_t graph{}; + ASSERT_EQ(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal), cudaSuccess); + enqueue(); + ASSERT_EQ(cudaStreamEndCapture(stream, &graph), cudaSuccess); + + std::size_t numNodes = 0; + ASSERT_EQ(cudaGraphGetNodes(graph, nullptr, &numNodes), cudaSuccess); + std::vector nodes(numNodes); + ASSERT_EQ(cudaGraphGetNodes(graph, nodes.data(), &numNodes), cudaSuccess); + std::size_t kernelNodes = 0; + std::vector actualGridZ; + for (auto const node : nodes) + { + cudaGraphNodeType nodeType{}; + ASSERT_EQ(cudaGraphNodeGetType(node, &nodeType), cudaSuccess); + if (nodeType != cudaGraphNodeTypeKernel) + { + continue; + } + ++kernelNodes; + cudaKernelNodeParams nodeParams{}; + ASSERT_EQ(cudaGraphKernelNodeGetParams(node, &nodeParams), cudaSuccess); + EXPECT_EQ(nodeParams.gridDim.y, 2U * numLayers); + actualGridZ.push_back(nodeParams.gridDim.z); + } + std::sort(actualGridZ.begin(), actualGridZ.end()); + std::sort(expectedGridZ.begin(), expectedGridZ.end()); + EXPECT_EQ(kernelNodes, expectedGridZ.size()); + EXPECT_EQ(actualGridZ, expectedGridZ); + ASSERT_EQ(cudaGraphDestroy(graph), cudaSuccess); + }; + + expectWholePageKernels( + [&] { invokeNvfp4ColdPageEncode(offloadPages.data(), offloadPages.size(), plan, coldPages.data(), stream); }); + expectWholePageKernels( + [&] { invokeNvfp4ColdPageDecode(onboardPages.data(), onboardPages.size(), plan, coldPages.data(), stream); }); +} + +TEST(Nvfp4ColdPageWholePageTest, TwoHundredFiftySevenPagesUseExactlyTwoWholePageKernelsPerDirection) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + if (!tensorrt_llm::common::isSM100Family()) + { + GTEST_SKIP() << "NVFP4 cold-page kernels require an SM100-family GPU"; + } + expectWholePageLaunchTopology(257, {1, 256}); +} + +TEST(Nvfp4ColdPageWholePageTest, FourThousandNinetySixPagesUseSixteenWholePageKernelsPerDirection) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + if (!tensorrt_llm::common::isSM100Family()) + { + GTEST_SKIP() << "NVFP4 cold-page kernels require an SM100-family GPU"; + } + expectWholePageLaunchTopology(4096, std::vector(16, 256)); +} + +class Nvfp4ColdPageTailTest : public testing::TestWithParam +{ +}; + +TEST_P(Nvfp4ColdPageTailTest, InactiveRowsDoNotAffectTheValidPrefix) +{ + runPartialPageTailIsolation(GetParam()); +} + +INSTANTIATE_TEST_SUITE_P( + AllRuntimeTypes, Nvfp4ColdPageTailTest, testing::Values(RawKind::kFloat16, RawKind::kBfloat16, RawKind::kFp8)); + +TEST(Nvfp4ColdPageValidationTest, EmptyBatchIsAnAsyncNoOp) +{ + invokeNvfp4ColdPageEncode(nullptr, 0U, Nvfp4ColdPageTestMetadata{}, nullptr, nullptr); + invokeNvfp4ColdPageDecode(nullptr, 0U, Nvfp4ColdPageTestMetadata{}, nullptr, nullptr); +} + +} // namespace diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index a111319795a3..1b80c8beeb9b 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -6202,6 +6202,7 @@ def launchTestJobs(pipeline, testFilter, globalVars) "DGX_H100-4_GPUs-PyTorch-Others-2": ["auto:dgx-h100-x4", "l0_dgx_h100", 2, 2, 4], "DGX_H100-4_GPUs-PyTorch-Ray-1": ["auto:dgx-h100-x4", "l0_dgx_h100", 1, 1, 4], "DGX_H100-4_GPUs-PyTorch-Post-Merge-1": ["auto:dgx-h100-x4", "l0_dgx_h100", 1, 1, 4], + "DGX_B200-CPP-1": ["auto:dgx-b200-flex", "l0_b200", 1, 1, 1, 1, true], "DGX_B200-PyTorch-1": ["auto:dgx-b200-flex", "l0_b200", 1, 9, 1, 1, true], "DGX_B200-PyTorch-2": ["auto:dgx-b200-flex", "l0_b200", 2, 9, 1, 1, true], "DGX_B200-PyTorch-3": ["auto:dgx-b200-flex", "l0_b200", 3, 9, 1, 1, true], diff --git a/tensorrt_llm/_torch/kv_cache_compression/quantization_for_cold_page/__init__.py b/tensorrt_llm/_torch/kv_cache_compression/quantization_for_cold_page/__init__.py new file mode 100644 index 000000000000..c2d2e98e9de8 --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/quantization_for_cold_page/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. diff --git a/tensorrt_llm/_torch/kv_cache_compression/quantization_for_cold_page/nvfp4_quantization.py b/tensorrt_llm/_torch/kv_cache_compression/quantization_for_cold_page/nvfp4_quantization.py new file mode 100644 index 000000000000..429a904336b1 --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/quantization_for_cold_page/nvfp4_quantization.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""NVFP4 cold-page layout, scales, metadata, and kernel dispatch.""" + +import json +import math +import os +import re +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Sequence + +import torch + +from tensorrt_llm.quantization.modelopt_config import ( + is_modelopt_quant_config, + read_modelopt_quant_config, +) + +from ...pyexecutor.resource_manager import DataType +from .quantization_for_cold_page import ColdPageQuantizationCompression + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import ColdPageQuantizationCompressionConfig + +_LayerScales = tuple[tuple[float, float], tuple[float, float]] + +_IDENTITY_NVFP4_SCALES: _LayerScales = ((1.0, 1.0), (1.0, 1.0)) +_MODEL_OPT_LANGUAGE_KV_SCALE_KEY = re.compile( + r"^model(?:\.language_model)?\.layers\.(?P\d+)\.self_attn\." + r"(?P[kv])_proj\.(?P=kind)_scale$" +) +_COLD_PAGE_ALIGNMENT = 16 +_ELEMENTS_PER_BYTE = 2 +_ELEMENTS_PER_SCALE = 16 +_ELEMENTS_PER_HALF_GROUP = 8 +_MAX_HALF_GROUPS_PER_TILE = 2048 +_MAX_BUFFERS_PER_LAUNCH = 256 +_WIDE_FIELDS = 6 +_INTEGER_FIELDS = 5 +_SCALE_FIELDS = 4 +_NVFP4_TRANSFORM = 0 +_LOSSLESS_TRANSFORM = 1 + + +@dataclass(frozen=True) +class _Nvfp4Scales: + nvfp4_orig_quant: float + nvfp4_quant_orig: float + fp8_orig_quant: float = 1.0 + fp8_quant_orig: float = 1.0 + + +@dataclass(frozen=True) +class _Nvfp4BufferLayout: + role: str + scales: _Nvfp4Scales | None = None + + +@dataclass(frozen=True) +class _Nvfp4LayerLayout: + layer_id: int + num_kv_heads: int + tokens_per_page: int + head_dim: int + buffers: tuple[_Nvfp4BufferLayout, ...] + + +@dataclass(frozen=True) +class _Nvfp4ColdPageMetadata: + """Python-owned launch metadata for one KVCM lifecycle.""" + + wide: torch.Tensor + integers: torch.Tensor + scales: torch.Tensor + num_buffers: int + max_half_groups_per_tile: int + cold_page_bytes: int + + +@dataclass +class _Nvfp4ColdPageCodecState: + """NVFP4 state owned by one target, draft, or retry codec.""" + + layer_layouts: dict[int, _Nvfp4LayerLayout] + layer_ids: tuple[int, ...] + runtime_type: int + lifecycle_metadata: tuple[_Nvfp4ColdPageMetadata, ...] = field(init=False) + + +def _load_modelopt_nvfp4_scales( + checkpoint_path: str | None, +) -> dict[int, _LayerScales]: + """Load optional ModelOpt NVFP4 K/V global scales by model layer.""" + + if checkpoint_path is None or os.environ.get("TRTLLM_LOAD_KV_SCALES", "1") != "1": + return {} + + checkpoint_dir = Path(checkpoint_path) + weight_files = sorted(checkpoint_dir.glob("*.safetensors")) + ordinary_files = [path for path in weight_files if "consolidated" not in path.name] + weight_files = ordinary_files or weight_files + if not weight_files: + raise FileNotFoundError( + f"No safetensors files in ModelOpt scale checkpoint {checkpoint_dir}" + ) + + metadata_path = checkpoint_dir / "hf_quant_config.json" + if metadata_path.exists(): + metadata = json.loads(metadata_path.read_text()) + else: + config_path = checkpoint_dir / "config.json" + metadata = ( + json.loads(config_path.read_text()).get("quantization_config") + if config_path.exists() + else None + ) + if not is_modelopt_quant_config(metadata): + return {} + if read_modelopt_quant_config(metadata).get("kv_cache_quant_algo") != "NVFP4": + return {} + + from safetensors import safe_open + + values: dict[int, dict[str, list[float]]] = {} + for file_path in weight_files: + with safe_open(str(file_path), framework="pt", device="cpu") as checkpoint: + for tensor_name in checkpoint.keys(): + match = _MODEL_OPT_LANGUAGE_KV_SCALE_KEY.fullmatch(tensor_name) + if match is None: + continue + value = float(checkpoint.get_tensor(tensor_name).reshape([]).item()) + if not math.isfinite(value) or value <= 0.0: + raise ValueError( + f"ModelOpt KV scale {file_path}:{tensor_name} must be finite and positive" + ) + layer_values = values.setdefault(int(match.group("layer_id")), {"k": [], "v": []}) + layer_values[match.group("kind")].append(value) + + result: dict[int, _LayerScales] = {} + for layer_id, layer_values in values.items(): + k_values, v_values = layer_values["k"], layer_values["v"] + if not k_values or not v_values: + raise ValueError(f"ModelOpt KV scales for layer {layer_id} must contain both K and V") + quant_orig = (max(k_values), max(v_values)) + orig_quant = (1.0 / quant_orig[0], 1.0 / quant_orig[1]) + stored_scales = torch.tensor( + (*orig_quant, *quant_orig), dtype=torch.float32, device="cpu" + ).tolist() + if any(not math.isfinite(value) or value <= 0.0 for value in stored_scales): + raise ValueError( + f"ModelOpt KV scales for layer {layer_id} are not representable as float32" + ) + result[layer_id] = ( + (stored_scales[0], stored_scales[1]), + (stored_scales[2], stored_scales[3]), + ) + return result + + +class Nvfp4ColdPageQuantizationCompression(ColdPageQuantizationCompression): + """NVFP4 layout, calibration metadata, and CUDA dispatch.""" + + def __init__(self, config: "ColdPageQuantizationCompressionConfig") -> None: + super().__init__(config) + self._model_scales = _load_modelopt_nvfp4_scales(config.scale_checkpoint_path) + + def build_codec_state( + self, + cache_config: object, + *, + runtime_dtype: DataType, + pp_layers: Sequence[int], + num_kv_heads_per_layer: Sequence[int], + head_dim_per_layer: Sequence[int], + is_draft: bool = False, + ) -> _Nvfp4ColdPageCodecState: + from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig + + runtime_type = { + DataType.HALF: 0, + DataType.BF16: 1, + DataType.FP8: 2, + }.get(runtime_dtype) + attention_layers = [ + layer for layer in cache_config.layers if isinstance(layer, AttentionLayerConfig) + ] + if attention_layers and runtime_type is None: + raise RuntimeError( + "NVFP4 cold-page compression supports FP16, BF16, or FP8 " + f"Attention KV, not {runtime_dtype}" + ) + + layer_layouts = [] + for layer in attention_layers: + layer_id = int(layer.layer_id) + buffer_roles = {str(buffer.role) for buffer in layer.buffers} + if "key" not in buffer_roles: + raise NotImplementedError( + "NVFP4 cold-page compression requires an Attention key buffer" + ) + + compressed_roles = ("key", "value") if "value" in buffer_roles else ("key",) + if len(compressed_roles) == 2 and not is_draft: + orig_quant, quant_orig = self._model_scales.get( + int(pp_layers[layer_id]), _IDENTITY_NVFP4_SCALES + ) + else: + orig_quant, quant_orig = _IDENTITY_NVFP4_SCALES + + num_kv_heads = int(num_kv_heads_per_layer[layer_id]) + tokens_per_page = int(cache_config.tokens_per_block) + head_dim = int(head_dim_per_layer[layer_id]) + if head_dim <= 0 or head_dim % _ELEMENTS_PER_SCALE != 0: + raise ValueError( + f"NVFP4 cold pages require head_dim divisible by 16, got {head_dim}" + ) + buffer_layouts = [ + _Nvfp4BufferLayout( + role=role, + scales=_Nvfp4Scales(orig_quant[index], quant_orig[index]), + ) + for index, role in enumerate(compressed_roles) + ] + for buffer in layer.buffers: + role = str(buffer.role) + if role not in compressed_roles: + buffer_layouts.append(_Nvfp4BufferLayout(role=role)) + + layer_layouts.append( + _Nvfp4LayerLayout( + layer_id=layer_id, + num_kv_heads=num_kv_heads, + tokens_per_page=tokens_per_page, + head_dim=head_dim, + buffers=tuple(buffer_layouts), + ) + ) + + layouts_by_layer = {layout.layer_id: layout for layout in layer_layouts} + return _Nvfp4ColdPageCodecState( + layer_layouts=layouts_by_layer, + layer_ids=tuple(sorted(layouts_by_layer)), + runtime_type=runtime_type if runtime_type is not None else 0, + ) + + def build_lifecycle_metadata( + self, codec_state: _Nvfp4ColdPageCodecState, lifecycle: object + ) -> _Nvfp4ColdPageMetadata: + wide_rows: list[list[int]] = [] + integer_rows: list[list[int]] = [] + scale_rows: list[list[float]] = [] + cold_page_bytes = 0 + max_half_groups_per_tile = 0 + + for layer_id, hot_buffers in lifecycle.layers.items(): + layout = codec_state.layer_layouts[int(layer_id)] + expected_roles = {buffer.role for buffer in layout.buffers} + if set(hot_buffers) != expected_roles: + raise ValueError(f"Cold-page layer {layer_id} roles do not match its KVCM layout") + elements = layout.num_kv_heads * layout.tokens_per_page * layout.head_dim + element_bytes = 1 if codec_state.runtime_type == 2 else 2 + expected_raw_bytes = elements * element_bytes + half_groups = elements // _ELEMENTS_PER_HALF_GROUP + compressed_count = sum(buffer.scales is not None for buffer in layout.buffers) + packed_bytes = elements // _ELEMENTS_PER_BYTE + scale_bytes = elements // _ELEMENTS_PER_SCALE + layer_start = cold_page_bytes + scale_start = layer_start + compressed_count * packed_bytes + cursor = scale_start + compressed_count * scale_bytes + + compressed_index = 0 + for buffer in layout.buffers: + is_compressed = buffer.scales is not None + hot = hot_buffers[buffer.role] + raw_base = int(hot.raw_base) + raw_slot_bytes = int(hot.raw_slot_bytes) + raw_bytes = int(hot.raw_bytes) + if raw_base <= 0 or raw_bytes <= 0 or raw_bytes > raw_slot_bytes: + raise ValueError("Cold-page hot buffer has invalid address or size") + + if is_compressed: + data_offset = layer_start + compressed_index * packed_bytes + scale_offset = scale_start + compressed_index * scale_bytes + compressed_index += 1 + if raw_bytes != expected_raw_bytes: + raise ValueError("Hot buffer size does not match NVFP4 geometry") + if raw_base % 16 or raw_slot_bytes % 16: + raise ValueError( + "NVFP4 hot address and Slot stride must be 16-byte aligned" + ) + max_half_groups_per_tile = max( + max_half_groups_per_tile, + min(half_groups, _MAX_HALF_GROUPS_PER_TILE), + ) + else: + data_offset = cursor + scale_offset = 0 + cursor += raw_bytes + + wide_rows.append( + [ + raw_base, + raw_slot_bytes, + raw_bytes, + data_offset, + scale_offset, + 0, + ] + ) + integer_rows.append( + [ + 0, + _NVFP4_TRANSFORM if is_compressed else _LOSSLESS_TRANSFORM, + layout.num_kv_heads if is_compressed else 0, + layout.tokens_per_page if is_compressed else 0, + layout.head_dim if is_compressed else 0, + ] + ) + buffer_scales = buffer.scales if is_compressed else _Nvfp4Scales(1.0, 1.0) + scale_rows.append( + [ + buffer_scales.nvfp4_orig_quant, + buffer_scales.nvfp4_quant_orig, + buffer_scales.fp8_orig_quant, + buffer_scales.fp8_quant_orig, + ] + ) + layer_end = ( + (cursor + _COLD_PAGE_ALIGNMENT - 1) // _COLD_PAGE_ALIGNMENT * _COLD_PAGE_ALIGNMENT + ) + wide_rows[-1][5] = cursor + integer_rows[-1][0] = layer_end - cursor + cold_page_bytes = layer_end + + num_buffers = len(wide_rows) + if not 0 < num_buffers <= _MAX_BUFFERS_PER_LAUNCH: + raise ValueError( + f"NVFP4 cold-page lifecycle has {num_buffers} buffers; " + f"the maximum is {_MAX_BUFFERS_PER_LAUNCH}" + ) + padding = _MAX_BUFFERS_PER_LAUNCH - num_buffers + return _Nvfp4ColdPageMetadata( + wide=torch.tensor( + wide_rows + [[0] * _WIDE_FIELDS for _ in range(padding)], + dtype=torch.int64, + device="cpu", + ), + integers=torch.tensor( + integer_rows + [[0] * _INTEGER_FIELDS for _ in range(padding)], + dtype=torch.int32, + device="cpu", + ), + scales=torch.tensor( + scale_rows + [[0.0] * _SCALE_FIELDS for _ in range(padding)], + dtype=torch.float32, + device="cpu", + ), + num_buffers=num_buffers, + max_half_groups_per_tile=max_half_groups_per_tile, + cold_page_bytes=cold_page_bytes, + ) + + def encode_cold_pages( + self, + codec_state: _Nvfp4ColdPageCodecState, + lifecycle_index: int, + cold_base: int, + page_indices: int, + num_pages: int, + stream: int, + ) -> None: + from tensorrt_llm.bindings.internal import kv_cache_compression as native + + metadata = codec_state.lifecycle_metadata[lifecycle_index] + native.nvfp4_cold_page_encode( + page_indices, + num_pages, + metadata.wide.data_ptr(), + metadata.integers.data_ptr(), + metadata.scales.data_ptr(), + metadata.num_buffers, + metadata.max_half_groups_per_tile, + metadata.cold_page_bytes, + codec_state.runtime_type, + cold_base, + stream, + ) + + def decode_cold_pages( + self, + codec_state: _Nvfp4ColdPageCodecState, + lifecycle_index: int, + cold_base: int, + page_indices: int, + num_pages: int, + stream: int, + ) -> None: + from tensorrt_llm.bindings.internal import kv_cache_compression as native + + metadata = codec_state.lifecycle_metadata[lifecycle_index] + native.nvfp4_cold_page_decode( + page_indices, + num_pages, + metadata.wide.data_ptr(), + metadata.integers.data_ptr(), + metadata.scales.data_ptr(), + metadata.num_buffers, + metadata.max_half_groups_per_tile, + metadata.cold_page_bytes, + codec_state.runtime_type, + cold_base, + stream, + ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/quantization_for_cold_page/quantization_for_cold_page.py b/tensorrt_llm/_torch/kv_cache_compression/quantization_for_cold_page/quantization_for_cold_page.py new file mode 100644 index 000000000000..d1c49f2bf0e8 --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/quantization_for_cold_page/quantization_for_cold_page.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Common runtime pipeline for cold-page quantization.""" + +from typing import Any, Sequence + +from ...pyexecutor.resource_manager import DataType, KVCacheCompressionManager + + +class ColdPageQuantizationCompression(KVCacheCompressionManager): + """Common codec registration and callbacks for cold-page quantizers.""" + + uses_iteration_lifecycle = False + provides_cold_page_codec = True + + def create_cold_page_codec( + self, + cache_config: object, + *, + runtime_dtype: DataType, + pp_layers: Sequence[int], + num_kv_heads_per_layer: Sequence[int], + head_dim_per_layer: Sequence[int], + is_draft: bool = False, + ) -> object: + """Create one native codec with state isolated to this KVCM.""" + + from tensorrt_llm.bindings.internal import kv_cache_compression as native + + codec_state = self.build_codec_state( + cache_config, + runtime_dtype=runtime_dtype, + pp_layers=pp_layers, + num_kv_heads_per_layer=num_kv_heads_per_layer, + head_dim_per_layer=head_dim_per_layer, + is_draft=is_draft, + ) + return native.create_python_cold_page_codec(self, codec_state) + + def configure(self, codec_state: Any, lifecycles: Sequence[object]) -> Sequence[object]: + """Resolve hot buffers and publish each lifecycle's cold-page size.""" + + from tensorrt_llm.bindings.internal import kv_cache_compression as native + + codec_state.lifecycle_metadata = tuple( + self.build_lifecycle_metadata(codec_state, lifecycle) for lifecycle in lifecycles + ) + properties = [] + for metadata in codec_state.lifecycle_metadata: + lifecycle = native.ColdPageLifecycleProperties() + lifecycle.cold_page_bytes = metadata.cold_page_bytes + lifecycle.page_index_location = native.ColdPageIndexLocation.HOST + properties.append(lifecycle) + return properties + + def build_codec_state( + self, + cache_config: object, + *, + runtime_dtype: DataType, + pp_layers: Sequence[int], + num_kv_heads_per_layer: Sequence[int], + head_dim_per_layer: Sequence[int], + is_draft: bool = False, + ) -> object: + """Build the format-specific state owned by one native codec.""" + raise NotImplementedError + + def build_lifecycle_metadata(self, codec_state: object, lifecycle: object) -> object: + """Resolve one KVCM lifecycle into format-specific launch metadata.""" + raise NotImplementedError diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 8462a4bd8cc7..468c26eff699 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -150,12 +150,10 @@ class TriAttentionCompressionManager(KVCacheCompressionManager): def __init__( self, config: "TriAttentionKvCacheCompressionConfig", - kv_cache_manager: KVCacheManagerV2, - draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, *, pretrained_config: "PretrainedConfig", ) -> None: - super().__init__(config, kv_cache_manager, draft_kv_cache_manager) + super().__init__(config) self.budget = config.budget self.beta = config.beta self.eviction_mode = config.eviction_mode @@ -168,6 +166,14 @@ def __init__( self._load_calibration() self._prepared_generation_batch: Optional["ScheduledRequests"] = None + + def bind_kv_cache_managers( + self, + kv_cache_manager: KVCacheManagerV2, + draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, + ) -> None: + """Finalize state whose geometry is owned by the constructed KVCMs.""" + super().bind_kv_cache_managers(kv_cache_manager, draft_kv_cache_manager) # Manager-lifetime constants. self._num_extra_kv_tokens = int(kv_cache_manager.num_extra_kv_tokens) self._protected_tail_capacity = ( @@ -420,7 +426,7 @@ def _evict_due_requests( # this request (pre-launch) instead of failing the batch. continue draft_cache = None - if self.draft_kv_cache_manager is not None: + if self.has_independent_draft_kv_cache: # A missing draft cache is a wiring bug: keep the precise KeyError. draft_cache = self.draft_kv_cache_manager.kv_cache_map[request_id] if not draft_cache.is_active: @@ -600,7 +606,7 @@ def cumulative_offsets(move_counts: List[int]) -> List[int]: if self._swa_window is not None: swa_offsets = cumulative_offsets([self._swa_window + tail for tail in tails]) draft_offsets = None - if self.draft_kv_cache_manager is not None: + if self.has_independent_draft_kv_cache: draft_offsets = cumulative_offsets( [self.budget + self._draft_protected_tail_capacity] * len(eviction_requests) ) @@ -681,7 +687,7 @@ def _initialize_eviction_state(self) -> None: """Create manager-lifetime state once.""" target_layout = self._create_kv_layout() draft_layout = ( - self._create_kv_layout(draft=True) if self.draft_kv_cache_manager is not None else None + self._create_kv_layout(draft=True) if self.has_independent_draft_kv_cache else None ) self._target_layout = target_layout self._draft_layout = draft_layout diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b4422c084463..084da74711ed 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -15,7 +15,7 @@ import copy import dataclasses import os -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Union import torch @@ -39,6 +39,7 @@ supports_native_fp8_lora) from tensorrt_llm.logger import logger from tensorrt_llm.mapping import CpType, Mapping +from tensorrt_llm.quantization import QuantAlgo from ..attention_backend import get_sparse_attn_kv_cache_manager from ..hostfunc import set_low_latency_dispatch @@ -77,9 +78,6 @@ SimpleUnifiedScheduler) from .seq_slot_manager import SeqSlotManager -if TYPE_CHECKING: - import transformers - GB = 1 << 30 @@ -1361,7 +1359,8 @@ def _create_kv_cache_manager( self, model_engine: PyTorchModelEngine, estimating_kv_cache: bool = False, - kv_cache_config_override: Optional[KvCacheConfig] = None + kv_cache_config_override: Optional[KvCacheConfig] = None, + cold_page_codec_provider: Optional[object] = None, ) -> KVCacheManager: mapping = self._mapping assert model_engine.model.model_config.is_generation, "Only construct KV cache for generation models." @@ -1399,6 +1398,7 @@ def _create_kv_cache_manager( execution_stream=self._execution_stream, layer_mask=spec_dec_layer_mask, is_disagg=self._is_disagg, + cold_page_codec_provider=cold_page_codec_provider, ) if not self._skip_est: @@ -1515,6 +1515,7 @@ def _create_one_model_draft_kv_cache_manager( max_seq_len: int, estimating_kv_cache: bool = False, kv_cache_config_override: Optional[KvCacheConfig] = None, + cold_page_codec_provider: Optional[object] = None, ) -> Optional[KVCacheManager]: """ Create a KV cache manager for draft model layers in one-model mode @@ -1586,6 +1587,7 @@ def _create_one_model_draft_kv_cache_manager( layer_mask=spec_dec_layer_mask, num_layers=num_draft_layers, is_disagg=self._is_disagg, + cold_page_codec_provider=cold_page_codec_provider, ) def _get_target_and_draft_cache_costs( @@ -2048,10 +2050,23 @@ def build_managers(self, budget_attr, self_kv_cache_config, draft_kv_cache_config)) + compression_config = self._llm_args.kv_cache_compression_config + compression_manager = create_kv_cache_compression_manager( + compression_config, + model_engine=self._model_engine, + kv_cache_config=self_kv_cache_config, + estimating_kv_cache=estimating_kv_cache and not self._skip_est, + ) + cold_page_codec_provider = ( + compression_manager if compression_manager is not None + and compression_manager.provides_cold_page_codec else None) + kv_cache_manager = self._create_kv_cache_manager( self._model_engine, estimating_kv_cache, - kv_cache_config_override=self_kv_cache_config) + kv_cache_config_override=self_kv_cache_config, + cold_page_codec_provider=cold_page_codec_provider, + ) # Carry the fp8 context-MLA workspace admission cap (computed in configure_kv_cache_capacity) onto # the real KV manager so the scheduler reads it directly instead of re-deriving from pool layout. @@ -2089,7 +2104,8 @@ def build_managers(self, draft_kv_cache_manager = self._create_one_model_draft_kv_cache_manager( original_max_seq_len, estimating_kv_cache, - kv_cache_config_override=draft_build_kv_cache_config) + kv_cache_config_override=draft_build_kv_cache_config, + cold_page_codec_provider=cold_page_codec_provider) # Encoder-decoder cross-attention pool cross_kv_cache_manager = None @@ -2103,9 +2119,17 @@ def build_managers(self, ResourceManagerType.DRAFT_KV_CACHE_MANAGER] = draft_kv_cache_manager resources[ ResourceManagerType.CROSS_KV_CACHE_MANAGER] = cross_kv_cache_manager + if (compression_manager is not None + and compression_manager.uses_iteration_lifecycle): + resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] = ( + compression_manager) def teardown_managers(self, resources: Dict) -> None: """Clean up KV caches for model, draft model, and cross pool.""" + compression_manager = resources.pop( + ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER, None) + if compression_manager is not None: + compression_manager.shutdown() resources[ResourceManagerType.KV_CACHE_MANAGER].shutdown() del resources[ResourceManagerType.KV_CACHE_MANAGER] draft_kv_cache_manager = resources[ @@ -2217,11 +2241,18 @@ def _create_kv_cache_manager( num_kv_heads: Optional[Union[int, List[int]]] = None, head_dim: Optional[int] = None, kv_cache_type=None, - is_disagg: bool = False) -> KVCacheManager: + is_disagg: bool = False, + cold_page_codec_provider: Optional[object] = None) -> KVCacheManager: """ Returns: A KVCacheManager instance for the given model engine or model config """ + if cold_page_codec_provider is not None and not issubclass( + kv_cache_manager_cls, KVCacheManagerV2): + raise ValueError( + "Cold-page quantization requires the resolved KV cache manager " + f"to be KVCacheManagerV2; selected {kv_cache_manager_cls.__name__}") + if (estimating_kv_cache and issubclass(kv_cache_manager_cls, KVCacheManagerV2) and kv_cache_config.pool_ratio is None @@ -2348,6 +2379,8 @@ def _create_kv_cache_manager( manager_extra_kwargs = {} if issubclass(kv_cache_manager_cls, KVCacheManagerV2): manager_extra_kwargs["enable_stats"] = enable_kv_cache_stats + manager_extra_kwargs[ + "cold_page_codec_provider"] = cold_page_codec_provider if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): manager_extra_kwargs["is_disagg"] = is_disagg @@ -2752,6 +2785,21 @@ def validate_kv_cache_compression_compatibility( spec_config: Optional[SpeculativeConfig], ) -> None: """Reject unsupported KV-cache compression feature combinations.""" + if config.algorithm == "quantization_for_cold_page": + from tensorrt_llm.runtime.kv_cache_manager_v2 import _BACKEND + + if _BACKEND == "python": + raise ValueError( + "Cold-page quantization requires the C++ KVCacheManagerV2 backend" + ) + if config.quant == "nvfp4" and not is_sm_100f(): + raise RuntimeError( + "NVFP4 cold-page quantization requires an SM100-family device " + "(SM100 or SM103).") + elif config.algorithm == "triattention" and not is_sm_100f(): + raise RuntimeError( + "TriAttention requires an SM100-family device (SM100 or SM103).") + if kv_cache_config.enable_block_reuse and not config.supports_block_reuse(): raise ValueError( f"KV-cache compression algorithm {config.algorithm!r} does not " @@ -2760,44 +2808,71 @@ def validate_kv_cache_compression_compatibility( if spec_config is None: return if not config.supports_speculative_decoding(): + guidance = ("; TriAttention requires eviction_mode='union'" + if config.algorithm == "triattention" else "") raise ValueError( f"KV-cache compression algorithm {config.algorithm!r} does not " - "support speculative decoding with its current configuration; " - "TriAttention requires eviction_mode='union'") + "support speculative decoding with its current configuration" + f"{guidance}") mode = spec_config.spec_dec_mode - if not (mode.is_mtp_one_model() or mode.is_eagle3_one_model()): + if config.algorithm == "quantization_for_cold_page": + supported = mode.is_mtp_eagle_one_model() or mode.is_eagle3_one_model() + guidance = "one-model MTP-EAGLE or EAGLE3" + else: + supported = mode.is_mtp_one_model() or mode.is_eagle3_one_model() + guidance = "one-model MTP or EAGLE3" + if not supported: raise ValueError( f"KV-cache compression does not support speculative decoding " - f"mode {mode.name}; use one-model MTP or EAGLE3") + f"mode {mode.name}; use {guidance}") def create_kv_cache_compression_manager( - config: KvCacheCompressionConfig, - kv_cache_manager: KVCacheManagerV2, - draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, - pretrained_config: Optional["transformers.PretrainedConfig"] = None, + config: Optional[KvCacheCompressionConfig], + *, + model_engine: PyTorchModelEngine, + kv_cache_config: KvCacheConfig, + estimating_kv_cache: bool = False, ) -> Optional[KVCacheCompressionManager]: - """Build the KV-cache compression manager for ``config.algorithm``, or return - None if no algorithm matches. + """Validate, select, and construct the configured manager before KVCM.""" + if config is None: + return None + if model_engine.mapping.has_cp_helix(): + # TODO: Revisit after KVCC validates HELIX-sharded Page ownership and migration. + raise ValueError( + "KV-cache compression does not support HELIX context parallelism.") + + if config.algorithm == "quantization_for_cold_page": + if config.quant != "nvfp4": + raise NotImplementedError( + f"Unsupported cold-page quantization format {config.quant!r}") + if estimating_kv_cache: + return None + quant_config = model_engine.model.model_config.quant_config + if (quant_config is not None and getattr( + quant_config, "kv_cache_quant_algo", None) == QuantAlgo.NVFP4): + logger.info( + "Skipping cold-page NVFP4 quantization because the active KV " + "cache already uses NVFP4; KVCM will migrate it losslessly.") + return None + + validate_kv_cache_compression_compatibility(config, kv_cache_config, + model_engine.spec_config) + from ..kv_cache_compression.quantization_for_cold_page.nvfp4_quantization import \ + Nvfp4ColdPageQuantizationCompression + + return Nvfp4ColdPageQuantizationCompression(config) - Called from ``create_py_executor`` and registered as a resource manager, - like the KV cache manager itself. Concrete algorithms add a dispatch branch - here. Feature compatibility is checked before resource-manager construction. - """ if config.algorithm == "triattention": - if not is_sm_100f(): - raise RuntimeError( - "TriAttention requires an SM100-family device (SM100 or SM103)." - ) + validate_kv_cache_compression_compatibility(config, kv_cache_config, + model_engine.spec_config) # TriAttention imports CuTe/CUTLASS; keep normal executor startup lazy. from ..kv_cache_compression.triattention.triattention import \ TriAttentionCompressionManager return TriAttentionCompressionManager( config, - kv_cache_manager, - draft_kv_cache_manager=draft_kv_cache_manager, - pretrained_config=pretrained_config, + pretrained_config=model_engine.model.model_config.pretrained_config, ) logger.warning( @@ -3067,24 +3142,13 @@ def create_py_executor_instance( resources[ResourceManagerType.SEQ_SLOT_MANAGER] = SeqSlotManager( max_num_sequences) - # Register the compression manager (if one is configured) with the other - # managers, before building ResourceManager, so it is part of the manager - # set from the start. Reads its own config, not the sparse-attention one. - kv_cache_compression_config = getattr(llm_args, - "kv_cache_compression_config", None) - if kv_cache_compression_config is not None: - draft_kv_cache_manager = resources.get( - ResourceManagerType.DRAFT_KV_CACHE_MANAGER) - compression_manager = create_kv_cache_compression_manager( - kv_cache_compression_config, - kv_cache_manager, - draft_kv_cache_manager=draft_kv_cache_manager, - pretrained_config=model_engine.model.model_config.pretrained_config, + compression_manager = resources.get( + ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER) + if compression_manager is not None: + compression_manager.bind_kv_cache_managers( + resources[ResourceManagerType.KV_CACHE_MANAGER], + resources.get(ResourceManagerType.DRAFT_KV_CACHE_MANAGER), ) - if compression_manager is not None: - resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] = ( - compression_manager) - resource_manager = ResourceManager(resources) # KV cache manager runs last (others may depend on it), except the @@ -3097,7 +3161,8 @@ def create_py_executor_instance( if cross_kv_cache_manager is not None: resource_manager.resource_managers.move_to_end( ResourceManagerType.CROSS_KV_CACHE_MANAGER, last=True) - # Compression is the final reconciler after every native KV manager. + # Iteration-driven compression is the final reconciler after every native + # KV manager. Cold-page quantization runs only at native storage migration. if (ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER in resource_manager.resource_managers): resource_manager.resource_managers.move_to_end( @@ -3537,14 +3602,6 @@ def _adjust_torch_mem_fraction(): def validate_feature_combination(llm_args, model_engine): # Validate the flags for features' combination - compression_config = llm_args.kv_cache_compression_config - if compression_config is not None: - validate_kv_cache_compression_compatibility( - compression_config, - llm_args.kv_cache_config, - model_engine.spec_config, - ) - def init_feature_status(llm_args) -> Dict[str, bool]: assert isinstance( llm_args, TorchLlmArgs diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 3eebc9539a0a..e3bfd3c2db4c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -819,6 +819,7 @@ def __init__( enable_stats: bool = False, num_reserved_index_slots: int = 1, is_estimating_kv_cache: bool = False, + cold_page_codec_provider: Optional[object] = None, **kwargs, ) -> None: self.mapping = mapping @@ -1134,14 +1135,34 @@ def append_to_kv_heads_per_layer( isinstance(tier, HostCacheTierConfig) for tier in config.cache_tiers ) + def create_cold_page_codec(cache_config: object) -> Optional[object]: + if cold_page_codec_provider is None: + return None + return cold_page_codec_provider.create_cold_page_codec( + cache_config, + runtime_dtype=self.dtype, + pp_layers=self.pp_layers, + num_kv_heads_per_layer=self.num_kv_heads_per_layer, + head_dim_per_layer=self.head_dim_per_layer, + is_draft=self.is_draft, + ) + candidate: Optional[KVCacheManagerPy] = None if not has_host_cache_tier: - candidate = KVCacheManagerPy(config, event_manager=self.event_manager) + candidate = KVCacheManagerPy( + config, + event_manager=self.event_manager, + cold_page_codec=create_cold_page_codec(config), + ) else: init_error: Optional[Exception] = None local_init_status = _KVCacheManagerInitStatus.KEEP_HOST try: - candidate = KVCacheManagerPy(config, event_manager=self.event_manager) + candidate = KVCacheManagerPy( + config, + event_manager=self.event_manager, + cold_page_codec=create_cold_page_codec(config), + ) except Exception as error: if isinstance(error, (CuError, KVCacheOutOfMemoryError)): local_init_status = _KVCacheManagerInitStatus.USE_NO_HOST @@ -1177,7 +1198,11 @@ def append_to_kv_heads_per_layer( if not isinstance(tier, HostCacheTierConfig) ], ) - candidate = KVCacheManagerPy(config, event_manager=self.event_manager) + candidate = KVCacheManagerPy( + config, + event_manager=self.event_manager, + cold_page_codec=create_cold_page_codec(config), + ) except Exception as error: fallback_error = error.with_traceback(None) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 8c0f4baf6488..bd8016d25efd 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2773,30 +2773,28 @@ def _free_blocks(self, block_list: list): class KVCacheCompressionManager(BaseResourceManager): - """Framework-level base class for all KV-cache compression managers. - - Inherits :class:`BaseResourceManager` so PyExecutor's main loop - auto-invokes ``prepare_resources`` / ``update_resources`` / - ``free_resources`` each iteration without any PyExecutor code changes; the - base implementations below translate those callbacks into the lifecycle - hooks. - - Concrete compression methods subclass this directly. The hooks default to - no-op; subclasses override what they need. The manager never inherits from - any cache manager because this layer decides *how* the physical KV is used, - not *what* physical KV exists. Subclasses hold ``KVCacheManagerV2`` as a tool. - - A subclass compacts through the ``KVCacheManagerV2`` it holds and records - the evicted count on ``LlmRequest.py_num_compressed_tokens``; the model - engine subtracts that count when building ``num_cached_tokens_per_seq``. + """Framework base for KV-cache compression methods in PyExecutor. + + Iteration-driven methods receive ResourceManager callbacks, while + storage-bound methods provide a cold-page codec during cache construction. + Subclasses coordinate through KVCacheManagerV2 without owning its pools, + mappings, or migration lifecycle. """ - def __init__( + uses_iteration_lifecycle = True + provides_cold_page_codec = False + + def __init__(self, config: "KvCacheCompressionConfig") -> None: + self.config = config + self.kv_cache_manager: Optional["KVCacheManagerV2"] = None + self.draft_kv_cache_manager: Optional["KVCacheManagerV2"] = None + + def bind_kv_cache_managers( self, - config: "KvCacheCompressionConfig", kv_cache_manager: "KVCacheManagerV2", draft_kv_cache_manager: Optional["KVCacheManagerV2"] = None, - ): + ) -> None: + """Bind the target and optional draft KVCMs after their construction.""" from .kv_cache_manager_v2 import KVCacheManagerV2 if not isinstance(kv_cache_manager, KVCacheManagerV2): @@ -2808,16 +2806,52 @@ def __init__( self.kv_cache_manager = kv_cache_manager self.draft_kv_cache_manager = draft_kv_cache_manager kv_cache_manager.kv_compression_manages_history = ( - config.changes_physical_kv_length) + self.config.changes_physical_kv_length) if draft_kv_cache_manager is not None: - # The draft cache is compacted together with the target. draft_kv_cache_manager.kv_compression_manages_history = ( - config.changes_physical_kv_length) + self.config.changes_physical_kv_length) @property def has_independent_draft_kv_cache(self) -> bool: return self.draft_kv_cache_manager is not None + def create_cold_page_codec( + self, + cache_config: object, + *, + runtime_dtype: DataType, + pp_layers: Sequence[int], + num_kv_heads_per_layer: Sequence[int], + head_dim_per_layer: Sequence[int], + is_draft: bool = False, + ) -> Optional[object]: + """Create a native cold-page codec when the algorithm provides one.""" + return None + + def encode_cold_pages( + self, + codec_state: object, + lifecycle_index: int, + cold_base: int, + page_indices: int, + num_pages: int, + stream: int, + ) -> None: + """Encode one complete KVCM migration batch into cold storage.""" + raise NotImplementedError + + def decode_cold_pages( + self, + codec_state: object, + lifecycle_index: int, + cold_base: int, + page_indices: int, + num_pages: int, + stream: int, + ) -> None: + """Decode one complete KVCM migration batch from cold storage.""" + raise NotImplementedError + # ================================================================== # # KV-cache lifecycle hooks (5, in temporal order). # # Subclasses override what they need; all default to no-op. # diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index aff2e9736977..e4520f2aadfa 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -8,9 +8,10 @@ # yapf: disable from .llm_args import (AttentionDpConfig, AutoDecodingConfig, BatchingType, BlockReuseConfig, CacheTransceiverConfig, CalibConfig, - CapacitySchedulerPolicy, ContextChunkingPolicy, - CudaGraphConfig, DecodeCudaGraphConfig, - DeepSeekSparseAttentionConfig, + CapacitySchedulerPolicy, + ColdPageQuantizationCompressionConfig, + ContextChunkingPolicy, CudaGraphConfig, + DecodeCudaGraphConfig, DeepSeekSparseAttentionConfig, DeepSeekV4SparseAttentionConfig, DFlashDecodingConfig, DraftTargetDecodingConfig, DSparkDecodingConfig, DynamicBatchConfig, Eagle3DecodingConfig, @@ -91,6 +92,7 @@ 'MiniMaxM3SparseAttentionConfig', 'SchedulingParams', 'SkipSoftmaxAttentionConfig', + 'ColdPageQuantizationCompressionConfig', 'TriAttentionKvCacheCompressionConfig', 'PrometheusMetricsConfig', 'PrefillCudaGraphBackend', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 1ffb613eb660..b81e3da9d97e 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3785,9 +3785,10 @@ class KvCacheCompressionConfig(StrictBaseModel): algorithm (e.g. periodic token eviction) alongside KVCacheManagerV2. Kept separate from SparseAttentionConfig by design -- compression changes - which KV is stored, not the attention computation. The manager is registered - as a resource manager in create_py_executor (_util.py), like the KV cache - manager itself. Concrete algorithms subclass this and add their parameters. + which KV is stored, not the attention computation. Iteration-driven methods + use the resource-manager cycle; storage-bound managers provide a native + codec that KVCacheManagerV2 retains and invokes. Concrete algorithms + subclass this and add their parameters. """ changes_physical_kv_length: ClassVar[bool] = False @@ -3806,6 +3807,37 @@ def supports_speculative_decoding(self) -> bool: return False +_KV_CACHE_COMPRESSION_ALGORITHM_TELEMETRY = TelemetryField.categorical( + "quantization_for_cold_page", "triattention") + + +class ColdPageQuantizationCompressionConfig(KvCacheCompressionConfig): + """Quantize Host and Disk KV pages without changing the active GPU cache.""" + + algorithm: Literal["quantization_for_cold_page"] = Field( + default="quantization_for_cold_page", + telemetry=False, + ) + quant: Literal["nvfp4"] = Field( + default="nvfp4", + description="Quantization format stored in the compressed cache tier.") + scale_checkpoint_path: Optional[str] = Field( + default=None, + min_length=1, + telemetry=False, + description= + "Optional local ModelOpt NVFP4 checkpoint directory supplying per-layer " + "K/V global scales. Omit it to use identity global scales.") + + def supports_block_reuse(self) -> bool: + """Block reuse is unchanged because token identity is preserved.""" + return True + + def supports_speculative_decoding(self) -> bool: + """Target and draft KVCMs encode their own cold pages independently.""" + return True + + class TriAttentionKvCacheCompressionConfig(KvCacheCompressionConfig): """TriAttention KV-cache compression: periodic decode-time eviction. @@ -3816,7 +3848,10 @@ class TriAttentionKvCacheCompressionConfig(KvCacheCompressionConfig): changes_physical_kv_length: ClassVar[bool] = True - algorithm: Literal["triattention"] = "triattention" + algorithm: Literal["triattention"] = Field( + default="triattention", + telemetry=_KV_CACHE_COMPRESSION_ALGORITHM_TELEMETRY, + ) eviction_mode: Literal["union", "per_head", "per_layer_perhead"] = Field( default="union", description= @@ -3858,7 +3893,8 @@ def supports_speculative_decoding(self) -> bool: KvCacheCompressionConfigType: TypeAlias = Annotated[ - Union[TriAttentionKvCacheCompressionConfig], + Union[ColdPageQuantizationCompressionConfig, + TriAttentionKvCacheCompressionConfig], Field(discriminator="algorithm"), ] diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py index c95649c19e17..4678caac9792 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py @@ -259,7 +259,10 @@ def __init__( self, config: KVCacheManagerConfig, event_manager: "KVCacheEventManager | None" = None, + cold_page_codec: object | None = None, ) -> None: + if cold_page_codec is not None: + raise NotImplementedError("Cold-page codecs require the C++ KVCacheManagerV2 backend") init_cuda_once() config = deepcopy(config) self._init_config = config diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 27ee0153bf3b..0382a52be805 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -615,10 +615,11 @@ }, { "allowed_values": [ + "quantization_for_cold_page", "triattention" ], "annotation": "Literal['triattention']", - "converter": "", + "converter": "allowlist", "kind": "categorical", "path": "kv_cache_compression_config.algorithm" }, @@ -654,6 +655,15 @@ "kind": "value", "path": "kv_cache_compression_config.normalize_scores" }, + { + "allowed_values": [ + "nvfp4" + ], + "annotation": "Literal['nvfp4']", + "converter": "", + "kind": "categorical", + "path": "kv_cache_compression_config.quant" + }, { "allowed_values": [], "annotation": "", diff --git a/tests/integration/defs/cpp/conftest.py b/tests/integration/defs/cpp/conftest.py index 75d79fc463c5..a48b71e10a61 100644 --- a/tests/integration/defs/cpp/conftest.py +++ b/tests/integration/defs/cpp/conftest.py @@ -173,6 +173,47 @@ def build_google_tests(request, build_type): ) +@pytest.fixture(scope="session") +def build_kv_cache_compression_tests(request, build_type): + """Build only the standalone NVFP4 cold-page kernel gtest. + + The binary uses NO_TLLM_LINKAGE, so this skips the full-library + build that build_google_tests pays for. + """ + cuda_arch = f"{request.param}-real" + + _logger.info(f"Using CUDA arch: {cuda_arch}") + + build_trt_llm( + build_type=build_type, + cuda_architectures=cuda_arch, + job_count=12, + use_ccache=True, + generator="Ninja", + nixl_root="/opt/nvidia/nvda_nixl", + skip_building_wheel=True, + configure_only=True, + ) + + build_dir = _cpp.find_build_dir(build_type) + _cpp.run_command( + [ + "cmake", + "--build", + str(build_dir), + "--config", + build_type, + "--parallel", + "12", + "--target", + "nvfp4ColdPageKernelsTest", + ], + cwd=build_dir, + env=_os.environ, + timeout=1800, + ) + + @pytest.fixture(scope="function", autouse=True) def keep_log_files(build_dir): """Backup previous cpp test results when run multiple ctest invocations.""" diff --git a/tests/integration/defs/cpp/test_unit_tests.py b/tests/integration/defs/cpp/test_unit_tests.py index 730ebbf389ed..9b3f8ead0bf0 100644 --- a/tests/integration/defs/cpp/test_unit_tests.py +++ b/tests/integration/defs/cpp/test_unit_tests.py @@ -35,3 +35,23 @@ def test_unit_tests(build_google_tests, test_group, build_dir, lora_setup): env=cpp_env, timeout=2700, parallel=parallel) + + +@pytest.mark.parametrize("build_kv_cache_compression_tests", ["80", "100"], + indirect=True) +def test_kv_cache_compression_unit_tests(build_kv_cache_compression_tests, + build_dir): + + xml_name = "results-unit-tests-kv_cache_compression.xml" + + # Run the binary directly: the lightweight fixture builds only this gtest, + # so a ctest directory scan would trip over unbuilt neighbors. + _cpp.run_command( + [ + f"{build_dir}/tests/unit_tests/kernels/nvfp4ColdPageKernelsTest", + f"--gtest_output=xml:{build_dir}/{xml_name}", + ], + cwd=build_dir, + env={**_os.environ}, + timeout=2700, + ) diff --git a/tests/integration/test_lists/test-db/l0_a30.yml b/tests/integration/test_lists/test-db/l0_a30.yml index 57a27879065a..3114a0a81a59 100644 --- a/tests/integration/test_lists/test-db/l0_a30.yml +++ b/tests/integration/test_lists/test-db/l0_a30.yml @@ -40,6 +40,7 @@ l0_a30: tests: # ------------- CPP tests --------------- - cpp/test_unit_tests.py::test_unit_tests[batch_manager-80] + - cpp/test_unit_tests.py::test_kv_cache_compression_unit_tests[80] - condition: ranges: system_gpu_count: diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 4a217b4294c5..26b1510a5b76 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -1,5 +1,21 @@ version: 0.0.1 l0_b200: +- condition: + ranges: + system_gpu_count: + gte: 1 + lte: 1 + wildcards: + gpu: + - '*b100*' + - '*b200*' + linux_distribution_name: ubuntu* + terms: + stage: pre_merge + backend: cpp + tests: + # ------------- CPP tests --------------- + - cpp/test_unit_tests.py::test_kv_cache_compression_unit_tests[100] - condition: ranges: system_gpu_count: diff --git a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/test_kv_cache_budget_split.py index 477e459b3e2c..f9675e56e375 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/executor/test_kv_cache_budget_split.py @@ -64,6 +64,7 @@ def _make_creator( c._speculative_config = None c._mapping = Mock() c._model_engine = Mock() + c._llm_args = SimpleNamespace(kv_cache_compression_config=None) c._kv_cache_manager_cls = Mock() c._kv_cache_manager_cls.get_cache_size_per_token = Mock( diff --git a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py index 4f628d6fbb94..644fe34c0508 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -1,26 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Unit tests for the KV-cache compression manager framework -(``KVCacheCompressionManager`` in ``resource_manager.py``) — the -``BaseResourceManager``-based single-manager design. - -Covers: -- :class:`KVCacheCompressionManager` contract: the four lifecycle hooks - default to no-op, zero resource counts, and it inherits - :class:`BaseResourceManager` (so PyExecutor auto-drives it once registered). -- The resource-manager API -> lifecycle-hook translation, gated on PyExecutor's - own signals: ``prepare_resources`` fires ``on_request_init`` on each - request's first prefill chunk (``is_first_context_chunk``); - ``update_resources`` fires ``on_context_step_end`` once with the - ``context_requests_last_chunk`` list + one ``on_generation_step_end`` per - iteration; ``free_resources`` fires ``on_request_finish``. -- :func:`create_kv_cache_compression_manager` factory. - -The base class lives in ``resource_manager.py`` (it is a resource manager, not a -sparse-attention backend); the ``create_kv_cache_compression_manager`` factory -lives in ``_util.py`` next to ``_create_kv_cache_manager``. -""" +"""Tests for the KV-cache compression manager lifecycle and factory.""" from types import SimpleNamespace from typing import ClassVar @@ -36,7 +17,12 @@ ResourceManager, ResourceManagerType, ) -from tensorrt_llm.llmapi.llm_args import KvCacheCompressionConfig +from tensorrt_llm.llmapi.llm_args import ( + ColdPageQuantizationCompressionConfig, + KvCacheCompressionConfig, + TriAttentionKvCacheCompressionConfig, +) +from tensorrt_llm.quantization import QuantAlgo # ---------------------------------------------------------------------- # # Mock infra: in-memory managers / requests (avoid touching V2 / model). # @@ -48,7 +34,8 @@ class _RecordingMixin: translation without real algorithm side-effects.""" def __init__(self, kv_cache_manager, record_list, name="m"): - super().__init__(_compression_config(), kv_cache_manager) + super().__init__(_compression_config()) + self.bind_kv_cache_managers(kv_cache_manager) self._record_list = record_list self._name = name @@ -57,7 +44,7 @@ def _record(self, hook_name: str): class _MockCompressionManager(_RecordingMixin, KVCacheCompressionManager): - """Mock manager that records the four lifecycle hooks.""" + """Mock manager that records iteration lifecycle hooks.""" def on_request_init(self, request): self._record("on_request_init") @@ -85,6 +72,25 @@ def _compression_config() -> KvCacheCompressionConfig: return KvCacheCompressionConfig(algorithm="test") +def _factory_model_engine( + *, + pretrained_config: object | None = None, + quant_config: object | None = None, + spec_config: object | None = None, + helix: bool = False, +) -> SimpleNamespace: + return SimpleNamespace( + mapping=SimpleNamespace(has_cp_helix=lambda: helix), + spec_config=spec_config, + model=SimpleNamespace( + model_config=SimpleNamespace( + pretrained_config=pretrained_config, + quant_config=quant_config, + ) + ), + ) + + def _v2_manager(*, is_draft: bool): from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 @@ -126,8 +132,9 @@ def test_inherits_base_resource_manager(self): # So PyExecutor's main loop auto-invokes prepare/update/free_resources. assert issubclass(KVCacheCompressionManager, BaseResourceManager) - def test_four_hooks_default_noop(self, fake_kv_cache_manager): - m = KVCacheCompressionManager(_compression_config(), fake_kv_cache_manager) + def test_lifecycle_hooks_default_noop(self, fake_kv_cache_manager): + m = KVCacheCompressionManager(_compression_config()) + m.bind_kv_cache_managers(fake_kv_cache_manager) assert m.on_request_init(MagicMock()) is None assert m.on_context_step_end([MagicMock()]) is None assert m.on_generation_step_begin(MagicMock()) is None @@ -137,12 +144,14 @@ def test_four_hooks_default_noop(self, fake_kv_cache_manager): def test_hooks_accept_extra_kwargs(self, fake_kv_cache_manager): # **kwargs lets the framework pass new args later without breaking # existing overrides. - m = KVCacheCompressionManager(_compression_config(), fake_kv_cache_manager) + m = KVCacheCompressionManager(_compression_config()) + m.bind_kv_cache_managers(fake_kv_cache_manager) assert m.on_request_init(MagicMock(), future_arg=1) is None assert m.on_generation_step_end(MagicMock(), future_arg=1) is None def test_resource_counts_are_zero(self, fake_kv_cache_manager): - m = KVCacheCompressionManager(_compression_config(), fake_kv_cache_manager) + m = KVCacheCompressionManager(_compression_config()) + m.bind_kv_cache_managers(fake_kv_cache_manager) # The manager owns no physical resources (the V2 cache manager does), # so it must not gate the scheduler. assert m.get_max_resource_count() == 0 @@ -155,7 +164,8 @@ def test_physical_length_change_marks_target_and_draft_v2(self): draft = _v2_manager(is_draft=True) config = _PhysicalLengthChangingConfig(algorithm="test") - manager = KVCacheCompressionManager(config, target, draft) + manager = KVCacheCompressionManager(config) + manager.bind_kv_cache_managers(target, draft) assert manager.kv_cache_manager is target assert manager.draft_kv_cache_manager is draft @@ -166,9 +176,11 @@ def test_physical_length_change_marks_target_and_draft_v2(self): def test_rejects_non_v2_ownership(self): config = _compression_config() with pytest.raises(TypeError, match="requires KVCacheManagerV2"): - KVCacheCompressionManager(config, MagicMock()) + KVCacheCompressionManager(config).bind_kv_cache_managers(MagicMock()) with pytest.raises(TypeError, match="requires KVCacheManagerV2"): - KVCacheCompressionManager(config, _v2_manager(is_draft=False), MagicMock()) + KVCacheCompressionManager(config).bind_kv_cache_managers( + _v2_manager(is_draft=False), MagicMock() + ) def test_request_field_defaults_to_zero(self): """LlmRequest carries the compression count (the manager's only @@ -276,42 +288,41 @@ def test_free_fires_finish(self, fake_kv_cache_manager): class TestFactory: - def test_returns_none_when_no_algorithm_registered(self, fake_kv_cache_manager): - # Framework-only: no concrete algorithm ships, so any config -> None. + def test_returns_none_when_no_algorithm_registered(self) -> None: cfg = MagicMock() cfg.algorithm = "made_up_method" - assert create_kv_cache_compression_manager(cfg, fake_kv_cache_manager) is None + assert ( + create_kv_cache_compression_manager( + cfg, + model_engine=_factory_model_engine(), + kv_cache_config=SimpleNamespace(enable_block_reuse=False), + ) + is None + ) - def test_warns_for_unregistered_algorithm(self, fake_kv_cache_manager): + def test_warns_for_unregistered_algorithm(self) -> None: cfg = MagicMock() cfg.algorithm = "made_up_method" with patch.object(util_mod, "logger") as mock_logger: - create_kv_cache_compression_manager(cfg, fake_kv_cache_manager) - mock_logger.warning.assert_called_once() - - def test_factory_accepts_independent_draft_manager(self): - cfg = MagicMock() - cfg.algorithm = "made_up_method" - target = _v2_manager(is_draft=False) - draft = _v2_manager(is_draft=True) - - assert ( create_kv_cache_compression_manager( cfg, - target, - draft_kv_cache_manager=draft, + model_engine=_factory_model_engine(), + kv_cache_config=SimpleNamespace(enable_block_reuse=False), ) - is None - ) + mock_logger.warning.assert_called_once() - def test_triattention_requires_sm100_family(self, fake_kv_cache_manager): + def test_triattention_requires_sm100_family(self): cfg = MagicMock() cfg.algorithm = "triattention" with ( patch.object(util_mod, "is_sm_100f", return_value=False), pytest.raises(RuntimeError, match="SM100-family"), ): - create_kv_cache_compression_manager(cfg, fake_kv_cache_manager) + util_mod.validate_kv_cache_compression_compatibility( + cfg, + SimpleNamespace(enable_block_reuse=False), + None, + ) def test_capabilities_default_false(self): config = KvCacheCompressionConfig(algorithm="offload") @@ -319,7 +330,8 @@ def test_capabilities_default_false(self): assert config.changes_physical_kv_length is False assert config.supports_block_reuse() is False assert config.supports_speculative_decoding() is False - m = KVCacheCompressionManager(config, target) + m = KVCacheCompressionManager(config) + m.bind_kv_cache_managers(target) assert target.kv_compression_manages_history is False assert not hasattr(m, "spec_config") @@ -343,6 +355,185 @@ def test_spec_gate_uses_config_capability(self): ) +@pytest.mark.cpu_only +class TestKvCacheCreatorLifecycle: + def test_estimation_still_creates_triattention_manager(self) -> None: + config = SimpleNamespace(algorithm="triattention") + pretrained_config = object() + expected_manager = MagicMock( + provides_cold_page_codec=False, + uses_iteration_lifecycle=True, + ) + creator = object.__new__(util_mod.KvCacheCreator) + creator._skip_est = False + creator._max_seq_len = 1024 + creator._kv_cache_config = SimpleNamespace(host_cache_size=None, disk_cache_size=None) + creator._llm_args = SimpleNamespace(kv_cache_compression_config=config) + creator._model_engine = _factory_model_engine(pretrained_config=pretrained_config) + creator._draft_model_engine = None + creator._is_encoder_decoder = MagicMock(return_value=False) + creator._should_create_separate_draft_kv_cache = MagicMock(return_value=False) + target_manager = object() + creator._create_kv_cache_manager = MagicMock(return_value=target_manager) + + with patch.object( + util_mod, + "create_kv_cache_compression_manager", + return_value=expected_manager, + ) as factory: + resources = {} + creator.build_managers(resources, estimating_kv_cache=True) + + factory.assert_called_once_with( + config, + model_engine=creator._model_engine, + kv_cache_config=creator._kv_cache_config, + estimating_kv_cache=True, + ) + assert resources[ResourceManagerType.KV_CACHE_MANAGER] is target_manager + assert resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] is expected_manager + expected_manager.bind_kv_cache_managers.assert_not_called() + + def test_teardown_pops_and_shuts_down_compression_manager(self) -> None: + creator = object.__new__(util_mod.KvCacheCreator) + compression_manager = MagicMock() + resources = { + ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER: compression_manager, + ResourceManagerType.KV_CACHE_MANAGER: MagicMock(), + ResourceManagerType.DRAFT_KV_CACHE_MANAGER: None, + } + + creator.teardown_managers(resources) + + compression_manager.shutdown.assert_called_once_with() + assert ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER not in resources + + +@pytest.mark.cpu_only +def test_executor_binds_iteration_manager_after_extra_resource_registration() -> None: + class _StopAfterBind(Exception): + pass + + target_manager = object() + draft_manager = object() + compression_manager = MagicMock() + resources = { + ResourceManagerType.KV_CACHE_MANAGER: target_manager, + ResourceManagerType.DRAFT_KV_CACHE_MANAGER: draft_manager, + } + llm_args = SimpleNamespace( + enable_low_latency_host_dispatch=False, + extra_resource_managers={ + ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER: compression_manager, + }, + ) + model_engine = SimpleNamespace(spec_config=None) + + with ( + patch.object(util_mod, "set_low_latency_dispatch"), + patch.object(util_mod, "ResourceManager", side_effect=_StopAfterBind), + pytest.raises(_StopAfterBind), + ): + util_mod.create_py_executor_instance( + dist=None, + resources=resources, + mapping=SimpleNamespace(), + llm_args=llm_args, + ctx_chunk_config=None, + model_engine=model_engine, + start_worker=False, + sampler=None, + drafter=None, + max_num_sequences=1, + ) + + compression_manager.bind_kv_cache_managers.assert_called_once_with( + target_manager, draft_manager + ) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "provides_cold_page_codec", + (True, False), + ids=("cold-codec", "iteration-manager"), +) +def test_build_routes_compression_manager_by_capabilities( + provides_cold_page_codec: bool, +) -> None: + creator = object.__new__(util_mod.KvCacheCreator) + creator._skip_est = False + creator._max_seq_len = 1024 + creator._kv_cache_config = SimpleNamespace() + compression_config = SimpleNamespace(algorithm="triattention") + pretrained_config = object() + creator._llm_args = SimpleNamespace(kv_cache_compression_config=compression_config) + creator._model_engine = _factory_model_engine(pretrained_config=pretrained_config) + creator._draft_model_engine = None + creator._kv_connector_manager = None + creator._is_kv_cache_manager_v2 = True + creator._fp8_ctx_mla_kv_len_cap = None + creator._is_encoder_decoder = MagicMock(return_value=False) + creator._should_create_separate_draft_kv_cache = MagicMock(return_value=True) + creator._needs_gpu_kv_cache_budget_split = MagicMock(return_value=False) + + build_order = [] + compression_manager = SimpleNamespace( + provides_cold_page_codec=provides_cold_page_codec, + uses_iteration_lifecycle=not provides_cold_page_codec, + bind_kv_cache_managers=MagicMock(side_effect=lambda *_args: build_order.append("bind")), + ) + target_config = object() + draft_config = object() + target_manager = SimpleNamespace() + draft_manager = object() + creator._split_kv_cache_budget_for_draft = MagicMock( + side_effect=[(target_config, draft_config), (target_config, draft_config)] + ) + creator._create_kv_cache_manager = MagicMock( + side_effect=lambda *_args, **_kwargs: build_order.append("target") or target_manager + ) + creator._create_one_model_draft_kv_cache_manager = MagicMock( + side_effect=lambda *_args, **_kwargs: build_order.append("draft") or draft_manager + ) + + resources = {} + with patch.object( + util_mod, + "create_kv_cache_compression_manager", + side_effect=lambda *_args, **_kwargs: build_order.append("factory") or compression_manager, + ) as factory: + creator.build_managers(resources) + + expected_codec_provider = compression_manager if provides_cold_page_codec else None + factory.assert_called_once_with( + compression_config, + model_engine=creator._model_engine, + kv_cache_config=target_config, + estimating_kv_cache=False, + ) + assert ( + creator._create_kv_cache_manager.call_args.kwargs["cold_page_codec_provider"] + is expected_codec_provider + ) + assert ( + creator._create_one_model_draft_kv_cache_manager.call_args.kwargs[ + "cold_page_codec_provider" + ] + is expected_codec_provider + ) + assert resources[ResourceManagerType.KV_CACHE_MANAGER] is target_manager + assert resources[ResourceManagerType.DRAFT_KV_CACHE_MANAGER] is draft_manager + if provides_cold_page_codec: + assert ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER not in resources + compression_manager.bind_kv_cache_managers.assert_not_called() + assert build_order == ["factory", "target", "draft"] + else: + assert resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] is compression_manager + compression_manager.bind_kv_cache_managers.assert_not_called() + assert build_order == ["factory", "target", "draft"] + + # ---------------------------------------------------------------------- # # 4. Canonical names live in resource_manager, not in the sparse module # # ---------------------------------------------------------------------- # @@ -372,6 +563,48 @@ def test_names_not_in_sparse_module(self): class TestCompressionCompatibility: + @pytest.mark.cpu_only + @pytest.mark.parametrize( + "config", + ( + ColdPageQuantizationCompressionConfig(), + TriAttentionKvCacheCompressionConfig(calibration_path="triattention-calibration.pt"), + ), + ) + def test_helix_is_rejected( + self, + config: KvCacheCompressionConfig, + ) -> None: + model_engine = SimpleNamespace( + mapping=SimpleNamespace(has_cp_helix=lambda: True), + spec_config=None, + model=SimpleNamespace(model_config=SimpleNamespace(quant_config=None)), + ) + with pytest.raises(ValueError, match="HELIX"): + create_kv_cache_compression_manager( + config, + model_engine=model_engine, + kv_cache_config=SimpleNamespace(enable_block_reuse=False), + ) + + @pytest.mark.cpu_only + def test_helix_is_rejected_before_redundant_cold_quantization(self) -> None: + model_engine = SimpleNamespace( + mapping=SimpleNamespace(has_cp_helix=lambda: True), + spec_config=None, + model=SimpleNamespace( + model_config=SimpleNamespace( + quant_config=SimpleNamespace(kv_cache_quant_algo=QuantAlgo.NVFP4) + ) + ), + ) + with pytest.raises(ValueError, match="HELIX"): + create_kv_cache_compression_manager( + ColdPageQuantizationCompressionConfig(), + model_engine=model_engine, + kv_cache_config=SimpleNamespace(enable_block_reuse=False), + ) + def test_raises_when_reuse_on(self): config = _compression_config() with pytest.raises(ValueError, match="block reuse"): diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/test_kv_cache_estimation.py index 038f602fd163..4199229af91c 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/test_kv_cache_estimation.py @@ -1026,8 +1026,13 @@ def test_separate_one_model_draft_normalizes_target_pool_ratio() -> None: return_value=Mock(), ) as create_manager, ): - creator._create_one_model_draft_kv_cache_manager(creator._max_seq_len) + codec_provider = object() + creator._create_one_model_draft_kv_cache_manager( + creator._max_seq_len, + cold_page_codec_provider=codec_provider, + ) draft_config = create_manager.call_args.kwargs["kv_cache_config"] assert draft_config.pool_ratio == [1.0] + assert create_manager.call_args.kwargs["cold_page_codec_provider"] is codec_provider assert creator._kv_cache_config.pool_ratio == target_pool_ratio diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py index a12893b2c52a..d2636a492161 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py @@ -117,6 +117,8 @@ def _make_manager_for_cache_tier_test( impl_side_effect: list[object], *, add_secondary_gpu_tier: bool = False, + cold_page_codec_provider: object | None = None, + is_draft: bool = False, mapping: Mapping | None = None, ) -> tuple[KVCacheManagerV2, Mock]: impl_constructor = Mock(side_effect=impl_side_effect) @@ -180,7 +182,9 @@ def build_cache_config( mapping=mapping, dtype=DataType.HALF, vocab_size=16, + is_draft=is_draft, execution_stream=Mock(), + cold_page_codec_provider=cold_page_codec_provider, ) return manager, impl_constructor @@ -505,6 +509,50 @@ def test_host_init_fallback_drops_only_host_tier(tmp_path) -> None: ] +@pytest.mark.cpu_only +def test_host_init_fallback_recreates_cold_codec_and_keeps_disk(tmp_path) -> None: + impl = Mock() + codecs = [object(), object()] + codec_provider = Mock() + codec_provider.create_cold_page_codec.side_effect = codecs + manager, impl_constructor = _make_manager_for_cache_tier_test( + KvCacheConfig( + max_gpu_total_bytes=16 << 20, + host_cache_size=16 << 20, + disk_cache_size=16 << 20, + disk_cache_path=str(tmp_path), + ), + [_CacheTierInitError("host tier init failed"), impl], + cold_page_codec_provider=codec_provider, + ) + + assert manager.can_evict + assert codec_provider.create_cold_page_codec.call_count == 2 + assert impl_constructor.call_count == 2 + assert impl_constructor.call_args_list[0].kwargs["cold_page_codec"] is codecs[0] + assert impl_constructor.call_args_list[1].kwargs["cold_page_codec"] is codecs[1] + fallback_tiers = impl_constructor.call_args_list[1].args[0].cache_tiers + assert [type(tier) for tier in fallback_tiers] == [ + GpuCacheTierConfig, + DiskCacheTierConfig, + ] + + +@pytest.mark.cpu_only +def test_cold_codec_provider_receives_draft_role() -> None: + impl = Mock() + codec_provider = Mock() + codec_provider.create_cold_page_codec.return_value = object() + _make_manager_for_cache_tier_test( + KvCacheConfig(max_gpu_total_bytes=16 << 20), + [impl], + cold_page_codec_provider=codec_provider, + is_draft=True, + ) + + assert codec_provider.create_cold_page_codec.call_args.kwargs["is_draft"] is True + + def test_extra_tokens_are_in_context_capacity() -> None: config = _make_cache_config_for_test( KvCacheConfig(avg_seq_len=264), diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 1ab526d83d40..a823dfd02386 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -348,11 +348,12 @@ def make_triattention(**overrides): ) with mock.patch.object(TriAttentionCompressionManager, "_initialize_eviction_state"): - return TriAttentionCompressionManager( + manager = TriAttentionCompressionManager( make_tri_config(**overrides), - make_fake_v2(), pretrained_config=make_test_pretrained_config(), ) + manager.bind_kv_cache_managers(make_fake_v2()) + return manager def make_eviction_request( diff --git a/tests/unittest/_torch/kv_cache_compression/test_quantization_for_cold_page.py b/tests/unittest/_torch/kv_cache_compression/test_quantization_for_cold_page.py new file mode 100644 index 000000000000..3d965a438241 --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_quantization_for_cold_page.py @@ -0,0 +1,918 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Control-plane tests for NVFP4 cold-page compression.""" + +import json +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch +from safetensors.torch import save_file + +from tensorrt_llm._torch.kv_cache_compression.quantization_for_cold_page.nvfp4_quantization import ( + Nvfp4ColdPageQuantizationCompression, + _load_modelopt_nvfp4_scales, +) +from tensorrt_llm._torch.kv_cache_compression.quantization_for_cold_page.quantization_for_cold_page import ( + ColdPageQuantizationCompression, +) +from tensorrt_llm._torch.pyexecutor import _util as util_mod +from tensorrt_llm._torch.pyexecutor.resource_manager import DataType +from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode +from tensorrt_llm._torch.speculative.utils import update_spec_config_from_model_config +from tensorrt_llm.llmapi.llm_args import ColdPageQuantizationCompressionConfig, MTPDecodingConfig +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization import QuantAlgo +from tensorrt_llm.runtime import kv_cache_manager_v2 as runtime_v2_mod +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, + BufferConfig, + SsmLayerConfig, +) + +pytestmark = pytest.mark.cpu_only + + +def _manager(scale_checkpoint_path=None): + config = ColdPageQuantizationCompressionConfig( + scale_checkpoint_path=( + str(scale_checkpoint_path) if scale_checkpoint_path is not None else None + ) + ) + return Nvfp4ColdPageQuantizationCompression(config) + + +def _factory_model_engine( + *, active_kv_quant: object | None = None, helix: bool = False +) -> SimpleNamespace: + return SimpleNamespace( + mapping=SimpleNamespace(has_cp_helix=lambda: helix), + spec_config=None, + model=SimpleNamespace( + model_config=SimpleNamespace( + quant_config=active_kv_quant, + pretrained_config=object(), + ) + ), + ) + + +def _cache_config(*layers): + configs = [] + for layer_id, kind in layers: + layer_type = SsmLayerConfig if kind == "ssm" else AttentionLayerConfig + roles = ("ssm_state", "conv_state") if kind == "ssm" else ("key", "value") + configs.append( + layer_type( + layer_id=layer_id, + buffers=[BufferConfig(role=role, size=128) for role in roles], + ) + ) + return SimpleNamespace(tokens_per_block=64, layers=tuple(configs)) + + +def _native() -> tuple[SimpleNamespace, MagicMock]: + codec = MagicMock() + module = SimpleNamespace( + ColdPageLifecycleProperties=lambda: SimpleNamespace(), + ColdPageIndexLocation=SimpleNamespace(HOST="host"), + create_python_cold_page_codec=MagicMock(return_value=codec), + nvfp4_cold_page_encode=MagicMock(), + nvfp4_cold_page_decode=MagicMock(), + ) + return module, codec + + +def _provider(native: SimpleNamespace) -> object: + return native.create_python_cold_page_codec.call_args.args[0] + + +def _codec_state(native: SimpleNamespace) -> object: + return native.create_python_cold_page_codec.call_args.args[1] + + +def _layouts(native: SimpleNamespace) -> list[object]: + return list(_codec_state(native).layer_layouts.values()) + + +def _configure_lifecycle(native: SimpleNamespace, layer_bytes: dict[int, dict[str, int]]) -> object: + address = 0x10000 + layers = {} + for layer_id, roles in layer_bytes.items(): + hot = {} + for role, raw_bytes in roles.items(): + hot[role] = SimpleNamespace( + raw_base=address, + raw_slot_bytes=(raw_bytes + 15) // 16 * 16, + raw_bytes=raw_bytes, + ) + address += 0x10000 + layers[layer_id] = hot + provider = _provider(native) + codec_state = _codec_state(native) + provider.configure(codec_state, [SimpleNamespace(layers=layers)]) + return codec_state.lifecycle_metadata[0] + + +def _configure_default_lifecycle(native: SimpleNamespace, raw_bytes: int) -> object: + return _configure_lifecycle(native, {0: {"key": raw_bytes, "value": raw_bytes}}) + + +def _write_quant_metadata(directory, algorithm="NVFP4"): + metadata = { + "producer": {"name": "modelopt"}, + "quantization": {"kv_cache_quant_algo": algorithm}, + } + (directory / "hf_quant_config.json").write_text(json.dumps(metadata)) + + +def _write_scales(directory, scales_by_layer, *, filename="model.safetensors", prefix="model"): + _write_quant_metadata(directory) + tensors = {} + for layer_id, (k_scale, v_scale) in scales_by_layer.items(): + base = f"{prefix}.layers.{layer_id}.self_attn" + tensors[f"{base}.k_proj.k_scale"] = torch.as_tensor(k_scale, dtype=torch.float32) + tensors[f"{base}.v_proj.v_scale"] = torch.as_tensor(v_scale, dtype=torch.float32) + save_file(tensors, str(directory / filename)) + + +def _validate_compression(mode: object | None = None) -> None: + spec_config = None if mode is None else SimpleNamespace(spec_dec_mode=mode) + with patch.object(util_mod, "is_sm_100f", return_value=True): + util_mod.validate_kv_cache_compression_compatibility( + ColdPageQuantizationCompressionConfig(), + SimpleNamespace(enable_block_reuse=False), + spec_config, + ) + + +def test_optional_modelopt_scales_map_pp_layers_and_default_missing_layers(tmp_path): + native, codec = _native() + _write_scales( + tmp_path, + {10: (0.5, 0.25)}, + filename="model-00001-of-00002.safetensors", + ) + _write_scales( + tmp_path, + {4: (0.125, 0.0625), 2: (0.75, 0.5)}, + filename="model-00002-of-00002.safetensors", + prefix="model.language_model", + ) + + with patch("tensorrt_llm.bindings.internal.kv_cache_compression", new=native): + result = _manager(tmp_path).create_cold_page_codec( + _cache_config((0, "attention"), (1, "attention"), (2, "attention")), + runtime_dtype=DataType.BF16, + pp_layers=(10, 4, 32), + num_kv_heads_per_layer=(8, 8, 8), + head_dim_per_layer=(128, 128, 128), + ) + + assert result is codec + layouts = _layouts(native) + assert [layout.layer_id for layout in layouts] == [0, 1, 2] + assert _codec_state(native).runtime_type == 1 + assert [ + ( + tuple(buffer.scales.nvfp4_orig_quant for buffer in layout.buffers), + tuple(buffer.scales.nvfp4_quant_orig for buffer in layout.buffers), + ) + for layout in layouts + ] == [ + ((2.0, 4.0), (0.5, 0.25)), + ((8.0, 16.0), (0.125, 0.0625)), + ((1.0, 1.0), (1.0, 1.0)), + ] + metadata = _configure_lifecycle( + native, + {layer_id: {"key": 131072, "value": 131072} for layer_id in range(3)}, + ) + assert metadata.scales[:6].tolist() == [ + [2.0, 0.5, 1.0, 1.0], + [4.0, 0.25, 1.0, 1.0], + [8.0, 0.125, 1.0, 1.0], + [16.0, 0.0625, 1.0, 1.0], + [1.0, 1.0, 1.0, 1.0], + [1.0, 1.0, 1.0, 1.0], + ] + + +def test_draft_codec_does_not_reuse_target_modelopt_scales(tmp_path) -> None: + native, _ = _native() + _write_scales(tmp_path, {10: (0.5, 0.25)}) + manager = _manager(tmp_path) + + with patch("tensorrt_llm.bindings.internal.kv_cache_compression", new=native): + manager.create_cold_page_codec( + _cache_config((0, "attention")), + runtime_dtype=DataType.BF16, + pp_layers=(10,), + num_kv_heads_per_layer=(8,), + head_dim_per_layer=(128,), + ) + target_layout = _layouts(native)[0] + manager.create_cold_page_codec( + _cache_config((0, "attention")), + runtime_dtype=DataType.BF16, + pp_layers=(10,), + num_kv_heads_per_layer=(8,), + head_dim_per_layer=(128,), + is_draft=True, + ) + + draft_layout = _layouts(native)[0] + assert [buffer.scales.nvfp4_orig_quant for buffer in target_layout.buffers] == [ + 2.0, + 4.0, + ] + assert [buffer.scales.nvfp4_orig_quant for buffer in draft_layout.buffers] == [ + 1.0, + 1.0, + ] + assert [buffer.scales.nvfp4_quant_orig for buffer in draft_layout.buffers] == [ + 1.0, + 1.0, + ] + + +def test_omitted_scale_checkpoint_uses_identity_and_keeps_kv_geometry(): + native, _ = _native() + cache_config = _cache_config((0, "attention")) + cache_config.tokens_per_block = 5 + + with patch("tensorrt_llm.bindings.internal.kv_cache_compression", new=native): + _manager().create_cold_page_codec( + cache_config, + runtime_dtype=DataType.HALF, + pp_layers=(10,), + num_kv_heads_per_layer=(4,), + head_dim_per_layer=(128,), + ) + + layout = _layouts(native)[0] + assert [buffer.role for buffer in layout.buffers] == ["key", "value"] + assert layout.num_kv_heads == 4 + assert layout.tokens_per_page == 5 + assert layout.head_dim == 128 + assert [buffer.scales.nvfp4_orig_quant for buffer in layout.buffers] == [ + 1.0, + 1.0, + ] + assert [buffer.scales.nvfp4_quant_orig for buffer in layout.buffers] == [ + 1.0, + 1.0, + ] + metadata = _configure_default_lifecycle(native, raw_bytes=5120) + assert metadata.cold_page_bytes == 2880 + assert metadata.wide[:2, 3].tolist() == [0, 1280] + assert metadata.wide[:2, 4].tolist() == [2560, 2720] + assert metadata.integers[:2, 0].tolist() == [0, 0] + + +def test_mha_layout_is_k_v_then_scales_and_layer_padding() -> None: + native, _ = _native() + cache_config = SimpleNamespace( + tokens_per_block=5, + layers=( + AttentionLayerConfig( + layer_id=0, + buffers=[ + BufferConfig(role="key", size=320), + BufferConfig(role="value", size=320), + ], + ), + ), + ) + + with patch("tensorrt_llm.bindings.internal.kv_cache_compression", new=native): + _manager().create_cold_page_codec( + cache_config, + runtime_dtype=DataType.HALF, + pp_layers=(0,), + num_kv_heads_per_layer=(1,), + head_dim_per_layer=(32,), + ) + + metadata = _configure_default_lifecycle(native, raw_bytes=320) + assert metadata.cold_page_bytes == 192 + assert metadata.wide[:2, 3].tolist() == [0, 80] + assert metadata.wide[:2, 4].tolist() == [160, 170] + assert metadata.wide[1, 5].item() == 180 + assert metadata.integers[:2, 0].tolist() == [0, 12] + + +def test_provider_creates_independent_state_per_kv_cache_manager() -> None: + native, _ = _native() + codecs = (object(), object()) + native.create_python_cold_page_codec.side_effect = codecs + provider = _manager() + + with patch("tensorrt_llm.bindings.internal.kv_cache_compression", new=native): + results = tuple( + provider.create_cold_page_codec( + _cache_config((0, "attention")), + runtime_dtype=DataType.BF16, + pp_layers=(layer_id,), + num_kv_heads_per_layer=(8,), + head_dim_per_layer=(128,), + ) + for layer_id in (0, 32) + ) + + assert results == codecs + assert native.create_python_cold_page_codec.call_count == 2 + calls = native.create_python_cold_page_codec.call_args_list + assert all(call.args[0] is provider for call in calls) + target_state, draft_state = (call.args[1] for call in calls) + assert target_state is not draft_state + assert target_state.layer_ids == draft_state.layer_ids == (0,) + + +def test_provider_forwards_a_4096_page_batch_through_one_native_call() -> None: + native, _ = _native() + encode = native.nvfp4_cold_page_encode + decode = native.nvfp4_cold_page_decode + + with patch("tensorrt_llm.bindings.internal.kv_cache_compression", new=native): + _manager().create_cold_page_codec( + _cache_config((0, "attention")), + runtime_dtype=DataType.BF16, + pp_layers=(0,), + num_kv_heads_per_layer=(1,), + head_dim_per_layer=(16,), + ) + provider = _provider(native) + codec_state = _codec_state(native) + hot = { + role: SimpleNamespace( + raw_base=0x1000 + index * 0x1000, + raw_slot_bytes=4096, + raw_bytes=2048, + ) + for index, role in enumerate(("key", "value")) + } + properties = provider.configure(codec_state, [SimpleNamespace(layers={0: hot})]) + provider.encode_cold_pages(codec_state, 0, 0x3000, 0x4000, 4096, 0x5000) + provider.decode_cold_pages(codec_state, 0, 0x3000, 0x4000, 4096, 0x5000) + + assert properties[0].cold_page_bytes == 1152 + assert properties[0].page_index_location == "host" + metadata = codec_state.lifecycle_metadata[0] + for operation in (encode, decode): + operation.assert_called_once() + arguments = operation.call_args.args + assert arguments == ( + 0x4000, + 4096, + metadata.wide.data_ptr(), + metadata.integers.data_ptr(), + metadata.scales.data_ptr(), + 2, + 128, + 1152, + 1, + 0x3000, + 0x5000, + ) + + +def test_codec_state_metadata_stays_on_cpu_with_non_cpu_default_device() -> None: + native, _ = _native() + with patch("tensorrt_llm.bindings.internal.kv_cache_compression", new=native): + _manager().create_cold_page_codec( + _cache_config((0, "attention")), + runtime_dtype=DataType.BF16, + pp_layers=(0,), + num_kv_heads_per_layer=(1,), + head_dim_per_layer=(16,), + ) + + with torch.device("meta"): + metadata = _configure_default_lifecycle(native, raw_bytes=2048) + for tensor, dtype, shape in ( + (metadata.wide, torch.int64, (256, 6)), + (metadata.integers, torch.int32, (256, 5)), + (metadata.scales, torch.float32, (256, 4)), + ): + assert tensor.device.type == "cpu" + assert tensor.dtype == dtype + assert tensor.shape == shape + assert tensor.is_contiguous() + + +def test_provider_rejects_invalid_resolved_hot_buffers() -> None: + native, _ = _native() + with patch("tensorrt_llm.bindings.internal.kv_cache_compression", new=native): + _manager().create_cold_page_codec( + _cache_config((0, "attention")), + runtime_dtype=DataType.BF16, + pp_layers=(0,), + num_kv_heads_per_layer=(1,), + head_dim_per_layer=(16,), + ) + + provider = _provider(native) + codec_state = _codec_state(native) + + def hot(raw_base: int = 0x1000, raw_bytes: int = 2048) -> SimpleNamespace: + return SimpleNamespace( + raw_base=raw_base, + raw_slot_bytes=2048, + raw_bytes=raw_bytes, + ) + + with pytest.raises(ValueError, match="roles do not match"): + provider.configure( + codec_state, + [SimpleNamespace(layers={0: {"key": hot(), "value": hot(), "extra": hot()}})], + ) + with pytest.raises(ValueError, match="size does not match"): + provider.configure( + codec_state, [SimpleNamespace(layers={0: {"key": hot(raw_bytes=32), "value": hot()}})] + ) + with pytest.raises(ValueError, match="16-byte aligned"): + provider.configure( + codec_state, + [SimpleNamespace(layers={0: {"key": hot(raw_base=0x1001), "value": hot()}})], + ) + + +def test_provider_rejects_more_than_256_lifecycle_buffers() -> None: + native, _ = _native() + layers = tuple( + AttentionLayerConfig( + layer_id=layer_id, + buffers=[ + BufferConfig(role="key", size=32), + BufferConfig(role="value", size=32), + ], + ) + for layer_id in range(129) + ) + cache_config = SimpleNamespace(tokens_per_block=1, layers=layers) + with patch("tensorrt_llm.bindings.internal.kv_cache_compression", new=native): + _manager().create_cold_page_codec( + cache_config, + runtime_dtype=DataType.BF16, + pp_layers=tuple(range(129)), + num_kv_heads_per_layer=(1,) * 129, + head_dim_per_layer=(16,) * 129, + ) + + with pytest.raises(ValueError, match="maximum is 256"): + _configure_lifecycle( + native, + {layer_id: {"key": 32, "value": 32} for layer_id in range(129)}, + ) + + +def test_unsupported_quant_is_rejected_before_manager_construction() -> None: + config = SimpleNamespace( + algorithm="quantization_for_cold_page", + quant="future-format", + scale_checkpoint_path="/not/a/checkpoint", + ) + with pytest.raises(NotImplementedError, match="future-format"): + util_mod.create_kv_cache_compression_manager( + config, + model_engine=_factory_model_engine(), + kv_cache_config=SimpleNamespace(enable_block_reuse=False), + ) + + +def test_scale_loader_matches_hf_shard_and_consolidated_policy(tmp_path): + _write_scales(tmp_path, {7: (0.5, 0.25)}, filename="model.safetensors") + _write_scales( + tmp_path, + {7: (0.125, 0.0625)}, + filename="consolidated.00.safetensors", + ) + assert _load_modelopt_nvfp4_scales(str(tmp_path))[7] == ( + (2.0, 4.0), + (0.5, 0.25), + ) + + consolidated_only = tmp_path / "consolidated-only" + consolidated_only.mkdir() + _write_scales( + consolidated_only, + {9: (0.125, 0.0625)}, + filename="consolidated.00.safetensors", + ) + assert _load_modelopt_nvfp4_scales(str(consolidated_only))[9] == ( + (8.0, 16.0), + (0.125, 0.0625), + ) + + +def test_scale_loader_reduces_duplicate_shards_like_native_qkv_loader(tmp_path): + _write_scales(tmp_path, {7: (0.25, 0.125)}, filename="model-00001.safetensors") + _write_scales( + tmp_path, + {7: (0.5, 0.25)}, + filename="model-00002.safetensors", + prefix="model.language_model", + ) + assert _load_modelopt_nvfp4_scales(str(tmp_path))[7] == ( + (2.0, 4.0), + (0.5, 0.25), + ) + + +def test_scale_loader_ignores_multimodal_towers_with_the_same_layer_id(tmp_path): + _write_quant_metadata(tmp_path) + tensors = { + "model.language_model.layers.7.self_attn.k_proj.k_scale": torch.tensor(0.5), + "model.language_model.layers.7.self_attn.v_proj.v_scale": torch.tensor(0.25), + "model.vision_tower.encoder.layers.7.self_attn.k_proj.k_scale": torch.tensor(4.0), + "model.vision_tower.encoder.layers.7.self_attn.v_proj.v_scale": torch.tensor(2.0), + "model.audio_tower.layers.7.self_attn.k_proj.k_scale": torch.tensor(8.0), + "model.audio_tower.layers.7.self_attn.v_proj.v_scale": torch.tensor(4.0), + } + save_file(tensors, str(tmp_path / "model.safetensors")) + + assert _load_modelopt_nvfp4_scales(str(tmp_path))[7] == ( + (2.0, 4.0), + (0.5, 0.25), + ) + + +def test_trtllm_load_kv_scales_zero_uses_identity(tmp_path, monkeypatch): + _write_scales(tmp_path, {7: (0.5, 0.25)}) + monkeypatch.setenv("TRTLLM_LOAD_KV_SCALES", "0") + assert _load_modelopt_nvfp4_scales(str(tmp_path)) == {} + + +def test_non_nvfp4_checkpoint_scales_are_not_reused(tmp_path): + _write_scales(tmp_path, {7: (0.5, 0.25)}) + _write_quant_metadata(tmp_path, "FP8") + assert _load_modelopt_nvfp4_scales(str(tmp_path)) == {} + + +def test_unquantized_checkpoint_uses_identity_scales(tmp_path): + save_file({"model.weight": torch.ones(1)}, str(tmp_path / "model.safetensors")) + assert _load_modelopt_nvfp4_scales(str(tmp_path)) == {} + + +def test_explicit_scale_checkpoint_requires_safetensors(tmp_path): + with pytest.raises(FileNotFoundError, match="No safetensors files"): + _load_modelopt_nvfp4_scales(str(tmp_path)) + + +@pytest.mark.parametrize("present_kind", ["k", "v"]) +def test_scale_checkpoint_requires_kv_pair(tmp_path, present_kind): + _write_scales(tmp_path, {7: (0.5, 0.5)}) + base = "model.layers.7.self_attn" + name = f"{base}.{present_kind}_proj.{present_kind}_scale" + save_file({name: torch.tensor(0.5)}, str(tmp_path / "model.safetensors")) + with pytest.raises(ValueError, match="both K and V"): + _load_modelopt_nvfp4_scales(str(tmp_path)) + + +def test_scale_checkpoint_requires_float32_reciprocals(tmp_path) -> None: + smallest_subnormal = torch.tensor(1e-45, dtype=torch.float32).item() + _write_scales(tmp_path, {7: (smallest_subnormal, smallest_subnormal)}) + with pytest.raises(ValueError, match="representable as float32"): + _load_modelopt_nvfp4_scales(str(tmp_path)) + + +def test_hybrid_codec_skips_ssm_layers_and_ssm_only_rank_is_lossless(tmp_path): + native, codec = _native() + _write_scales(tmp_path, {4: (0.5, 0.25)}) + with patch("tensorrt_llm.bindings.internal.kv_cache_compression", new=native): + _manager(tmp_path).create_cold_page_codec( + _cache_config((0, "ssm"), (1, "attention")), + runtime_dtype=DataType.BF16, + pp_layers=(10, 4), + num_kv_heads_per_layer=(0, 8), + head_dim_per_layer=(128, 128), + ) + layout = _layouts(native)[0] + assert layout.layer_id == 1 + assert [buffer.scales.nvfp4_orig_quant for buffer in layout.buffers] == [ + 2.0, + 4.0, + ] + + result = _manager().create_cold_page_codec( + _cache_config((0, "ssm")), + runtime_dtype=DataType.INT8, + pp_layers=(10,), + num_kv_heads_per_layer=(0,), + head_dim_per_layer=(128,), + ) + + assert result is codec + assert _codec_state(native).layer_ids == () + + +def test_mla_key_only_layout_with_index_key_uses_identity_scales(tmp_path): + native, codec = _native() + _write_scales(tmp_path, {10: (0.5, 0.25)}) + cache_config = SimpleNamespace( + tokens_per_block=64, + layers=( + AttentionLayerConfig( + layer_id=0, + buffers=[ + BufferConfig(role="key", size=64 * 576 * 2), + BufferConfig(role="index_key", size=64 * 132), + ], + ), + ), + ) + with patch("tensorrt_llm.bindings.internal.kv_cache_compression", new=native): + result = _manager(tmp_path).create_cold_page_codec( + cache_config, + runtime_dtype=DataType.BF16, + pp_layers=(10,), + num_kv_heads_per_layer=(1,), + head_dim_per_layer=(576,), + ) + + assert result is codec + layout = _layouts(native)[0] + assert layout.layer_id == 0 + assert [buffer.role for buffer in layout.buffers] == ["key", "index_key"] + assert [buffer.scales is not None for buffer in layout.buffers] == [True, False] + assert _codec_state(native).runtime_type == 1 + assert layout.num_kv_heads == 1 + assert layout.tokens_per_page == 64 + assert layout.head_dim == 576 + scales = layout.buffers[0].scales + assert scales.nvfp4_orig_quant == scales.nvfp4_quant_orig == 1.0 + assert layout.buffers[1].scales is None + metadata = _configure_lifecycle(native, {0: {"key": 64 * 576 * 2, "index_key": 64 * 132}}) + assert metadata.cold_page_bytes == 29184 + assert metadata.wide[:2, 3].tolist() == [0, 20736] + assert metadata.wide[:2, 4].tolist() == [18432, 0] + + +def test_mla_all_non_latent_roles_are_explicit_lossless_spans() -> None: + native, _ = _native() + cache_config = SimpleNamespace( + tokens_per_block=5, + layers=( + AttentionLayerConfig( + layer_id=0, + buffers=[ + BufferConfig(role="key", size=320), + BufferConfig(role="index_key", size=68), + BufferConfig(role="rope_state", size=7), + ], + ), + ), + ) + + with patch("tensorrt_llm.bindings.internal.kv_cache_compression", new=native): + _manager().create_cold_page_codec( + cache_config, + runtime_dtype=DataType.BF16, + pp_layers=(0,), + num_kv_heads_per_layer=(1,), + head_dim_per_layer=(32,), + ) + + layout = _layouts(native)[0] + assert [buffer.role for buffer in layout.buffers] == [ + "key", + "index_key", + "rope_state", + ] + assert [buffer.scales is not None for buffer in layout.buffers] == [True, False, False] + metadata = _configure_lifecycle(native, {0: {"key": 320, "index_key": 68, "rope_state": 7}}) + assert metadata.wide[:3, 3].tolist() == [0, 90, 158] + assert metadata.wide[2, 5].item() == 165 + assert metadata.integers[2, 0].item() == 11 + assert metadata.cold_page_bytes == 176 + + +def test_lossless_layout_uses_resolved_hot_buffer_bytes() -> None: + native, _ = _native() + cache_config = SimpleNamespace( + tokens_per_block=4, + layers=( + AttentionLayerConfig( + layer_id=0, + buffers=[ + BufferConfig(role="key", size=128), + BufferConfig(role="index_key", size=3), + ], + ), + ), + ) + + with patch("tensorrt_llm.bindings.internal.kv_cache_compression", new=native): + _manager().create_cold_page_codec( + cache_config, + runtime_dtype=DataType.BF16, + pp_layers=(0,), + num_kv_heads_per_layer=(1,), + head_dim_per_layer=(16,), + ) + + metadata = _configure_lifecycle(native, {0: {"key": 128, "index_key": 6}}) + assert metadata.wide[:2, 3].tolist() == [0, 36] + assert metadata.wide[1, 5].item() == 42 + assert metadata.integers[1, 0].item() == 6 + assert metadata.cold_page_bytes == 48 + + +@pytest.mark.parametrize( + ("owns_index", "expected_layers", "expected_buffers", "expected_bytes"), + [ + ([True] * 61, 61, 122, 1_780_224), + ( + [layer < 3 or (layer >= 6 and layer % 4 == 2) for layer in range(78)], + 78, + 99, + 1_794_816, + ), + ], + ids=("deepseek-v3.2", "glm-5.2"), +) +def test_mla_model_layouts_are_built_in_python( + owns_index: list[bool], + expected_layers: int, + expected_buffers: int, + expected_bytes: int, +) -> None: + native, _ = _native() + layers = [] + for layer_id, has_index in enumerate(owns_index): + buffers = [BufferConfig(role="key", size=64 * 576 * 2)] + if has_index: + buffers.append(BufferConfig(role="index_key", size=64 * (128 + 4))) + layers.append(AttentionLayerConfig(layer_id=layer_id, buffers=buffers)) + cache_config = SimpleNamespace(tokens_per_block=64, layers=tuple(layers)) + + with patch("tensorrt_llm.bindings.internal.kv_cache_compression", new=native): + _manager().create_cold_page_codec( + cache_config, + runtime_dtype=DataType.BF16, + pp_layers=tuple(range(expected_layers)), + num_kv_heads_per_layer=(1,) * expected_layers, + head_dim_per_layer=(576,) * expected_layers, + ) + + layouts = _layouts(native) + assert len(layouts) == expected_layers + assert sum(len(layout.buffers) for layout in layouts) == expected_buffers + layer_bytes = { + layer_id: { + "key": 64 * 576 * 2, + **({"index_key": 64 * (128 + 4)} if has_index else {}), + } + for layer_id, has_index in enumerate(owns_index) + } + assert _configure_lifecycle(native, layer_bytes).cold_page_bytes == expected_bytes + + +def test_fp8_runtime_uses_modelopt_nvfp4_scales(tmp_path): + native, _ = _native() + _write_scales(tmp_path, {10: (0.5, 0.25)}) + with patch("tensorrt_llm.bindings.internal.kv_cache_compression", new=native): + _manager(tmp_path).create_cold_page_codec( + _cache_config((0, "attention")), + runtime_dtype=DataType.FP8, + pp_layers=(10,), + num_kv_heads_per_layer=(8,), + head_dim_per_layer=(128,), + ) + + layout = _layouts(native)[0] + assert _codec_state(native).runtime_type == 2 + assert [buffer.scales.nvfp4_orig_quant for buffer in layout.buffers] == [ + 2.0, + 4.0, + ] + assert [buffer.scales.nvfp4_quant_orig for buffer in layout.buffers] == [ + 0.5, + 0.25, + ] + assert all( + buffer.scales.fp8_orig_quant == buffer.scales.fp8_quant_orig == 1.0 + for buffer in layout.buffers + ) + + +def test_runtime_admission_is_checked_before_manager_creation(monkeypatch) -> None: + monkeypatch.setattr(runtime_v2_mod, "_BACKEND", "python") + with pytest.raises(ValueError, match=r"require.*C\+\+ KVCacheManagerV2"): + _validate_compression() + + monkeypatch.setattr(runtime_v2_mod, "_BACKEND", "cpp") + monkeypatch.setattr(util_mod, "is_sm_100f", lambda: False) + with pytest.raises(RuntimeError, match="requires an SM100-family device"): + util_mod.create_kv_cache_compression_manager( + ColdPageQuantizationCompressionConfig(), + model_engine=_factory_model_engine(), + kv_cache_config=SimpleNamespace(enable_block_reuse=False), + ) + + monkeypatch.setattr(util_mod, "is_sm_100f", lambda: True) + assert isinstance( + util_mod.create_kv_cache_compression_manager( + ColdPageQuantizationCompressionConfig(), + model_engine=_factory_model_engine(), + kv_cache_config=SimpleNamespace(enable_block_reuse=False), + ), + Nvfp4ColdPageQuantizationCompression, + ) + + +def test_speculative_admission_accepts_verified_one_model_modes(monkeypatch) -> None: + monkeypatch.setattr(runtime_v2_mod, "_BACKEND", "cpp") + + _validate_compression(SpeculativeDecodingMode.EAGLE3_ONE_MODEL) + _validate_compression(SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL) + for mode in ( + SpeculativeDecodingMode.MTP, + SpeculativeDecodingMode.MTP_EAGLE, + SpeculativeDecodingMode.EAGLE3, + SpeculativeDecodingMode.DFLASH, + ): + with pytest.raises(ValueError, match="one-model MTP-EAGLE or EAGLE3"): + _validate_compression(mode) + + +def test_qwen35_mtp3_resolves_to_supported_one_model_mode(monkeypatch) -> None: + monkeypatch.setattr(runtime_v2_mod, "_BACKEND", "cpp") + + spec_config = MTPDecodingConfig(max_draft_len=3) + update_spec_config_from_model_config( + spec_config, + SimpleNamespace(mtp_num_hidden_layers=1), + ) + + assert spec_config.spec_dec_mode is SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL + assert spec_config.max_draft_len == 3 + with patch.object(util_mod, "is_sm_100f", return_value=True): + util_mod.validate_kv_cache_compression_compatibility( + ColdPageQuantizationCompressionConfig(), + SimpleNamespace(enable_block_reuse=False), + spec_config, + ) + + +def test_cold_manager_is_disabled_for_estimation_and_active_nvfp4(monkeypatch) -> None: + monkeypatch.setattr(runtime_v2_mod, "_BACKEND", "cpp") + monkeypatch.setattr(util_mod, "is_sm_100f", lambda: True) + + def build( + *, + estimating: bool = False, + skip_est: bool = False, + active_kv_quant: object | None = None, + ) -> tuple[dict, util_mod.KvCacheCreator]: + creator = object.__new__(util_mod.KvCacheCreator) + creator._skip_est = skip_est + creator._max_seq_len = 1024 + creator._kv_cache_config = SimpleNamespace( + host_cache_size=None, + disk_cache_size=None, + enable_block_reuse=False, + ) + creator._llm_args = SimpleNamespace( + kv_cache_compression_config=ColdPageQuantizationCompressionConfig() + ) + creator._model_engine = _factory_model_engine(active_kv_quant=active_kv_quant) + creator._draft_model_engine = None + creator._kv_connector_manager = None + creator._fp8_ctx_mla_kv_len_cap = None + creator._is_encoder_decoder = MagicMock(return_value=False) + creator._should_create_separate_draft_kv_cache = MagicMock(return_value=False) + creator._create_kv_cache_manager = MagicMock(return_value=SimpleNamespace()) + creator.configure_kv_cache_capacity = MagicMock() + resources = {} + creator.build_managers(resources, estimating_kv_cache=estimating) + return resources, creator + + resources, creator = build() + manager = creator._create_kv_cache_manager.call_args.kwargs["cold_page_codec_provider"] + assert isinstance(manager, Nvfp4ColdPageQuantizationCompression) + assert isinstance(manager, ColdPageQuantizationCompression) + assert manager.provides_cold_page_codec + assert not manager.uses_iteration_lifecycle + assert util_mod.ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER not in resources + _, estimation_creator = build(estimating=True) + assert ( + estimation_creator._create_kv_cache_manager.call_args.kwargs["cold_page_codec_provider"] + is None + ) + _, skip_est_creator = build(estimating=True, skip_est=True) + assert isinstance( + skip_est_creator._create_kv_cache_manager.call_args.kwargs["cold_page_codec_provider"], + Nvfp4ColdPageQuantizationCompression, + ) + with patch.object(util_mod.logger, "info") as log: + active_resources, active_creator = build( + active_kv_quant=QuantConfig(kv_cache_quant_algo=QuantAlgo.NVFP4) + ) + assert util_mod.ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER not in active_resources + assert ( + active_creator._create_kv_cache_manager.call_args.kwargs["cold_page_codec_provider"] is None + ) + log.assert_called_once() diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py index 6824e5d1ce83..b67ab51b2b90 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -91,7 +91,13 @@ def test_speculative_admission_gates_raise(gate, match): eviction_mode="per_head" if gate == "union_only_per_head" else "union", ) - with pytest.raises(ValueError, match=match): + with ( + mock.patch( + "tensorrt_llm._torch.pyexecutor._util.is_sm_100f", + return_value=True, + ), + pytest.raises(ValueError, match=match), + ): validate_kv_cache_compression_compatibility( config, SimpleNamespace(enable_block_reuse=False), diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 92a8bda7ed4d..0dcc4430b04b 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -69,16 +69,24 @@ def _make_hf_config(**values): return SimpleNamespace(get_text_config=lambda: text_config) +def _factory_model_engine(pretrained_config: object) -> SimpleNamespace: + return SimpleNamespace( + mapping=SimpleNamespace(has_cp_helix=lambda: False), + spec_config=None, + model=SimpleNamespace( + model_config=SimpleNamespace( + pretrained_config=pretrained_config, + quant_config=None, + ) + ), + ) + + class TestConfigAndFactory: - def test_factory_allows_block_reuse_and_propagates_config_fields(self): + def test_factory_allows_block_reuse_and_propagates_config_fields(self) -> None: # The factory contract is independent of GPU-owned persistent buffers. fake_v2 = _make_fake_v2(enable_block_reuse=True) cfg = _make_tri_config(budget=32, beta=16, eviction_mode="per_head") - validate_kv_cache_compression_compatibility( - cfg, - SimpleNamespace(enable_block_reuse=True), - None, - ) with ( mock.patch( "tensorrt_llm._torch.pyexecutor._util.is_sm_100f", @@ -90,9 +98,10 @@ def test_factory_allows_block_reuse_and_propagates_config_fields(self): ): mgr = create_kv_cache_compression_manager( cfg, - kv_cache_manager=fake_v2, - pretrained_config=_make_test_pretrained_config(), + model_engine=_factory_model_engine(_make_test_pretrained_config()), + kv_cache_config=SimpleNamespace(enable_block_reuse=True), ) + mgr.bind_kv_cache_managers(fake_v2) assert isinstance(mgr, TriAttentionCompressionManager) assert mgr.budget == 32 assert mgr.beta == 16 @@ -286,7 +295,7 @@ def test_prepare_does_not_evict_and_update_runs_final_hook_once(self): @staticmethod def _make_due_decode_request(seq_len, *, num_extra_kv_tokens=0, kv_reserve_draft_tokens=0): # The growth and protected-tail capacity constants snapshot the - # manager at construction, so the reserve widths are set up front. + # manager at binding, so the reserve widths are set up front. request = _make_request( 7, py_prompt_len=1024, @@ -299,9 +308,9 @@ def _make_due_decode_request(seq_len, *, num_extra_kv_tokens=0, kv_reserve_draft with mock.patch.object(TriAttentionCompressionManager, "_initialize_eviction_state"): mgr = TriAttentionCompressionManager( _make_tri_config(budget=8), - fake_v2, pretrained_config=_make_test_pretrained_config(), ) + mgr.bind_kv_cache_managers(fake_v2) cache = SimpleNamespace( capacity=seq_len, history_length=1024, @@ -441,14 +450,12 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): def test_one_model_draft_co_compression_is_accepted(self, spec_mode): draft_manager = _make_fake_v2(is_draft=True) with mock.patch.object(TriAttentionCompressionManager, "_initialize_eviction_state"): - TriAttentionCompressionManager( + manager = TriAttentionCompressionManager( _make_tri_config(budget=8), - _make_fake_v2(), - draft_kv_cache_manager=draft_manager, pretrained_config=_make_test_pretrained_config(), ) + manager.bind_kv_cache_managers(_make_fake_v2(), draft_manager) - from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_compatibility from tensorrt_llm.llmapi.llm_args import Eagle3DecodingConfig, MTPDecodingConfig spec_config = ( @@ -461,11 +468,15 @@ def test_one_model_draft_co_compression_is_accepted(self, spec_mode): ) ) - validate_kv_cache_compression_compatibility( - _make_tri_config(budget=8), - SimpleNamespace(enable_block_reuse=False), - spec_config, - ) + with mock.patch( + "tensorrt_llm._torch.pyexecutor._util.is_sm_100f", + return_value=True, + ): + validate_kv_cache_compression_compatibility( + _make_tri_config(budget=8), + SimpleNamespace(enable_block_reuse=False), + spec_config, + ) class TestFixedScoreMetadata: diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index d74aa6054cd9..cfb82e4b9791 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -276,7 +276,7 @@ methods: default: null status: prototype kv_cache_compression_config: - annotation: Union[tensorrt_llm.llmapi.llm_args.TriAttentionKvCacheCompressionConfig, NoneType] + annotation: Union[tensorrt_llm.llmapi.llm_args.ColdPageQuantizationCompressionConfig, tensorrt_llm.llmapi.llm_args.TriAttentionKvCacheCompressionConfig, NoneType] default: null status: prototype otlp_traces_endpoint: diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 0487b403d11e..2bd9840c81bb 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -3461,8 +3461,46 @@ def test_no_custom_init_methods(self): @pytest.mark.cpu_only def test_kv_cache_compression_config_dispatches_by_algorithm(): - from tensorrt_llm.llmapi.llm_args import \ - TriAttentionKvCacheCompressionConfig + from tensorrt_llm.llmapi.llm_args import ( + ColdPageQuantizationCompressionConfig, + TriAttentionKvCacheCompressionConfig) + + cold_config = TorchLlmArgs( + model="/tmp/dummy_model", + kv_cache_compression_config={ + "algorithm": "quantization_for_cold_page", + "quant": "nvfp4", + }, + ).kv_cache_compression_config + + assert isinstance(cold_config, ColdPageQuantizationCompressionConfig) + assert cold_config.model_dump() == { + "algorithm": "quantization_for_cold_page", + "quant": "nvfp4", + "scale_checkpoint_path": None, + } + assert not cold_config.changes_physical_kv_length + assert cold_config.supports_block_reuse() + assert cold_config.supports_speculative_decoding() + + cold_config_with_scales = TorchLlmArgs( + model="/tmp/dummy_model", + kv_cache_compression_config={ + "algorithm": "quantization_for_cold_page", + "quant": "nvfp4", + "scale_checkpoint_path": "/tmp/nvfp4-kv-scales", + }, + ).kv_cache_compression_config + assert cold_config_with_scales.scale_checkpoint_path == "/tmp/nvfp4-kv-scales" + + with pytest.raises(ValidationError): + TorchLlmArgs( + model="/tmp/dummy_model", + kv_cache_compression_config={ + "algorithm": "quantization_for_cold_page", + "quant": "fp8", + }, + ) config_dict = yaml.safe_load(""" kv_cache_compression_config: diff --git a/tests/unittest/usage/test_llmapi_config_telemetry_docs.py b/tests/unittest/usage/test_llmapi_config_telemetry_docs.py index 059fad4224ef..e6e8f68b7c83 100644 --- a/tests/unittest/usage/test_llmapi_config_telemetry_docs.py +++ b/tests/unittest/usage/test_llmapi_config_telemetry_docs.py @@ -286,6 +286,67 @@ def test_build_capture_manifest_matches_committed_golden(): _assert_committed_manifest_current(golden_manifest()) +def test_kv_cache_compression_discriminator_captures_both_algorithms() -> None: + """The shared allowlist captures either compression discriminator.""" + from tensorrt_llm.llmapi.llm_args import ( + ColdPageQuantizationCompressionConfig, + TorchLlmArgs, + TriAttentionKvCacheCompressionConfig, + ) + from tensorrt_llm.usage.llmapi_config import ( + build_capture_manifest, + collect_llm_api_config_payloads, + ) + + entry = next( + item + for item in build_capture_manifest(TorchLlmArgs) + if item.path == "kv_cache_compression_config.algorithm" + ) + assert repr(entry.annotation) == "typing.Literal['triattention']" + assert entry.converter == "allowlist" + assert set(entry.allowed_values) == { + "quantization_for_cold_page", + "triattention", + } + cold_field = ColdPageQuantizationCompressionConfig.model_fields["algorithm"] + assert cold_field.json_schema_extra["telemetry"] == {"exclude": True} + tri_field = TriAttentionKvCacheCompressionConfig.model_fields["algorithm"] + assert tri_field.json_schema_extra["telemetry"] == { + "kind": "categorical", + "converter": "allowlist", + "allowed_values": ["quantization_for_cold_page", "triattention"], + } + + private_paths = ( + "/private/modelopt-scales", + "/private/triattention.pt", + ) + configs = ( + ColdPageQuantizationCompressionConfig( + scale_checkpoint_path=private_paths[0], + ), + TriAttentionKvCacheCompressionConfig( + calibration_path=private_paths[1], + ), + ) + for config in configs: + args = TorchLlmArgs( + model="/model", + kv_cache_compression_config=config, + ) + config_json, metadata_json = collect_llm_api_config_payloads(args) + captured = json.loads(config_json) + metadata = json.loads(metadata_json) + assert captured["kv_cache_compression_config.algorithm"] == config.algorithm + assert "kv_cache_compression_config.scale_checkpoint_path" not in captured + assert "kv_cache_compression_config.calibration_path" not in captured + for private_path in private_paths: + assert private_path not in config_json + assert private_path not in metadata_json + assert metadata["capture_succeeded"] is True + + def test_load_generator_does_not_leak_sys_modules(): """_load_generator must not leak its temporary module into sys.modules.