Skip to content

[None][feat] Integrate FP4 indexer for DSA on Blackwell #13340

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
longlee0622 merged 1 commit into from
May 7, 2026
Merged

[None][feat] Integrate FP4 indexer for DSA on Blackwell #13340

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion 3rdparty/fetch_content.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
{
"name": "deepgemm",
"git_repository": "https://github.com/deepseek-ai/DeepGEMM",
"git_tag": "4ff3f54d9b7ed3129e4f36f9871232ea7ecab86b",
"git_tag": "c491439ed5966833d56883ca302b6f72e74f8105",
"git_submodules_recurse": true,
"source_subdir": "dont-add-this-project-with-add-subdirectory"
},
Expand Down
31 changes: 26 additions & 5 deletions cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -760,7 +760,7 @@ class WindowBlockManager
std::shared_ptr<kv_connector::KvCacheConnectorManager> kvCacheConnectorManager,
radix_block_tree::UnifiedBlockTree& lookupTree, std::shared_ptr<kvc::BaseLoopbackAgent> loopbackAgent = nullptr,
bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128,
SizeType32 indexerKCacheIndexHeadDim = 0,
SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false,
std::optional<LinearAttentionMetadata> linearAttentionMetadata = std::nullopt,
SizeType32 numPlaceholderBlocks = 0);

Expand All @@ -781,6 +781,11 @@ class WindowBlockManager
return mIndexerKCacheIndexHeadDim;
}

[[nodiscard]] bool getIndexerKCacheUseFp4() const
{
return mIndexerKCacheUseFp4;
}

void allocatePools(bool useUvm);

void releasePools();
Expand Down Expand Up @@ -1292,6 +1297,9 @@ class WindowBlockManager
SizeType32 mIndexerKCacheQuantBlockSize;
// Index head dim for indexer K cache
SizeType32 mIndexerKCacheIndexHeadDim;
// Whether the indexer K cache stores FP4-packed data (half the byte count
// per token vs. FP8). Drives the createIndexerKCachePools() formula.
bool mIndexerKCacheUseFp4{false};

std::optional<LinearAttentionMetadata> mLinearAttentionMetadata;
};
Expand All @@ -1314,6 +1322,7 @@ class BlockManager
std::shared_ptr<kv_connector::KvCacheConnectorManager> kvCacheConnectorManager = nullptr,
std::optional<kvc::BaseAgentConfig> agentConfig = std::nullopt, bool enableIndexerKCache = false,
SizeType32 indexerKCacheQuantBlockSize = 128, SizeType32 indexerKCacheIndexHeadDim = 0,
bool indexerKCacheUseFp4 = false,
std::optional<LinearAttentionMetadata> linearAttentionMetadata = std::nullopt);

[[nodiscard]] bool isEnableIndexerKCache() const
Expand All @@ -1331,6 +1340,11 @@ class BlockManager
return mIndexerKCacheIndexHeadDim;
}

[[nodiscard]] bool getIndexerKCacheUseFp4() const
{
return mIndexerKCacheUseFp4;
}

[[nodiscard]] bool isEnablePartialReuse() const
{
return mWindowBlockManagers.begin()->second.isEnablePartialReuse();
Expand Down Expand Up @@ -1733,6 +1747,7 @@ class BlockManager
bool mIsEnableIndexerKCache{false};
SizeType32 mIndexerKCacheQuantBlockSize{0};
SizeType32 mIndexerKCacheIndexHeadDim{0};
bool mIndexerKCacheUseFp4{false};
std::optional<LinearAttentionMetadata> mLinearAttentionMetadata;
};

Expand Down Expand Up @@ -1854,6 +1869,7 @@ class BaseKVCacheManager
[[nodiscard]] virtual bool isEnableIndexerKCache() const = 0;
[[nodiscard]] virtual SizeType32 getIndexerKCacheIndexHeadDim() const = 0;
[[nodiscard]] virtual SizeType32 getIndexerKCacheQuantBlockSize() const = 0;
[[nodiscard]] virtual bool getIndexerKCacheUseFp4() const = 0;

// void removeToken(SizeType32 seqSlotIdx);
virtual void rewindKVCache(LlmRequest::RequestIdType requestId, SizeType32 rewindLengths) = 0;
Expand Down Expand Up @@ -2014,7 +2030,7 @@ class KVCacheManager : public BaseKVCacheManager
bool copyOnpartialReuse = true,
std::shared_ptr<kv_connector::KvCacheConnectorManager> kvCacheConnectorManager = nullptr,
bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128,
SizeType32 indexerKCacheIndexHeadDim = 0,
SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false,
Comment thread
dc3671 marked this conversation as resolved.
std::optional<LinearAttentionMetadata> linearAttentionMetadata = std::nullopt);

KVCacheManager(std::vector<SizeType32> const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock,
Expand All @@ -2027,7 +2043,7 @@ class KVCacheManager : public BaseKVCacheManager
bool copyOnpartialReuse = true,
std::shared_ptr<kv_connector::KvCacheConnectorManager> kvCacheConnectorManager = nullptr,
bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128,
SizeType32 indexerKCacheIndexHeadDim = 0,
SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false,
std::optional<LinearAttentionMetadata> linearAttentionMetadata = std::nullopt);

KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock,
Expand All @@ -2040,7 +2056,7 @@ class KVCacheManager : public BaseKVCacheManager
bool copyOnpartialReuse = true,
std::shared_ptr<kv_connector::KvCacheConnectorManager> kvCacheConnectorManager = nullptr,
bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128,
SizeType32 indexerKCacheIndexHeadDim = 0,
SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false,
std::optional<LinearAttentionMetadata> linearAttentionMetadata = std::nullopt);

KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock,
Expand All @@ -2049,7 +2065,7 @@ class KVCacheManager : public BaseKVCacheManager
int64_t stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false,
CacheType cacheType = CacheType::kSELF, bool enablePartialReuse = true, bool copyOnpartialReuse = true,
bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128,
SizeType32 indexerKCacheIndexHeadDim = 0,
SizeType32 indexerKCacheIndexHeadDim = 0, bool indexerKCacheUseFp4 = false,
std::optional<LinearAttentionMetadata> linearAttentionMetadata = std::nullopt);

~KVCacheManager() override = default;
Expand Down Expand Up @@ -2238,6 +2254,11 @@ class KVCacheManager : public BaseKVCacheManager
return mBlockManager.getIndexerKCacheQuantBlockSize();
}

[[nodiscard]] bool getIndexerKCacheUseFp4() const override
{
return mBlockManager.getIndexerKCacheUseFp4();
}

void removeToken(LlmRequest::RequestIdType requestId);
void rewindKVCache(LlmRequest::RequestIdType requestId, SizeType32 rewindLengths) override;

Expand Down
26 changes: 21 additions & 5 deletions cpp/include/tensorrt_llm/executor/dataTransceiverState.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -53,7 +53,7 @@ class CacheState final
std::vector<SizeType32> const& attentionLayerNumPerPP, nvinfer1::DataType dataType,
AttentionType attentionType = AttentionType::kDEFAULT, int kvFactor = 2, bool enableBlockReuse = false,
bool enablePartialReuse = false, bool hasIndexerKCache = false, SizeType32 indexerDimPerHead = 0,
SizeType32 indexerKCacheQuantBlockSize = 128)
SizeType32 indexerKCacheQuantBlockSize = 128, bool indexerKCacheUseFp4 = false)
: mModelConfig(std::move(modelConfig))
Comment thread
dc3671 marked this conversation as resolved.
, mParallelConfig{worldConfig.getTensorParallelism(), worldConfig.getPipelineParallelism(),
worldConfig.getContextParallelism(), worldConfig.enableAttentionDP(), worldConfig.getTensorParallelRank(),
Expand All @@ -66,14 +66,16 @@ class CacheState final
mHasIndexerKCache = hasIndexerKCache;
mIndexerDimPerHead = indexerDimPerHead;
mIndexerKCacheQuantBlockSize = indexerKCacheQuantBlockSize;
mIndexerKCacheUseFp4 = indexerKCacheUseFp4;
}

CacheState(std::vector<SizeType32> nbKvHeadPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock,
SizeType32 tensorParallelism, SizeType32 pipelineParallelism, SizeType32 contextParallelism,
std::vector<SizeType32> const& attentionLayerNumPerPP, nvinfer1::DataType dataType,
AttentionType attentionType = AttentionType::kDEFAULT, int kvFactor = 2, bool enableAttentionDP = false,
int DPrank = 0, int DPsize = 0, bool enableBlockReuse = false, bool enablePartialReuse = false,
bool hasIndexerKCache = false, SizeType32 indexerDimPerHead = 0, SizeType32 indexerKCacheQuantBlockSize = 128)
bool hasIndexerKCache = false, SizeType32 indexerDimPerHead = 0, SizeType32 indexerKCacheQuantBlockSize = 128,
bool indexerKCacheUseFp4 = false)
: mModelConfig{std::move(nbKvHeadPerLayer), sizePerHead, tokensPerBlock}
, mParallelConfig{tensorParallelism, pipelineParallelism, contextParallelism, enableAttentionDP, DPrank, DPsize,
attentionLayerNumPerPP}
Expand All @@ -85,14 +87,16 @@ class CacheState final
mHasIndexerKCache = hasIndexerKCache;
mIndexerDimPerHead = indexerDimPerHead;
mIndexerKCacheQuantBlockSize = indexerKCacheQuantBlockSize;
mIndexerKCacheUseFp4 = indexerKCacheUseFp4;
}

CacheState(SizeType32 nbAttentionLayers, SizeType32 nbKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock,
SizeType32 tensorParallelism, SizeType32 pipelineParallelism, SizeType32 contextParallelism,
std::vector<SizeType32> const& attentionLayerNumPerPP, nvinfer1::DataType dataType,
AttentionType attentionType = AttentionType::kDEFAULT, int kvFactor = 2, bool enableAttentionDP = false,
int DPrank = 0, int DPsize = 0, bool enableBlockReuse = false, bool enablePartialReuse = false,
bool hasIndexerKCache = false, SizeType32 indexerDimPerHead = 0, SizeType32 indexerKCacheQuantBlockSize = 128)
bool hasIndexerKCache = false, SizeType32 indexerDimPerHead = 0, SizeType32 indexerKCacheQuantBlockSize = 128,
bool indexerKCacheUseFp4 = false)
: mModelConfig{std::vector(nbAttentionLayers, nbKvHeads), sizePerHead, tokensPerBlock}
, mParallelConfig{tensorParallelism, pipelineParallelism, contextParallelism, enableAttentionDP, DPrank, DPsize,
attentionLayerNumPerPP}
Expand All @@ -104,13 +108,18 @@ class CacheState final
mHasIndexerKCache = hasIndexerKCache;
mIndexerDimPerHead = indexerDimPerHead;
mIndexerKCacheQuantBlockSize = indexerKCacheQuantBlockSize;
mIndexerKCacheUseFp4 = indexerKCacheUseFp4;
}

[[nodiscard]] bool operator==(kv_cache::CacheState const& other) const noexcept
{
return mModelConfig == other.mModelConfig && mParallelConfig == other.mParallelConfig
&& mAttentionConfig == other.mAttentionConfig && mDataType == other.mDataType
&& mRnnCacheState == other.mRnnCacheState;
&& mRnnCacheState == other.mRnnCacheState && mEnableBlockReuse == other.mEnableBlockReuse
&& mEnablePartialReuse == other.mEnablePartialReuse && mHasIndexerKCache == other.mHasIndexerKCache
&& mIndexerDimPerHead == other.mIndexerDimPerHead
&& mIndexerKCacheQuantBlockSize == other.mIndexerKCacheQuantBlockSize
&& mIndexerKCacheUseFp4 == other.mIndexerKCacheUseFp4;
}
Comment thread
dc3671 marked this conversation as resolved.

struct ModelConfig
Expand Down Expand Up @@ -247,6 +256,11 @@ class CacheState final
return mIndexerKCacheQuantBlockSize;
}

[[nodiscard]] bool getIndexerKCacheUseFp4() const
{
return mIndexerKCacheUseFp4;
}

// =========================================================================
// RNN/Mamba cache state (optional, present only for hybrid models)
// =========================================================================
Expand Down Expand Up @@ -308,6 +322,7 @@ class CacheState final
sstring << "hasIndexerKCache:" << mHasIndexerKCache << "\n";
sstring << "indexerDimPerHead:" << mIndexerDimPerHead << "\n";
sstring << "indexerKCacheQuantBlockSize:" << mIndexerKCacheQuantBlockSize << "\n";
sstring << "indexerKCacheUseFp4:" << mIndexerKCacheUseFp4 << "\n";
if (mRnnCacheState.has_value())
{
auto const& rnn = mRnnCacheState.value();
Expand Down Expand Up @@ -337,6 +352,7 @@ class CacheState final
bool mHasIndexerKCache{false};
SizeType32 mIndexerDimPerHead{0};
SizeType32 mIndexerKCacheQuantBlockSize{128};
bool mIndexerKCacheUseFp4{false};
// RNN/Mamba cache state (optional, for hybrid models)
std::optional<RnnCacheState> mRnnCacheState;
};
Expand Down
11 changes: 6 additions & 5 deletions cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
Expand Down Expand Up @@ -147,10 +147,11 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa
{
kvFactor = 1;
}
mCacheState = std::make_unique<executor::kv_cache::CacheState>(cacheStateModelCfg, worldConfig,
attentionLayerNumPerPP, dataType, attentionType, kvFactor, cacheManager->isEnableBlockReuse(),
cacheManager->isEnablePartialReuse(), cacheManager->isEnableIndexerKCache(),
cacheManager->getIndexerKCacheIndexHeadDim(), cacheManager->getIndexerKCacheQuantBlockSize());
mCacheState
= std::make_unique<executor::kv_cache::CacheState>(cacheStateModelCfg, worldConfig, attentionLayerNumPerPP,
dataType, attentionType, kvFactor, cacheManager->isEnableBlockReuse(), cacheManager->isEnablePartialReuse(),
cacheManager->isEnableIndexerKCache(), cacheManager->getIndexerKCacheIndexHeadDim(),
cacheManager->getIndexerKCacheQuantBlockSize(), cacheManager->getIndexerKCacheUseFp4());

if (mCacheState->getParallelConfig().mEnableAttentionDP)
{
Expand Down
Loading
Loading