[TRTLLM-16288][feat] Integrate sparse KV offload for DeepSeek-V4 Flash and GLM5.2 on KVCM2

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batch.cpp

@@ -106,8 +106,9 @@ Batch::Batch(KvCacheManager& manager, CUstream stream, std::vector<IndexTableSpe
     // A row that no member occupies must still read as "no such block" rather than as slot
     // zero, since the fetch has no other way to tell them apart.
     std::ranges::fill(_buf<int32_t>(BufName::kPageTables), kBadPageIndex.value());
-    // Watermarks of a free row are zero: a request with nothing eligible.
-    for (auto const name : {BufName::kHistory, BufName::kCapacity})
+    // Watermarks of a free row are zero: a request with nothing eligible. The eviction
+    // clock starts there too, which publish() advances off before any resolve reads it.
+    for (auto const name : {BufName::kHistory, BufName::kCapacity, BufName::kSparseClock})
     {
         std::ranges::fill(_buf<int32_t>(name), 0);
     }
@@ -1096,6 +1097,7 @@ void Batch::_layoutBufs(size_t scratchElements) noexcept
         {BufName::kPageTables, sizeof(int32_t) * _rowElements() * groups, alignof(int32_t)},
         {BufName::kHistory, sizeof(int32_t) * rows, alignof(int32_t)},
         {BufName::kCapacity, sizeof(int32_t) * rows, alignof(int32_t)},
+        {BufName::kSparseClock, sizeof(int32_t), alignof(int32_t)},
         {BufName::kLeases, sizeof(SparseLeaseIndex) * rows, alignof(SparseLeaseIndex)},
         {BufName::kScratch, sizeof(int32_t) * scratchElements, alignof(int32_t)},
     }};
@@ -1115,6 +1117,11 @@ SparseLeaseIndex const* Batch::_leaseMirror() const noexcept
     return reinterpret_cast<SparseLeaseIndex const*>(sparseLeases());
 }
 
+int32_t const* Batch::_sparseStepAgeMirror() const noexcept
+{
+    return reinterpret_cast<int32_t const*>(_mirrorAt(BufName::kSparseClock));
+}
+
 size_t Batch::_mirrorBytes() const noexcept
 {
     // From the histories through the leases: the page tables ahead of them are read from the
@@ -1220,8 +1227,7 @@ BatchRow Batch::addKvCache(KvCache& kvCache, std::optional<BatchRow> requestedRo
     else
     {
         row = *requestedRow;
-        TLLM_CHECK_WITH_INFO(
-            row != kNoBatchRow && row < mMaxBatchSize, "Batch::addKvCache: row is beyond max_batch_size");
+        TLLM_CHECK_WITH_INFO(row >= 0 && row < mMaxBatchSize, "Batch::addKvCache: row is beyond max_batch_size");
         if (mRows[row] != nullptr)
         {
             throw LogicError("Batch::addKvCache: row is already occupied");
@@ -1471,6 +1477,16 @@ void Batch::publish()
     }
     mFlushCountersDirty = false;
 
+    // The sparse eviction clock, which the resolves of this step then read. It wraps into
+    // SparseEntry's timestamp field, and a rebase every kRebaseStep steps is what keeps a
+    // timestamp from falling far enough behind the wrap to read as a recent one.
+    //
+    // Staged rather than passed by value, with the watermarks and for the same reason: a
+    // captured resolve replays against whatever the device holds, not against what the host
+    // held when the graph was recorded.
+    mSparseStepAge = (mSparseStepAge + 1) & SparseEntry::kTimestampMask;
+    _buf<int32_t>(BufName::kSparseClock).front() = static_cast<int32_t>(mSparseStepAge);
+
     // Refusing rather than waiting, because by this point a wait would be too late to help.
     // Everything that writes the staged rows has to be separated from the copy reading them:
     // the members' own mutations, which happened before this call, and the watermarks and
@@ -1677,6 +1693,7 @@ void Batch::prepareSparseWriteIndices(
     resolve.inFlightSlot = scratch.inFlightSlot;
     resolve.maxInFlightUnits = scratch.maxInFlightUnits;
     resolve.missCount = scratch.missCount;
+    resolve.stepAge = _sparseStepAgeMirror();
     // No selection on this path, so there is nothing to resolve into a slot.
     resolve.selected = nullptr;
     resolve.outSlot = nullptr;
@@ -1865,6 +1882,7 @@ void Batch::resolveSparseSelection(std::vector<BufferId> const& buffers, int32_t
     resolve.maxInFlightUnits = scratch.maxInFlightUnits;
     resolve.work = _sparseFetchWork(scratch, span.selectionSlot);
     resolve.missCount = scratch.missCount;
+    resolve.stepAge = _sparseStepAgeMirror();
     // Every row adds its own hits and misses as it resolves them.
     resolve.hitMiss = scratch.hitMiss;
     launchSparseResolve(resolve, stream);

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/batch.h

@@ -719,6 +719,7 @@ class Batch
         kPageTables = 0,
         kHistory,
         kCapacity,
+        kSparseClock,
         kLeases,
         kScratch,
         kCount
@@ -799,6 +800,9 @@ class Batch
     //! The same array where a kernel reads it, which is what the launches are handed.
     [[nodiscard]] SparseLeaseIndex const* _leaseMirror() const noexcept;
 
+    //! The staged eviction clock where a resolve reads it, for the same reason.
+    [[nodiscard]] int32_t const* _sparseStepAgeMirror() const noexcept;
+
     //! Bytes of the mirror: the staged block from the histories through the leases, which is
     //! everything a captured kernel reads. The page tables ahead of them are read from the
     //! staging, by the conversion, so carrying them across would serve nothing.
@@ -880,6 +884,19 @@ class Batch
     //! because the step's clear must not touch it.
     CudaUniqPtr<SparseHitMiss> mHitMiss;
 
+    //! The eviction clock every resolve of this Batch reads. Advanced in publish(), which
+    //! is the one place a step begins: a step resolves once per layer, so advancing per
+    //! resolve would age an entry by the layer count, and advancing from the watermarks
+    //! would quantise the clock by a life cycle's unit size.
+    //!
+    //! Shared by the life cycles. A single clock is what lets one Batch-wide advance serve
+    //! them, and nothing compares an entry of one against an entry of another.
+    //!
+    //! The host counter only; publish() stages it into BufName::kSparseClock, which is what
+    //! a resolve reads. See SparseResolveArgs::stepAge for why a kernel cannot take it by
+    //! value.
+    uint32_t mSparseStepAge = 0;
+
     //! Device storage behind every life cycle's per-BufferId copy descriptors.
     CudaUniqPtr<SparseCopyDesc> mCopyDesc;
 

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp

@@ -263,6 +263,34 @@ CudaUniqPtr<> deviceCopy(std::vector<int32_t> const& values)
     return owned;
 }
 
+CudaUniqPtr<> deviceCopyWords(std::vector<uint32_t> const& values)
+{
+    CUstream const stream = nullptr;
+    auto owned = allocateDeviceMemory(sizeof(uint32_t) * values.size());
+    if (values.empty())
+    {
+        return owned;
+    }
+    cuCheck(cuMemcpyHtoDAsync(
+        reinterpret_cast<CUdeviceptr>(owned.get()), values.data(), sizeof(uint32_t) * values.size(), stream));
+    cuCheck(cuStreamSynchronize(stream));
+    return owned;
+}
+
+std::vector<uint32_t> readBackWords(CudaUniqPtr<> const& buffer, size_t count)
+{
+    CUstream const stream = nullptr;
+    std::vector<uint32_t> out(count);
+    if (count == 0)
+    {
+        return out;
+    }
+    cuCheck(
+        cuMemcpyDtoHAsync(out.data(), reinterpret_cast<CUdeviceptr>(buffer.get()), sizeof(uint32_t) * count, stream));
+    cuCheck(cuStreamSynchronize(stream));
+    return out;
+}
+
 std::vector<int32_t> readBack(CudaUniqPtr<> const& buffer, size_t count)
 {
     CUstream const stream = nullptr;
@@ -280,22 +308,30 @@ std::vector<int32_t> readBack(CudaUniqPtr<> const& buffer, size_t count)
 } // namespace
 
 KvCacheIntrospection::SparseResolveProbe KvCacheIntrospection::sparseResolve(
-    std::vector<int32_t> const& residentOrdinal, std::vector<int32_t> const& lruOrder, int numLayers,
+    std::vector<int32_t> const& residentOrdinal, std::vector<int32_t> const& entryAge, int numLayers,
     int entriesPerLayer, std::vector<int32_t> const& selected, int selectionWidth,
     std::vector<int32_t> const& historyUnits, std::vector<int32_t> const& capacityUnits, int maxInFlightUnits,
     std::vector<int32_t> const& inactiveRows)
 {
     CUstream const stream = nullptr;
+    // Far enough into the clock that every age the probe accepts has a timestamp of its own
+    // below it, so the state it builds needs no wrap.
+    auto const stepAge = SparseEntry::kMaxAge;
     auto const numRows = static_cast<int>(historyUnits.size());
     auto const perRowEntries = static_cast<size_t>(numLayers) * entriesPerLayer;
     auto const entryCount = perRowEntries * static_cast<size_t>(numRows);
     auto const perRowSelection = static_cast<size_t>(numLayers) * selectionWidth;
     auto const selectionCount = perRowSelection * static_cast<size_t>(numRows);
     TLLM_CHECK_WITH_INFO(maxInFlightUnits >= 0, "sparseResolve: maxInFlightUnits counts units and cannot be negative");
     TLLM_CHECK_WITH_INFO(numRows != 0 && capacityUnits.size() == historyUnits.size()
-            && residentOrdinal.size() == entryCount && lruOrder.size() == entryCount
+            && residentOrdinal.size() == entryCount && (entryAge.empty() || entryAge.size() == entryCount)
             && selected.size() == selectionCount,
         "sparseResolve: input sizes must match the declared shape");
+    for (auto const age : entryAge)
+    {
+        TLLM_CHECK_WITH_INFO(age >= 0 && static_cast<uint32_t>(age) <= SparseEntry::kMaxAge,
+            "sparseResolve: an entry age must be within the range the eviction clock keeps distinguishable");
+    }
 
     auto selection = deviceCopy(selected);
     auto outSlot = deviceCopy(std::vector<int32_t>(selectionCount, 0));
@@ -310,19 +346,19 @@ KvCacheIntrospection::SparseResolveProbe KvCacheIntrospection::sparseResolve(
     auto const inFlightCount = countCount * static_cast<size_t>(maxInFlightUnits);
     auto inFlight = deviceCopy(std::vector<int32_t>(inFlightCount, kNoSparseSlot));
 
-    // One metadata block per lease, its resident ordinals then its recency order, which is
-    // how the arena lays a lease out. The probe takes and returns each array whole, so they
-    // are interleaved on the way in and separated on the way out.
-    auto const leaseInts = 2 * perRowEntries;
-    std::vector<int32_t> metadataHost(leaseInts * static_cast<size_t>(numRows));
-    for (int row = 0; row < numRows; ++row)
+    // One metadata block per lease: the two caller arrays packed into the entry words the
+    // arena actually holds. The probe takes and returns them separately, since an ordinal
+    // and an age are what a test states.
+    std::vector<uint32_t> metadataHost(entryCount);
+    for (size_t i = 0; i < entryCount; ++i)
     {
-        auto const from = perRowEntries * static_cast<size_t>(row);
-        auto const to = leaseInts * static_cast<size_t>(row);
-        std::copy_n(residentOrdinal.begin() + from, perRowEntries, metadataHost.begin() + to);
-        std::copy_n(lruOrder.begin() + from, perRowEntries, metadataHost.begin() + to + perRowEntries);
+        auto const ordinal = residentOrdinal[i];
+        auto const age = entryAge.empty() ? 0U : static_cast<uint32_t>(entryAge[i]);
+        metadataHost[i] = ordinal < 0
+            ? SparseEntry::kFreeWord
+            : SparseEntry::pack(SparseEntry::usedByte(stepAge - age, false), static_cast<uint32_t>(ordinal));
     }
-    auto metadata = deviceCopy(metadataHost);
+    auto metadata = deviceCopyWords(metadataHost);
 
     // Row r takes lease r, except an inactive row, which holds none for the kernel to skip.
     std::vector<int32_t> leaseHost(static_cast<size_t>(numRows));
@@ -333,11 +369,13 @@ KvCacheIntrospection::SparseResolveProbe KvCacheIntrospection::sparseResolve(
                                                                                              : row;
     }
     auto leases = deviceCopy(leaseHost);
+    // Device-resident, as the Batch publishes it: the resolve reads the clock rather than
+    // taking it by value.
+    auto stepAgeBuffer = deviceCopy(std::vector<int32_t>(1, static_cast<int32_t>(stepAge)));
 
     SparseArenaView arena;
     arena.metadataBase = metadata.get();
-    arena.leaseStride = sizeof(int32_t) * leaseInts;
-    arena.lruOffset = static_cast<int32_t>(perRowEntries);
+    arena.leaseStride = sizeof(uint32_t) * perRowEntries;
     arena.numLayers = numLayers;
     // One metadata row per layer: this probe drives a model where every layer selects for
     // itself, which is what a resolve sees whenever no selection_id is shared.
@@ -380,21 +418,24 @@ KvCacheIntrospection::SparseResolveProbe KvCacheIntrospection::sparseResolve(
         args.work.capacity = numRows * selectionWidth;
         args.work.numLists = 1;
         args.missCount = reinterpret_cast<int32_t*>(missCount.get()) + countOffset;
+        args.stepAge = reinterpret_cast<int32_t const*>(stepAgeBuffer.get());
 
         launchSparseResolve(args, stream);
     }
     cuCheck(cuStreamSynchronize(stream));
 
     SparseResolveProbe result;
-    auto const metadataOut = readBack(metadata, leaseInts * static_cast<size_t>(numRows));
+    auto const metadataOut = readBackWords(metadata, entryCount);
     result.residentOrdinal.resize(entryCount);
-    result.lruOrder.resize(entryCount);
-    for (int row = 0; row < numRows; ++row)
+    result.entryAge.resize(entryCount);
+    result.entryHit.resize(entryCount);
+    for (size_t i = 0; i < entryCount; ++i)
     {
-        auto const from = leaseInts * static_cast<size_t>(row);
-        auto const to = perRowEntries * static_cast<size_t>(row);
-        std::copy_n(metadataOut.begin() + from, perRowEntries, result.residentOrdinal.begin() + to);
-        std::copy_n(metadataOut.begin() + from + perRowEntries, perRowEntries, result.lruOrder.begin() + to);
+        auto const word = metadataOut[i];
+        auto const ordinal = SparseEntry::ordinal(word);
+        result.residentOrdinal[i] = ordinal;
+        result.entryAge[i] = ordinal < 0 ? -1 : static_cast<int32_t>(SparseEntry::age(word, stepAge));
+        result.entryHit[i] = ordinal < 0 ? 0 : static_cast<int32_t>(SparseEntry::wasHit(word));
     }
     result.outSlot = readBack(outSlot, selectionCount);
     result.missRow = readBack(missRow, selectionCount);
@@ -475,7 +516,6 @@ std::vector<uint8_t> KvCacheIntrospection::sparseUnitCopy(std::vector<uint8_t> c
     // This probe drives the addressing, not the metadata.
     arena.metadataBase = nullptr;
     arena.leaseStride = 0;
-    arena.lruOffset = 0;
     arena.numLayers = 0;
     arena.numSelections = 0;
     arena.entriesPerSelection = 0;
@@ -552,11 +592,13 @@ KvCacheIntrospection::SparseFlushList KvCacheIntrospection::sparseFlushList(std:
     auto outCount = deviceCopy(std::vector<int32_t>(static_cast<size_t>(numLayers), 0));
 
     // The scan spans the batch, so the hook stands one row in for a single request, holding
-    // one lease laid out as the arena lays one: its resident ordinals, then its recency
-    // order, which the scan does not read.
-    std::vector<int32_t> metadataHost(2 * entryCount, 0);
-    std::copy(residentOrdinal.begin(), residentOrdinal.end(), metadataHost.begin());
-    auto metadata = deviceCopy(metadataHost);
+    // one lease laid out as the arena lays one. The scan reads only the ordinals, so the
+    // age each word carries is left at zero.
+    std::vector<uint32_t> metadataHost(entryCount);
+    std::transform(residentOrdinal.begin(), residentOrdinal.end(), metadataHost.begin(),
+        [](int32_t ordinal)
+        { return ordinal < 0 ? SparseEntry::kFreeWord : SparseEntry::pack(0, static_cast<uint32_t>(ordinal)); });
+    auto metadata = deviceCopyWords(metadataHost);
     auto leases = deviceCopy(std::vector<int32_t>(1, 0));
 
     // A layer per metadata row, in order: the probe drives unshared selections.
@@ -566,8 +608,7 @@ KvCacheIntrospection::SparseFlushList KvCacheIntrospection::sparseFlushList(std:
 
     SparseArenaView arena;
     arena.metadataBase = metadata.get();
-    arena.leaseStride = sizeof(int32_t) * 2 * entryCount;
-    arena.lruOffset = static_cast<int32_t>(entryCount);
+    arena.leaseStride = sizeof(uint32_t) * entryCount;
     arena.numLayers = numLayers;
     arena.numSelections = numLayers;
     arena.entriesPerSelection = entriesPerLayer;

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h

@@ -110,13 +110,18 @@ class KvCacheIntrospection
     static std::pair<std::vector<int32_t>, std::vector<int32_t>> batchWatermarks(Batch const& batch);
 
     // Run the sparse resolve kernel over caller-supplied state, so the
-    // probe index, victim choice and recency rebuild can be exercised without a request.
+    // probe index and victim choice can be exercised without a request.
     // Synchronises. Inputs are [layer][entry] and [layer][selectionWidth] row-major.
     struct SparseResolveProbe
     {
         std::vector<int32_t> residentOrdinal; //!< [layer][entry] after the resolve
-        std::vector<int32_t> lruOrder;        //!< [layer][entry] after the resolve
-        std::vector<int32_t> outSlot;         //!< [layer][width] entry per selection
+        //! [layer][entry] steps since each entry's last use, and whether that use found it
+        //! already resident. Together they are the eviction order: an entry is given away
+        //! before an entry of the same age that was hit, and before any younger one. -1 and
+        //! false where the entry is free.
+        std::vector<int32_t> entryAge;
+        std::vector<int32_t> entryHit;
+        std::vector<int32_t> outSlot; //!< [layer][width] entry per selection
         //! The compacted work list, one per layer: [layer][rows * width], of which the
         //! first missTotal[layer] are valid, in whatever order the CTAs raced to.
         std::vector<int32_t> missRow;
@@ -127,8 +132,11 @@ class KvCacheIntrospection
         std::vector<int32_t> inFlightSlot; //!< [layer][maxInFlightUnits]
     };
 
+    //! @param entryAge [layer][entry] steps since each entry's last use, as the state the
+    //! resolve starts from. Empty leaves every entry at zero, which makes the eviction order
+    //! among them arbitrary.
     static SparseResolveProbe sparseResolve(std::vector<int32_t> const& residentOrdinal,
-        std::vector<int32_t> const& lruOrder, int numLayers, int entriesPerLayer, std::vector<int32_t> const& selected,
+        std::vector<int32_t> const& entryAge, int numLayers, int entriesPerLayer, std::vector<int32_t> const& selected,
         int selectionWidth, std::vector<int32_t> const& historyUnits, std::vector<int32_t> const& capacityUnits,
         int maxInFlightUnits, std::vector<int32_t> const& inactiveRows);
 

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp

@@ -190,7 +190,7 @@ void KvCacheManager::_computeSparseCacheShape()
 
         // Layers of one life cycle must currently declare the same sparse_cache. The arena
         // is per life cycle and its metadata is one stride: a selection's rows are found at
-        // `residentOrdinal + slot * entriesPerSelection`, so a single entry count serves
+        // `entry + slot * entriesPerSelection`, so a single entry count serves
         // them all, and cellsPerBuffer follows from maxCachedUnits and unitsPerPage alike.
         // See SPARSE_ATTENTION_DESIGN.md.
         auto& shape = mSparseCacheShape[lcId];