From 3a5cfd753a4877cb25e3793d817b18b6df36604d Mon Sep 17 00:00:00 2001 From: Thor Johnsen Date: Wed, 2 Sep 2026 17:17:08 -0700 Subject: [PATCH] [None][fix] Exempt SSM pages from the page-index release assertion SharedPageLock::releasePageIndex() asserts that the index it just cleared matches the page's slot. SSM pages are locked with kBadBlockOrdinal because they have no per-block index slot, so KvCache::updateBasePageIndex() tracks nothing for them and reports kBadPageIndex -- which never equals a real slot. The assertion therefore fails for every SSM page release, so any hybrid attention/SSM sequence aborts when its KvCache closes, as soon as TLLM_DEBUG_MODE=1 makes TLLM_CHECK_DEBUG live. The Python original carries the exemption this port dropped (_page.py, SharedPageLock.unlock): the expected value is BAD_PAGE_INDEX when the ordinal is BAD_BLOCK_ORDINAL. Restore it. Verified on GB200 with a throwaway gtest that opens a hybrid attention/SSM KvCache, commits two blocks and closes it. Against the unfixed library it aborts under TLLM_DEBUG_MODE=1 with Assertion failed: oldBaseIndex == slotIdToPageIndexValue(page()->slotId()) (kv_cache_manager_v2/page.cpp:395) SharedPageLock::releasePageIndex() <- SharedPageLock::unlock() <- ~SharedPageLock() <- KvCache::_clearBlocks() <- KvCache::close() and passes with it. The nine existing kvCacheManagerV2* gtests (71 cases) pass with and without TLLM_DEBUG_MODE=1 after the change. That test is not included here because it cannot enable the mode it needs: DebugConfig::isCheckDebugEnabled() caches the environment in a function-local static that is first read during libtensorrt_llm.so's static initialisation, so neither SetUpTestSuite() nor a priority-101 constructor in the test binary runs early enough -- both were tried and both let the unfixed library pass. It would take an ENVIRONMENT property on the ctest target, which no test in the tree uses today. Nothing in cpp/tests currently runs under TLLM_DEBUG_MODE=1, which is why this defect went unnoticed; wiring up a debug-mode lane is worth doing but wants its own change. Signed-off-by: Thor Johnsen --- .../batch_manager/kv_cache_manager_v2/page.cpp | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) 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 ca500682b1ff..898358121cb4 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 @@ -391,8 +391,13 @@ void SharedPageLock::releasePageIndex() { int oldBaseIndex = mUser.kvCache->updateBasePageIndex(mUser.beamIndex, mUser.ordinal, mUser.lifeCycle, kBadPageIndex.value()); - // Mirrors Python assertion: old base index must match this page's slot ID. - TLLM_CHECK_DEBUG(oldBaseIndex == slotIdToPageIndexValue(page()->slotId())); + // Mirrors Python assertion: old base index must match this page's slot ID -- except for SSM + // pages, which are locked with kBadBlockOrdinal because they have no per-block index slot. + // updateBasePageIndex() tracks nothing for those and reports kBadPageIndex, so comparing + // against the page's slot would fail for every SSM page (_page.py SharedPageLock.unlock() + // carries the same exemption; the port dropped it). + TLLM_CHECK_DEBUG(oldBaseIndex + == (mUser.ordinal == kBadBlockOrdinal ? kBadPageIndex.value() : slotIdToPageIndexValue(page()->slotId()))); (void) oldBaseIndex; }