diff --git a/benchmarks/cpp/disaggServerBenchmark.cpp b/benchmarks/cpp/disaggServerBenchmark.cpp index 057e7898afeb..bc3a7a2659fd 100644 --- a/benchmarks/cpp/disaggServerBenchmark.cpp +++ b/benchmarks/cpp/disaggServerBenchmark.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * 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"); @@ -543,7 +543,7 @@ texec::Request makeExecutorContextRequest(Sample const& sample, SizeType32 const std::nullopt, // logitsPostProcessorName std::nullopt, // logitsPostProcessor encoderInputTokenIds.has_value() ? encoderInputTokenIds : std::nullopt, - std::nullopt); // cacheSaltID + std::nullopt); // cacheSalt request.setRequestType(tensorrt_llm::executor::RequestType::REQUEST_TYPE_CONTEXT_ONLY); return request; } diff --git a/benchmarks/cpp/gptManagerBenchmark.cpp b/benchmarks/cpp/gptManagerBenchmark.cpp index b4f0948c1155..287cbba343ce 100644 --- a/benchmarks/cpp/gptManagerBenchmark.cpp +++ b/benchmarks/cpp/gptManagerBenchmark.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * 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"); @@ -838,7 +838,7 @@ texec::Request makeExecutorRequest(Sample const& sample, SizeType32 const& beamW std::nullopt, // logitsPostProcessorName std::nullopt, // logitsPostProcessor encoderInputTokenIds.has_value() ? encoderInputTokenIds : std::nullopt, - std::nullopt); // cacheSaltID + std::nullopt); // cacheSalt } void benchmarkExecutor(std::optional const& decoderEngineDir, diff --git a/cpp/include/tensorrt_llm/batch_manager/blockKey.h b/cpp/include/tensorrt_llm/batch_manager/blockKey.h index 002b4356c869..920212845331 100644 --- a/cpp/include/tensorrt_llm/batch_manager/blockKey.h +++ b/cpp/include/tensorrt_llm/batch_manager/blockKey.h @@ -29,7 +29,6 @@ using VecTokens = std::vector; using UniqueToken = tensorrt_llm::runtime::UniqueToken; using VecUniqueTokens = tensorrt_llm::runtime::VecUniqueTokens; using LoraTaskIdType = tensorrt_llm::runtime::LoraTaskIdType; -using CacheSaltIDType = tensorrt_llm::runtime::CacheSaltIDType; using MmKey = tensorrt_llm::executor::MmKey; //! \brief Generate the multimodal extra keys for a single KV cache block. @@ -49,7 +48,8 @@ struct BlockKey // Extra keys for multimodal data (similar to VLLM's approach) // Each extra key is a pair of (mm_hash, start_offset_in_block) std::vector extraKeys; - std::optional cacheSaltID = std::nullopt; + // Cache salt string. Used as part of the block key so blocks from different salts do not match. + std::optional cacheSalt = std::nullopt; BlockKey() = default; @@ -64,12 +64,12 @@ struct BlockKey } explicit BlockKey(bool usesExtraIds, std::optional loraTaskId, VecUniqueTokens uniqueTokens, - std::vector extraKeys = {}, std::optional cacheSaltID = std::nullopt) + std::vector extraKeys = {}, std::optional cacheSalt = std::nullopt) : usesExtraIds{usesExtraIds} , loraTaskId{loraTaskId} , uniqueTokens{std::move(uniqueTokens)} , extraKeys{std::move(extraKeys)} - , cacheSaltID{cacheSaltID} + , cacheSalt{std::move(cacheSalt)} { } @@ -86,7 +86,7 @@ struct BlockKey } //! \brief Count the number of leading tokens that match between this key and \p other. - //! \details Returns 0 immediately when loraTaskId, extraKeys, or cacheSaltID differ, because those fields must + //! \details Returns 0 immediately when loraTaskId, extraKeys, or cacheSalt differ, because those fields must //! match exactly before token content is considered. //! \param other The key to compare against. //! \return Number of leading uniqueTokens that are identical in both keys. @@ -94,7 +94,7 @@ struct BlockKey { SizeType32 numMatched{0}; if (usesExtraIds == other.usesExtraIds && loraTaskId == other.loraTaskId && extraKeys == other.extraKeys - && cacheSaltID == other.cacheSaltID) + && cacheSalt == other.cacheSalt) { auto [matchEnd, otherMatchEnd] = std::mismatch( uniqueTokens.begin(), uniqueTokens.end(), other.uniqueTokens.begin(), other.uniqueTokens.end()); diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index d3966adf2f20..c665f7a8df95 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -82,7 +82,6 @@ using UniqueToken = tensorrt_llm::runtime::UniqueToken; using VecUniqueTokens = tensorrt_llm::runtime::VecUniqueTokens; using LoraTaskIdType = tensorrt_llm::runtime::LoraTaskIdType; using BlocksPerWindow = std::map>; -using CacheSaltIDType = tensorrt_llm::runtime::CacheSaltIDType; using MmKey = tensorrt_llm::executor::MmKey; using WindowSizeType = SizeType32; diff --git a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h index 263a15b50970..886147a09c73 100644 --- a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h +++ b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h @@ -108,7 +108,6 @@ class GenericLlmRequest using MillisecondsType = std::chrono::milliseconds; using TimePoint = std::chrono::time_point; using Duration = std::chrono::time_point::duration; - using CacheSaltIDType = runtime::CacheSaltIDType; GenericLlmRequest(RequestIdType requestId, SizeType32 maxNewTokens, std::shared_ptr const& inputTokens, runtime::SamplingConfig const& samplingConfig, bool isStreaming, std::optional endId = std::nullopt, @@ -147,11 +146,12 @@ class GenericLlmRequest std::optional languageAdapterUid = std::nullopt, std::optional allottedTimeMs = std::nullopt, std::optional const& contextPhaseParams = std::nullopt, - std::optional cacheSaltID = std::nullopt, std::optional arrivalTime = std::nullopt, + std::optional arrivalTime = std::nullopt, std::optional>> agent_hierarchy = std::nullopt, std::optional>> multimodalItemRunCuOffsets = std::nullopt, std::optional>> multimodalRunPositions = std::nullopt, - std::optional>> multimodalRunLengths = std::nullopt) + std::optional>> multimodalRunLengths = std::nullopt, + std::optional cacheSalt = std::nullopt) : mRequestId(requestId) , mPromptLen(inputTokens->size()) , mMaxNewTokens(maxNewTokens) @@ -213,7 +213,7 @@ class GenericLlmRequest , mGuidedDecodingParams(std::move(guidedDecodingParams)) , mLanguageAdapterUid(languageAdapterUid) , mAllottedTimeMs(allottedTimeMs) - , mCacheSaltID(cacheSaltID) + , mCacheSalt(std::move(cacheSalt)) , mAgentHierarchy(std::move(agent_hierarchy)) { if (mEncoderTokens.has_value() || encoderInputFeatures.has_value()) @@ -242,7 +242,7 @@ class GenericLlmRequest executor::PriorityType priority = executor::Request::kDefaultPriority, SizeType32 numReturnSequences = 1, std::optional languageAdapterUid = std::nullopt, std::optional const& contextPhaseParams = std::nullopt, - std::optional cacheSaltID = std::nullopt) + std::optional cacheSalt = std::nullopt) : mRequestId(requestId) , mPromptLen(inputTokens.size()) , mMaxNewTokens(maxNewTokens) @@ -283,7 +283,7 @@ class GenericLlmRequest , mContextPhaseParams(contextPhaseParams) , mNumReturnSequences(numReturnSequences) , mLanguageAdapterUid(languageAdapterUid) - , mCacheSaltID(cacheSaltID) + , mCacheSalt(std::move(cacheSalt)) { if (mEncoderTokens.has_value()) { @@ -323,7 +323,7 @@ class GenericLlmRequest , mGuidedDecodingParams(req.getGuidedDecodingParams()) , mLanguageAdapterUid(req.getLanguageAdapterUid()) , mAllottedTimeMs(req.getAllottedTimeMs()) - , mCacheSaltID(req.getCacheSaltID()) + , mCacheSalt(req.getCacheSalt()) { if (req.getRequestType() == executor::RequestType::REQUEST_TYPE_GENERATION_ONLY) { @@ -1897,9 +1897,9 @@ class GenericLlmRequest return mLanguageAdapterUid; } - [[nodiscard]] std::optional getCacheSaltID() const + [[nodiscard]] std::optional getCacheSalt() const { - return mCacheSaltID; + return mCacheSalt; } std::vector getLanguageAdapterRouting( @@ -2196,8 +2196,8 @@ class GenericLlmRequest bool mUseDraftModel{false}; - // Cache salt id for each request. - std::optional mCacheSaltID{std::nullopt}; + // Cache salt string. Used in BlockKey hashing/matching and surfaced in KV cache events. + std::optional mCacheSalt{std::nullopt}; std::optional>> mAgentHierarchy{std::nullopt}; @@ -2394,11 +2394,12 @@ class LlmRequest : public GenericLlmRequest std::optional languageAdapterUid = std::nullopt, std::optional allottedTimeMs = std::nullopt, std::optional const& contextPhaseParams = std::nullopt, - std::optional cacheSaltID = std::nullopt, std::optional arrivalTime = std::nullopt, + std::optional arrivalTime = std::nullopt, std::optional>> agent_hierarchy = std::nullopt, std::optional> multimodalItemRunCuOffsets = std::nullopt, std::optional> multimodalRunPositions = std::nullopt, - std::optional> multimodalRunLengths = std::nullopt) + std::optional> multimodalRunLengths = std::nullopt, + std::optional cacheSalt = std::nullopt) : Base(requestId, maxNewTokens, std::make_shared>(std::move(inputTokens)), samplingConfig, isStreaming, endId, padId, std::move(embeddingBias), std::move(badWordsList), std::move(stopWordsList), @@ -2431,8 +2432,8 @@ class LlmRequest : public GenericLlmRequest inputTokenExtraIds ? std::make_optional(std::make_shared(std::move(*inputTokenExtraIds))) : std::optional>(std::nullopt), numReturnSequences, std::move(eagleConfig), skipCrossAttnBlocks, returnPerfMetrics, - std::move(guidedDecodingParams), languageAdapterUid, allottedTimeMs, contextPhaseParams, cacheSaltID, - arrivalTime, std::move(agent_hierarchy), + std::move(guidedDecodingParams), languageAdapterUid, allottedTimeMs, contextPhaseParams, arrivalTime, + std::move(agent_hierarchy), multimodalItemRunCuOffsets.has_value() ? std::make_shared>(std::move(multimodalItemRunCuOffsets.value())) : std::optional>>(std::nullopt), @@ -2441,7 +2442,8 @@ class LlmRequest : public GenericLlmRequest : std::optional>>(std::nullopt), multimodalRunLengths.has_value() ? std::make_shared>(std::move(multimodalRunLengths.value())) - : std::optional>>(std::nullopt)) + : std::optional>>(std::nullopt), + std::move(cacheSalt)) { } diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index 1f625d57084c..f716bef6e3cc 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/cpp/include/tensorrt_llm/executor/executor.h @@ -714,8 +714,9 @@ class Request /// @param allottedTimeMs The allotted time in milliseconds after which the request is cancelled with a timedOut /// finish reason. The request may exceed this time slightly, but at most by 1 forward pass (in pipeline parallelism /// that may involve multiple micro-batches). A request can be timed-out before ever being scheduled. - /// @param cacheSaltID Salt ID for KV cache blocks to limit the kv cache reuse to the requests with the same string. /// @param disaggRequestId Disaggregated request ID. + /// @param cacheSalt Optional cache salt string. If provided, KV cache blocks are tagged so reuse is limited to + /// requests with the same salt. The string is also surfaced in KV cache events. Defaults to std::nullopt. Request(VecTokens inputTokenIds, SizeType32 maxTokens, bool streaming = false, SamplingConfig const& samplingConfig = SamplingConfig(), OutputConfig const& outputConfig = OutputConfig(), std::optional const& endId = std::nullopt, std::optional const& padId = std::nullopt, @@ -743,8 +744,7 @@ class Request std::optional guidedDecodingParams = std::nullopt, std::optional languageAdapterUid = std::nullopt, std::optional allottedTimeMs = std::nullopt, - std::optional cacheSaltID = std::nullopt, - std::optional disaggRequestId = std::nullopt); + std::optional disaggRequestId = std::nullopt, std::optional cacheSalt = std::nullopt); /// @brief This logits postprocessor name will dispatch to the batched logits postprocessor static auto constexpr kBatchedPostProcessorName = "batched"; @@ -792,7 +792,7 @@ class Request [[nodiscard]] std::optional getGuidedDecodingParams() const; [[nodiscard]] std::optional getLanguageAdapterUid() const; [[nodiscard]] std::optional getAllottedTimeMs() const; - [[nodiscard]] std::optional getCacheSaltID() const; + [[nodiscard]] std::optional getCacheSalt() const; [[nodiscard]] std::optional> getAdditionalOutputNames() const; [[nodiscard]] std::optional getDisaggRequestId() const; @@ -829,7 +829,7 @@ class Request void setGuidedDecodingParams(GuidedDecodingParams const& guidedDecodingParams); void setLanguageAdapterUid(SizeType32 languageAdapterUid); void setAllottedTimeMs(MillisecondsType allottedTimeMs); - void setCacheSaltID(CacheSaltIDType cacheSaltID); + void setCacheSalt(std::optional cacheSalt); void setDisaggRequestId(IdType disaggRequestId); private: @@ -1729,13 +1729,14 @@ struct KVCacheStoredBlockData KVCacheStoredBlockData(IdType blockHash, tensorrt_llm::runtime::VecUniqueTokens tokens, std::optional loraId, SizeType32 cacheLevel, SizeType32 priority, - std::vector mmKeys = {}) + std::vector mmKeys = {}, std::optional cacheSalt = std::nullopt) : blockHash{blockHash} , tokens{std::move(tokens)} , loraId{loraId} , cacheLevel{cacheLevel} , priority{priority} , mmKeys{std::move(mmKeys)} + , cacheSalt{std::move(cacheSalt)} { } @@ -1751,6 +1752,8 @@ struct KVCacheStoredBlockData SizeType32 priority; /// @brief The multimodal keys of the block std::vector mmKeys; + /// @brief The original cache salt string of the block, if any + std::optional cacheSalt; }; struct KVCacheStoredData diff --git a/cpp/include/tensorrt_llm/executor/types.h b/cpp/include/tensorrt_llm/executor/types.h index 0800865df7f1..2e6051291629 100644 --- a/cpp/include/tensorrt_llm/executor/types.h +++ b/cpp/include/tensorrt_llm/executor/types.h @@ -59,7 +59,6 @@ using RandomSeedType = std::uint64_t; using VecLogProbs = std::vector; using StreamPtr = std::shared_ptr; using MillisecondsType = std::chrono::milliseconds; -using CacheSaltIDType = std::uint64_t; using LogitsPostProcessor = std::function)>; using LogitsPostProcessorMap = std::unordered_map; diff --git a/cpp/include/tensorrt_llm/runtime/common.h b/cpp/include/tensorrt_llm/runtime/common.h index 7a3079d0bd75..2cda8821c133 100644 --- a/cpp/include/tensorrt_llm/runtime/common.h +++ b/cpp/include/tensorrt_llm/runtime/common.h @@ -44,7 +44,6 @@ using TokenIdType = std::int32_t; using LoraTaskIdType = std::uint64_t; using TokenExtraIdType = std::uint64_t; using VecTokenExtraIds = std::vector; -using CacheSaltIDType = std::uint64_t; struct UniqueToken { diff --git a/cpp/tensorrt_llm/batch_manager/blockKey.cpp b/cpp/tensorrt_llm/batch_manager/blockKey.cpp index 33092a5a37fa..e8125b0106f4 100644 --- a/cpp/tensorrt_llm/batch_manager/blockKey.cpp +++ b/cpp/tensorrt_llm/batch_manager/blockKey.cpp @@ -334,7 +334,7 @@ std::vector buildBlockKeys( currentTokenIdx += uniqueTokens.size(); blockKeys.emplace_back(llmRequest.getInputTokensExtraIds().has_value(), llmRequest.getLoraTaskId(), - std::move(uniqueTokens), std::move(extraKeys), llmRequest.getCacheSaltID()); + std::move(uniqueTokens), std::move(extraKeys), llmRequest.getCacheSalt()); } return blockKeys; } @@ -342,7 +342,7 @@ std::vector buildBlockKeys( bool BlockKey::operator==(BlockKey const& other) const noexcept { return (usesExtraIds == other.usesExtraIds && loraTaskId == other.loraTaskId && uniqueTokens == other.uniqueTokens - && extraKeys == other.extraKeys && cacheSaltID == other.cacheSaltID); + && extraKeys == other.extraKeys && cacheSalt == other.cacheSalt); } BlockKey BlockKey::shorten(int newNumTokens) const @@ -364,10 +364,10 @@ size_t BlockKeyHasher::hash(BlockKey const& blockKey, std::size_t parentHash) no // Constants provide very good distribution - each input bit affects each output bit with ~50% probability. size_t seed = blockKey.uniqueTokens.size() ^ parentHash * UINT64_C(0xbf58476d1ce4e5b9); - if (parentHash == 0 && blockKey.cacheSaltID) + if (parentHash == 0 && blockKey.cacheSalt) { - // Only hashing the cache salt ID for the first block in the sequence - uint64_t c = blockKey.cacheSaltID.value(); + // Only mix the cache salt into the hash for the first block in the sequence. + uint64_t c = static_cast(std::hash{}(blockKey.cacheSalt.value())); seed = hash64Mix(c, seed); } diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheEventManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheEventManager.cpp index c8f6ddd474f4..2a986adf310f 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheEventManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheEventManager.cpp @@ -105,7 +105,8 @@ void KVCacheEventManager::enqueueStoredEvent(std::vector const& blocks for (auto const& block : blocks) { data.blocks.emplace_back(block->getHash(), block->getUniqueTokens(), block->getBlockKey().loraTaskId, - block->isPrimary() ? kPrimaryLevel : kSecondaryLevel, block->getPriority(), block->getExtraKeys()); + block->isPrimary() ? kPrimaryLevel : kSecondaryLevel, block->getPriority(), block->getExtraKeys(), + block->getBlockKey().cacheSalt); } enqueueEvent({mEventId++, data, windowSize, mAttentionDpRank}); diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 0d3a8f164027..93e7772da123 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -2076,7 +2076,7 @@ std::shared_ptr WindowBlockManager::findBlocksInReuseTreeByBlockKe for (auto const& blockedUniqueTokensList : blockedUniqueTokens) { blockKeys.emplace_back(blockKey.usesExtraIds, blockKey.loraTaskId, blockedUniqueTokensList, blockKey.extraKeys, - blockKey.cacheSaltID); + blockKey.cacheSalt); } return searchReuseTree(blockKeys); } @@ -4453,7 +4453,7 @@ std::vector KVCacheManager::commitAndGetBlockHashesForRequest( bool const usesExtraIds = llmRequest.getInputTokensExtraIds().has_value(); auto const loraTaskId = llmRequest.getLoraTaskId(); - auto const cacheSaltID = llmRequest.getCacheSaltID(); + auto const cacheSalt = llmRequest.getCacheSalt(); std::vector hashes; hashes.reserve(static_cast(limit)); @@ -4469,7 +4469,7 @@ std::vector KVCacheManager::commitAndGetBlockHashesForRequest( SizeType32 const tokenEnd = tokenStart + tokensPerBlock; auto extraKeys = generateBlockHashExtraKeys(llmRequest, tokenStart, tokenEnd); VecUniqueTokens blockTokens(uniqueTokens.begin() + tokenStart, uniqueTokens.begin() + tokenEnd); - BlockKey blockKey(usesExtraIds, loraTaskId, std::move(blockTokens), std::move(extraKeys), cacheSaltID); + BlockKey blockKey(usesExtraIds, loraTaskId, std::move(blockTokens), std::move(extraKeys), cacheSalt); block->setBlockKey(blockKey, /*isFull=*/true); // setHash() chains through mPrevBlockInSeq, which was wired in addBlockToBeam. The // loop walks blocks in allocation order, so by the time we reach block b its diff --git a/cpp/tensorrt_llm/executor/request.cpp b/cpp/tensorrt_llm/executor/request.cpp index 5ac62d3fcb64..e32045892ba7 100644 --- a/cpp/tensorrt_llm/executor/request.cpp +++ b/cpp/tensorrt_llm/executor/request.cpp @@ -40,8 +40,8 @@ Request::Request(VecTokens inputTokenIds, SizeType32 maxTokens, bool streaming, std::optional encoderOutputLength, std::optional crossAttentionMask, SizeType32 numReturnSequences, std::optional eagleConfig, std::optional skipCrossAttnBlocks, std::optional guidedDecodingParams, std::optional languageAdapterUid, - std::optional allottedTimeMs, std::optional cacheSaltID, - std::optional disaggRequestId) + std::optional allottedTimeMs, std::optional disaggRequestId, + std::optional cacheSalt) : mImpl(std::make_unique(std::move(inputTokenIds), maxTokens, streaming, samplingConfig, outputConfig, endId, padId, std::move(positionIds), std::move(badWords), std::move(stopWords), std::move(embeddingBias), std::move(externalDraftTokensConfig), std::move(pTuningConfig), std::move(multimodalInput), @@ -50,7 +50,7 @@ Request::Request(VecTokens inputTokenIds, SizeType32 maxTokens, bool streaming, std::move(encoderInputTokenIds), clientId, returnAllGeneratedTokens, priority, type, std::move(contextPhaseParams), std::move(encoderInputFeatures), encoderOutputLength, crossAttentionMask, numReturnSequences, eagleConfig, skipCrossAttnBlocks, std::move(guidedDecodingParams), languageAdapterUid, - allottedTimeMs, cacheSaltID, disaggRequestId)) + allottedTimeMs, disaggRequestId, std::move(cacheSalt))) { } @@ -249,9 +249,9 @@ std::optional Request::getLanguageAdapterUid() const return mImpl->getLanguageAdapterUid(); } -std::optional Request::getCacheSaltID() const +std::optional Request::getCacheSalt() const { - return mImpl->getCacheSaltID(); + return mImpl->getCacheSalt(); } std::optional Request::getDisaggRequestId() const @@ -424,9 +424,9 @@ void Request::setLanguageAdapterUid(SizeType32 languageAdapterUid) mImpl->setLanguageAdapterUid(languageAdapterUid); } -void Request::setCacheSaltID(CacheSaltIDType cacheSaltID) +void Request::setCacheSalt(std::optional cacheSalt) { - mImpl->setCacheSaltID(cacheSaltID); + mImpl->setCacheSalt(std::move(cacheSalt)); } void Request::setDisaggRequestId(IdType disaggRequestId) diff --git a/cpp/tensorrt_llm/executor/requestImpl.h b/cpp/tensorrt_llm/executor/requestImpl.h index 281f81d462a7..55610885b1ae 100644 --- a/cpp/tensorrt_llm/executor/requestImpl.h +++ b/cpp/tensorrt_llm/executor/requestImpl.h @@ -32,6 +32,21 @@ class Request::Impl { public: + //! Maximum allowed length of a cache salt string. Cache salts are copied into every BlockKey and emitted + //! with KV cache events, so unbounded strings would inflate memory and serialization cost proportional to + //! the number of blocks. + static constexpr std::size_t kMaxCacheSaltLength{256}; + + static std::optional validateCacheSalt(std::optional cacheSalt) + { + if (cacheSalt.has_value() && cacheSalt->size() > kMaxCacheSaltLength) + { + TLLM_THROW("cacheSalt length (%zu) exceeds the maximum supported length (%zu).", cacheSalt->size(), + kMaxCacheSaltLength); + } + return cacheSalt; + } + Impl(VecTokens inputTokenIds, SizeType32 maxNewTokens, bool streaming, SamplingConfig const& samplingConfig, OutputConfig outputConfig, std::optional const& endId, std::optional const& padId, std::optional> positionIds, std::optional> badWords, @@ -48,7 +63,7 @@ class Request::Impl std::optional crossAttentionMask, SizeType32 numReturnSequences, std::optional eagleConfig, std::optional skipCrossAttnBlocks, std::optional guidedDecodingParams, std::optional languageAdapterUid, std::optional allottedTimeMs, - std::optional cacheSaltID, std::optional disaggRequestId) + std::optional disaggRequestId, std::optional cacheSalt = std::nullopt) : mInputTokenIds(std::move(inputTokenIds)) , mMaxNewTokens(maxNewTokens) , mStreaming(streaming) @@ -85,7 +100,7 @@ class Request::Impl , mGuidedDecodingParams(std::move(guidedDecodingParams)) , mLanguageAdapterUid(languageAdapterUid) , mAllottedTimeMs(allottedTimeMs) - , mCacheSaltID(cacheSaltID) + , mCacheSalt(validateCacheSalt(std::move(cacheSalt))) , mDisaggRequestId(disaggRequestId) { validate(); @@ -298,9 +313,9 @@ class Request::Impl return mLanguageAdapterUid; } - [[nodiscard]] std::optional getCacheSaltID() const + [[nodiscard]] std::optional getCacheSalt() const { - return mCacheSaltID; + return mCacheSalt; } [[nodiscard]] std::optional getDisaggRequestId() const @@ -482,9 +497,9 @@ class Request::Impl mLanguageAdapterUid = languageAdapterUid; } - void setCacheSaltID(CacheSaltIDType cacheSaltID) + void setCacheSalt(std::optional cacheSalt) { - mCacheSaltID = cacheSaltID; + mCacheSalt = validateCacheSalt(std::move(cacheSalt)); } void setDisaggRequestId(IdType disaggRequestId) @@ -565,8 +580,8 @@ class Request::Impl lambda(mGuidedDecodingParams); lambda(mLanguageAdapterUid); lambda(mAllottedTimeMs ? std::make_optional(mAllottedTimeMs->count()) : std::nullopt); - lambda(mCacheSaltID); lambda(mDisaggRequestId); + lambda(mCacheSalt); } VecTokens mInputTokenIds; @@ -605,7 +620,7 @@ class Request::Impl std::optional mGuidedDecodingParams; std::optional mLanguageAdapterUid; std::optional mAllottedTimeMs; - std::optional mCacheSaltID; + std::optional mCacheSalt; std::optional mDisaggRequestId; }; diff --git a/cpp/tensorrt_llm/executor/serialization.cpp b/cpp/tensorrt_llm/executor/serialization.cpp index e4c325423126..ce081e10c603 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -885,8 +885,8 @@ Request Serialization::deserializeRequest(std::istream& is) auto allottedTimeMs = allottedTimeInt ? std::optional(std::chrono::milliseconds(*allottedTimeInt)) : std::nullopt; - auto cacheSaltID = su::deserialize>(is); auto disaggRequestId = su::deserialize>(is); + auto cacheSalt = su::deserialize>(is); return Request(std::move(inputTokenIds), maxNewTokens, streaming, samplingConfig, outputConfig, endId, padId, std::move(positionIds), std::move(badWords), std::move(stopWords), std::move(embeddingBias), @@ -896,7 +896,7 @@ Request Serialization::deserializeRequest(std::istream& is) std::move(encoderInputTokenIds), clientId, returnAllGeneratedTokens, priority, requestType, std::move(contextPhaseParams), std::move(encoderInputFeatures), encoderOutputLength, std::move(crossAttentionMask), numReturnSequences, std::move(eagleConfig), std::move(skipCrossAttnBlocks), - std::move(guidedDecodingParams), languageAdapterUid, allottedTimeMs, cacheSaltID, disaggRequestId); + std::move(guidedDecodingParams), languageAdapterUid, allottedTimeMs, disaggRequestId, std::move(cacheSalt)); } void Serialization::serialize(Request const& request, std::ostream& os) @@ -2517,6 +2517,7 @@ size_t Serialization::serializedSize(KVCacheStoredBlockData const& data) totalSize += su::serializedSize(data.cacheLevel); totalSize += su::serializedSize(data.priority); totalSize += su::serializedSize(data.mmKeys); + totalSize += su::serializedSize(data.cacheSalt); return totalSize; } @@ -2528,6 +2529,7 @@ void Serialization::serialize(KVCacheStoredBlockData const& data, std::ostream& su::serialize(data.cacheLevel, os); su::serialize(data.priority, os); su::serialize(data.mmKeys, os); + su::serialize(data.cacheSalt, os); } KVCacheStoredBlockData Serialization::deserializeKVCacheStoredBlockData(std::istream& is) @@ -2538,8 +2540,9 @@ KVCacheStoredBlockData Serialization::deserializeKVCacheStoredBlockData(std::ist auto cacheLevel = su::deserialize(is); auto priority = su::deserialize(is); auto mmKeys = su::deserialize>(is); + auto cacheSalt = su::deserialize>(is); - return KVCacheStoredBlockData{blockHash, tokens, loraId, cacheLevel, priority, mmKeys}; + return KVCacheStoredBlockData{blockHash, tokens, loraId, cacheLevel, priority, mmKeys, cacheSalt}; } // KVcacheRemovedData @@ -2686,7 +2689,7 @@ size_t Serialization::serializedSize(tensorrt_llm::batch_manager::kv_cache_manag totalSize += su::serializedSize(key.uniqueTokens); // std::vector where MmKey is pair, SizeType32> totalSize += su::serializedSize(key.extraKeys); - totalSize += su::serializedSize(key.cacheSaltID); + totalSize += su::serializedSize(key.cacheSalt); return totalSize; } @@ -2696,7 +2699,7 @@ void Serialization::serialize(tensorrt_llm::batch_manager::kv_cache_manager::Blo su::serialize(key.loraTaskId, os); su::serialize(key.uniqueTokens, os); su::serialize(key.extraKeys, os); - su::serialize(key.cacheSaltID, os); + su::serialize(key.cacheSalt, os); } tensorrt_llm::batch_manager::kv_cache_manager::BlockKey Serialization::deserializeBlockKey(std::istream& is) @@ -2705,13 +2708,13 @@ tensorrt_llm::batch_manager::kv_cache_manager::BlockKey Serialization::deseriali auto loraTaskId = su::deserialize>(is); auto uniqueTokens = su::deserialize>(is); auto extraKeys = su::deserialize>(is); - auto cacheSaltID = su::deserialize>(is); + auto cacheSalt = su::deserialize>(is); tensorrt_llm::batch_manager::kv_cache_manager::BlockKey key; key.usesExtraIds = usesExtraIds; key.loraTaskId = std::move(loraTaskId); key.uniqueTokens = std::move(uniqueTokens); key.extraKeys = std::move(extraKeys); - key.cacheSaltID = std::move(cacheSaltID); + key.cacheSalt = std::move(cacheSalt); return key; } diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp index 5447f13c60e2..086cdf7547c2 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp @@ -202,7 +202,7 @@ void initBindings(nb::module_& m) .def_prop_ro("llm_request_type", &GenLlmReq::getLlmRequestType) .def_prop_ro("parent_request_id", &GenLlmReq::getParentRequestId) .def_prop_ro("is_child", &GenLlmReq::isChild) - .def_prop_ro("cache_salt_id", &GenLlmReq::getCacheSaltID) + .def_prop_ro("cache_salt", &GenLlmReq::getCacheSalt) .def_prop_ro("kv_cache_retention_config", &GenLlmReq::getKvCacheRetentionConfig) .def_prop_ro("multimodal_hashes", [](GenLlmReq& self) @@ -347,12 +347,12 @@ void initBindings(nb::module_& m) std::optional language_adapter_uid, std::optional allotted_time_ms, std::optional context_phase_params, - std::optional cache_salt_id, std::optional arrival_time, std::optional>> agent_hierarchy, std::optional> multimodal_item_run_cu_offsets, std::optional> multimodal_run_positions, - std::optional> multimodal_run_lengths) + std::optional> multimodal_run_lengths, + std::optional cache_salt) { auto makeOptionalTensor = [](std::optional const& atTensor, bool unsqueeze = false) { @@ -392,9 +392,9 @@ void initBindings(nb::module_& m) encoder_input_tokens, return_encoder_output, client_id, priority, encoder_input_features_tensor_ptr, encoder_output_length, cross_attention_mask_tensor_ptr, llm_request_type, input_token_extra_ids, num_return_sequences, eagle_config, skip_cross_attn_blocks_tensor_ptr, return_perf_metrics, - guided_decoding_params, language_adapter_uid, allotted_time_ms, context_phase_params, cache_salt_id, - arrival_time, std::move(agent_hierarchy), multimodal_item_run_cu_offsets, multimodal_run_positions, - multimodal_run_lengths}; + guided_decoding_params, language_adapter_uid, allotted_time_ms, context_phase_params, arrival_time, + std::move(agent_hierarchy), multimodal_item_run_cu_offsets, multimodal_run_positions, + multimodal_run_lengths, std::move(cache_salt)}; }, nb::arg("request_id"), nb::arg("max_new_tokens"), nb::arg("input_tokens"), nb::arg("sampling_config"), nb::arg("is_streaming"), nb::arg("end_id") = std::nullopt, nb::arg("pad_id") = std::nullopt, @@ -420,10 +420,10 @@ void initBindings(nb::module_& m) nb::arg("eagle_config") = std::nullopt, nb::arg("skip_cross_attn_blocks") = std::nullopt, nb::arg("return_perf_metrics") = false, nb::arg("guided_decoding_params") = std::nullopt, nb::arg("language_adapter_uid") = std::nullopt, nb::arg("allotted_time_ms") = std::nullopt, - nb::arg("context_phase_params") = std::nullopt, nb::arg("cache_salt_id") = std::nullopt, - nb::arg("arrival_time") = std::nullopt, nb::arg("agent_hierarchy") = std::nullopt, - nb::arg("multimodal_item_run_cu_offsets") = std::nullopt, - nb::arg("multimodal_run_positions") = std::nullopt, nb::arg("multimodal_run_lengths") = std::nullopt) + nb::arg("context_phase_params") = std::nullopt, nb::arg("arrival_time") = std::nullopt, + nb::arg("agent_hierarchy") = std::nullopt, nb::arg("multimodal_item_run_cu_offsets") = std::nullopt, + nb::arg("multimodal_run_positions") = std::nullopt, nb::arg("multimodal_run_lengths") = std::nullopt, + nb::arg("cache_salt") = std::nullopt) .def("check_token_id_range", &tb::LlmRequest::checkTokenIdRange, nb::arg("vocab_size")) .def(nb::init()) .def("validate", &tb::LlmRequest::validate, nb::arg("max_input_len"), nb::arg("max_seq_len"), diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.cpp index 796909bd419a..21f9ea39823b 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.cpp @@ -127,11 +127,11 @@ std::shared_ptr LlmRequest::toTrtLlm() const mLanguageAdapterUid, // mAllottedTimeMs, // mContextPhaseParams, // - mCacheSaltID, // mPerfMetrics.timingMetrics.arrivalTime, // mAgentHierarchy, // mMultimodalItemRunCuOffsets, // mMultimodalRunPositions, // - mMultimodalRunLengths // + mMultimodalRunLengths, // + mCacheSalt // ); } diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.h b/cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.h index 967870c8177c..387d915ab4aa 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.h +++ b/cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.h @@ -86,7 +86,7 @@ class LlmRequest : public tb::GenericLlmRequest std::optional languageAdapterUid = std::nullopt, std::optional allottedTimeMs = std::nullopt, std::optional const& contextPhaseParams = std::nullopt, - std::optional cacheSaltID = std::nullopt, std::optional arrivalTime = std::nullopt, + std::optional arrivalTime = std::nullopt, std::optional>> agent_hierarchy = std::nullopt, std::optional> multimodalItemRunCuOffsets = std::nullopt, std::optional> multimodalRunPositions = std::nullopt, @@ -155,7 +155,6 @@ class LlmRequest : public tb::GenericLlmRequest languageAdapterUid, // allottedTimeMs, // contextPhaseParams, // - cacheSaltID, // arrivalTime, // std::move(agent_hierarchy), // multimodalItemRunCuOffsets.has_value() diff --git a/cpp/tensorrt_llm/nanobind/executor/bindings.cpp b/cpp/tensorrt_llm/nanobind/executor/bindings.cpp index fbec513de3a1..b0ad31b7347e 100644 --- a/cpp/tensorrt_llm/nanobind/executor/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/bindings.cpp @@ -230,6 +230,7 @@ void initBindings(nb::module_& m) .def_ro("lora_id", &tle::KVCacheStoredBlockData::loraId) .def_ro("cache_level", &tle::KVCacheStoredBlockData::cacheLevel) .def_ro("priority", &tle::KVCacheStoredBlockData::priority) + .def_ro("cache_salt", &tle::KVCacheStoredBlockData::cacheSalt) .def_prop_ro("mm_keys", [](tle::KVCacheStoredBlockData const& self) { diff --git a/cpp/tensorrt_llm/nanobind/executor/request.cpp b/cpp/tensorrt_llm/nanobind/executor/request.cpp index 6cffe7740c13..b502370504d1 100644 --- a/cpp/tensorrt_llm/nanobind/executor/request.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/request.cpp @@ -613,7 +613,7 @@ void initRequestBindings(nb::module_& m) self.getClientId(), self.getReturnAllGeneratedTokens(), self.getPriority(), self.getRequestType(), self.getContextPhaseParams(), self.getEncoderInputFeatures(), self.getEncoderOutputLength(), self.getCrossAttentionMask(), self.getEagleConfig(), self.getSkipCrossAttnBlocks(), - self.getGuidedDecodingParams(), self.getCacheSaltID(), self.getDisaggRequestId()); + self.getGuidedDecodingParams(), self.getDisaggRequestId(), self.getCacheSalt()); }; auto requestSetstate = [](tle::Request& self, nb::tuple const& state) { @@ -642,7 +642,7 @@ void initRequestBindings(nb::module_& m) nb::cast>(state[29]), 1, nb::cast>(state[30]), nb::cast>(state[31]), nb::cast>(state[32]), std::nullopt, std::nullopt, - nb::cast>(state[33]), nb::cast>(state[34])); + nb::cast>(state[33]), nb::cast>(state[34])); }; nb::class_ request(m, "Request", nb::dynamic_attr()); @@ -683,8 +683,8 @@ void initRequestBindings(nb::module_& m) std::optional, // guidedDecodingParams std::optional, // languageAdapterUid std::optional, // allottedTimeMs - std::optional, // cacheSaltID - std::optional // disaggRequestId + std::optional, // disaggRequestId + std::optional // cacheSalt >(), // clang-format off nb::arg("input_token_ids"), @@ -724,9 +724,9 @@ void initRequestBindings(nb::module_& m) nb::arg("guided_decoding_params") = nb::none(), nb::arg("language_adapter_uid") = nb::none(), nb::arg("allotted_time_ms") = nb::none(), - nb::arg("cache_salt_id") = nb::none(), - nb::arg("disagg_request_id") = nb::none() - ) // clang-format on + nb::arg("disagg_request_id") = nb::none(), + nb::arg("cache_salt") = nb::none() + ) // clang-format on .def_prop_ro("input_token_ids", &tle::Request::getInputTokenIds) .def_prop_ro("max_tokens", &tle::Request::getMaxTokens) .def_prop_rw("streaming", &tle::Request::getStreaming, &tle::Request::setStreaming) @@ -768,7 +768,7 @@ void initRequestBindings(nb::module_& m) .def_prop_rw( "guided_decoding_params", &tle::Request::getGuidedDecodingParams, &tle::Request::setGuidedDecodingParams) .def_prop_rw("allotted_time_ms", &tle::Request::getAllottedTimeMs, &tle::Request::setAllottedTimeMs) - .def_prop_rw("cache_salt_id", &tle::Request::getCacheSaltID, &tle::Request::setCacheSaltID) + .def_prop_rw("cache_salt", &tle::Request::getCacheSalt, &tle::Request::setCacheSalt) .def_prop_rw("context_phase_params", &tle::Request::getContextPhaseParams, &tle::Request::setContextPhaseParams) .def_prop_rw("disagg_request_id", &tle::Request::getDisaggRequestId, &tle::Request::setDisaggRequestId) .def_prop_rw("priority", &tle::Request::getPriority, &tle::Request::setPriority) diff --git a/cpp/tests/unit_tests/batch_manager/agentTreeTest.cpp b/cpp/tests/unit_tests/batch_manager/agentTreeTest.cpp index 32e0b3e0ea17..5104f2a88a71 100644 --- a/cpp/tests/unit_tests/batch_manager/agentTreeTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/agentTreeTest.cpp @@ -64,7 +64,7 @@ class AgentTreeTest : public ::testing::Test std::nullopt, std::nullopt, false, std::nullopt, false, std::nullopt, false, std::nullopt, tensorrt_llm::executor::Request::kDefaultPriority, std::nullopt, std::nullopt, std::nullopt, tb::LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, 1, std::nullopt, std::nullopt, - false, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, agentHierarchy); + false, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, agentHierarchy); } LlmRequestPtr createAgentDeepResearchRequest(SizeType32 nodeId, SizeType32 requestId) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index a5d566f3b408..a7493f7019a8 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -2094,12 +2094,11 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdAndLoraTaskIdTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); } -TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) +TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltTest) { - // Test that cache_salt_id prevents KV cache reuse between requests with same tokens - // but different cache_salt_id values. + // Test that cache_salt prevents KV cache reuse between requests with same tokens + // but different cache_salt values. using VecTokenExtraIds = LlmRequest::VecTokenExtraIds; - using CacheSaltIDType = LlmRequest::CacheSaltIDType; auto constexpr numLayers = 12; auto constexpr numKvHeads = 6; @@ -2135,7 +2134,7 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) auto const inputLength = static_cast(inputTokens->size()); /////////////////////////////////////////////////////////////////////////// - // Test Case 1: Request without cache_salt_id + // Test Case 1: Request without cache_salt LlmRequest::RequestIdType requestId{0}; auto llmRequest0 = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, @@ -2143,8 +2142,8 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt, std::nullopt, false, std::nullopt, false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt, std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences, std::nullopt, - std::nullopt, false, std::nullopt, std::nullopt, std::nullopt, std::nullopt, - std::nullopt); // No cache_salt_id + std::nullopt, false, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + std::nullopt); // No cache_salt GenerationRequest seq0{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()}; @@ -2177,21 +2176,22 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); /////////////////////////////////////////////////////////////////////////// - // Test Case 2: Request with same tokens but with cache_salt_id = 12345 + // Test Case 2: Request with same tokens but with cache_salt = "tenant-A" requestId = 1; - CacheSaltIDType cacheSaltId1{12345}; + std::string const cacheSalt1{"tenant-A"}; auto llmRequest1 = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt, std::nullopt, false, std::nullopt, false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt, std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences, std::nullopt, - std::nullopt, false, std::nullopt, std::nullopt, std::nullopt, std::nullopt, - cacheSaltId1); // With cache_salt_id = 12345 + std::nullopt, false, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + std::nullopt, std::nullopt, std::nullopt, + cacheSalt1); // With cache_salt = "tenant-A" GenerationRequest seq1{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()}; - // Should NOT reuse blocks despite same tokens, because cache_salt_id is different + // Should NOT reuse blocks despite same tokens, because cache_salt is different auto promptLen1 = llmRequest1->getNumTokens(beamIdx); auto numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock()); auto prepopulatedPromptLen1 = blockManager @@ -2215,7 +2215,7 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); /////////////////////////////////////////////////////////////////////////// - // Test Case 3: Request with same tokens and same cache_salt_id = 12345 + // Test Case 3: Request with same tokens and same cache_salt = "tenant-A" requestId = 2; auto llmRequest2 = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, @@ -2223,12 +2223,13 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt, std::nullopt, false, std::nullopt, false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt, std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences, std::nullopt, - std::nullopt, false, std::nullopt, std::nullopt, std::nullopt, std::nullopt, - cacheSaltId1); // Same cache_salt_id = 12345 + std::nullopt, false, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + std::nullopt, std::nullopt, std::nullopt, + cacheSalt1); // Same cache_salt = "tenant-A" GenerationRequest seq2{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()}; - // SHOULD reuse blocks because both tokens and cache_salt_id match + // SHOULD reuse blocks because both tokens and cache_salt match auto promptLen2 = llmRequest2->getNumTokens(beamIdx); auto numContextBlocks2 = tc::ceilDiv(promptLen2, blockManager.getTokensPerBlock()); auto prepopulatedPromptLen2 = blockManager @@ -2252,21 +2253,22 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); /////////////////////////////////////////////////////////////////////////// - // Test Case 4: Request with same tokens but different cache_salt_id = 67890 + // Test Case 4: Request with same tokens but different cache_salt = "tenant-B" requestId = 3; - CacheSaltIDType cacheSaltId2{67890}; + std::string const cacheSalt2{"tenant-B"}; auto llmRequest3 = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt, std::nullopt, false, std::nullopt, false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt, std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences, std::nullopt, - std::nullopt, false, std::nullopt, std::nullopt, std::nullopt, std::nullopt, - cacheSaltId2); // Different cache_salt_id = 67890 + std::nullopt, false, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + std::nullopt, std::nullopt, std::nullopt, + cacheSalt2); // Different cache_salt = "tenant-B" GenerationRequest seq3{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()}; - // Should NOT reuse blocks from any previous request because cache_salt_id is different + // Should NOT reuse blocks from any previous request because cache_salt is different auto promptLen3 = llmRequest3->getNumTokens(beamIdx); auto numContextBlocks3 = tc::ceilDiv(promptLen3, blockManager.getTokensPerBlock()); auto prepopulatedPromptLen3 = blockManager @@ -2284,7 +2286,7 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks); /////////////////////////////////////////////////////////////////////////// - // Test Case 5: Request without cache_salt_id again + // Test Case 5: Request without cache_salt again requestId = 4; auto llmRequest4 = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, @@ -2292,12 +2294,12 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) std::nullopt, std::nullopt, std::nullopt, std::nullopt, false, false, false, std::nullopt, std::nullopt, false, std::nullopt, false, std::nullopt, false, std::nullopt, 0.5, std::nullopt, std::nullopt, std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, std::nullopt, numReturnSequences, std::nullopt, - std::nullopt, false, std::nullopt, std::nullopt, std::nullopt, std::nullopt, - std::nullopt); // No cache_salt_id + std::nullopt, false, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + std::nullopt); // No cache_salt GenerationRequest seq4{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()}; - // Should reuse blocks from request0 (blocks 0,1) because both have no cache_salt_id + // Should reuse blocks from request0 (blocks 0,1) because both have no cache_salt auto promptLen4 = llmRequest4->getNumTokens(beamIdx); auto numContextBlocks4 = tc::ceilDiv(promptLen4, blockManager.getTokensPerBlock()); auto prepopulatedPromptLen4 = blockManager diff --git a/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp b/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp index d80d0be456b4..98569bfc2aa0 100644 --- a/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp +++ b/cpp/tests/unit_tests/executor/serializeUtilsTest.cpp @@ -1097,7 +1097,7 @@ TEST(SerializeUtilsTest, BlockKeyWithExtras) VecUniqueTokens uniqueTokens{UniqueToken{10, 100}, UniqueToken{20, 200}}; std::optional loraTaskId = LoraTaskIdType{42}; - // Note: cacheSaltID is intentionally not set since it is not serialized + // Note: cacheSalt is intentionally not set; round-tripping with it set is covered separately. BlockKey key(true, loraTaskId, uniqueTokens, extraKeys); testSerializeDeserialize(key); diff --git a/examples/cpp/executor/executorExampleKvEvents.cpp b/examples/cpp/executor/executorExampleKvEvents.cpp index a48cbdfa9769..ea1923294382 100644 --- a/examples/cpp/executor/executorExampleKvEvents.cpp +++ b/examples/cpp/executor/executorExampleKvEvents.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * 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"); @@ -50,13 +50,14 @@ struct RuntimeOptions struct KVCacheBlock { KVCacheBlock(size_t hash, int cacheLevel, int priority, std::optional loraId = std::nullopt, - std::shared_ptr prevBlock = nullptr); + std::shared_ptr prevBlock = nullptr, std::optional cacheSalt = std::nullopt); size_t hash; int cacheLevel; int priority; std::optional loraId; + std::optional cacheSalt; std::shared_ptr prevBlock; std::unordered_map> nextBlocks; @@ -196,12 +197,13 @@ RuntimeOptions parseArgs(int argc, char* argv[]) return runtimeOpts; } -KVCacheBlock::KVCacheBlock( - size_t hash, int cacheLevel, int priority, std::optional loraId, std::shared_ptr prevBlock) +KVCacheBlock::KVCacheBlock(size_t hash, int cacheLevel, int priority, std::optional loraId, + std::shared_ptr prevBlock, std::optional cacheSalt) : hash{hash} , cacheLevel{cacheLevel} , priority{priority} , loraId{loraId} + , cacheSalt{std::move(cacheSalt)} , prevBlock{prevBlock} , nextBlocks{} { @@ -255,7 +257,7 @@ void RadixTree::pollEvents() TLLM_CHECK(block.tokens.size() > 0); auto thisBlock = std::make_shared( - block.blockHash, block.cacheLevel, block.priority, block.loraId, prevBlock); + block.blockHash, block.cacheLevel, block.priority, block.loraId, prevBlock, block.cacheSalt); blockTable[block.blockHash] = thisBlock; // Link the parent to the new block diff --git a/examples/llm-api/llm_kv_cache_connector.py b/examples/llm-api/llm_kv_cache_connector.py index 0e6aa3d83aa9..882478993e5a 100644 --- a/examples/llm-api/llm_kv_cache_connector.py +++ b/examples/llm-api/llm_kv_cache_connector.py @@ -192,7 +192,7 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput): len(block_ids)): if len(chunks[block_pos]) == self.block_size: hashed_tokens = self._hash_tokens(chunks[block_pos], - req.cache_salt_id) + req.cache_salt) file_path = self._file_path(hashed_tokens) @@ -202,11 +202,10 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput): return metadata - def _hash_tokens(self, tokens: list[int], - cache_salt_id: Optional[int]) -> int: - # cache_salt_id must participate in the hash so that requests carrying + def _hash_tokens(self, tokens: list[int], cache_salt: Optional[str]) -> int: + # cache_salt must participate in the hash so that requests carrying # different salts (or no salt) cannot collide on the same cache file. - return abs(hash((cache_salt_id, tuple(tokens)))) + return abs(hash((cache_salt, tuple(tokens)))) def _file_path(self, hash_value: int) -> Path: return Path(self.cache_folder) / f"{hash_value}.pt" @@ -238,7 +237,7 @@ def get_num_new_matched_tokens( for chunk in remaining_chunks: # Only do full blocks. if len(chunk) == self.block_size: - hashed_tokens = self._hash_tokens(chunk, request.cache_salt_id) + hashed_tokens = self._hash_tokens(chunk, request.cache_salt) file_path = self._file_path(hashed_tokens) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index f5034256e142..99da2a42265c 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -82,9 +82,9 @@ class RequestData: # Per-request cache salt that the KV cache manager uses to isolate reuse # between requests carrying different salts. Connectors that key cached # content on token sequences (e.g. by hashing tokens to a file path or - # remote object id) MUST mix cache_salt_id into their identifiers, + # remote object id) MUST mix cache_salt into their identifiers, # otherwise blocks from a different salt could be incorrectly reused. - cache_salt_id: Optional[int] = None + cache_salt: Optional[str] = None # A class to store some basic data regarding all inflight requests. @@ -361,7 +361,7 @@ def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManag num_scheduled_tokens, block_hashes=block_hashes, priorities=priorities, - cache_salt_id=req.cache_salt_id, + cache_salt=req.cache_salt, ) diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 16db68c79f9d..311f0b63d5e4 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1095,7 +1095,7 @@ def executor_request_to_llm_request( priority=executor_request.priority, llm_request_type=llm_request_type, context_phase_params=executor_request.context_phase_params, - cache_salt_id=executor_request.cache_salt_id, + cache_salt=executor_request.cache_salt, arrival_time=getattr(executor_request, "py_arrival_time", None), py_multimodal_data=getattr(executor_request, "py_multimodal_data", None), diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index d8e948d15689..c67bec2f22b9 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -3,6 +3,7 @@ import copy import enum +import hashlib import math import os from abc import ABC, abstractmethod @@ -2996,11 +2997,10 @@ def prepare_context(self, req: LlmRequest) -> bool: all_tokens, req, end=len(all_tokens) - 1) else: tokens = None - kv_cache = self._create_kv_cache( - req.py_request_id, - req.lora_task_id, - tokens, - cache_salt_id=req.cache_salt_id) + kv_cache = self._create_kv_cache(req.py_request_id, + req.lora_task_id, + tokens, + cache_salt=req.cache_salt) if kv_cache is None: return False kv_cache.cuda_stream = self._stream.cuda_stream @@ -3123,11 +3123,10 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): for req in scheduled_batch.context_requests: kv_cache = self.kv_cache_map.get(req.py_request_id) if kv_cache is None: - kv_cache = self._create_kv_cache( - req.py_request_id, - req.lora_task_id, - None, - cache_salt_id=req.cache_salt_id) + kv_cache = self._create_kv_cache(req.py_request_id, + req.lora_task_id, + None, + cache_salt=req.cache_salt) kv_cache.stop_committing() if not self._resume_and_restore(req.py_request_id, kv_cache): raise RuntimeError( @@ -3295,7 +3294,7 @@ def release_resources(current_request: LlmRequest, if prepare_resource: # Dummy/warmup request. ``stop_committing()`` below blocks all # writes to the radix tree, so the choice of branch does not - # affect committed state. ``cache_salt_id`` is left defaulted + # affect committed state. ``cache_salt`` is left defaulted # to None to avoid coupling synthetic data to any salted branch. kv_cache = self._create_kv_cache(req.py_request_id, req.lora_task_id, input_tokens) @@ -3674,7 +3673,7 @@ def _create_kv_cache(self, request_id: int, lora_task_id: int | None, input_tokens: Sequence[TokenIdExt] | None, - cache_salt_id: int | None = None): + cache_salt: str | None = None): assert request_id not in self.kv_cache_map, f"KV cache for request {request_id} already exists" if self.index_mapper.num_free_slots() == 0: logger.warning( @@ -3683,8 +3682,14 @@ def _create_kv_cache(self, "Skipping KV cache creation; request will retry next iteration.", request_id, self.index_mapper.size(), self.index_mapper.size()) return None + # ReuseScope.salt is int|None; derive a deterministic int from the + # cache_salt string so the same string yields the same reuse namespace + # across processes (matches C++ blockKey hashing on cacheSalt). + salt_int = (int.from_bytes( + hashlib.sha256(cache_salt.encode("utf-8")).digest()[:8], "little") + if cache_salt is not None else None) kv_cache = self.impl.create_kv_cache( - ReuseScope(lora_id=lora_task_id, salt=cache_salt_id), + ReuseScope(lora_id=lora_task_id, salt=salt_int), input_tokens, ) self.kv_cache_map[request_id] = kv_cache diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 841f396104f5..955a0ab55082 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -1212,6 +1212,8 @@ def _stored_block_to_json(data): for token in data.tokens ], # "lora_id": data.lora_id, # TODO (shreyasm): enable serialization of lora_id + "cache_salt": + data.cache_salt, "cache_level": data.cache_level, "priority": diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 2c6b19cb247a..fda27eb8fe8c 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -594,8 +594,8 @@ def _deduce_max_tokens(request: GenerationRequest, kv_cache_retention_config=request.kv_cache_retention_config, context_phase_params=context_phase_params, type=request_type, - cache_salt_id=request.cache_salt_id, disagg_request_id=disagg_request_id, + cache_salt=request.cache_salt, priority=request.priority) executor_request.py_original_end_id = request.sampling_params.end_id executor_request.py_num_logprobs = request.sampling_params.logprobs diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index a5f5efea6427..5ea904531f22 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -134,7 +134,7 @@ def generate_async( postproc_params: Optional[PostprocParams] = None, multimodal_params: Optional[MultimodalParams] = None, scheduling_params: Optional[SchedulingParams] = None, - cache_salt_id: Optional[int] = None, + cache_salt: Optional[str] = None, arrival_time: Optional[float] = None, priority: float = DEFAULT_REQUEST_PRIORITY, ) -> GenerationResult: @@ -162,7 +162,7 @@ def generate_async( trace_headers=trace_headers, multimodal_params=multimodal_params, scheduling_params=scheduling_params, - cache_salt_id=cache_salt_id, + cache_salt=cache_salt, arrival_time=arrival_time, priority=priority) result = self.submit(request) @@ -180,6 +180,7 @@ def generate( prompt_adapter_request: Optional[Union[ PromptAdapterRequest, List[PromptAdapterRequest]]] = None, disaggregated_params: Optional[DisaggregatedParams] = None, + cache_salt: Optional[Union[str, List[Optional[str]]]] = None, ) -> Union[GenerationResult, List[GenerationResult]]: """Generate output for the given prompt token ids in the synchronous mode. Synchronous generation accepts either single prompt or batched prompts. @@ -205,6 +206,7 @@ def generate( pa_req = prompt_adapter_request[i] else: pa_req = prompt_adapter_request + cs = cache_salt[i] if isinstance(cache_salt, list) else cache_salt future = self.generate_async( p, sampling_params=sp, @@ -212,7 +214,8 @@ def generate( lora_request=lora_req, prompt_adapter_request=pa_req, streaming=False, - disaggregated_params=disaggregated_params) + disaggregated_params=disaggregated_params, + cache_salt=cs) futures.append(future) for future in futures: diff --git a/tensorrt_llm/executor/request.py b/tensorrt_llm/executor/request.py index adbc2358b372..43ea11706e54 100644 --- a/tensorrt_llm/executor/request.py +++ b/tensorrt_llm/executor/request.py @@ -89,6 +89,8 @@ def local_path(self): class GenerationRequest: + # Mirrors C++ Request::Impl::kMaxCacheSaltLength + MAX_CACHE_SALT_LEN: int = 256 def __init__( self, @@ -105,7 +107,7 @@ def __init__( postproc_params: Optional[PostprocParams] = None, multimodal_params: Optional[MultimodalParams] = None, scheduling_params: Optional[SchedulingParams] = None, - cache_salt_id: Optional[int] = None, + cache_salt: Optional[str] = None, arrival_time: Optional[float] = None, priority: float = DEFAULT_REQUEST_PRIORITY, ): @@ -134,7 +136,21 @@ def __init__( self.disaggregated_params = disaggregated_params self.trace_headers = trace_headers self.scheduling_params = scheduling_params - self.cache_salt_id = cache_salt_id + if cache_salt is not None: + if not isinstance(cache_salt, str): + raise TypeError( + f"cache_salt must be str or None, got {type(cache_salt).__name__}" + ) + # The C++ side validates against UTF-8 byte length, so do the same here + # (Python `len()` would count Unicode code points, which can pass this + # guard but fail at C++ dispatch for non-ASCII salts). + cache_salt_byte_len = len(cache_salt.encode("utf-8")) + if cache_salt_byte_len > self.MAX_CACHE_SALT_LEN: + raise ValueError( + f"cache_salt UTF-8 byte length ({cache_salt_byte_len}) " + f"exceeds the maximum supported length " + f"({self.MAX_CACHE_SALT_LEN}).") + self.cache_salt = cache_salt self.arrival_time = arrival_time if not (0.0 <= priority <= 1.0): raise ValueError( diff --git a/tensorrt_llm/inputs/__init__.py b/tensorrt_llm/inputs/__init__.py index 4f2d4cc7e99b..3b9a5d51d052 100644 --- a/tensorrt_llm/inputs/__init__.py +++ b/tensorrt_llm/inputs/__init__.py @@ -24,8 +24,7 @@ async_load_audio, async_load_image, async_load_video, convert_image_mode, default_multimodal_input_loader, encode_base64_content_from_url, encode_base64_image, - get_cache_salt_id, load_base64_image_embeds, load_image, - load_video) + load_base64_image_embeds, load_image, load_video) # yapf: enable @@ -69,7 +68,6 @@ "encode_base64_image", "load_image", "load_video", - "get_cache_salt_id", "compute_retained_tokens_count", "compute_retained_tokens_from_tubelet_budget", "compute_retention_mask", diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index 59483e625790..eaa98cad44ef 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -27,8 +27,7 @@ _safe_aiohttp_get, _safe_request_get) from tensorrt_llm.inputs.media_io import \ convert_image_mode as convert_image_mode -from tensorrt_llm.inputs.multimodal import (MultimodalServerConfig, - default_hasher) +from tensorrt_llm.inputs.multimodal import MultimodalServerConfig from tensorrt_llm.inputs.multimodal_data import \ BaseModalityData as BaseModalityData from tensorrt_llm.inputs.multimodal_data import VideoData as VideoData @@ -894,14 +893,3 @@ def convert_to_conversation_message( inputs.append(input) return inputs - - -def get_cache_salt_id(cache_salt: str) -> int: - b = cache_salt.encode("utf-8") - h = default_hasher(b).digest(length=8) - cache_salt_id = int.from_bytes(h, "little", signed=False) - if cache_salt_id < 0 or cache_salt_id >= (1 << 64): - raise ValueError( - f"cache_salt_id must be in [0, 2**64 - 1], got {cache_salt_id}.") - - return cache_salt_id diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index d8c471d8624e..f99aa1593c9f 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -36,7 +36,7 @@ from ..executor.utils import (RequestError, create_mpi_comm_session, get_spawn_proxy_process_env) from ..inputs import (PromptInputs, create_input_processor, - create_input_processor_with_hash, get_cache_salt_id, + create_input_processor_with_hash, maybe_compute_mm_embed_cumsum, prompt_inputs) from ..logger import logger from ..sampling_params import SamplingParams @@ -476,8 +476,6 @@ def generate_async( sampling_params = self._prepare_sampling_params(sampling_params) - cache_salt_id = get_cache_salt_id( - cache_salt) if cache_salt is not None else None # With pytorch backend, py_executor has logic to handle max_tokens of 1, # so set to 1 to avoid allocating unnecessary KV cache blocks for single request # TODO: Also support for trt backend @@ -520,7 +518,7 @@ def generate_async( postproc_params=_postproc_params, multimodal_params=multimodal_params, scheduling_params=scheduling_params, - cache_salt_id=cache_salt_id, + cache_salt=cache_salt, arrival_time=arrival_time, priority=priority, ) diff --git a/tests/unittest/llmapi/test_llm_kv_cache_events.py b/tests/unittest/llmapi/test_llm_kv_cache_events.py index 5831c5b7a57a..ee002905d26c 100644 --- a/tests/unittest/llmapi/test_llm_kv_cache_events.py +++ b/tests/unittest/llmapi/test_llm_kv_cache_events.py @@ -109,6 +109,9 @@ def test_kv_cache_event_data_serialization(): # Verify mm_keys field exists (empty for text-only requests) assert "mm_keys" in serialized_event[0]["data"]["blocks"][0] assert serialized_event[0]["data"]["blocks"][0]["mm_keys"] == [] + # Verify cache_salt field exists (None for unsalted requests) + assert "cache_salt" in serialized_event[0]["data"]["blocks"][0] + assert serialized_event[0]["data"]["blocks"][0]["cache_salt"] is None req2 = create_llm_request(1, [1, 2, 3, 4, 5]) kv_cache_manager.impl.add_sequence_batch( @@ -779,7 +782,7 @@ def test_mm_keys_in_stored_events(): events = llm.get_kv_cache_events(5) - # Find stored events and verify mm_keys field + # Find stored events and verify mm_keys and cache_salt fields for event in events: if event and event["data"]["type"] == "stored": blocks = event["data"]["blocks"] @@ -789,6 +792,86 @@ def test_mm_keys_in_stored_events(): assert isinstance(block["mm_keys"], list) # For text-only requests, mm_keys should be empty assert block["mm_keys"] == [] + # cache_salt should be present (None for unsalted requests) + assert "cache_salt" in block + assert block["cache_salt"] is None + + +def test_cache_salt_in_stored_events(): + """Test that cache_salt string is preserved in stored block events.""" + llm = create_llm() + sampling_params = SamplingParams(max_tokens=6, temperature=0.01) + prompt = "Hello, my name is" + + _ = llm.generate(prompt, + sampling_params=sampling_params, + cache_salt="tenant-A") + + events = llm.get_kv_cache_events(5) + + # Find stored events and verify cache_salt field + found_stored = False + for event in events: + if event and event["data"]["type"] == "stored": + found_stored = True + blocks = event["data"]["blocks"] + for block in blocks: + assert "cache_salt" in block + assert block["cache_salt"] == "tenant-A" + + assert found_stored, "No stored events found" + + +def test_cache_salt_max_length_validation(): + """cache_salt longer than MAX_CACHE_SALT_LEN UTF-8 bytes is rejected.""" + from tensorrt_llm.executor.request import GenerationRequest + + max_len = GenerationRequest.MAX_CACHE_SALT_LEN + sampling_params = SamplingParams() + + # ASCII salt at the limit is accepted. + GenerationRequest(prompt_token_ids=[1, 2, 3], + sampling_params=sampling_params, + cache_salt="a" * max_len) + + # ASCII salt one byte over the limit is rejected. + with pytest.raises(ValueError, match="cache_salt UTF-8 byte length"): + GenerationRequest(prompt_token_ids=[1, 2, 3], + sampling_params=sampling_params, + cache_salt="a" * (max_len + 1)) + + # Non-ASCII salt: each character is 3 UTF-8 bytes. A salt whose + # `len()` is below the limit but whose UTF-8 byte count exceeds it + # must be rejected (this is the case Python's len()-based check missed). + char_count = (max_len // 3) + 1 # len() is well below max_len + salt = "中" * char_count # Chinese character, 3 UTF-8 bytes each + assert len(salt) <= max_len + assert len(salt.encode("utf-8")) > max_len + with pytest.raises(ValueError, match="cache_salt UTF-8 byte length"): + GenerationRequest(prompt_token_ids=[1, 2, 3], + sampling_params=sampling_params, + cache_salt=salt) + + +def test_non_ascii_cache_salt_in_stored_events(): + """Test that a non-ASCII cache_salt string is preserved in stored block events.""" + llm = create_llm() + sampling_params = SamplingParams(max_tokens=6, temperature=0.01) + prompt = "Hello, my name is" + salt = "tenant-中文" # mixed ASCII + Chinese + + _ = llm.generate(prompt, sampling_params=sampling_params, cache_salt=salt) + + events = llm.get_kv_cache_events(5) + + found_stored = False + for event in events: + if event and event["data"]["type"] == "stored": + found_stored = True + for block in event["data"]["blocks"]: + assert block.get("cache_salt") == salt + + assert found_stored, "No stored events found" def test_expected_kv_cache_events():