[brb-nv]

Description

The blockKey contruction in WindowBlockManager::findBlocksInReuseTreeByBlockKey is getting very slow with ISL. This change fixes that. Main insight is that push_back involves copy construction copying of uniqueTokens - which gets overwritten immediately anyway.

Let's say N is num_prefix_blocks and T is num_prefix_tokens.

for (auto const& blockedUniqueTokensList : blockedUniqueTokens)   // N iterations
{
    blockKeys.push_back(blockKey);                                 // <- This line is O(T)
    blockKeys.back().uniqueTokens = blockedUniqueTokensList;      // overwrites it, O(1)
}

Test Coverage

$ ./cpp/build/tests/unit_tests/batch_manager/kvCacheManagerTest --gtest_filter='*FindBlocksInReuseTreeByBlockKeys*'

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #50068 [ run ] triggered by Bot. Commit: 4842da9 Link to invocation

[tensorrt-cicd]

PR_Github #50068 [ run ] completed with state SUCCESS. Commit: 4842da9
/LLM/main/L0_MergeRequest_PR pipeline #39624 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[chuangz0]

agent suggests use the following code to replace fromReuseTree in getBlockRangeForSending

    // Equivalent to fromReuseTree but O(K): the request's cacheBlockIds already records reused-prefix and
    // newly-allocated suffix in chronological order, so the trailing slice avoids the reuse-tree walk and mutex.
    auto const windowSize = cacheManager->getBlockManager().getWindowSizesMetadata().begin()->first;
    auto blockRange = BlockRange::fromAllBlockIds(*cacheManager, llmRequest.mRequestId);
    auto const& allBlockIds = blockRange.getBlockIdsPerWindow().at(windowSize);
    int32_t const numBlocksToSend = indexFromEnd + 1;
    TLLM_CHECK_WITH_INFO(static_cast<int32_t>(allBlockIds.size()) >= numBlocksToSend,
        "sender sequence has %zu blocks but receiver requested last %d", allBlockIds.size(), numBlocksToSend);
    std::vector<SizeType32> tailBlockIds(allBlockIds.end() - numBlocksToSend, allBlockIds.end());
    blockRange.setBlockIdsForWindow(windowSize, std::move(tailBlockIds));
    return blockRange;

@Tabrizian Do you think it reasonable? the overhead should be smaller.

[VALLIS-NERIA]

LGTM

[Tabrizian]

@chuangz0 The ctx request is terminated here so blocks for that request may not be available anymore. This was also added as a part of refactor to remove request ID and allow arbitrary KVCache transfer.

[brb-nv]

Merging as this has been verified to help e2e perf by downstream user. If any refactor is needed, we can take as a follow-up on main @chuangz0 @Tabrizian .

I believe this function has been used for a while. Fixing it with minimal changes in this MR seems like a reasonable thing to do.