diff --git a/.gitignore b/.gitignore index 6c115ad2a40c..8754313d88fb 100644 --- a/.gitignore +++ b/.gitignore @@ -81,6 +81,7 @@ ad-test-workspace/ */tllm_debug/** *.patch !cpp/tensorrt_llm/deep_ep/*.patch +!3rdparty/patches/*.patch !3rdparty/vendor_patches/*.patch examples/disaggregated/slurm/benchmark/logs/ scripts/attribution/data/checksum_to_paths.yml diff --git a/3rdparty/fetch_content.json b/3rdparty/fetch_content.json index 25f64fbda4f4..4db7c8262044 100644 --- a/3rdparty/fetch_content.json +++ b/3rdparty/fetch_content.json @@ -39,10 +39,9 @@ { "name": "deepgemm", "git_repository": "https://github.com/deepseek-ai/DeepGEMM", - "git_tag": "8b1392b978f5a03c828dd1711090d7fb50958b8a", + "git_tag": "b64107f2b9599ca76445b7f62eedf66bae1d095b", "git_submodules_recurse": true, - "source_subdir": "dont-add-this-project-with-add-subdirectory", - "patch_file": "patches/deepgemm_configurable_barrier_timeout.patch" + "source_subdir": "dont-add-this-project-with-add-subdirectory" }, { "name": "eigen", @@ -57,6 +56,12 @@ "git_submodules_recurse": true, "source_subdir": "dont-add-this-project-with-add-subdirectory" }, + { + "name": "fmt", + "git_repository": "https://github.com/fmtlib/fmt", + "git_tag": "553ec11ec06fbe0beebfbb45f9dc3c9eabd83d28", + "source_subdir": "dont-add-this-project-with-add-subdirectory" + }, { "name": "googlebenchmark", "git_repository": "https://github.com/google/benchmark", diff --git a/3rdparty/patches/deepgemm_configurable_barrier_timeout.patch b/3rdparty/patches/deepgemm_configurable_barrier_timeout.patch deleted file mode 100644 index b9d89bcf6d4c..000000000000 --- a/3rdparty/patches/deepgemm_configurable_barrier_timeout.patch +++ /dev/null @@ -1,82 +0,0 @@ ---- a/csrc/jit/compiler.hpp -+++ b/csrc/jit/compiler.hpp -@@ -56,12 +56,18 @@ public: - "--ptxas-options=--register-usage-level=10", - get_env("DG_JIT_CPP_STANDARD", 20)); - if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PTXAS_VERBOSE", 0) or get_env("DG_JIT_PTXAS_CHECK", 0)) - flags += " --ptxas-options=--verbose,--warn-on-local-memory-usage"; - if (get_env("DG_JIT_WITH_LINEINFO", 0)) - flags += " -Xcompiler -rdynamic -lineinfo"; -+ -+ // Allow overriding the in-kernel barrier timeout, see `comm/barrier.cuh` -+ // NOTES: `flags` is part of the kernel signature, so changing this value -+ // naturally invalidates the JIT cache -+ if (const auto barrier_timeout_seconds = get_env("DG_JIT_BARRIER_TIMEOUT_SECONDS", 0); barrier_timeout_seconds > 0) -+ flags += fmt::format(" -DDG_BARRIER_TIMEOUT_SECONDS={}", barrier_timeout_seconds); - } - - virtual ~Compiler() = default; - - std::filesystem::path make_tmp_dir() const { - return make_dirs(cache_dir_path / "tmp"); ---- a/deep_gemm/include/deep_gemm/comm/barrier.cuh -+++ b/deep_gemm/include/deep_gemm/comm/barrier.cuh -@@ -5,14 +5,22 @@ - #include - #include - #include - - namespace deep_gemm::comm { - --// 60s timeout, at 2 GHz --constexpr int64_t kNumTimeoutCycles = 60ll * 2000000000ll; -+// Barrier timeout in seconds, overridable at JIT compile time. The host-side compiler -+// turns the `DG_JIT_BARRIER_TIMEOUT_SECONDS` environment variable into this definition, -+// so slow-but-healthy rendezvous (e.g. large expert-parallel groups whose first -+// collective follows a long warm-up) can be distinguished from a genuine hang. -+#ifndef DG_BARRIER_TIMEOUT_SECONDS -+#define DG_BARRIER_TIMEOUT_SECONDS 60 -+#endif -+ -+// Timeout in cycles, at 2 GHz -+constexpr int64_t kNumTimeoutCycles = static_cast(DG_BARRIER_TIMEOUT_SECONDS) * 2000000000ll; - - CUTLASS_DEVICE void cluster_sync_with_relaxed_arrive() { - // Perform cluster_sync with `barrier.cluster.arrive.relaxed` - // This is slightly faster than `cute::cluster_sync` but has weaker memory ordering guarantee - cute::cluster_arrive_relaxed(); - cute::cluster_wait(); -@@ -31,14 +39,14 @@ CUTLASS_DEVICE void grid_sync(const layout::Workspace& workspace, - count_ptr, sm_idx == 0 ? (kFinishSumTag - (kNumSMs - 1)) : 1); - uint32_t new_value; - const auto start_clock = clock64(); - do { - new_value = ptx::ld_acq(count_ptr); - if (clock64() - start_clock >= kNumTimeoutCycles) { -- printf("DeepGEMM grid sync timeout: sm=%u, thread=%u, grid_sync_idx=%u, old=%u, current=%u, expected_tag=%u\n", -- sm_idx, thread_idx, kGridSyncIndex, old_value, new_value, old_value ^ kFinishSumTag); -+ printf("DeepGEMM grid sync timeout (%ds): sm=%u, thread=%u, grid_sync_idx=%u, old=%u, current=%u, expected_tag=%u\n", -+ DG_BARRIER_TIMEOUT_SECONDS, sm_idx, thread_idx, kGridSyncIndex, old_value, new_value, old_value ^ kFinishSumTag); - DG_DEVICE_ASSERT(false and "Grid sync timeout"); - } - } while (((new_value ^ old_value) & kFinishSumTag) == 0); - } - sync_scope(); - } -@@ -72,14 +80,14 @@ CUTLASS_DEVICE void nvlink_barrier(const layout::Workspace& workspace, - if (thread_idx == 0) { - ptx::red_add(counter_ptr, 1); - const int target = signal_sign ? 0 : static_cast(kNumRanks); - const auto start_clock = clock64(); - while (ptx::ld_acq_sys(signal_ptr) != target) { - if (clock64() - start_clock >= kNumTimeoutCycles) { -- printf("DeepGEMM NVLink barrier timeout: rank=%d, counter=%d, signal=%d, target=%d, phase=%d, sign=%d, tag=%d\n", -- sym_buffer.rank_idx, *counter_ptr, ptx::ld_acq_sys(signal_ptr), target, signal_phase, signal_sign, kTag); -+ printf("DeepGEMM NVLink barrier timeout (%ds): rank=%d, counter=%d, signal=%d, target=%d, phase=%d, sign=%d, tag=%d\n", -+ DG_BARRIER_TIMEOUT_SECONDS, sym_buffer.rank_idx, *counter_ptr, ptx::ld_acq_sys(signal_ptr), target, signal_phase, signal_sign, kTag); - DG_DEVICE_ASSERT(false and "NVLink barrier timeout"); - } - } - } - } - diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp index d38b7ec2d7ef..1b015fc48af6 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp @@ -482,7 +482,8 @@ std::vector> Block::clearStaleBlocksAfterPageUnlink( // If this is a full-attention block or a sink block: evict subtree. // Mirrors Python: pages = remove_subtree(self) - if (alc && (!alc->windowSize.has_value() || block.ordinal() < BlockOrdinal{alc->numSinkBlocks})) + if (alc && alc->requiresPrefixCoverage() + && (!alc->windowSize.has_value() || block.ordinal() < BlockOrdinal{alc->numSinkBlocks})) { pruneStart = block.prev; detachedBlocks.push_back(removeSubtree(block)); @@ -887,6 +888,10 @@ std::vector BlockRadixTree::pruneMatch( bool shortened = false; for (auto [lcId, attn] : attnLcs) { + if (!attn->requiresPrefixCoverage()) + { + continue; + } auto const staleRange = attn->getStaleRange(numTok, mTokensPerBlock); int const staleBeg = staleRange.beg.value(); int const staleEnd = staleRange.end.value(); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h index bb40369610e3..a6cf8530f21b 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h @@ -137,6 +137,19 @@ inline void validateNoDuplicateBufferRoles(std::vector const& buff } // namespace detail +//! Controls whether an attention lifecycle participates in prefix matching and page reuse. +enum class AttentionReusePolicy +{ + //! Commit reusable pages and require live-window coverage when matching a prefix. + REQUIRED = 0, + //! Keep request-local working pages; never publish or reuse them through the radix tree. + PRIVATE = 1, + //! Commit reusable pages without constraining the REQUIRED prefix match. The caller selects + //! available groups at resume; missing or unselected groups use fresh pages for recomputation. + //! These pages may be evicted independently without removing the REQUIRED prefix. + OPTIONAL = 2, +}; + struct AttentionLayerConfig { static constexpr LayerType type = LayerType::ATTENTION; @@ -150,6 +163,9 @@ struct AttentionLayerConfig // nullopt or 0 = no sink tokens. std::optional numSinkTokens; + // Private working pages never constrain prefix matching or enter the tree. + AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; + [[nodiscard]] std::optional windowSize() const noexcept { return slidingWindowSize; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index 3a75bbaf2a9e..5a0aaba48fa1 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -102,6 +102,7 @@ KvCache::KvCache(KvCacheManager& manager, ReuseScope reuseScope, std::optional(mHistoryLength)); @@ -150,6 +151,128 @@ StorageManager* KvCache::storageManager() const return &mManager->storage(); } +void KvCache::_collectReuseCandidates() +{ + auto const end = numCommittedTokens(); + auto const& lifeCycles = mManager->lifeCycles(); + mOptionalCandidates.resize(lifeCycles.size()); + for (LifeCycleId lc{0}; lc < lifeCycles.size(); ++lc) + { + auto const* attention = std::get_if(&lifeCycles[lc]); + auto const policy = attention ? attention->reusePolicy : AttentionReusePolicy::REQUIRED; + ReuseGroupStatus status{lc, policy, end, end > 0 && policy == AttentionReusePolicy::REQUIRED, {}}; + if (status.complete) + { + if (!attention) + { + status.coverage.emplace_back(end, end); // Exact recurrent-state endpoint. + } + else + { + auto stale = getStaleRange(lifeCycles[lc], end, mTokensPerBlock); + if (stale.beg > BlockOrdinal{0}) + { + status.coverage.emplace_back(0, std::min(end, stale.beg.value() * mTokensPerBlock)); + } + if (stale.end.value() * mTokensPerBlock < end) + { + status.coverage.emplace_back(stale.end.value() * mTokensPerBlock, end); + } + } + } + // OPTIONAL checks coverage at the REQUIRED endpoint without shortening it. + // Use the ordinary SWA live ranges, including sinks and multi-page windows. + if (attention && attention->allowsCheckpoint() && end > 0 && end % mTokensPerBlock == 0 + && attention->windowSize) + { + auto const stale = getStaleRange(lifeCycles[lc], end, mTokensPerBlock); + auto& candidates = mOptionalCandidates[lc]; + status.complete = true; + auto collectRange = [&](BlockOrdinal beg, BlockOrdinal limit) + { + for (auto ordinal = beg; ordinal < limit; ++ordinal) + { + auto const& block = mBlocks[ordinal].treeBlock; + auto* page = block && !block->isOrphan() ? block->getPage(lc) : nullptr; + if (!page || page->numTokensInBlock < mTokensPerBlock) + { + status.complete = false; + return; + } + candidates.emplace_back(ordinal, page->hold()); + } + if (beg < limit) + { + status.coverage.emplace_back(beg.value() * mTokensPerBlock, limit.value() * mTokensPerBlock); + } + }; + collectRange(BlockOrdinal{0}, stale.beg); + if (status.complete) + { + collectRange(stale.end, BlockOrdinal{end / mTokensPerBlock}); + } + if (!status.complete) + { + candidates.clear(); + status.coverage.clear(); + } + } + mReuseStatus.push_back(std::move(status)); + } +} + +bool KvCache::_needsFreshPrefixPages(LifeCycleId lc) const +{ + return !requiresPrefixCoverage(mManager->lifeCycles()[lc]) + && std::find(mSelectedOptionalGroups.begin(), mSelectedOptionalGroups.end(), lc) + == mSelectedOptionalGroups.end(); +} + +void KvCache::_selectOptionalGroups(std::vector const& groups) +{ + // Validate everything before touching the previous attempt's mappings. + for (auto lc : groups) + { + if (lc < LifeCycleId{0} || lc >= mOptionalCandidates.size() + || mReuseStatus[lc.value()].policy != AttentionReusePolicy::OPTIONAL || !mReuseStatus[lc.value()].complete) + { + throw std::invalid_argument("Selected OPTIONAL group has no reusable state at the claimed prefix"); + } + } + for (auto lc : mSelectedOptionalGroups) + { + for (auto const& [ordinal, holder] : mOptionalCandidates[lc]) + { + mBlocks[ordinal].pages[kDefaultBeamIndex][lc] = std::monostate{}; + } + } + for (auto lc : groups) + { + auto const stale = _getStaleRange(mHistoryLength, mManager->lifeCycles()[lc]); + for (auto const& [ordinal, holder] : mOptionalCandidates[lc]) + { + if (!stale.contains(ordinal)) + { + mBlocks[ordinal].pages[kDefaultBeamIndex][lc] = holder; + } + } + } + mSelectedOptionalGroups = groups; + std::sort(mSelectedOptionalGroups.begin(), mSelectedOptionalGroups.end()); + mSelectedOptionalGroups.erase( + std::unique(mSelectedOptionalGroups.begin(), mSelectedOptionalGroups.end()), mSelectedOptionalGroups.end()); + for (LifeCycleId lc{0}; lc < mOptionalCandidates.size(); ++lc) + { + if (std::find(groups.begin(), groups.end(), lc) == groups.end() + && mReuseStatus[lc.value()].policy == AttentionReusePolicy::OPTIONAL) + { + mOptionalCandidates[lc].clear(); + mReuseStatus[lc.value()].complete = false; + mReuseStatus[lc.value()].coverage.clear(); + } + } +} + std::vector KvCache::_activePages() const { std::vector result; @@ -172,6 +295,10 @@ std::vector KvCache::_activePages() const // Attention lifecycle: yield non-stale blocks (sink + window). LifeCycle const& lc = lcs.getLifeCycle(lcId); + if (mNeverResumed && _needsFreshPrefixPages(lcId)) + { + continue; // Allocated by first resume; private pages are never claimed from the tree. + } auto staleRange = _getStaleRange(mHistoryLength, lc); BlockOrdinal staleBeg = staleRange.beg; BlockOrdinal staleEnd = staleRange.end; @@ -378,7 +505,8 @@ void KvCache::activate() } } -bool KvCache::resume(std::optional stream, std::optional isDecoding) +bool KvCache::resume(std::optional stream, std::optional isDecoding, + std::optional> optionalReuseGroups) { KVCM2_API_GUARD(); TLLM_CHECK(mStatus == Status::SUSPENDED); @@ -397,6 +525,14 @@ bool KvCache::resume(std::optional stream, std::optional isDecod throw std::invalid_argument("Cannot return a decoding cache to prefill"); auto restorePhase = FuncGuard([&]() { mIsDecoding = oldIsDecoding; }); mIsDecoding = isDecoding.value_or(mIsDecoding); + if (mNeverResumed) + { + _selectOptionalGroups(optionalReuseGroups.value_or(std::vector{})); + } + else if (optionalReuseGroups.has_value()) + { + throw std::invalid_argument("OPTIONAL reuse selection is only valid before the first resume"); + } // Check utilization against threshold. auto const utilizations = mManager->storage().getUtilization(kHotLevel); @@ -410,6 +546,26 @@ bool KvCache::resume(std::optional stream, std::optional isDecod auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); LifeCycleId numLc = storageMgr.numLifeCycles(); + // Snapshot source levels before allocation or activation can migrate the selected pages. + ReusedBlocksByLevelByLifeCycle optionalReusedByLevel; + if (mNeverResumed && _shouldRecordStats()) + { + for (auto lc : mSelectedOptionalGroups) + { + auto& byLevel = optionalReusedByLevel[lc]; + byLevel.full = CountsByLevel(storageMgr.numCacheLevels(), 0); + byLevel.partial = CountsByLevel(storageMgr.numCacheLevels(), 0); + auto const stale = _getStaleRange(mHistoryLength, mManager->lifeCycles()[lc]); + for (auto const& [ordinal, holder] : mOptionalCandidates[lc]) + { + if (!stale.contains(ordinal)) + { + ++byLevel.full.at(holder->page->cacheLevel); + } + } + } + } + // Pre-allocate GPU slots for deferred copies (partial blocks + SSM) and scratch slots // before locking, so we never end up in a state where pages are locked but we can't allocate. TypedVec> deferredSlots(numLc); @@ -432,6 +588,22 @@ bool KvCache::resume(std::optional stream, std::optional isDecod } } }); + TypedVec> privateOrdinals(numLc); + TypedVec> privateSlots(numLc); + auto releasePrivateSlots = FuncGuard( + [&]() + { + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + for (auto& slot : privateSlots[lc]) + { + if (slot.hasValidSlot()) + { + storageMgr.releaseSlot(lc, kHotLevel, std::move(slot)); + } + } + } + }); // Compute scratch slot deltas UNCONDITIONALLY (mirrors Python: _take_excess_scratch_slots // is called outside _never_resumed). @@ -448,8 +620,24 @@ bool KvCache::resume(std::optional stream, std::optional isDecod for (LifeCycleId lc{0}; lc < numLc; ++lc) { bool isSsm = ssmLcId.has_value() && lc == *ssmLcId; - if (isSsm || hasPartial) + auto const& lifeCycle = mManager->lifeCycles()[lc]; + if (_needsFreshPrefixPages(lc)) + { + auto const stale = _getStaleRange(mHistoryLength, lifeCycle); + for (BlockOrdinal ord{0}; ord < stale.beg; ++ord) + { + privateOrdinals[lc].push_back(ord); + } + for (BlockOrdinal ord = stale.end; ord < mBlocks.size(); ++ord) + { + privateOrdinals[lc].push_back(ord); + } + numSlotsNeeded[lc] += static_cast(privateOrdinals[lc].size()); + } + else if (isSsm || hasPartial) + { numSlotsNeeded[lc] += 1; + } } } @@ -486,7 +674,14 @@ bool KvCache::resume(std::optional stream, std::optional isDecod if (!tmpSlots[lc].empty()) { // Mirrors Python: `if self._never_resumed and (... SsmLifeCycle or has_partial):` - bool needsDeferred = mNeverResumed && ((ssmLcId.has_value() && lc == *ssmLcId) || hasPartial); + bool const isPrivate = _needsFreshPrefixPages(lc); + bool needsDeferred + = mNeverResumed && !isPrivate && ((ssmLcId.has_value() && lc == *ssmLcId) || hasPartial); + for (size_t i = 0; i < privateOrdinals[lc].size(); ++i) + { + privateSlots[lc].push_back(std::move(tmpSlots[lc].back())); + tmpSlots[lc].pop_back(); + } if (needsDeferred) { // Python uses pop() here: reserve one slot for deferred copy, then treat the rest as scratch. @@ -524,6 +719,35 @@ bool KvCache::resume(std::optional stream, std::optional isDecod mStatus = Status::ACTIVE; auto rollbackActivation = FuncGuard([&]() { _deactivate(); }); + // Private recovery pages have no reusable source. Their contents are + // initialized by the caller before inference, after all allocations succeed. + std::vector privateReadyEvents; + for (auto const& slots : privateSlots) + { + for (auto const& slot : slots) + { + privateReadyEvents.push_back(slot.readyEvent.handle()); + } + } + streamWaitEvents(reinterpret_cast(cudaStream()), std::move(privateReadyEvents)); + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + for (size_t i = 0; i < privateSlots[lc].size(); ++i) + { + auto const ordinal = privateOrdinals[lc][i]; + auto page = makeShared(*this, ordinal, lc, kHotLevel, kDefaultBeamIndex); + page->setSlot(privateSlots[lc][i]); + mBlocks[ordinal].pages[kDefaultBeamIndex][lc] + = page->lock(*this, kDefaultBeamIndex, ordinal, lc, /*skipWait=*/true); + if (mPendingStats.recordAllocationRange(lc, ordinal, ordinal + 1, + /*beamWidth=*/1, /*countAsMissed=*/true, /*countAsGeneration=*/false, _shouldRecordManagerStats(), + _shouldRecordRequestStats())) + { + mManager->markStatsDirty(id); + } + } + } + // Deferred copy: for partial blocks and SSM, copy from now-locked source pages // to pre-allocated GPU slots, then unlock sources and replace with new pages. if (mNeverResumed) @@ -637,6 +861,19 @@ bool KvCache::resume(std::optional stream, std::optional isDecod mBlocks[lastOrdinal].treeBlock = nullptr; } + if (mNeverResumed && _shouldRecordStats()) + { + for (auto const& [lc, byLevel] : optionalReusedByLevel) + { + int const reused = static_cast(countsByLevelTotal(byLevel.full)); + if (mPendingStats.recordReuse( + lc, reused, 0, byLevel, _shouldRecordManagerStats(), _shouldRecordRequestStats())) + { + mManager->markStatsDirty(id); + } + } + } + mOptionalCandidates.clear(); // Deferred copies survive a failed decode admission and must not be repeated on retry. mNeverResumed = false; if (mIsDecoding) @@ -719,6 +956,7 @@ void KvCache::suspend() KVCM2_API_GUARD(); auto const apiLock = mManager->lockExclusive(); TLLM_CHECK_DEBUG(mStatus == Status::ACTIVE); + mOptionalCandidates.clear(); TLLM_CHECK_DEBUG(_checkSanity()); TLLM_CHECK_DEBUG(!mFinishEvent.has_value()); @@ -776,6 +1014,7 @@ void KvCache::close() if (mStatus == Status::CLOSED) return; + mOptionalCandidates.clear(); discardPendingStats(); stopCommitting(); TLLM_CHECK_DEBUG(_checkSanity()); @@ -1221,7 +1460,10 @@ void KvCache::_snapshotPartialBlockToTree(BlockOrdinal ordinal, bool commitSsm) std::vector attachedLcs; for (auto const& [lcIdx, attn] : mManager->lifeCycles().attentionLifeCycles()) { - (void) attn; + if (!attn->requiresPrefixCoverage()) + { + continue; + } auto& bp = beamBlock[lcIdx]; if (blockPageIsNull(bp) || treeBlock->pageCoverage(lcIdx) >= numTokens) { @@ -1874,8 +2116,13 @@ std::vector KvCache::_unlockStaleBlocks(int newHistoryLeng { auto& sb = mBlocks[ord]; bool isCommitted = sb.isCommitted(); - bool holdForCommit - = !mManager->commitMinSnapshot() && !isCommitted && (mCommitState == CommitState::ALLOWED); + // A partial tail can expire pages still needed at the preceding + // aligned checkpoint. Keep its entire live window until publication. + int const alignedEnd = newHistoryLength / mTokensPerBlock * mTokensPerBlock; + bool const keepAlignedCheckpoint = alc.allowsCheckpoint() + && ord < BlockOrdinal{alignedEnd / mTokensPerBlock} && !_getStaleRange(alignedEnd, lc).contains(ord); + bool holdForCommit = !alc.isPrivate() && (!mManager->commitMinSnapshot() || keepAlignedCheckpoint) + && !isCommitted && (mCommitState == CommitState::ALLOWED); for (BeamIndex bi{0}; bi < sb.pages.size(); ++bi) { @@ -1923,8 +2170,10 @@ TypedVec KvCache::_takeUncommittedPage( TypedVec result(numLc, TakenPage{nullptr, false}); for (LifeCycleId lc{0}; lc < numLc; ++lc) { - if (skipLc.has_value() && lc == *skipLc) + if ((skipLc.has_value() && lc == *skipLc) || isPrivateLifeCycle(mManager->lifeCycles()[lc])) + { continue; + } auto& bp = sb.pages[beamIdx][lc]; if (auto* lock = std::get_if(&bp)) { @@ -2042,6 +2291,10 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) { if (ssmLcId.has_value() && lc == *ssmLcId) continue; + if (isPrivateLifeCycle(mManager->lifeCycles()[lc])) + { + continue; + } auto& bp = sb.pages[kDefaultBeamIndex][lc]; if (blockPageIsNull(bp)) continue; @@ -2580,6 +2833,10 @@ void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match) if (ssmLcId.has_value() && lcId == *ssmLcId) continue; + if (!requiresPrefixCoverage(allLc[lcId])) + { + continue; + } auto staleRange = getStaleRange(allLc[lcId], numTokens, mTokensPerBlock); BlockOrdinal staleStart = staleRange.beg; BlockOrdinal staleEnd = staleRange.end; @@ -2746,8 +3003,19 @@ SharedPtr const& KvCache::_getTreeBlock(BlockOrdinal ordinal) const auto page = blockPageGetPage(beamBlock[lcId]); // committed->block may differ from `ret`: a page can be moved to a longer // sibling block, or replaced there by one with larger token coverage. - auto committed = dynamicPointerCast(page); - TLLM_CHECK(committed); + auto const* attention = std::get_if(&mManager->lifeCycles()[lcId]); + if (attention && attention->allowsCheckpoint()) + { + TLLM_CHECK(dynamicPointerCast(page) || dynamicPointerCast(page)); + } + else if (!requiresPrefixCoverage(mManager->lifeCycles()[lcId])) + { + TLLM_CHECK(dynamicPointerCast(page)); + } + else + { + TLLM_CHECK(dynamicPointerCast(page)); + } } } } @@ -2804,6 +3072,12 @@ bool KvCache::_checkSanity() const auto const& staleRange = staleRanges[lc]; auto const& scratchRange = scratchRangesVec[lc]; + bool const isPrivate = !requiresPrefixCoverage(lcs[lc]); + if (_needsFreshPrefixPages(lc) && mNeverResumed) + { + TLLM_CHECK_DEBUG(blockPageIsNull(bp)); + continue; + } if (scratchRange.contains(ordinal)) { @@ -2812,7 +3086,7 @@ bool KvCache::_checkSanity() const } else if (staleRange.beg <= ordinal && ordinal < staleRange.end) { - if (isCommitted || mCommitState != CommitState::ALLOWED) + if (isPrivateLifeCycle(lcs[lc]) || isCommitted || mCommitState != CommitState::ALLOWED) { TLLM_CHECK_DEBUG(blockPageIsNull(bp)); } @@ -2841,7 +3115,17 @@ bool KvCache::_checkSanity() const if (!blockPageIsNull(bp)) { auto page = blockPageGetPage(bp); - TLLM_CHECK_DEBUG(isCommitted == (dynamicPointerCast(page) != nullptr)); + auto const* attention = std::get_if(&lcs[lc]); + auto committed = dynamicPointerCast(page); + if (attention && attention->allowsCheckpoint() && committed) + { + TLLM_CHECK_DEBUG(isCommitted && committed->numTokensInBlock > 0 + && committed->numTokensInBlock <= mTokensPerBlock); + } + else + { + TLLM_CHECK_DEBUG((isCommitted && !isPrivate) == (committed != nullptr)); + } } } } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h index 90fe45a5fc49..8444d3049f72 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h @@ -88,8 +88,8 @@ struct SeqBlock bool ret = treeBlock != nullptr; if (TLLM_UNLIKELY(gDebug)) { - // When committed: must have one canonical beam, and all non-null - // pages must be CommittedPage. + // Committed blocks have one canonical beam; publication concerns reusable pages. + // Request-private state stays uncommitted. if (ret) { TLLM_CHECK(pages.size() == BeamIndex{1}); @@ -98,7 +98,7 @@ struct SeqBlock if (!blockPageIsNull(bp)) { auto pg = blockPageGetPage(bp); - TLLM_CHECK(!pg || dynamicPointerCast(pg)); + TLLM_CHECK(!pg || dynamicPointerCast(pg) || pg->isPrivate()); } } else @@ -204,6 +204,25 @@ class PageStorageSnapshot std::vector mReadyEvents; }; +// Availability at the claimed prefix, indexed by LayerGroupId. OPTIONAL +// candidates are held until the first successful resume or close. +// Describes the initial prefix claim, not subsequent request progress. +struct ReuseGroupStatus +{ + // Lifecycle group returned by getLayerGroupId(layerId), not a physical pool index. + LayerGroupId groupId; + // REQUIRED constrains the prefix match; OPTIONAL is selected at resume; PRIVATE is never reused. + AttentionReusePolicy policy; + // Exclusive token endpoint of the prefix claimed by REQUIRED groups, shared by all groups. + int endpoint; + // All state needed by this group at endpoint is available; does not imply GPU residency or selection. + // False for an empty prefix, PRIVATE groups, or missing/discarded OPTIONAL candidates. + bool complete; + // Retained token ranges [begin, end), including block-rounded SWA windows and sinks. + // SSM uses (endpoint, endpoint) for its exact state; empty when complete is false. + std::vector> coverage; +}; + // --------------------------------------------------------------------------- // KvCache — manages the per-sequence KV cache state. // Mirrors Python's _KVCache. @@ -243,7 +262,8 @@ class KvCache : public std::enable_shared_from_this // isDecoding defaults to the current phase. A cache starts in prefill and cannot return to it // after decode admission. Set true when admitting a suspended request directly to decode. // Returns false if utilization too high or out of memory. - bool resume(std::optional stream = std::nullopt, std::optional isDecoding = std::nullopt); + bool resume(std::optional stream = std::nullopt, std::optional isDecoding = std::nullopt, + std::optional> optionalReuseGroups = std::nullopt); // Enter decode only after prefill has submitted its final KV accesses. Offloads complete sparse // history, deferring pages still needed on GPU by another owner. Retries deferred pages even @@ -405,6 +425,11 @@ class KvCache : public std::enable_shared_from_this //! survive, so only the tail stops counting as a local hit. void dropPartialBlockCachedTokenAttribution(); + std::vector const& reuseStatus() const noexcept + { + return mReuseStatus; + } + // Internal diagnostic: prefix supported by the attention pages alone, // before recurrent-state (SSM) snapshot pruning shortened the reuse. int numReusableTokensBeforeHybridPruning() const noexcept @@ -782,6 +807,14 @@ class KvCache : public std::enable_shared_from_this PendingStats mPendingStats; + std::vector mReuseStatus; + TypedVec>>> mOptionalCandidates; + std::vector mSelectedOptionalGroups; + void _collectReuseCandidates(); + void _selectOptionalGroups(std::vector const& groups); + // PRIVATE and unselected OPTIONAL groups need fresh pages for the claimed prefix. + bool _needsFreshPrefixPages(LifeCycleId lc) const; + // SWA scratch slot support. bool mEnableSwaScratchReuse = false; TypedVec> mScratchSlots; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp index 9f4a4f0eec29..a4c57cacea07 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp @@ -41,7 +41,8 @@ LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock) else { bool const isSparse = !cfg.buffers.empty() && cfg.buffers.front().isSparse; - return AttnLifeCycle::make(cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, isSparse); + return AttnLifeCycle::make( + cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, isSparse, cfg.reusePolicy); } }, layer); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h index c62d9b40d2c2..bdf550e6fcc1 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h @@ -41,6 +41,22 @@ struct AttnLifeCycle std::optional windowSize; // nullopt = no sliding window int numSinkBlocks = 0; // divUp(numSinkTokens, tokensPerBlock) bool isSparse = false; + AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; + + bool isPrivate() const noexcept + { + return reusePolicy == AttentionReusePolicy::PRIVATE; + } + + bool requiresPrefixCoverage() const noexcept + { + return reusePolicy == AttentionReusePolicy::REQUIRED; + } + + bool allowsCheckpoint() const noexcept + { + return reusePolicy == AttentionReusePolicy::OPTIONAL; + } HalfOpenRange getStaleRange(int historyLength, int tokensPerBlock) const { @@ -58,7 +74,8 @@ struct AttnLifeCycle bool operator==(AttnLifeCycle const& o) const noexcept { - return windowSize == o.windowSize && numSinkBlocks == o.numSinkBlocks && isSparse == o.isSparse; + return windowSize == o.windowSize && numSinkBlocks == o.numSinkBlocks && isSparse == o.isSparse + && reusePolicy == o.reusePolicy; } bool operator<(AttnLifeCycle const& o) const noexcept @@ -71,18 +88,22 @@ struct AttnLifeCycle { return numSinkBlocks < o.numSinkBlocks; } - return isSparse < o.isSparse; + if (isSparse != o.isSparse) + { + return isSparse < o.isSparse; + } + return reusePolicy < o.reusePolicy; } - static AttnLifeCycle make( - std::optional ws, std::optional numSinkTokens, int tokensPerBlock, bool isSparse = false) + static AttnLifeCycle make(std::optional ws, std::optional numSinkTokens, int tokensPerBlock, + bool isSparse = false, AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED) { TLLM_CHECK_DEBUG(tokensPerBlock > 0); TLLM_CHECK_DEBUG(!ws.has_value() || *ws > 0); TLLM_CHECK_DEBUG(!numSinkTokens.has_value() || *numSinkTokens >= 0); TLLM_CHECK_DEBUG((!numSinkTokens.has_value() || *numSinkTokens == 0) || ws.has_value()); int sinkBlocks = divUp(numSinkTokens.value_or(0), tokensPerBlock); - return AttnLifeCycle{ws, sinkBlocks, isSparse}; + return AttnLifeCycle{ws, sinkBlocks, isSparse, reusePolicy}; } }; @@ -114,6 +135,18 @@ struct SsmLifeCycle // --------------------------------------------------------------------------- using LifeCycle = std::variant; +inline bool isPrivateLifeCycle(LifeCycle const& lc) +{ + auto const* attn = std::get_if(&lc); + return attn && attn->isPrivate(); +} + +inline bool requiresPrefixCoverage(LifeCycle const& lc) +{ + auto const* attn = std::get_if(&lc); + return !attn || attn->requiresPrefixCoverage(); +} + // Free function: compute stale range via std::visit. inline HalfOpenRange getStaleRange(LifeCycle const& lc, int historyLength, int tokensPerBlock) { diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp index 7ada3a2e97e7..2fb89b143ea6 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp @@ -42,6 +42,11 @@ Page::Page(StorageManager* mgr, LifeCycleId lc, CacheLevel level, Priority prio) { } +bool Page::isPrivate() const +{ + return !isCommitted() && !requiresPrefixCoverage(manager->lifeCycles()[lifeCycle]); +} + Page::~Page() { KVCM2_POISON_ON_EXCEPT( @@ -259,8 +264,11 @@ PageHolder::~PageHolder() // page with a larger recorded token count) is unreachable for reuse, so keeping it // in the eviction LRU would just pin a slot until memory pressure hits. auto* cp = static_cast(page.get()); - if (cp->block == nullptr || cp->block->isOrphan() || !cp->block->holdsPage(*cp)) + if (page->scheduledForEviction() + && (cp->block == nullptr || cp->block->isOrphan() || !cp->block->holdsPage(*cp))) + { manager->excludeFromEviction(*page); + } } else { @@ -531,6 +539,15 @@ void SharedPageLock::acquirePageIndex() void SharedPageLock::releasePageIndex() { + if (mUser.ordinal == kBadBlockOrdinal) + { + // SSM state has one live page per beam and no attention page-index + // entry. updateBasePageIndex deliberately returns BAD for this ordinal. + auto const& lc = page()->manager->lifeCycles().getLifeCycle(mUser.lifeCycle); + auto const* attention = std::get_if(&lc); + TLLM_CHECK_DEBUG(std::holds_alternative(lc) || (attention && attention->allowsCheckpoint())); + return; + } int oldBaseIndex = mUser.kvCache->updateBasePageIndex(mUser.beamIndex, mUser.ordinal, mUser.lifeCycle, kBadPageIndex.value()); // SSM pages use kBadBlockOrdinal, for which updateBasePageIndex returns kBadPageIndex. @@ -717,7 +734,9 @@ void ScratchSlotLock::unlock() { TLLM_CHECK_DEBUG(mSlot.hasValidSlot()); mSlot.readyEvent = mOwner->finishEvent(); - mOwner->storageManager()->releaseSlot(mLifeCycle, kHotLevel, std::move(mSlot)); + // Slot's implicit move leaves its optional ID engaged. Transfer ownership + // explicitly so destruction cannot release this physical slot a second time. + mOwner->storageManager()->releaseSlot(mLifeCycle, kHotLevel, detachSlot()); TLLM_CHECK_DEBUG(!mSlot.hasValidSlot()); } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h index c60e8eafc411..9ecd538fe079 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h @@ -62,6 +62,8 @@ class Page : public Slot, public EnableSharedFromThis virtual bool isCommitted() const = 0; + bool isPrivate() const; + PageStatus status() const noexcept; bool scheduledForEviction() const noexcept diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp index bdb429ba9074..c9e3718b14b9 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp @@ -206,15 +206,17 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) slotGroups.push_back(std::move(var)); } - // Keep sparse and dense lifecycles in separate GPU pools even when slot sizes match. + // Distinct sparsity and reuse policies need separate GPU pools and eviction queues. struct PoolGroupKey { + AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; std::vector slotSizes; bool isSparse = false; bool operator<(PoolGroupKey const& other) const { - return std::tie(slotSizes, isSparse) < std::tie(other.slotSizes, other.isSparse); + return std::tie(reusePolicy, slotSizes, isSparse) + < std::tie(other.reusePolicy, other.slotSizes, other.isSparse); } }; @@ -224,7 +226,8 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) auto const sizes = sg.slotSizeList(); auto const* attn = std::get_if(®istry.getLifeCycle(sg.lifeCycleId)); bool const isSparse = attn != nullptr && attn->isSparse; - poolGroups[{.slotSizes = sizes.raw(), .isSparse = isSparse}].push_back(std::move(sg)); + auto const policy = attn ? attn->reusePolicy : AttentionReusePolicy::REQUIRED; + poolGroups[{.reusePolicy = policy, .slotSizes = sizes.raw(), .isSparse = isSparse}].push_back(std::move(sg)); } StorageConfig out; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp index 94ca946bba55..dc2a7519ffaa 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp @@ -17,6 +17,7 @@ #include "kv_cache_manager_v2/storageManager.h" #include "kv_cache_manager_v2/batchedPageCopy.h" +#include "kv_cache_manager_v2/blockRadixTree.h" #include "kv_cache_manager_v2/common.h" #include "kv_cache_manager_v2/copyEngine.h" #include "kv_cache_manager_v2/exceptions.h" @@ -397,12 +398,14 @@ StorageManager::StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfi TypedVec coldGrouping(numLifeCycles()); TypedVec coldSlotDescList; - std::map coldGroupByPageBytes; + std::map, PoolGroupIndex> coldGroupByPageBytes; for (LifeCycleId lifeCycle{0}; lifeCycle < numLifeCycles(); ++lifeCycle) { size_t const coldPageBytes = coldPageBytesByLifeCycle[lifeCycle]; + auto const* attn = std::get_if(&mLifeCycles[lifeCycle]); + auto const policy = attn ? attn->reusePolicy : AttentionReusePolicy::REQUIRED; auto [it, inserted] = coldGroupByPageBytes.emplace( - coldPageBytes, PoolGroupIndex{static_cast(coldSlotDescList.size().value())}); + std::make_pair(policy, coldPageBytes), PoolGroupIndex{static_cast(coldSlotDescList.size().value())}); PoolGroupIndex const coldPgIdx = it->second; if (inserted) { @@ -744,7 +747,14 @@ bool StorageManager::isEvictable(Page const& page, std::optional lev { PageStatus s = page.status(); CacheLevel lvl = level.value_or(page.cacheLevel); - return (s == PageStatus::DROPPABLE && page.isCommitted()) || (s == PageStatus::HELD && lvl < numCacheLevels() - 1); + if (s == PageStatus::DROPPABLE && page.isCommitted()) + { + // A migration batch can outlive checkpoint invalidation. Do not put + // detached bundle members back in an LRU that would retain their slots. + auto const& committed = static_cast(page); + return committed.block && !committed.block->isOrphan() && committed.block->holdsPage(committed); + } + return s == PageStatus::HELD && lvl < numCacheLevels() - 1; } // --------------------------------------------------------------------------- @@ -1079,10 +1089,7 @@ void StorageManager::_prepareFreeSlots(TypedVecstatus() != PageStatus::HELD) - { - lvl.controller.scheduleForEviction(*page); - } + scheduleForEviction(*page); } } } @@ -1943,7 +1950,16 @@ TypedVec StorageManager::projectPoolGroupRatio( TypedVec StorageManager::ratioFromBatch(BatchDesc const& batch, int tokensPerBlock, std::optional const& swaScratchReuse, size_t granularity) const { - auto const numSlots = computeSlotsForBatch(batch, tokensPerBlock, swaScratchReuse); + auto numSlots = computeSlotsForBatch(batch, tokensPerBlock, swaScratchReuse); + // A soft preference, not a feasibility floor: keep one optional endpoint + // per typical request in addition to its private working pages. + for (auto const& [lc, attention] : mLifeCycles.attentionLifeCycles()) + { + if (attention->allowsCheckpoint()) + { + numSlots[lc] += slotCountValueFromSize(batch.kvCaches.size()); + } + } return normalizeToRatio(slotsToBytes(numSlots, granularity)); } diff --git a/cpp/tensorrt_llm/deep_gemm/CMakeLists.txt b/cpp/tensorrt_llm/deep_gemm/CMakeLists.txt index bedc7cad6f62..46292e954440 100644 --- a/cpp/tensorrt_llm/deep_gemm/CMakeLists.txt +++ b/cpp/tensorrt_llm/deep_gemm/CMakeLists.txt @@ -1,3 +1,17 @@ +# Copyright (c) 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. You may obtain a copy of +# the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. + add_custom_target(deep_gemm) if(WIN32) @@ -12,10 +26,11 @@ set(DEEP_GEMM_SOURCE_DIR ${CMAKE_BINARY_DIR}/_deps/deepgemm-src) get_filename_component(DEEP_GEMM_SOURCE_DIR ${DEEP_GEMM_SOURCE_DIR} ABSOLUTE) # Check if submodules are initialized -if(NOT EXISTS ${DEEP_GEMM_SOURCE_DIR}/third-party/cutlass/include) +if(NOT EXISTS ${DEEP_GEMM_SOURCE_DIR}/third-party/cutlass/include + OR NOT EXISTS ${DEEP_GEMM_SOURCE_DIR}/third-party/deep_jit/include) message( FATAL_ERROR - "DeepGEMM submodules not initialized. Something went wrong with FetchContent.\n" + "DeepGEMM submodules (cutlass, deep_jit) not initialized. Something went wrong with FetchContent.\n" ) endif() @@ -83,7 +98,8 @@ foreach(SOURCE_FILE ${DEEP_GEMM_ALL_FILES}) PROPERTY CMAKE_CONFIGURE_DEPENDS ${SOURCE_FILE}) endforeach() -# Copy third-party includes (cutlass and fmt) to the include directory +# Copy third-party includes (cutlass) to the include directory. DeepJIT is +# header-only and host-side only, so it is not needed at kernel JIT time. set(DEEP_GEMM_INCLUDE_DEST ${DEEP_GEMM_PYTHON_DEST}/include) file(MAKE_DIRECTORY ${DEEP_GEMM_INCLUDE_DEST}) file(COPY ${DEEP_GEMM_SOURCE_DIR}/third-party/cutlass/include/cute @@ -95,8 +111,31 @@ file(COPY ${DEEP_GEMM_SOURCE_DIR}/third-party/cutlass/include/cutlass find_library(TORCH_PYTHON_LIB torch_python REQUIRED HINTS ${TORCH_INSTALL_PREFIX}/lib) +# Use the standard formatter when available. Older C++20 standard libraries +# build host-only copies against header-only fmt, without modifying the fetched +# sources or changing the CUDA kernels included in the wheel. +include(CheckCXXSourceCompiles) +check_cxx_source_compiles( + "#include + int main() { return std::format(\"{}\", 42) != \"42\"; }" + DEEP_GEMM_HAVE_STD_FORMAT) +set(DEEP_GEMM_CSRC_DIR ${DEEP_GEMM_SOURCE_DIR}/csrc) +set(DEEP_JIT_INCLUDE_DIR ${DEEP_GEMM_SOURCE_DIR}/third-party/deep_jit/include) +if(NOT DEEP_GEMM_HAVE_STD_FORMAT) + message( + STATUS "std::format unavailable; DeepGEMM host code uses header-only fmt") + FetchContent_MakeAvailable(fmt) + include(${CMAKE_CURRENT_SOURCE_DIR}/formatCompatibility.cmake) + set(DEEP_GEMM_COMPAT_DIR ${CMAKE_CURRENT_BINARY_DIR}/format_compat) + trtllm_copy_with_fmt(${DEEP_GEMM_CSRC_DIR} ${DEEP_GEMM_COMPAT_DIR}/csrc) + trtllm_copy_with_fmt(${DEEP_JIT_INCLUDE_DIR} + ${DEEP_GEMM_COMPAT_DIR}/deep_jit/include) + set(DEEP_GEMM_CSRC_DIR ${DEEP_GEMM_COMPAT_DIR}/csrc) + set(DEEP_JIT_INCLUDE_DIR ${DEEP_GEMM_COMPAT_DIR}/deep_jit/include) +endif() + # Build deep_gemm_cpp_tllm extension (matching deep_gemm's setup.py) -set(DEEP_GEMM_SOURCES ${DEEP_GEMM_SOURCE_DIR}/csrc/python_api.cpp) +set(DEEP_GEMM_SOURCES ${DEEP_GEMM_CSRC_DIR}/python_api.cpp) pybind11_add_module(deep_gemm_cpp_tllm ${DEEP_GEMM_SOURCES}) set_target_properties( @@ -115,13 +154,33 @@ target_compile_options(deep_gemm_cpp_tllm PRIVATE ${TORCH_CXX_FLAGS} -O3 -fPIC # Extension name definition target_compile_definitions(deep_gemm_cpp_tllm PRIVATE TORCH_EXTENSION_NAME=deep_gemm_cpp_tllm) +if(NOT DEEP_GEMM_HAVE_STD_FORMAT) + target_compile_definitions(deep_gemm_cpp_tllm PRIVATE FMT_HEADER_ONLY=1) + target_include_directories( + deep_gemm_cpp_tllm PRIVATE ${fmt_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}) +endif() # Include directories matching deep_gemm setup.py target_include_directories( deep_gemm_cpp_tllm - PRIVATE ${CUDA_INCLUDE_DIRS} ${DEEP_GEMM_SOURCE_DIR}/deep_gemm/include - ${DEEP_GEMM_SOURCE_DIR}/third-party/cutlass/include - ${DEEP_GEMM_SOURCE_DIR}/third-party/fmt/include) + PRIVATE ${CUDA_INCLUDE_DIRS} ${CUDA_TOOLKIT_ROOT_DIR}/include/cccl + ${DEEP_GEMM_SOURCE_DIR}/deep_gemm/include ${DEEP_JIT_INCLUDE_DIR} + ${DEEP_GEMM_SOURCE_DIR}/third-party/cutlass/include) + +# DeepJIT resolves libdwfl through dlopen at runtime but includes the elfutils +# header for its types. Hosts without the elfutils development headers use the +# minimal declaration shim under deep_jit_shim/ instead. +include(CheckIncludeFileCXX) +check_include_file_cxx("elfutils/libdwfl.h" DEEP_GEMM_HAVE_LIBDWFL_HEADER) +if(NOT DEEP_GEMM_HAVE_LIBDWFL_HEADER) + message( + STATUS + "elfutils/libdwfl.h not found; DeepJIT builds against the libdwfl declaration shim" + ) + target_include_directories(deep_gemm_cpp_tllm + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/deep_jit_shim) +endif() # Link libraries (matching deep_gemm setup.py: cuda, cudart + torch) target_link_libraries( diff --git a/cpp/tensorrt_llm/deep_gemm/deep_jit_shim/elfutils/libdwfl.h b/cpp/tensorrt_llm/deep_gemm/deep_jit_shim/elfutils/libdwfl.h new file mode 100644 index 000000000000..5058a21b29d7 --- /dev/null +++ b/cpp/tensorrt_llm/deep_gemm/deep_jit_shim/elfutils/libdwfl.h @@ -0,0 +1,72 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * Minimal declarations of the libdwfl (elfutils) entry points that DeepJIT's + * exception backtraces resolve through dlopen("libdw.so.1"). DeepJIT only + * includes for these types and prototypes and never links + * libdw, so build hosts without the elfutils development headers use this + * shim instead; when the real header is installed it takes precedence. + */ + +#pragma once + +#include +#include + +#ifdef __cplusplus +extern "C" +{ +#endif + + typedef uint64_t Dwarf_Addr; + typedef uint64_t Dwarf_Word; + typedef uint64_t GElf_Addr; + typedef uint32_t GElf_Word; + + typedef struct Elf Elf; + typedef struct Dwfl Dwfl; + typedef struct Dwfl_Module Dwfl_Module; + typedef struct Dwfl_Line Dwfl_Line; + typedef struct Dwfl_Shim_Elf64_Shdr GElf_Shdr; + + typedef struct + { + int (*find_elf)( + Dwfl_Module* mod, void** userdata, char const* modname, Dwarf_Addr base, char** file_name, Elf** elfp); + int (*find_debuginfo)(Dwfl_Module* mod, void** userdata, char const* modname, Dwarf_Addr base, + char const* file_name, char const* debuglink_file, GElf_Word debuglink_crc, char** debuginfo_file_name); + int (*section_address)(Dwfl_Module* mod, void** userdata, char const* modname, Dwarf_Addr base, + char const* secname, GElf_Word shndx, GElf_Shdr const* shdr, Dwarf_Addr* addr); + char** debuginfo_path; + } Dwfl_Callbacks; + + Dwfl* dwfl_begin(Dwfl_Callbacks const* callbacks); + int dwfl_linux_proc_report(Dwfl* dwfl, pid_t pid); + int dwfl_report_end(Dwfl* dwfl, + int (*removed)(Dwfl_Module* mod, void* userdata, char const* name, Dwarf_Addr low_addr, void* arg), void* arg); + Dwfl_Module* dwfl_addrmodule(Dwfl* dwfl, Dwarf_Addr address); + char const* dwfl_module_addrname(Dwfl_Module* mod, GElf_Addr address); + Dwfl_Line* dwfl_module_getsrc(Dwfl_Module* mod, Dwarf_Addr address); + char const* dwfl_lineinfo( + Dwfl_Line* line, Dwarf_Addr* addr, int* linep, int* colp, Dwarf_Word* mtime, Dwarf_Word* length); + int dwfl_linux_proc_find_elf( + Dwfl_Module* mod, void** userdata, char const* modname, Dwarf_Addr base, char** file_name, Elf** elfp); + +#ifdef __cplusplus +} +#endif diff --git a/cpp/tensorrt_llm/deep_gemm/formatCompatibility.cmake b/cpp/tensorrt_llm/deep_gemm/formatCompatibility.cmake new file mode 100644 index 000000000000..49be60025f3f --- /dev/null +++ b/cpp/tensorrt_llm/deep_gemm/formatCompatibility.cmake @@ -0,0 +1,53 @@ +# Copyright (c) 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. You may obtain a copy of +# the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. + +function(trtllm_copy_with_fmt source_dir destination_dir) + file( + GLOB_RECURSE source_files CONFIGURE_DEPENDS + LIST_DIRECTORIES FALSE + "${source_dir}/*") + file(REMOVE_RECURSE "${destination_dir}") + foreach(source_file IN LISTS source_files) + file(RELATIVE_PATH relative_path "${source_dir}" "${source_file}") + file(READ "${source_file}" content) + if(relative_path STREQUAL "jit_kernels/impls/runtime_utils.hpp") + # fmt's hexadecimal floats include a 0x prefix; std::format's do not. Keep + # the JIT float literals identical to the standard formatter. + set(float_format + "std::format(\"{}0x{:a}f\", std::signbit(v) ? \"-\" : \"\", std::abs(v))" + ) + string(FIND "${content}" "${float_format}" float_format_offset) + if(float_format_offset EQUAL -1) + message( + FATAL_ERROR + "DeepGEMM float literal formatter has changed; update the fmt compatibility conversion" + ) + endif() + string(REPLACE "${float_format}" + "::tensorrt_llm::deepGemm::formatFloatLiteral(v)" content + "${content}") + string(REPLACE "#include " + "#include \"formatCompatibility.h\"\n#include " + content "${content}") + endif() + string(REPLACE "#include " "#include " content + "${content}") + string(REPLACE "std::format(" "fmt::format(" content "${content}") + file(WRITE "${destination_dir}/${relative_path}" "${content}") + set_property( + DIRECTORY + APPEND + PROPERTY CMAKE_CONFIGURE_DEPENDS "${source_file}") + endforeach() +endfunction() diff --git a/cpp/tensorrt_llm/deep_gemm/formatCompatibility.h b/cpp/tensorrt_llm/deep_gemm/formatCompatibility.h new file mode 100644 index 000000000000..92648f32c06c --- /dev/null +++ b/cpp/tensorrt_llm/deep_gemm/formatCompatibility.h @@ -0,0 +1,45 @@ +/* + * Copyright (c) 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. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include + +namespace tensorrt_llm::deepGemm +{ + +// Format a finite float as a CUDA literal, matching std::format's hexadecimal +// representation, including signed zero and subnormal values. +inline std::string formatFloatLiteral(float const value) +{ + constexpr int kBufferSize = 32; + std::array buffer{}; + // Every finite float is represented exactly as a normal double (or zero), + // giving subnormal floats the same normalized exponent across libraries. + auto const result = std::to_chars( + buffer.data(), buffer.data() + buffer.size(), static_cast(std::abs(value)), std::chars_format::hex); + if (result.ec != std::errc{}) + { + throw std::system_error(std::make_error_code(result.ec)); + } + return std::string(std::signbit(value) ? "-0x" : "0x") + std::string(buffer.data(), result.ptr) + "f"; +} + +} // namespace tensorrt_llm::deepGemm diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu index 9379dbad7310..001772ed74eb 100644 --- a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu @@ -205,9 +205,10 @@ struct VecType<16> // Cache scale / pre-store quantization type for postProcessScatterKernel. // The store dtype is implied: kNone keeps the input dtype (elem_bytes -// controls bf16 vs fp32), kFP8* writes one byte per element, and -// kMXFP4Blockwise writes packed FP4 (two values per byte) plus per-32 -// UE8M0 scale bytes. +// controls bf16 vs fp32), kFP8* writes one byte per element, and the FP4 +// types write packed FP4 (two values per byte) plus one scale byte per +// block -- per-32 UE8M0 for kMXFP4Blockwise, per-16 E4M3 for +// kNVFP4Blockwise. enum class CacheScaleType { kNone = 0, @@ -308,8 +309,13 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo { using KvScoreElemT = typename std::conditional::type; using StateElemT = typename std::conditional::type; - // DeepSeek-V4 model configures compress_ratio to be 4 or 128. - static_assert(COMPRESS_RATIO == 4 || COMPRESS_RATIO == 128, "Unsupported COMPRESS_RATIO"); + // DeepSeek-V4 configures compress_ratio to 4 or 128; DeepSeek-V4.1 uses 2. + // (V4.1's other ratio, 1, needs no reduction at all and never reaches here.) + static_assert(COMPRESS_RATIO == 2 || COMPRESS_RATIO == 4 || COMPRESS_RATIO == 128, "Unsupported COMPRESS_RATIO"); + // Overlapping pooling windows are a V4 ratio-4 property, not a property of + // small ratios: V4.1's ratio-2 compressor keeps one window in flight and + // sizes its state (max_batch, ratio, head_dim) with no overlap half + // (model.py:451-456), so CR=2 takes the non-overlap path like CR=128. constexpr bool IS_OVERLAP = (COMPRESS_RATIO == 4); constexpr int ELEM_BYTES_FOR_VEC = (KV_SCORE_ELEM_BYTES > STATE_ELEM_BYTES) ? KV_SCORE_ELEM_BYTES : STATE_ELEM_BYTES; @@ -663,9 +669,9 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo // HD — HEAD_DIM in {128, 512} // KV_EB — kv_score element bytes in {2 (bf16), 4 (fp32)} // STATE_EB — paged state element bytes in {2 (bf16), 4 (fp32)} -// CR — COMPRESS_RATIO in {4, 128} +// CR — COMPRESS_RATIO in {2 (V4.1), 4, 128} // NRW — NUM_RED_WARPS — 4 when CR=128 (multi-warp Phase 3 reduction -// hides DRAM latency for the heavier R=128 chunk); 1 when CR=4. +// hides DRAM latency for the heavier R=128 chunk); 1 when CR=2 or 4. // // Multi-warp SMEM budget (per block): 3 * NRW * ELEM_PER_BLOCK * sizeof(float). // HD=128: ELEM_PER_BLOCK=128 → 6 KB @@ -680,6 +686,8 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo // Master list. Order does not matter; the dispatcher walks linearly. // clang-format off #define FOREACH_DECODE_CONFIG(F) \ + /* CR=2 (DeepSeek-V4.1): single-warp, same reason as CR=4 and half the loop. */ \ + FOREACH_DECODE_DTYPE(F, 128, 2, 1) FOREACH_DECODE_DTYPE(F, 512, 2, 1) \ /* CR=4: single-warp only (small reduction; multi-warp would over-subscribe). */ \ FOREACH_DECODE_DTYPE(F, 128, 4, 1) FOREACH_DECODE_DTYPE(F, 512, 4, 1) \ /* CR=128: multi-warp Phase 3 reduction for every supported next_n. */ \ @@ -722,8 +730,8 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k int compress_ratio, int next_n, int kv_score_elem_bytes, int state_elem_bytes, int out_elem_bytes, int32_t const* new_tokens_per_seq, cudaStream_t stream) { - TLLM_CHECK_WITH_INFO( - compress_ratio == 4 || compress_ratio == 128, "pagedKvCompressLaunch only supports compress_ratio 4 or 128"); + TLLM_CHECK_WITH_INFO(compress_ratio == 2 || compress_ratio == 4 || compress_ratio == 128, + "pagedKvCompressLaunch only supports compress_ratio 2 (V4.1), 4 or 128, got %d", compress_ratio); TLLM_CHECK_WITH_INFO( (kv_score_elem_bytes == 2 || kv_score_elem_bytes == 4) && (state_elem_bytes == 2 || state_elem_bytes == 4), "pagedKvCompressLaunch only supports bf16/fp32 kv_score and paged state"); @@ -873,7 +881,8 @@ __global__ void prefillReductionKernel(void const* __restrict__ kv_score_raw, fl { using KvScoreElemT = typename std::conditional::type; using StateElemT = typename std::conditional::type; - static_assert(COMPRESS_RATIO == 4 || COMPRESS_RATIO == 128, "Unsupported COMPRESS_RATIO"); + // See pagedKvCompressKernel: V4.1 adds CR=2, which is non-overlapping. + static_assert(COMPRESS_RATIO == 2 || COMPRESS_RATIO == 4 || COMPRESS_RATIO == 128, "Unsupported COMPRESS_RATIO"); constexpr bool IS_OVERLAP = (COMPRESS_RATIO == 4); constexpr int ELEM_BYTES_FOR_VEC @@ -1228,8 +1237,10 @@ __global__ void prefillReductionKernel(void const* __restrict__ kv_score_raw, fl INST_PREFILL(HD, 2, 2, CR, NRW) \ INST_PREFILL(HD, 2, 4, CR, NRW) INST_PREFILL(HD, 4, 2, CR, NRW) INST_PREFILL(HD, 4, 4, CR, NRW) +INST_PREFILL_DTYPES(128, 2, 1) INST_PREFILL_DTYPES(128, 4, 1) INST_PREFILL_DTYPES(128, 128, 4) +INST_PREFILL_DTYPES(512, 2, 1) INST_PREFILL_DTYPES(512, 4, 1) INST_PREFILL_DTYPES(512, 128, 4) #undef INST_PREFILL_DTYPES @@ -1256,8 +1267,8 @@ void prefillReductionLaunch(void const* kv_score, float const* ape, void* paged_ int out_elem_bytes, cudaStream_t stream) { bool const overlap = (compress_ratio == 4); - TLLM_CHECK_WITH_INFO( - compress_ratio == 4 || compress_ratio == 128, "prefillReductionLaunch only supports compress_ratio 4 or 128"); + TLLM_CHECK_WITH_INFO(compress_ratio == 2 || compress_ratio == 4 || compress_ratio == 128, + "prefillReductionLaunch only supports compress_ratio 2 (V4.1), 4 or 128, got %d", compress_ratio); TLLM_CHECK_WITH_INFO( (kv_score_elem_bytes == 2 || kv_score_elem_bytes == 4) && (state_elem_bytes == 2 || state_elem_bytes == 4), "prefillReductionLaunch only supports bf16/fp32 kv_score and paged state"); @@ -1304,14 +1315,18 @@ void prefillReductionLaunch(void const* kv_score, float const* ape, void* paged_ if (head_dim == 512) { - if (compress_ratio == 4) + if (compress_ratio == 2) + DISPATCH_PREFILL_DTYPE(512, 2, 1); + else if (compress_ratio == 4) DISPATCH_PREFILL_DTYPE(512, 4, 1); else DISPATCH_PREFILL_DTYPE(512, 128, 4); } else { - if (compress_ratio == 4) + if (compress_ratio == 2) + DISPATCH_PREFILL_DTYPE(128, 2, 1); + else if (compress_ratio == 4) DISPATCH_PREFILL_DTYPE(128, 4, 1); else DISPATCH_PREFILL_DTYPE(128, 128, 4); @@ -1340,6 +1355,8 @@ void prefillReductionLaunch(void const* kv_score, float const* ape, void* paged_ // - kFP8Blockwise: one fp8 byte per value + one fp32 scale per 128 values. // - kMXFP4Blockwise: packed fp4 (two values per byte) + one ue8m0 byte // per 32 values. +// - kNVFP4Blockwise: packed fp4 (two values per byte) + one e4m3 byte +// per 16 values (V4.1 tech report §2.4.4). // // Pipeline (10 steps, all in fp32 registers): // 1. Vectorized load compressed token from kv_comp diff --git a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu index 47d7640e7e62..e5a307f25074 100644 --- a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu +++ b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu @@ -168,6 +168,13 @@ void dispatchDeepseekV4QNorm( // (kPairsPerLane-1) iterations cover the nope range and the last iteration // covers the rope range exactly. +// `applyNorm`: false makes the RMS scale an identity, for models that have no +// per-head query norm at all (V4.1 deleted V4's). The norm here carries no learned +// weight -- it is one per-row scalar -- so skipping it is exactly `normScale = 1`, +// and the warp reduction that produced it goes with it. Runtime rather than a fifth +// template parameter: the branch is warp-uniform, and the instantiation set is +// already 8-way (2 dtypes x kFuseRope x kWideVec). +// // `kFuseRope`: rotate the rope segment here too and store it FP8 in the same row. // Register-local -- `kRopePairs == kWarpSize`, so lane `l` owns rope pair `l` and // needs only one float2 of cos/sin. @@ -180,7 +187,7 @@ void dispatchDeepseekV4QNorm( template __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8_e4m3* __restrict__ quant_q_nope, T* __restrict__ q_pe_out, float const* __restrict__ quant_scale_qkv_ptr, int totalRows, - int quantQNopeRowStrideBytes, float eps, float2 const* __restrict__ cos_sin_cache = nullptr, + int quantQNopeRowStrideBytes, float eps, bool applyNorm, float2 const* __restrict__ cos_sin_cache = nullptr, int const* __restrict__ cache_seq_lens = nullptr, int num_heads = 0, int seq_len = 0, int const* __restrict__ cu_q_seqlens = nullptr, int num_seqs = 0, int num_heads_shift = -1, int seq_len_shift = -1) { @@ -272,8 +279,15 @@ __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8 } } - sumSquares = warpReduceSum(sumSquares); - float const normScale = rsqrtf(sumSquares / static_cast(kHeadDim) + eps); + // The per-lane squares above stay unconditional: they ride along with loads this + // path needs anyway, and branching inside the unrolled loop would cost the norming + // model. The cross-lane reduction is what is worth skipping. + float normScale = 1.0F; + if (applyNorm) + { + sumSquares = warpReduceSum(sumSquares); + normScale = rsqrtf(sumSquares / static_cast(kHeadDim) + eps); + } float const fp8Scale = normScale * quantScale; // Position depends on the token, which every lane of the warp shares, so this is @@ -410,8 +424,8 @@ __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8 template void dispatchDeepseekV4QNormFused(void const* input, void* quant_q_nope, void* q_pe_out, void const* quant_scale_qkv_ptr, int totalRows, int headDim, int nopeDim, int quantQNopeRowStrideBytes, float eps, - void const* cos_sin_cache, int const* cache_seq_lens, int num_heads, int seq_len, int const* cu_q_seqlens, - int num_seqs, cudaStream_t stream) + bool applyNorm, void const* cos_sin_cache, int const* cache_seq_lens, int num_heads, int seq_len, + int const* cu_q_seqlens, int num_seqs, cudaStream_t stream) { TLLM_CHECK_WITH_INFO(headDim == 512, "deepseekV4QNormFused only supports head_dim=512, got %d", headDim); TLLM_CHECK_WITH_INFO(nopeDim == 448, "deepseekV4QNormFused only supports nope_dim=448, got %d", nopeDim); @@ -457,7 +471,7 @@ void dispatchDeepseekV4QNormFused(void const* input, void* quant_q_nope, void* q tensorrt_llm::common::launchWithPdlWhenEnabled("deepseekV4QNormFused", deepseekV4QNormFusedKernel, grid, block, 0, stream, static_cast(input), static_cast<__nv_fp8_e4m3*>(quant_q_nope), static_cast(q_pe_out), - static_cast(quant_scale_qkv_ptr), totalRows, quantQNopeRowStrideBytes, eps, + static_cast(quant_scale_qkv_ptr), totalRows, quantQNopeRowStrideBytes, eps, applyNorm, static_cast(cos_sin_cache), cache_seq_lens, num_heads, seq_len, cu_q_seqlens, num_seqs, num_heads_shift, seq_len_shift); } @@ -466,7 +480,7 @@ void dispatchDeepseekV4QNormFused(void const* input, void* quant_q_nope, void* q tensorrt_llm::common::launchWithPdlWhenEnabled("deepseekV4QNormFused", deepseekV4QNormFusedKernel, grid, block, 0, stream, static_cast(input), static_cast<__nv_fp8_e4m3*>(quant_q_nope), static_cast(q_pe_out), - static_cast(quant_scale_qkv_ptr), totalRows, quantQNopeRowStrideBytes, eps, + static_cast(quant_scale_qkv_ptr), totalRows, quantQNopeRowStrideBytes, eps, applyNorm, static_cast(nullptr), static_cast(nullptr), 0, 0, static_cast(nullptr), 0, -1, -1); } @@ -504,8 +518,8 @@ void invokeDeepseekV4QNorm( void invokeDeepseekV4QNormFusedFp8(void const* input, void* quant_q_nope, void* q_pe_out, void const* quant_scale_qkv_ptr, int totalRows, int headDim, int nopeDim, int quantQNopeRowStrideBytes, - bool isBfloat16, float eps, void const* cos_sin_cache, int const* cache_seq_lens, int num_heads, int seq_len, - int const* cu_q_seqlens, int num_seqs, cudaStream_t stream) + bool isBfloat16, float eps, bool applyNorm, void const* cos_sin_cache, int const* cache_seq_lens, int num_heads, + int seq_len, int const* cu_q_seqlens, int num_seqs, cudaStream_t stream) { if (totalRows == 0) { @@ -515,14 +529,14 @@ void invokeDeepseekV4QNormFusedFp8(void const* input, void* quant_q_nope, void* if (isBfloat16) { dispatchDeepseekV4QNormFused<__nv_bfloat16>(input, quant_q_nope, q_pe_out, quant_scale_qkv_ptr, totalRows, - headDim, nopeDim, quantQNopeRowStrideBytes, eps, cos_sin_cache, cache_seq_lens, num_heads, seq_len, - cu_q_seqlens, num_seqs, stream); + headDim, nopeDim, quantQNopeRowStrideBytes, eps, applyNorm, cos_sin_cache, cache_seq_lens, num_heads, + seq_len, cu_q_seqlens, num_seqs, stream); } else { dispatchDeepseekV4QNormFused(input, quant_q_nope, q_pe_out, quant_scale_qkv_ptr, totalRows, headDim, - nopeDim, quantQNopeRowStrideBytes, eps, cos_sin_cache, cache_seq_lens, num_heads, seq_len, cu_q_seqlens, - num_seqs, stream); + nopeDim, quantQNopeRowStrideBytes, eps, applyNorm, cos_sin_cache, cache_seq_lens, num_heads, seq_len, + cu_q_seqlens, num_seqs, stream); } } diff --git a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h index a09b68f6f6cc..82fa43dc30c5 100644 --- a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h +++ b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h @@ -50,8 +50,8 @@ void invokeDeepseekV4QNorm( // the plain no-RoPE behaviour. void invokeDeepseekV4QNormFusedFp8(void const* input, void* quant_q_nope, void* q_pe_out, void const* quant_scale_qkv_ptr, int totalRows, int headDim, int nopeDim, int quantQNopeRowStrideBytes, - bool isBfloat16, float eps, void const* cos_sin_cache, int const* cache_seq_lens, int num_heads, int seq_len, - int const* cu_q_seqlens, int num_seqs, cudaStream_t stream); + bool isBfloat16, float eps, bool applyNorm, void const* cos_sin_cache, int const* cache_seq_lens, int num_heads, + int seq_len, int const* cu_q_seqlens, int num_seqs, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/fusedCatFp4.cu b/cpp/tensorrt_llm/kernels/fusedCatFp4.cu index d42568de0401..8cf81aac1871 100644 --- a/cpp/tensorrt_llm/kernels/fusedCatFp4.cu +++ b/cpp/tensorrt_llm/kernels/fusedCatFp4.cu @@ -33,12 +33,7 @@ namespace kernels namespace { -// Constants (fixed for DSV3.2 FP4 indexer: head_dim=128, per-block-32 quant). -// -// INV_FP4_E2M1_MAX / MIN_AMAX mirror DeepGEMM's per_token_cast_to_fp4 reference -// (use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) at -// tensorrt_llm/deep_gemm/utils/math.py. Keep in sync with that helper when -// bumping DeepGEMM — test_fused_cat_fp4_matches_deepgemm is the guard. +// Constants for the FP4 indexer: 128-wide rows with block-32 UE8M0 scales. constexpr int HEAD_DIM = 128; constexpr int WARP_SIZE = 32; constexpr int ELEMS_PER_THREAD = 4; // 128 / 32 = 4 elements per thread. @@ -57,13 +52,13 @@ union BF16x4 }; /// FP4 E2M1 quantize a single scaled value. -/// Returns a 4-bit code matching DeepGEMM's table: +/// Returns a 4-bit code with ties rounded toward zero: /// magnitudes {0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0} /// sign bit 3 set iff value < 0 and magnitude code != 0. __device__ __forceinline__ uint32_t quantizeFp4E2M1(float scaled) { float ax = fminf(fabsf(scaled), 6.0f); - // Strict `>` matches Triton reference and DeepGEMM (boundary values round down). + // Strict `>` selects the lower magnitude at every midpoint. uint32_t idx = static_cast( (ax > 0.25f) + (ax > 0.75f) + (ax > 1.25f) + (ax > 1.75f) + (ax > 2.5f) + (ax > 3.5f) + (ax > 5.0f)); uint32_t code = idx & 0x7u; diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh index b997621a8fcc..28fc74cfe6a1 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh +++ b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh @@ -802,7 +802,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 template + uint32_t kNumPmapThreads, uint32_t kNumSplits = 1, bool kFuseNorm = false, bool kLaggedPre = false> __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_allinone_tf32_pmap_gemm_atomic_impl(const uint32_t shape_m, const __grid_constant__ cute::TmaDescriptor tensor_map_residual, // residual_prev, bf16 @@ -825,7 +825,11 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) // norm_weight must be bf16 [HIDDEN]; norm_eps is the RMSNorm epsilon. // When kFuseNorm is false, norm_weight/norm_eps are ignored. __nv_bfloat16 const* __restrict__ norm_weight, float norm_eps, float rms_eps, float hc_pre_eps, - float hc_sinkhorn_eps, float hc_post_mult_value, uint32_t sinkhorn_repeat) + float hc_sinkhorn_eps, float hc_post_mult_value, uint32_t sinkhorn_repeat, + // kLaggedPre (DeepSeek-V4.1 wiring): layer_input is collapsed with the previous sublayer's + // gate pre_mix_ext [M, HC_MULT] and this token's own gate is written to pre_mix_out [M, HC_MULT]. + // Both are ignored when kLaggedPre is false (V4). + float const* __restrict__ pre_mix_ext = nullptr, float* __restrict__ pre_mix_out = nullptr) { #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) and (__CUDA_ARCH__ < 1100)) or defined(__CLION_IDE__) using Barrier = cutlass::arch::ClusterTransactionBarrier; @@ -1417,6 +1421,11 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) float v = y_row[lane_bf] * rstd * s0 + hc_base[lane_bf]; pre_mix_local = 1.0f / (1.0f + __expf(-v)) + hc_pre_eps; + if constexpr (kLaggedPre) + { + if (warp_in_team == 0) + pre_mix_out[static_cast(tok) * HC_MULT + lane_bf] = pre_mix_local; + } if (warp_in_team == 0) { @@ -1477,7 +1486,12 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) float pm[HC_MULT]; #pragma unroll for (uint32_t j = 0; j < HC_MULT; ++j) - pm[j] = __shfl_sync(0xffffffff, pre_mix_local, j); + { + if constexpr (kLaggedPre) + pm[j] = pre_mix_ext[static_cast(tok) * HC_MULT + j]; + else + pm[j] = __shfl_sync(0xffffffff, pre_mix_local, j); + } // Layer_input[tok, h] = sum_j pm[j] * residual_cur[tok, j, h]. // When WARPS_PER_TOK>1, warp_in_team 0..WARPS_PER_TOK-1 together cover diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu index c098a4203b0e..012ecd2b01e4 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu @@ -104,6 +104,7 @@ static constexpr uint32_t FHC_BLOCK_K = 64; #ifdef TRTLLM_MHC_ENABLE_FUSED_HC static constexpr uint32_t FHC_HIDDEN_FLASH = 4096; static constexpr uint32_t FHC_HIDDEN_PRO = 7168; +static constexpr uint32_t FHC_HIDDEN_V41 = 5120; // DeepSeek-V4.1-Flash static constexpr uint32_t FHC_BLOCK_M = 64; static constexpr uint32_t FHC_BLOCK_N = 32; static constexpr uint32_t FHC_SWIZZLE_CD = 128; @@ -120,12 +121,13 @@ static constexpr uint32_t FHC_NUM_PMAP_TH = 128; template