From 1e422bf6c17031dda7cfd63184408119c1bc500f Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Mon, 13 Apr 2026 16:04:04 +0000 Subject: [PATCH 1/2] Revert "[TRTLLM-11540][feat] Add EAGLE3 dynamic tree speculative decoding support (#12062)" This reverts commit 4ece13c0bf35553e5a339e1f2ce8ebc5cf9e99d8. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../speculativeDecoding/dynamicTreeKernels.cu | 356 ------ .../speculativeDecoding/dynamicTreeKernels.h | 98 -- .../kvCacheUpdateKernels.cu | 205 ---- .../kvCacheUpdateKernels.h | 15 - .../trtllmGenKernels/fmha/fmhaRunnerParams.h | 4 - .../fmha/prepareCustomMask.cu | 11 +- cpp/tensorrt_llm/kernels/xqaDispatcher.cpp | 1 - cpp/tensorrt_llm/thop/CMakeLists.txt | 1 - cpp/tensorrt_llm/thop/dynamicTreeOp.cpp | 229 ---- .../thop/parallelDecodeKVCacheUpdateOp.cpp | 64 - .../features/feature-combination-matrix.md | 43 +- docs/source/features/speculative-decoding.md | 44 +- docs/source/models/supported-models.md | 24 +- examples/llm-api/quickstart_advanced.py | 183 +-- .../_torch/attention_backend/interface.py | 5 +- .../_torch/attention_backend/sparse/dsa.py | 4 +- .../_torch/attention_backend/trtllm.py | 210 ++-- tensorrt_llm/_torch/modules/attention.py | 32 - .../_torch/pyexecutor/model_engine.py | 90 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 2 +- .../_torch/pyexecutor/py_executor_creator.py | 15 +- .../_torch/pyexecutor/resource_manager.py | 12 +- .../_torch/speculative/drafting_loops.py | 29 +- .../_torch/speculative/dynamic_tree_ops.py | 198 ---- tensorrt_llm/_torch/speculative/eagle3.py | 200 +--- .../_torch/speculative/eagle3_dynamic_tree.py | 1049 ----------------- .../_torch/speculative/spec_sampler_base.py | 37 +- .../_torch/speculative/spec_tree_manager.py | 361 ++---- tensorrt_llm/_torch/speculative/utils.py | 15 +- tensorrt_llm/llmapi/llm_args.py | 35 +- tests/integration/test_lists/waives.txt | 1 - .../_torch/modeling/test_modeling_llama.py | 88 +- ...test_draft_token_prepare_for_generation.py | 21 +- .../test_draft_token_tree_sampling.py | 4 +- .../_torch/speculative/test_eagle3.py | 91 -- .../_torch/thop/parallel/test_custom_ops.py | 3 - tests/unittest/others/test_kv_cache_update.py | 201 ---- 37 files changed, 414 insertions(+), 3567 deletions(-) delete mode 100644 cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu delete mode 100644 cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h delete mode 100644 cpp/tensorrt_llm/thop/dynamicTreeOp.cpp delete mode 100644 tensorrt_llm/_torch/speculative/dynamic_tree_ops.py delete mode 100644 tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu deleted file mode 100644 index 1dcfdde71671..000000000000 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu +++ /dev/null @@ -1,356 +0,0 @@ -/* - * Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. - * Portions Copyright (c) 2025 by SGLang team (original implementation). - * - * 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. - */ - -#include "dynamicTreeKernels.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" - -TRTLLM_NAMESPACE_BEGIN - -using namespace tensorrt_llm::runtime; - -namespace kernels::speculative_decoding -{ - -//! \param parentList [in] layer-wise parent indices [bs, topK*(depth-1)+1] -//! \param selectedIndex [in] resampled history buffer indices [bs, draftTokenNum-1] -//! \param treeMask [out] attention mask (which nodes each node can see) -//! \param positions [out] position id per node [bs, draftTokenNum] -//! \param retrieveIndex [out] tree node -> local index mapping [bs, draftTokenNum] -//! \param retrieveNextToken [out] first-child pointer [bs, draftTokenNum], -1=none -//! \param retrieveNextSibling [out] next-sibling pointer [bs, draftTokenNum], -1=none -//! \param topK top-K value per layer -//! \param depth max tree depth (number of draft layers) -//! \param draftTokenNum total tree nodes per batch (including root) -__global__ void buildDynamicTreeKernel(int64_t const* parentList, int64_t const* selectedIndex, int32_t* treeMask, - int32_t* positions, int32_t* retrieveIndex, int32_t* retrieveNextToken, int32_t* retrieveNextSibling, - SizeType32 topK, SizeType32 depth, SizeType32 draftTokenNum) -{ - int32_t bid = blockIdx.x; - int32_t tid = threadIdx.x; - - if (tid >= draftTokenNum) - { - return; - } - - // treeMask layout: [batchSize, draftTokenNum, draftTokenNum] (QLEN_ONLY mode) - int32_t tokenTreeIdx = draftTokenNum * draftTokenNum * bid + draftTokenNum * tid + 1; - - treeMask[tokenTreeIdx - 1] = 1; // self-attention diagonal - for (int32_t i = 0; i < draftTokenNum - 1; i++) - { - treeMask[tokenTreeIdx + i] = 0; - } - - int32_t position = 0; - - if (tid == 0) - { - positions[bid * draftTokenNum] = 0; - - // Reverse iteration: inserting at list head produces forward sibling order - for (int32_t i = draftTokenNum - 1; i > 0; --i) - { - retrieveIndex[bid * draftTokenNum + i] = i; - - int64_t parentTbIdx = selectedIndex[bid * (draftTokenNum - 1) + i - 1] / topK; - int32_t parentPosition = 0; - - if (parentTbIdx > 0) - { - int64_t parentTokenIdx = parentList[bid * (topK * (depth - 1) + 1) + parentTbIdx]; - for (; parentPosition < draftTokenNum; ++parentPosition) - { - if (selectedIndex[bid * (draftTokenNum - 1) + parentPosition] == parentTokenIdx) - { - ++parentPosition; // +1 because position 0 is root - break; - } - } - } - - if (parentPosition == draftTokenNum) - { - printf( - "WARNING: Invalid dynamic tree! Detected a token with no parent token selected. " - "Please check if the logprob has nan. The token will be ignored.\n"); - continue; - } - - if (retrieveNextToken[bid * draftTokenNum + parentPosition] == -1) - { - retrieveNextToken[bid * draftTokenNum + parentPosition] = i; - } - else - { - int32_t originNextToken = retrieveNextToken[bid * draftTokenNum + parentPosition]; - retrieveNextToken[bid * draftTokenNum + parentPosition] = i; - retrieveNextSibling[bid * draftTokenNum + i] = originNextToken; - } - } - retrieveIndex[bid * draftTokenNum] = 0; - } - else - { - // Walk up to root, setting treeMask ancestor bits and counting depth - int32_t curPosition = tid - 1; - while (position < depth + 1) - { - position += 1; - treeMask[tokenTreeIdx + curPosition] = 1; - - int64_t parentTbIdx = selectedIndex[bid * (draftTokenNum - 1) + curPosition] / topK; - if (parentTbIdx == 0) - { - break; - } - - int64_t tokenIdx = parentList[bid * (topK * (depth - 1) + 1) + parentTbIdx]; - for (curPosition = 0; curPosition < draftTokenNum; ++curPosition) - { - if (selectedIndex[bid * (draftTokenNum - 1) + curPosition] == tokenIdx) - { - break; - } - } - if (curPosition == draftTokenNum) - { - break; - } - } - positions[bid * draftTokenNum + tid] = position; - } -} - -//! Bit-packed variant of buildDynamicTreeKernel. -//! \param numInt32PerRow int32 count per treeMask row (buffer stride; >= ceil(draftTokenNum/32) if padded) -__global__ void buildDynamicTreeKernelPacked(int64_t const* parentList, int64_t const* selectedIndex, int32_t* treeMask, - int32_t* positions, int32_t* retrieveIndex, int32_t* retrieveNextToken, int32_t* retrieveNextSibling, - SizeType32 topK, SizeType32 depth, SizeType32 draftTokenNum, SizeType32 numInt32PerRow) -{ - int32_t bid = blockIdx.x; - int32_t tid = threadIdx.x; - - if (tid >= draftTokenNum) - { - return; - } - - int32_t rowBaseIdx = (bid * draftTokenNum + tid) * numInt32PerRow; - - treeMask[rowBaseIdx] = 1; // bit 0 = root, always visible - - int32_t position = 0; - - if (tid == 0) - { - positions[bid * draftTokenNum] = 0; - - for (int32_t i = draftTokenNum - 1; i > 0; --i) - { - retrieveIndex[bid * draftTokenNum + i] = i; - - int64_t parentTbIdx = selectedIndex[bid * (draftTokenNum - 1) + i - 1] / topK; - int32_t parentPosition = 0; - - if (parentTbIdx > 0) - { - int64_t parentTokenIdx = parentList[bid * (topK * (depth - 1) + 1) + parentTbIdx]; - for (; parentPosition < draftTokenNum; ++parentPosition) - { - if (selectedIndex[bid * (draftTokenNum - 1) + parentPosition] == parentTokenIdx) - { - ++parentPosition; - break; - } - } - } - - if (parentPosition == draftTokenNum) - { - printf("WARNING: Invalid dynamic tree! Detected a token with no parent token selected.\n"); - continue; - } - - if (retrieveNextToken[bid * draftTokenNum + parentPosition] == -1) - { - retrieveNextToken[bid * draftTokenNum + parentPosition] = i; - } - else - { - int32_t originNextToken = retrieveNextToken[bid * draftTokenNum + parentPosition]; - retrieveNextToken[bid * draftTokenNum + parentPosition] = i; - retrieveNextSibling[bid * draftTokenNum + i] = originNextToken; - } - } - retrieveIndex[bid * draftTokenNum] = 0; - } - else - { - int32_t curPosition = tid - 1; - while (position < depth + 1) - { - position += 1; - - int32_t bitPosition = curPosition + 1; // +1 because bit 0 is root - int32_t int32Idx = bitPosition / 32; - int32_t bitIdx = bitPosition % 32; - - if (int32Idx < numInt32PerRow) - { - atomicOr(&treeMask[rowBaseIdx + int32Idx], 1 << bitIdx); - } - - int64_t parentTbIdx = selectedIndex[bid * (draftTokenNum - 1) + curPosition] / topK; - if (parentTbIdx == 0) - { - break; - } - - int64_t tokenIdx = parentList[bid * (topK * (depth - 1) + 1) + parentTbIdx]; - for (curPosition = 0; curPosition < draftTokenNum; ++curPosition) - { - if (selectedIndex[bid * (draftTokenNum - 1) + curPosition] == tokenIdx) - { - break; - } - } - if (curPosition == draftTokenNum) - { - break; - } - } - positions[bid * draftTokenNum + tid] = position; - } -} - -void invokeBuildDynamicTree(int64_t const* parentList, int64_t const* selectedIndex, void* treeMask, int32_t* positions, - int32_t* retrieveIndex, int32_t* retrieveNextToken, int32_t* retrieveNextSibling, SizeType32 batchSize, - SizeType32 topK, SizeType32 depth, SizeType32 numDraftTokens, TreeMaskMode treeMaskMode, cudaStream_t stream, - SizeType32 numInt32PerRow) -{ - dim3 grid(batchSize); - dim3 block(numDraftTokens); - - if (treeMaskMode == TreeMaskMode::QLEN_ONLY_BITPACKING) - { - TLLM_CHECK_WITH_INFO( - numInt32PerRow > 0, "numInt32PerRow must be the packed treeMask row stride in int32s (from buffer shape)."); - buildDynamicTreeKernelPacked<<>>(parentList, selectedIndex, - static_cast(treeMask), positions, retrieveIndex, retrieveNextToken, retrieveNextSibling, topK, - depth, numDraftTokens, numInt32PerRow); - } - else - { - buildDynamicTreeKernel<<>>(parentList, selectedIndex, static_cast(treeMask), - positions, retrieveIndex, retrieveNextToken, retrieveNextSibling, topK, depth, numDraftTokens); - } - - sync_check_cuda_error(stream); -} - -//! \param predicts [out] accepted token ids + bonus token [bs * numDraftTokens] -//! \param acceptIndex [out] accepted path as local tree positions [bs, numSpeculativeTokens] -//! \param acceptTokenNum [out] number of accepted draft tokens per batch [bs] -//! \param candidates [in] candidate token id per tree node [bs, numDraftTokens] -//! \param retrieveIndex [in] tree node -> local index [bs, numDraftTokens] -//! \param retrieveNextToken [in] first-child pointer [bs, numDraftTokens], -1=none -//! \param retrieveNextSibling [in] next-sibling pointer [bs, numDraftTokens], -1=none -//! \param targetPredict [in] target model prediction per position [bs * numDraftTokens] -//! \param batchSize batch size -//! \param numSpeculativeTokens second dim of acceptIndex (>= max possible accepts + 1) -//! \param numDraftTokens total tree nodes per batch (including root) -__global__ void verifyDynamicTreeGreedyKernel(int64_t* predicts, int64_t* acceptIndex, int64_t* acceptTokenNum, - int64_t* acceptToken, int64_t const* candidates, int32_t const* retrieveIndex, int32_t const* retrieveNextToken, - int32_t const* retrieveNextSibling, int64_t const* targetPredict, bool const* treeValid, uint32_t batchSize, - uint32_t numSpeculativeTokens, uint32_t numDraftTokens) -{ - uint32_t bx = blockIdx.x; - uint32_t batchOffset = bx * numDraftTokens; - - // First-gen or dummy request: no valid tree, accept only the bonus token - if (treeValid != nullptr && !treeValid[bx]) - { - acceptTokenNum[bx] = 0; - acceptIndex[bx * numSpeculativeTokens] = 0; - acceptToken[bx * numSpeculativeTokens] = targetPredict[batchOffset]; - predicts[batchOffset] = targetPredict[batchOffset]; - return; - } - - int32_t lastAcceptedLocalIdx = retrieveIndex[batchOffset]; - acceptIndex[bx * numSpeculativeTokens] = lastAcceptedLocalIdx; - uint32_t numAcceptedTokens = 0; - int32_t curIndex = 0; - - // Root token: target prediction at root position - acceptToken[bx * numSpeculativeTokens] = targetPredict[batchOffset + lastAcceptedLocalIdx]; - - for (uint32_t j = 1; j < numSpeculativeTokens; ++j) - { - curIndex = retrieveNextToken[batchOffset + curIndex]; - - while (curIndex != -1) - { - int32_t draftLocalIdx = retrieveIndex[batchOffset + curIndex]; - int64_t draftTokenId = candidates[batchOffset + curIndex]; - int64_t targetTokenId = targetPredict[batchOffset + lastAcceptedLocalIdx]; - - if (draftTokenId == targetTokenId) - { - predicts[batchOffset + lastAcceptedLocalIdx] = targetTokenId; - ++numAcceptedTokens; - acceptIndex[bx * numSpeculativeTokens + numAcceptedTokens] = draftLocalIdx; - // Accepted token: target prediction at accepted draft position - acceptToken[bx * numSpeculativeTokens + numAcceptedTokens] = targetPredict[batchOffset + draftLocalIdx]; - lastAcceptedLocalIdx = draftLocalIdx; - break; - } - else - { - curIndex = retrieveNextSibling[batchOffset + curIndex]; - } - } - - if (curIndex == -1) - break; - } - - acceptTokenNum[bx] = numAcceptedTokens; - // Bonus token from target model at the last accepted position - predicts[batchOffset + lastAcceptedLocalIdx] = targetPredict[batchOffset + lastAcceptedLocalIdx]; -} - -void invokeVerifyDynamicTreeGreedy(int64_t* predicts, int64_t* acceptIndex, int64_t* acceptTokenNum, - int64_t* acceptToken, int64_t const* candidates, int32_t const* retrieveIndex, int32_t const* retrieveNextToken, - int32_t const* retrieveNextSibling, int64_t const* targetPredict, bool const* treeValid, SizeType32 batchSize, - SizeType32 numDraftTokens, SizeType32 numSpecStep, cudaStream_t stream) -{ - dim3 grid(batchSize); - dim3 block(1); - - verifyDynamicTreeGreedyKernel<<>>(predicts, acceptIndex, acceptTokenNum, acceptToken, - candidates, retrieveIndex, retrieveNextToken, retrieveNextSibling, targetPredict, treeValid, batchSize, - numSpecStep, numDraftTokens); - - sync_check_cuda_error(stream); -} - -} // namespace kernels::speculative_decoding - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h deleted file mode 100644 index c48b0f25005a..000000000000 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h +++ /dev/null @@ -1,98 +0,0 @@ -/* - * Copyright (c) 2024-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 "tensorrt_llm/common/config.h" -#include "tensorrt_llm/runtime/common.h" -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::speculative_decoding -{ - -//! \brief Tree mask mode for dynamic tree building -enum class TreeMaskMode : int32_t -{ - QLEN_ONLY = 1, //! Query length only (bool array) [bs, draftTokenNum, draftTokenNum] - QLEN_ONLY_BITPACKING = 2, //! Query length with bit packing (32 bits per int32) -}; - -//! \brief Build dynamic tree structure efficiently -//! Ported from SGLang's eagle_utils.cu for high-performance tree construction. -//! \param parentList [batchSize, topK * (depth - 1) + 1], on GPU. int64. -//! Parent indices for each token in the tree (layer-local relative indices). -//! \param selectedIndex [batchSize, numDraftTokens - 1], on GPU. int64. -//! Selected token indices (excluding root node). -//! \param treeMask output buffer [varies by mode], on GPU. int32. -//! Attention mask for tree structure. Shape depends on treeMaskMode. -//! \param positions output buffer [batchSize * numDraftTokens], on GPU. int32. -//! Position IDs for each draft token. -//! \param retrieveIndex output buffer [batchSize, numDraftTokens], on GPU. int32. -//! Local indices for retrieving tokens (0, 1, ..., numDraftTokens-1). -//! \param retrieveNextToken output buffer [batchSize, numDraftTokens], on GPU. int32. -//! Index of the first child token for each node. -//! \param retrieveNextSibling output buffer [batchSize, numDraftTokens], on GPU. int32. -//! Index of the next sibling token for each node. -//! \param batchSize runtime::SizeType32. Batch size. -//! \param topK runtime::SizeType32. Number of top-K tokens per node. -//! \param depth runtime::SizeType32. Tree depth. -//! \param numDraftTokens runtime::SizeType32. Total number of draft tokens. -//! \param treeMaskMode TreeMaskMode. Attention mask mode. -//! \param stream cuda stream -//! \param numInt32PerRow For QLEN_ONLY_BITPACKING: treeMask row stride in int32s (must be > 0). Ignored otherwise. -void invokeBuildDynamicTree(int64_t const* parentList, int64_t const* selectedIndex, void* treeMask, int32_t* positions, - int32_t* retrieveIndex, int32_t* retrieveNextToken, int32_t* retrieveNextSibling, runtime::SizeType32 batchSize, - runtime::SizeType32 topK, runtime::SizeType32 depth, runtime::SizeType32 numDraftTokens, TreeMaskMode treeMaskMode, - cudaStream_t stream, runtime::SizeType32 numInt32PerRow); - -//! \brief Verify dynamic tree using greedy strategy -//! Verifies draft tokens against target model predictions using tree traversal. -//! All index/token pointer parameters use int64. -//! \param predicts output buffer [seqLensSum], on GPU. int64. -//! Verified token predictions. -//! \param acceptIndex output buffer [batchSize, numSpecStep], on GPU. int64. -//! Indices of accepted tokens. -//! \param acceptTokenNum output buffer [batchSize], on GPU. int64. -//! Number of accepted tokens per request. -//! \param acceptToken output buffer [batchSize, numSpecStep], on GPU. int64. -//! Contiguous accepted token ids along the accepted path (root prediction, -//! accepted draft predictions, bonus token). Entry count per request = -//! acceptTokenNum + 1. -//! \param candidates [batchSize, numDraftTokens], on GPU. int64. -//! Candidate draft tokens. -//! \param retrieveIndex [batchSize, numDraftTokens], on GPU. int32. -//! Indices for retrieving tokens. -//! \param retrieveNextToken [batchSize, numDraftTokens], on GPU. int32. -//! Index of the first child token. -//! \param retrieveNextSibling [batchSize, numDraftTokens], on GPU. int32. -//! Index of the next sibling token. -//! \param targetPredict [batchSize, numDraftTokens], on GPU. int64. -//! Target model predictions. -//! \param batchSize runtime::SizeType32. Batch size. -//! \param numDraftTokens runtime::SizeType32. Total number of draft tokens. -//! \param numSpecStep runtime::SizeType32. Number of speculative steps. -//! \param stream cuda stream -void invokeVerifyDynamicTreeGreedy(int64_t* predicts, int64_t* acceptIndex, int64_t* acceptTokenNum, - int64_t* acceptToken, int64_t const* candidates, int32_t const* retrieveIndex, int32_t const* retrieveNextToken, - int32_t const* retrieveNextSibling, int64_t const* targetPredict, bool const* treeValid, - runtime::SizeType32 batchSize, runtime::SizeType32 numDraftTokens, runtime::SizeType32 numSpecStep, - cudaStream_t stream); - -} // namespace kernels::speculative_decoding - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.cu b/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.cu index 40fe23b8f8a8..8d1ca4530d87 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.cu +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.cu @@ -135,164 +135,8 @@ __global__ void updateKVCacheDraftTokenLocationBatchedKernel(std::array -__global__ void updateKVCacheDraftTokenLocationBatchedKernel2D(std::array kvCacheBuffers, - IndexType const* acceptedDraftTokensIndices2D, SizeType32 const* numAcceptedTokens, SizeType32 maxDraftLen, - SizeType32 const* pastKeyValueLengths, SizeType32 rewindDraftTokenCommonCount, - SizeType32 const* rewindDraftTokenSeparateAdjustments, SizeType32 const* seqSlotRemapping, - SizeType32 const* batchSlots, SizeType32 eltCountPerHead) -{ - auto const seqIdx = static_cast(blockIdx.x); - auto const headIdx = static_cast(blockIdx.y); - auto const layerIdx = static_cast(blockIdx.z); - auto const warpIdx = static_cast(threadIdx.x / 32); - auto const warpCount = static_cast(blockDim.x / 32); - auto const laneIdx = static_cast(threadIdx.x & 0x1f); - auto const seqDraftCount = max(numAcceptedTokens[seqIdx] - 1, 0); - auto const seqSlot = seqSlotRemapping == nullptr ? seqIdx : seqSlotRemapping[seqIdx]; - auto const maxEltCountPerMove = seqDraftCount > 0 - ? static_cast(kUpdateKVCacheKernelShmSize / sizeof(MoveEltType) / seqDraftCount) - : 0; - auto const eltCountPerMove = min(maxEltCountPerMove, eltCountPerHead); - if (seqDraftCount == 0 || eltCountPerMove == 0) - { - return; - } - auto const seqRow = seqIdx * maxDraftLen; - KVCacheBuffer& kvCacheBuffer = kvCacheBuffers[layerIdx]; - auto tokenStartIdx = pastKeyValueLengths[seqSlot] - rewindDraftTokenCommonCount; - if (rewindDraftTokenSeparateAdjustments != nullptr) - { - auto const batchSlot = batchSlots == nullptr ? seqIdx : batchSlots[seqIdx]; - tokenStartIdx -= rewindDraftTokenSeparateAdjustments[batchSlot]; - } - // Guard against negative tokenStartIdx (e.g. during CUDA graph warmup - // when kv_lens < rewindDraftTokenCommonCount). - if (tokenStartIdx < 0) - { - return; - } - __shared__ char loadSmemBuffer[kUpdateKVCacheKernelShmSize]; - auto* eltLoadSmemBuffer = reinterpret_cast(&loadSmemBuffer[0]); - for (SizeType32 startChannelOffset = 0; startChannelOffset < eltCountPerHead; startChannelOffset += eltCountPerMove) - { - SizeType32 eltCountCurrentMove = min(eltCountPerMove, eltCountPerHead - startChannelOffset); - for (SizeType32 tokenIdx = warpIdx; tokenIdx < seqDraftCount; tokenIdx += warpCount) - { - // Clamp negative indices (e.g. -1 fill during CUDA graph warmup) to 0 - // to prevent OOB access. Correctness is N/A for warmup data. - auto const tokenPos = max(acceptedDraftTokensIndices2D[seqRow + tokenIdx], 0); - auto* tokenSmemBuffer = eltLoadSmemBuffer + tokenIdx * eltCountCurrentMove; - auto const tokenKVPosition = tokenStartIdx + tokenPos; - auto* kPtr = reinterpret_cast(kvCacheBuffer.getKBlockPtr(seqSlot, tokenKVPosition)); - for (SizeType32 loadChannelIdx = laneIdx; loadChannelIdx < eltCountCurrentMove; loadChannelIdx += 32) - { - auto const channelIdx = loadChannelIdx + startChannelOffset; - auto const kvLocationIdx - = kvCacheBuffer.getKVLocalIdx(tokenKVPosition, headIdx, eltCountPerHead, channelIdx); - tokenSmemBuffer[loadChannelIdx] = kPtr[kvLocationIdx]; - } - } - __syncthreads(); - for (SizeType32 tokenIdx = warpIdx; tokenIdx < seqDraftCount; tokenIdx += warpCount) - { - auto const tokenPos = tokenIdx; - auto* tokenSmemBuffer = eltLoadSmemBuffer + tokenIdx * eltCountCurrentMove; - auto const tokenKVPosition = tokenStartIdx + tokenPos; - auto* kPtr = reinterpret_cast(kvCacheBuffer.getKBlockPtr(seqSlot, tokenKVPosition)); - for (SizeType32 loadChannelIdx = laneIdx; loadChannelIdx < eltCountCurrentMove; loadChannelIdx += 32) - { - auto const channelIdx = loadChannelIdx + startChannelOffset; - auto const kvLocationIdx - = kvCacheBuffer.getKVLocalIdx(tokenKVPosition, headIdx, eltCountPerHead, channelIdx); - kPtr[kvLocationIdx] = tokenSmemBuffer[loadChannelIdx]; - } - } - __syncthreads(); - for (SizeType32 tokenIdx = warpIdx; tokenIdx < seqDraftCount; tokenIdx += warpCount) - { - auto const tokenPos = max(acceptedDraftTokensIndices2D[seqRow + tokenIdx], 0); - auto* tokenSmemBuffer = eltLoadSmemBuffer + tokenIdx * eltCountCurrentMove; - auto const tokenKVPosition = tokenStartIdx + tokenPos; - auto* vPtr = reinterpret_cast(kvCacheBuffer.getVBlockPtr(seqSlot, tokenKVPosition)); - for (SizeType32 loadChannelIdx = laneIdx; loadChannelIdx < eltCountCurrentMove; loadChannelIdx += 32) - { - auto const channelIdx = loadChannelIdx + startChannelOffset; - auto const kvLocationIdx - = kvCacheBuffer.getKVLocalIdx(tokenKVPosition, headIdx, eltCountPerHead, channelIdx); - tokenSmemBuffer[loadChannelIdx] = vPtr[kvLocationIdx]; - } - } - __syncthreads(); - for (SizeType32 tokenIdx = warpIdx; tokenIdx < seqDraftCount; tokenIdx += warpCount) - { - auto const tokenPos = tokenIdx; - auto* tokenSmemBuffer = eltLoadSmemBuffer + tokenPos * eltCountCurrentMove; - auto const tokenKVPosition = tokenStartIdx + tokenPos; - auto* vPtr = reinterpret_cast(kvCacheBuffer.getVBlockPtr(seqSlot, tokenKVPosition)); - for (SizeType32 loadChannelIdx = laneIdx; loadChannelIdx < eltCountCurrentMove; loadChannelIdx += 32) - { - auto const channelIdx = loadChannelIdx + startChannelOffset; - auto const kvLocationIdx - = kvCacheBuffer.getKVLocalIdx(tokenKVPosition, headIdx, eltCountPerHead, channelIdx); - vPtr[kvLocationIdx] = tokenSmemBuffer[loadChannelIdx]; - } - } - __syncthreads(); - } -} } // namespace -template -void updateKVCacheDraftTokenLocationBatched2D(KVCacheBuffer const* kvCacheBuffers, - IndexType const* acceptedDraftTokensIndices2D, SizeType32 const* numAcceptedTokens, SizeType32 maxDraftLen, - SizeType32 const* pastKeyValueLengths, SizeType32 layerCount, SizeType32 seqCount, SizeType32 numKVHeads, - SizeType32 sizeInBytesPerKVHead, SizeType32 rewindDraftTokenCommonCount, - SizeType32 const* rewindDraftTokenSeparateAdjustments, SizeType32 const* seqSlotRemapping, - SizeType32 const* batchSlots, cudaStream_t stream) -{ - static_assert(MaxLayerCount * sizeof(KVCacheBuffer) <= 3072); - if (seqCount == 0 || layerCount == 0) - { - return; - } - SizeType32 alignedBytes = 16; - while (alignedBytes > 0 && (sizeInBytesPerKVHead % alignedBytes != 0)) - { - alignedBytes >>= 1; - } - TLLM_CHECK_WITH_INFO(alignedBytes > 0, "alignedByte should be positive"); - SizeType32 eltCountPerHead = sizeInBytesPerKVHead / alignedBytes; - dim3 grid(seqCount, numKVHeads, layerCount); - dim3 block(128, 1, 1); - std::array kvCacheBufferArray; - for (SizeType32 i = 0; i < layerCount; i++) - { - kvCacheBufferArray[i] = kvCacheBuffers[i]; - } - void (*pKernelFunc)(std::array, IndexType const*, SizeType32 const*, SizeType32, - SizeType32 const*, SizeType32, SizeType32 const*, SizeType32 const*, SizeType32 const*, SizeType32) - = nullptr; - switch (alignedBytes) - { - case 16: pKernelFunc = &updateKVCacheDraftTokenLocationBatchedKernel2D; break; - case 8: pKernelFunc = &updateKVCacheDraftTokenLocationBatchedKernel2D; break; - case 4: - pKernelFunc = &updateKVCacheDraftTokenLocationBatchedKernel2D; - break; - case 2: pKernelFunc = &updateKVCacheDraftTokenLocationBatchedKernel2D; break; - default: - TLLM_CHECK_WITH_INFO(alignedBytes == 1, "Strange alignedBytes"); - pKernelFunc = &updateKVCacheDraftTokenLocationBatchedKernel2D; - break; - } - pKernelFunc<<>>(kvCacheBufferArray, acceptedDraftTokensIndices2D, numAcceptedTokens, - maxDraftLen, pastKeyValueLengths, rewindDraftTokenCommonCount, rewindDraftTokenSeparateAdjustments, - seqSlotRemapping, batchSlots, eltCountPerHead); - sync_check_cuda_error(stream); -} - template void updateKVCacheDraftTokenLocationBatched(KVCacheBuffer const* kvCacheBuffers, SizeType32 const* seqAcceptedDraftTokenOffsets, IndexType const* packedAcceptedDraftTokensIndices, @@ -397,55 +241,6 @@ void updateKVCacheDraftTokenLocation(KVCacheBuffer const* kvCacheBuffers, } } -template -void updateKVCacheDraftTokenLocation2D(KVCacheBuffer const* kvCacheBuffers, - IndexType const* acceptedDraftTokensIndices2D, SizeType32 const* numAcceptedTokens, SizeType32 maxDraftLen, - SizeType32 const* pastKeyValueLengths, SizeType32 layerCount, SizeType32 seqCount, SizeType32 numKVHeads, - SizeType32 sizeInBytesPerKVHead, SizeType32 rewindDraftTokenCommonCount, - SizeType32 const* rewindDraftTokenSeparateAdjustments, SizeType32 const* seqSlotRemapping, - SizeType32 const* batchSlots, cudaStream_t stream) -{ - SizeType32 startLayer = 0; - static constexpr SizeType32 kMaxLayersPerIter = 32; - while (startLayer < layerCount) - { - SizeType32 microBatchLayerCount = std::min(layerCount - startLayer, kMaxLayersPerIter); - updateKVCacheDraftTokenLocationBatched2D(kvCacheBuffers + startLayer, - acceptedDraftTokensIndices2D, numAcceptedTokens, maxDraftLen, pastKeyValueLengths, microBatchLayerCount, - seqCount, numKVHeads, sizeInBytesPerKVHead, rewindDraftTokenCommonCount, - rewindDraftTokenSeparateAdjustments, seqSlotRemapping, batchSlots, stream); - startLayer += microBatchLayerCount; - } -} - -void updateKVBlockArrayDraftTokenLocation2D(IndexType const* acceptedDraftTokensIndices2D, - SizeType32 const* numAcceptedTokens, SizeType32 maxDraftLen, SizeType32 const* pastKeyValueLengths, - void* const* pointerArray, KVBlockArray::DataType* offsetArray, SizeType32 layerCount, SizeType32 seqCount, - SizeType32 numKVHeads, SizeType32 sizeInBytesPerKVHead, SizeType32 rewindDraftTokenCommonCount, - SizeType32 const* rewindDraftTokenSeparateAdjustments, SizeType32 const* seqSlotRemapping, - SizeType32 const* batchSlots, SizeType32 maxKVCacheLen, SizeType32 maxBlocksPerSeq, SizeType32 tokensPerBlock, - bool canUseOneMoreBlock, cudaStream_t stream) -{ - std::vector kvBlockArrays; - kvBlockArrays.reserve(layerCount); - auto const bytesPerToken = numKVHeads * sizeInBytesPerKVHead; - auto const bytesPerBlock = tokensPerBlock * bytesPerToken; - for (SizeType32 layerIdx = 0; layerIdx < layerCount; layerIdx++) - { - auto const layerOffset = layerIdx * 2 * bytesPerBlock; - auto* const primaryPoolPointer - = reinterpret_cast(reinterpret_cast(pointerArray[0]) + layerOffset); - auto* const secondaryPoolPointer - = reinterpret_cast(reinterpret_cast(pointerArray[1]) + layerOffset); - - kvBlockArrays.emplace_back(seqCount, maxBlocksPerSeq, tokensPerBlock, bytesPerToken, maxKVCacheLen, - maxKVCacheLen, 0, canUseOneMoreBlock, primaryPoolPointer, secondaryPoolPointer, offsetArray); - } - updateKVCacheDraftTokenLocation2D(kvBlockArrays.data(), acceptedDraftTokensIndices2D, numAcceptedTokens, - maxDraftLen, pastKeyValueLengths, layerCount, seqCount, numKVHeads, sizeInBytesPerKVHead, - rewindDraftTokenCommonCount, rewindDraftTokenSeparateAdjustments, seqSlotRemapping, batchSlots, stream); -} - void updateLinearKVCacheDraftTokenLocation(SizeType32 const* seqAcceptedDraftTokenOffsets, IndexType const* packedAcceptedDraftTokensIndices, SizeType32 const* pastKeyValueLengths, int8_t* const* pastKeyValueList, SizeType32 layerCount, SizeType32 seqCount, SizeType32 numKVHeads, diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.h b/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.h index fd3cbd53512f..f8551db9b75c 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.h +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.h @@ -208,21 +208,6 @@ void updateKVBlockArrayDraftTokenLocation(runtime::SizeType32 const* seqAccepted runtime::SizeType32 maxKVCacheLen, runtime::SizeType32 maxBlocksPerSeq, runtime::SizeType32 tokensPerBlock, bool canUseOneMoreBlock, cudaStream_t stream); -/*! - * Update Block KV cache using 2D indices tensor (CUDA graph compatible). - * Accepts a fixed-size 2D tensor where each row contains accepted draft token indices, - * padded with -1. The kernel computes draft count as max(numAcceptedTokens[i] - 1, 0). - * Context requests naturally have numAcceptedTokens=1 → 0 draft tokens. - */ -void updateKVBlockArrayDraftTokenLocation2D(IndexType const* acceptedDraftTokensIndices2D, - runtime::SizeType32 const* numAcceptedTokens, runtime::SizeType32 maxDraftLen, - runtime::SizeType32 const* pastKeyValueLengths, void* const* pointerArray, KVBlockArray::DataType* offsetArray, - runtime::SizeType32 layerCount, runtime::SizeType32 seqCount, runtime::SizeType32 numKVHeads, - runtime::SizeType32 sizeInBytesPerKVHead, runtime::SizeType32 rewindDraftTokenCommonCount, - runtime::SizeType32 const* rewindDraftTokenSeparateAdjustments, runtime::SizeType32 const* seqSlotRemapping, - runtime::SizeType32 const* batchSlots, runtime::SizeType32 maxKVCacheLen, runtime::SizeType32 maxBlocksPerSeq, - runtime::SizeType32 tokensPerBlock, bool canUseOneMoreBlock, cudaStream_t stream); - } // namespace kernels::speculative_decoding TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h index 1212b7ad057c..d286eab8df49 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h @@ -292,10 +292,6 @@ struct TllmGenFmhaRunnerParams int32_t mLayerIdx = 0; // Whether the spec-dec tree is used. bool mIsSpecDecTree = false; - // The max seqLenQ used as row stride for generalPackedCustoMaskPtr. - // When seqlensQPtr[i] < mPackedMaskMaxSeqLenQ, the packed mask tensor has - // row stride ceilDiv(mPackedMaskMaxSeqLenQ, 32) rather than ceilDiv(seqLenQ, 32). - int32_t mPackedMaskMaxSeqLenQ = 0; // set the attention mask type TllmGenFmhaRunnerParams& setAttentionMaskType(std::int8_t maskType) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/prepareCustomMask.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/prepareCustomMask.cu index 18ec12fd52c3..58e8a2c6c366 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/prepareCustomMask.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/prepareCustomMask.cu @@ -131,13 +131,6 @@ __global__ void prepareCustomMaskBuffersKernelForKeepsMmaAb( // The sequence length of tensor KV. int32_t const seqLenKv = seqLensKvPtr[batchIdx]; - // The packed mask tensor row stride: use mPackedMaskMaxSeqLenQ when available, - // because the packed mask tensor has shape [bs, maxSeqLenQ, ceilDiv(maxSeqLenQ, 32)] - // even when the actual seqLenQ is smaller (e.g., drafter layers in dynamic tree). - int32_t const packedMaskSeqLenQ - = runnerParams.mPackedMaskMaxSeqLenQ > 0 ? runnerParams.mPackedMaskMaxSeqLenQ : seqLenQ; - int32_t const packedMaskNumBlocks = ceilDiv(packedMaskSeqLenQ, 32); - // Calculate global Q token index (flattened across heads) int32_t const qTokensPerBlock = static_cast(blockDim.x); int32_t const flattenedQIdx = qGroupIdx * qTokensPerBlock + qThreadIdx; @@ -172,12 +165,12 @@ __global__ void prepareCustomMaskBuffersKernelForKeepsMmaAb( else { // Sparse region: check the input mask - // Input mask shape: [bs, packedMaskSeqLenQ, packedMaskNumBlocks] + // Input mask shape: [bs, seqLenQ, ceilDiv(seqLenQ, 32)] // The KV dimension in the mask corresponds to Q positions (tree mask) int32_t const qPosInTree = tokenIdxKv - firstSparseMaskOffsetKv; if (qPosInTree < seqLenQ) { - int32_t const qMaskBaseIdx = (batchIdx * packedMaskSeqLenQ + tokenIdxQ) * packedMaskNumBlocks; + int32_t const qMaskBaseIdx = (batchIdx * seqLenQ + tokenIdxQ) * ceilDiv(seqLenQ, 32); int32_t const packedMaskIdx = qMaskBaseIdx + (qPosInTree >> 5); int32_t const bitPos = qPosInTree & 0x1F; randomMask = (customMaskInputPtr[packedMaskIdx] >> bitPos) & 1; diff --git a/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp b/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp index 29a94017e738..af025394b836 100644 --- a/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp +++ b/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp @@ -499,7 +499,6 @@ void XqaDispatcher::runImpl( tllmRunnerParams.mLayerIdx = params.layer_idx; tllmRunnerParams.seqLensQPtr = params.spec_decoding_generation_lengths; tllmRunnerParams.generalPackedCustoMaskPtr = params.spec_decoding_packed_mask; - tllmRunnerParams.mPackedMaskMaxSeqLenQ = params.spec_decoding_max_generation_length; tllmRunnerParams.customMaskPtr = params.spec_decoding_bl_tree_mask; tllmRunnerParams.customMaskOffsetsPtr = params.spec_decoding_bl_tree_mask_offset; tllmRunnerParams.firstSparseMaskOffsetsKvPtr = params.spec_bl_tree_first_sparse_mask_offset_kv; diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index cfe79ca7493b..5d715a87fece 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -108,7 +108,6 @@ add_library( weightOnlyQuantGemm.cpp weightOnlyQuantOp.cpp specDecOp.cpp - dynamicTreeOp.cpp loraOp.cpp finegrained_mixed_dtype_gemm_thop.cpp tinygemm2.cpp diff --git a/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp b/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp deleted file mode 100644 index 2c70fcfc5cb1..000000000000 --- a/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp +++ /dev/null @@ -1,229 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h" -#include "tensorrt_llm/runtime/torchUtils.h" - -namespace th = torch; -namespace tk = tensorrt_llm::kernels::speculative_decoding; - -TRTLLM_NAMESPACE_BEGIN - -namespace torch_ext -{ - -//! \brief Build dynamic tree structure (in-place, writes to pre-allocated output buffers) -//! All index tensors use int64 to match PyTorch's default integer dtype. -void build_dynamic_tree_op(th::Tensor& parentList, th::Tensor& selectedIndex, th::Tensor& treeMask, - th::Tensor& positions, th::Tensor& retrieveIndex, th::Tensor& retrieveNextToken, th::Tensor& retrieveNextSibling, - int64_t topK, int64_t depth, int64_t numDraftTokens, int64_t treeMaskMode, int64_t numInt32PerRow) -{ - // Validate inputs - TORCH_CHECK(parentList.dim() == 2, "parentList must be 2D tensor"); - TORCH_CHECK(selectedIndex.dim() == 2, "selectedIndex must be 2D tensor"); - TORCH_CHECK(parentList.scalar_type() == torch::kInt64, "parentList must be int64 tensor"); - TORCH_CHECK(selectedIndex.scalar_type() == torch::kInt64, "selectedIndex must be int64 tensor"); - - int64_t batchSize = parentList.size(0); - TORCH_CHECK(selectedIndex.size(0) == batchSize, "Batch size mismatch"); - TORCH_CHECK(selectedIndex.size(1) == numDraftTokens - 1, "selectedIndex size mismatch"); - TORCH_CHECK(numDraftTokens <= 1024, "numDraftTokens (", numDraftTokens, ") exceeds CUDA block size limit of 1024"); - - auto device = parentList.device(); - auto stream = at::cuda::getCurrentCUDAStream(device.index()); - - // Reset output buffers - treeMask.zero_(); - positions.zero_(); - retrieveIndex.zero_(); - retrieveNextToken.fill_(-1); - retrieveNextSibling.fill_(-1); - - // Call kernel - tk::invokeBuildDynamicTree(parentList.data_ptr(), selectedIndex.data_ptr(), treeMask.data_ptr(), - positions.data_ptr(), retrieveIndex.data_ptr(), retrieveNextToken.data_ptr(), - retrieveNextSibling.data_ptr(), batchSize, topK, depth, numDraftTokens, - static_cast(treeMaskMode), stream, numInt32PerRow); -} - -//! \brief Verify dynamic tree using greedy strategy -std::tuple verify_dynamic_tree_greedy_op(th::Tensor& candidates, - th::Tensor& retrieveIndex, th::Tensor& retrieveNextToken, th::Tensor& retrieveNextSibling, - th::Tensor& targetPredict, th::Tensor& treeValid, int64_t numSpecStep) -{ - // Validate inputs - TORCH_CHECK(candidates.dim() == 2, "candidates must be 2D tensor"); - TORCH_CHECK(retrieveIndex.dim() == 2, "retrieveIndex must be 2D tensor"); - TORCH_CHECK(retrieveNextToken.dim() == 2, "retrieveNextToken must be 2D tensor"); - TORCH_CHECK(retrieveNextSibling.dim() == 2, "retrieveNextSibling must be 2D tensor"); - TORCH_CHECK(targetPredict.dim() == 2, "targetPredict must be 2D tensor"); - - int64_t batchSize = candidates.size(0); - int64_t numDraftTokens = candidates.size(1); - - TORCH_CHECK(retrieveIndex.size(0) == batchSize, "Batch size mismatch"); - TORCH_CHECK(retrieveNextToken.size(0) == batchSize, "Batch size mismatch"); - TORCH_CHECK(retrieveNextSibling.size(0) == batchSize, "Batch size mismatch"); - TORCH_CHECK(targetPredict.size(0) == batchSize, "Batch size mismatch"); - - TORCH_CHECK(retrieveIndex.size(1) == numDraftTokens, "retrieveIndex size mismatch"); - TORCH_CHECK(retrieveNextToken.size(1) == numDraftTokens, "retrieveNextToken size mismatch"); - TORCH_CHECK(retrieveNextSibling.size(1) == numDraftTokens, "retrieveNextSibling size mismatch"); - TORCH_CHECK(targetPredict.size(1) == numDraftTokens, "targetPredict size mismatch"); - - TORCH_CHECK(treeValid.dim() == 1, "treeValid must be 1D tensor"); - TORCH_CHECK(treeValid.size(0) >= batchSize, "treeValid buffer too small"); - TORCH_CHECK(treeValid.scalar_type() == torch::kBool, "treeValid must be bool tensor"); - - auto device = candidates.device(); - auto stream = at::cuda::getCurrentCUDAStream(device.index()); - - // Compute total sequence length sum for predicts buffer - int64_t seqLensSum = batchSize * numDraftTokens; - - // Allocate output tensors as int64 - auto predicts = torch::zeros({seqLensSum}, torch::TensorOptions().dtype(torch::kInt64).device(device)); - auto acceptIndex - = torch::zeros({batchSize, numSpecStep}, torch::TensorOptions().dtype(torch::kInt64).device(device)); - auto acceptTokenNum = torch::zeros({batchSize}, torch::TensorOptions().dtype(torch::kInt64).device(device)); - auto acceptToken - = torch::zeros({batchSize, numSpecStep}, torch::TensorOptions().dtype(torch::kInt64).device(device)); - - // Convert candidates and targetPredict to int64 if needed - auto candidatesInt64 = candidates.to(torch::kInt64); - auto targetPredictInt64 = targetPredict.to(torch::kInt64); - - tk::invokeVerifyDynamicTreeGreedy(predicts.data_ptr(), acceptIndex.data_ptr(), - acceptTokenNum.data_ptr(), acceptToken.data_ptr(), candidatesInt64.data_ptr(), - retrieveIndex.data_ptr(), retrieveNextToken.data_ptr(), - retrieveNextSibling.data_ptr(), targetPredictInt64.data_ptr(), treeValid.data_ptr(), - batchSize, numDraftTokens, numSpecStep, stream); - - return std::make_tuple(predicts, acceptIndex, acceptTokenNum, acceptToken); -} - -//! \brief In-place variant of verify_dynamic_tree_greedy_op. -//! Writes to pre-allocated output buffers instead of allocating new tensors. -void verify_dynamic_tree_greedy_out_op(th::Tensor& candidates, th::Tensor& retrieveIndex, th::Tensor& retrieveNextToken, - th::Tensor& retrieveNextSibling, th::Tensor& targetPredict, th::Tensor& predicts, th::Tensor& acceptIndex, - th::Tensor& acceptTokenNum, th::Tensor& acceptToken, th::Tensor& treeValid, int64_t numSpecStep) -{ - // Validate inputs - TORCH_CHECK(candidates.dim() == 2, "candidates must be 2D tensor"); - TORCH_CHECK(retrieveIndex.dim() == 2, "retrieveIndex must be 2D tensor"); - TORCH_CHECK(retrieveNextToken.dim() == 2, "retrieveNextToken must be 2D tensor"); - TORCH_CHECK(retrieveNextSibling.dim() == 2, "retrieveNextSibling must be 2D tensor"); - TORCH_CHECK(targetPredict.dim() == 2, "targetPredict must be 2D tensor"); - TORCH_CHECK(candidates.scalar_type() == torch::kInt64, "candidates must be int64 tensor"); - TORCH_CHECK(targetPredict.scalar_type() == torch::kInt64, "targetPredict must be int64 tensor"); - - int64_t batchSize = candidates.size(0); - int64_t numDraftTokens = candidates.size(1); - - TORCH_CHECK(retrieveIndex.size(0) == batchSize, "Batch size mismatch"); - TORCH_CHECK(retrieveNextToken.size(0) == batchSize, "Batch size mismatch"); - TORCH_CHECK(retrieveNextSibling.size(0) == batchSize, "Batch size mismatch"); - TORCH_CHECK(targetPredict.size(0) == batchSize, "Batch size mismatch"); - - TORCH_CHECK(retrieveIndex.size(1) == numDraftTokens, "retrieveIndex size mismatch"); - TORCH_CHECK(retrieveNextToken.size(1) == numDraftTokens, "retrieveNextToken size mismatch"); - TORCH_CHECK(retrieveNextSibling.size(1) == numDraftTokens, "retrieveNextSibling size mismatch"); - TORCH_CHECK(targetPredict.size(1) == numDraftTokens, "targetPredict size mismatch"); - - TORCH_CHECK(treeValid.dim() == 1, "treeValid must be 1D tensor"); - TORCH_CHECK(treeValid.size(0) >= batchSize, "treeValid buffer too small"); - TORCH_CHECK(treeValid.scalar_type() == torch::kBool, "treeValid must be bool tensor"); - - // Validate output buffers - TORCH_CHECK(predicts.scalar_type() == torch::kInt64, "predicts must be int64 tensor"); - TORCH_CHECK(acceptIndex.scalar_type() == torch::kInt64, "acceptIndex must be int64 tensor"); - TORCH_CHECK(acceptTokenNum.scalar_type() == torch::kInt64, "acceptTokenNum must be int64 tensor"); - TORCH_CHECK(acceptToken.scalar_type() == torch::kInt64, "acceptToken must be int64 tensor"); - TORCH_CHECK(predicts.size(0) >= batchSize * numDraftTokens, "predicts buffer too small"); - TORCH_CHECK(acceptIndex.size(0) >= batchSize && acceptIndex.size(1) >= numSpecStep, "acceptIndex buffer too small"); - TORCH_CHECK(acceptTokenNum.size(0) >= batchSize, "acceptTokenNum buffer too small"); - TORCH_CHECK(acceptToken.size(0) >= batchSize && acceptToken.size(1) >= numSpecStep, "acceptToken buffer too small"); - - auto stream = at::cuda::getCurrentCUDAStream(candidates.device().index()); - - // Zero output buffers in-place - predicts.zero_(); - acceptIndex.zero_(); - acceptTokenNum.zero_(); - acceptToken.zero_(); - - tk::invokeVerifyDynamicTreeGreedy(predicts.data_ptr(), acceptIndex.data_ptr(), - acceptTokenNum.data_ptr(), acceptToken.data_ptr(), candidates.data_ptr(), - retrieveIndex.data_ptr(), retrieveNextToken.data_ptr(), - retrieveNextSibling.data_ptr(), targetPredict.data_ptr(), treeValid.data_ptr(), - batchSize, numDraftTokens, numSpecStep, stream); -} - -} // namespace torch_ext - -TRTLLM_NAMESPACE_END - -//////////////////////////////////////////////////////////////////////////////////////////////////////////// - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def( - "build_dynamic_tree_op(Tensor parentList, Tensor selectedIndex, " - "Tensor(a!) treeMask, Tensor(b!) positions, Tensor(c!) retrieveIndex, " - "Tensor(d!) retrieveNextToken, Tensor(e!) retrieveNextSibling, " - "int topK, int depth, int numDraftTokens, int treeMaskMode, " - "int numInt32PerRow) -> ()"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("build_dynamic_tree_op", &tensorrt_llm::torch_ext::build_dynamic_tree_op); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////////////// - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def( - "verify_dynamic_tree_greedy_op(Tensor candidates, Tensor retrieveIndex, Tensor retrieveNextToken, " - "Tensor retrieveNextSibling, Tensor targetPredict, Tensor treeValid, int numSpecStep) -> (Tensor, Tensor, " - "Tensor, Tensor)"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("verify_dynamic_tree_greedy_op", &tensorrt_llm::torch_ext::verify_dynamic_tree_greedy_op); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////////////// - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def( - "verify_dynamic_tree_greedy_out_op(" - "Tensor candidates, Tensor retrieveIndex, Tensor retrieveNextToken, " - "Tensor retrieveNextSibling, Tensor targetPredict, " - "Tensor(a!) predicts, Tensor(b!) acceptIndex, Tensor(c!) acceptTokenNum, " - "Tensor(d!) acceptToken, Tensor treeValid, int numSpecStep) -> ()"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("verify_dynamic_tree_greedy_out_op", &tensorrt_llm::torch_ext::verify_dynamic_tree_greedy_out_op); -} diff --git a/cpp/tensorrt_llm/thop/parallelDecodeKVCacheUpdateOp.cpp b/cpp/tensorrt_llm/thop/parallelDecodeKVCacheUpdateOp.cpp index 3a25d903bf32..400cf81033ba 100644 --- a/cpp/tensorrt_llm/thop/parallelDecodeKVCacheUpdateOp.cpp +++ b/cpp/tensorrt_llm/thop/parallelDecodeKVCacheUpdateOp.cpp @@ -113,73 +113,9 @@ void updateKVCacheDraftTokenLocation(torch::Tensor seqAcceptedDraftTokenOffsetsT } } -void updateKVCacheDraftTokenLocation2D(torch::Tensor acceptedDraftTokensIndices2DTensor, - torch::Tensor numAcceptedTokensTensor, torch::Tensor pastKeyValueLengthsTensor, bool usePagedKVCache, - int64_t layerCount, int64_t numKVHeads, int64_t headSizeInBytes, int64_t rewindDraftTokenCount, - int64_t maxKVCacheLen, th::optional pointerArrayOpt = th::nullopt, - th::optional offsetArrayOpt = th::nullopt, th::optional maxBlocksPerSeqOpt = th::nullopt, - th::optional tokensPerBlockOpt = th::nullopt, th::optional stream_ptr = th::nullopt) -{ - TLLM_CHECK_WITH_INFO( - at::cuda::is_available(), "update_kv_cache_draft_token_location_2d should be called with cuda enabled."); - cudaStream_t stream; - if (stream_ptr.has_value()) - { - stream = reinterpret_cast(stream_ptr.value()); - } - else - { - stream = at::cuda::getCurrentCUDAStream(); - } - TLLM_CHECK_WITH_INFO(acceptedDraftTokensIndices2DTensor.dim() == 2 - && acceptedDraftTokensIndices2DTensor.scalar_type() == torch::kInt, - "accepted_draft_tokens_indices_2d tensor should be 2D int tensor."); - int seqCount = acceptedDraftTokensIndices2DTensor.size(0); - int maxDraftLen = acceptedDraftTokensIndices2DTensor.size(1); - - TLLM_CHECK_WITH_INFO(numAcceptedTokensTensor.dim() == 1 && numAcceptedTokensTensor.size(0) == seqCount - && numAcceptedTokensTensor.scalar_type() == torch::kInt, - "num_accepted_tokens tensor should be 1D int tensor with same length as seqCount."); - - TLLM_CHECK_WITH_INFO(pastKeyValueLengthsTensor.dim() == 1 && pastKeyValueLengthsTensor.size(0) == seqCount - && pastKeyValueLengthsTensor.scalar_type() == torch::kInt, - "past_key_value_lengths tensor should be 1D int tensor with same length as seqCount"); - - if (usePagedKVCache) - { - TLLM_CHECK_WITH_INFO( - pointerArrayOpt.has_value(), "pool_pointer_array should be set when using paged KV cache."); - TLLM_CHECK_WITH_INFO(offsetArrayOpt.has_value(), "block_offset_array should be set when using paged KV cache."); - TLLM_CHECK_WITH_INFO( - maxBlocksPerSeqOpt.has_value(), "max_blocks_per_seq should be set when using paged KV cache."); - TLLM_CHECK_WITH_INFO( - tokensPerBlockOpt.has_value(), "tokens_per_block should be set when using paged KV cache."); - - auto const& pointerArray = pointerArrayOpt.value(); - auto const& offsetArray = offsetArrayOpt.value(); - bool constexpr canUseOneMoreBlock{true}; - - tksd::updateKVBlockArrayDraftTokenLocation2D(acceptedDraftTokensIndices2DTensor.data_ptr(), - numAcceptedTokensTensor.data_ptr(), maxDraftLen, pastKeyValueLengthsTensor.data_ptr(), - reinterpret_cast(pointerArray.data_ptr()), - reinterpret_cast( - offsetArray.data_ptr()), - layerCount, seqCount, numKVHeads, headSizeInBytes, rewindDraftTokenCount, nullptr, nullptr, nullptr, - maxKVCacheLen, maxBlocksPerSeqOpt.value(), tokensPerBlockOpt.value(), canUseOneMoreBlock, stream); - } - else - { - TLLM_CHECK_WITH_INFO(false, "update_kv_cache_draft_token_location_2d only supports paged KV cache."); - } -} - } // namespace torch_ext TRTLLM_NAMESPACE_END static auto update_kv_cache_draft_token_location = torch::RegisterOperators( "tensorrt_llm::update_kv_cache_draft_token_location", &tensorrt_llm::torch_ext::updateKVCacheDraftTokenLocation); - -static auto update_kv_cache_draft_token_location_2d - = torch::RegisterOperators("tensorrt_llm::update_kv_cache_draft_token_location_2d", - &tensorrt_llm::torch_ext::updateKVCacheDraftTokenLocation2D); diff --git a/docs/source/features/feature-combination-matrix.md b/docs/source/features/feature-combination-matrix.md index 2d74fadc5598..7dbfdea2b07c 100644 --- a/docs/source/features/feature-combination-matrix.md +++ b/docs/source/features/feature-combination-matrix.md @@ -1,24 +1,23 @@ # Feature Combination Matrix -| Feature | Overlap Scheduler | CUDA Graph | Tensor Parallelism | Pipeline Parallelism | Expert Parallelism | Helix Parallelism | Attention Data Parallelism | Disaggregated Serving | Chunked Prefill | MTP | EAGLE-3(One Model Engine) — Linear | EAGLE-3(One Model Engine) — Dynamic | EAGLE-3(Two Model Engine) | Torch Sampler | TLLM C++ Sampler | KV Cache Reuse | Sliding Window Attention | Logits Post Processor | Guided Decoding | LoRA | -| ---------------------------------------- | ----------------- | ---------- | ------------------ | -------------------- | ------------------ | ----------------- | -------------------------- | --------------------- | --------------- | -------- | ---------------------------------- | ----------------------------------- | ------------------------- | ------------- | ---------------- | -------------- | ---------------------- | --------------------- | --------------- | -------- | -| Overlap Scheduler | --- | | | | | | | | | | | | | | | | | | | | -| CUDA Graph | Yes | --- | | | | | | | | | | | | | | | | | | | -| Tensor Parallelism | Yes | Yes | --- | | | | | | | | | | | | | | | | | | -| Pipeline Parallelism | Yes | Yes | Yes | --- | | | | | | | | | | | | | | | | | -| Expert Parallelism | Yes | Yes | Yes | Yes | --- | | | | | | | | | | | | | | | | -| Helix Parallelism | Untested | Yes | Yes | Yes | Yes | --- | | | | | | | | | | | | | | | -| Attention Data Parallelism | Yes | Yes | Yes | Yes | Yes | Known issues | --- | | | | | | | | | | | | | | -| Disaggregated Serving | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | | | | | | | | | | -| Chunked Prefill | Yes | Yes | Yes | Untested | Yes | Yes | Yes | Yes | --- | | | | | | | | | | | | -| MTP | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | --- | | | | | | | | | | | -| EAGLE-3(One Model Engine) — Linear | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | No | --- | | | | | | | | | | -| EAGLE-3(One Model Engine) — Dynamic | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | No | No | --- | | | | | | | | | -| EAGLE-3(Two Model Engine) | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | No | No | No | --- | | | | | | | | -| Torch Sampler | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | | | | -| TLLM C++ Sampler | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | No | No | No | --- | | | | | | -| KV Cache Reuse | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | | -| Sliding Window Attention | Yes | Yes | Yes | Yes | Yes | Untested | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | -| Logits Post Processor | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | Yes | No | No | No | No | Yes | Yes | Yes | Yes | --- | | | -| Guided Decoding | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | -| LoRA | Yes | Yes | Yes | Yes | Untested | Untested | Untested | Untested | Yes | Untested | Untested | Untested | Untested | Yes | Yes | Yes | Yes | Yes | Untested | --- | +| Feature | Overlap Scheduler | CUDA Graph | Tensor Parallelism | Pipeline Parallelism | Expert Parallelism | Helix Parallelism | Attention Data Parallelism | Disaggregated Serving | Chunked Prefill | MTP | EAGLE-3(One Model Engine) | EAGLE-3(Two Model Engine) | Torch Sampler | TLLM C++ Sampler | KV Cache Reuse | Sliding Window Attention | Logits Post Processor | Guided Decoding | LoRA | +| -------------------------- | ----------------- | ---------- | ------------------ | -------------------- | ------------------ | ----------------- | -------------------------- | --------------------- | --------------- | -------- | ------------------------- | ------------------------- | ------------- | ---------------- | -------------- | ---------------------- | --------------------- | --------------- | -------- | +| Overlap Scheduler | --- | | | | | | | | | | | | | | | | | | | +| CUDA Graph | Yes | --- | | | | | | | | | | | | | | | | | | +| Tensor Parallelism | Yes | Yes | --- | | | | | | | | | | | | | | | | | +| Pipeline Parallelism | Yes | Yes | Yes | --- | | | | | | | | | | | | | | | | +| Expert Parallelism | Yes | Yes | Yes | Yes | --- | | | | | | | | | | | | | | | +| Helix Parallelism | Untested | Yes | Yes | Yes | Yes | --- | | | | | | | | | | | | | | +| Attention Data Parallelism | Yes | Yes | Yes | Yes | Yes | Known issues | --- | | | | | | | | | | | | | +| Disaggregated Serving | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | | | | | | | | | +| Chunked Prefill | Yes | Yes | Yes | Untested | Yes | Yes | Yes | Yes | --- | | | | | | | | | | | +| MTP | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | --- | | | | | | | | | | +| EAGLE-3(One Model Engine) | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | No | --- | | | | | | | | | +| EAGLE-3(Two Model Engine) | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | No | No | --- | | | | | | | | +| Torch Sampler | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | | | | +| TLLM C++ Sampler | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | No | No | --- | | | | | | +| KV Cache Reuse | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | | +| Sliding Window Attention | Yes | Yes | Yes | Yes | Yes | Untested | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | +| Logits Post Processor | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | Yes | No | No | No | Yes | Yes | Yes | Yes | --- | | | +| Guided Decoding | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | +| LoRA | Yes | Yes | Yes | Yes | Untested | Untested | Untested | Untested | Yes | Untested | Untested | Untested | Yes | Yes | Yes | Yes | Yes | Untested | --- | diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index 5364d9b6fcd7..d1673deaffea 100644 --- a/docs/source/features/speculative-decoding.md +++ b/docs/source/features/speculative-decoding.md @@ -28,7 +28,7 @@ llm = LLM("/path/to/target_model", speculative_config=speculative_config, disabl ### EAGLE 3 The EAGLE 3 algorithm is described in the paper [EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test](https://arxiv.org/pdf/2503.01840). -By default, each request uses a single sequence (linear chain) of draft tokens with length `max_draft_len`. Optionally, dynamic tree draft generation can be enabled to improve acceptance rates — see [Dynamic Tree Mode](#dynamic-tree-mode) below. +TRT-LLM supports a modified version of the algorithm presented in the paper: tree structures for draft sequences are not supported. Instead, each request uses a single sequence of draft tokens with length `max_draft_len`. The following draft model checkpoints can be used for EAGLE 3: * Llama 3 variants: [use the checkpoints from the authors of the original EAGLE 3 paper](https://huggingface.co/yuhuili). @@ -50,36 +50,6 @@ llm = LLM(model, speculative_config=speculative_config) EAGLE 3 can be combined with the [Suffix Automaton enhancement](#suffix-automaton-sa-enhancement) for improved acceptance rates on repetitive content. See the SA section below for details. -#### Dynamic Tree Mode - -Dynamic tree mode enables tree-structured draft generation for EAGLE 3, where the drafter expands multiple candidate tokens at each layer instead of a single token. This can improve acceptance rates compared to linear drafting at the cost of additional compute per generation step. - -To enable dynamic tree mode, set `use_dynamic_tree=True` on the `Eagle3DecodingConfig` and provide the following parameters: - -* `use_dynamic_tree` (`bool`): Enables dynamic tree draft generation. Mutually exclusive with `eagle_choices` (static tree). -* `dynamic_tree_max_topK` (`int`): Maximum number of tokens to expand per node at each draft layer. -* `max_total_draft_tokens` (`int`, optional): Total draft token budget for the tree. Must satisfy `max_draft_len <= max_total_draft_tokens <= dynamic_tree_max_topK * max_draft_len`. Defaults to `dynamic_tree_max_topK * max_draft_len` if not set. -* `max_batch_size` (`int`): Required when `use_dynamic_tree=True` for pre-allocating dynamic tree CUDA buffers. - -```python -from tensorrt_llm.llmapi import Eagle3DecodingConfig - -speculative_config = Eagle3DecodingConfig( - max_draft_len=6, - speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", - use_dynamic_tree=True, - dynamic_tree_max_topK=10, - max_total_draft_tokens=60, - max_batch_size=4, -) - -llm = LLM("/path/to/target_model", speculative_config=speculative_config) -``` - -```{note} -Dynamic tree mode is currently **not supported** for models that use sliding window attention or MLA (Multi-Latent Attention), such as DeepSeek and gpt-oss models. -``` - ### NGram The NGram method is an implementation of [this Prompt Lookup Decoding algorithm](https://github.com/apoorvumang/prompt-lookup-decoding). @@ -229,18 +199,6 @@ speculative_config: speculative_model: /path/to/draft/model ``` -```yaml -# Dynamic tree mode -speculative_config: - decoding_type: Eagle3 - max_draft_len: 6 - speculative_model: /path/to/eagle3_model - use_dynamic_tree: true - dynamic_tree_max_topK: 10 - max_total_draft_tokens: 60 - max_batch_size: 4 -``` - ```yaml # SA combination: enable Suffix Automaton enhancement with any supported technique speculative_config: diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 52f5cda34ac8..c8d3a31ba628 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -43,18 +43,18 @@ The following is a table of supported models for the PyTorch backend: Note: Support for other models may vary. Features marked "N/A" are not applicable to the model architecture. -| Model Architecture/Feature | Overlap Scheduler | CUDA Graph | Attention Data Parallelism | Disaggregated Serving | Chunked Prefill | MTP | EAGLE-3(One Model Engine) — Linear | EAGLE-3(One Model Engine) — Dynamic | EAGLE-3(Two Model Engine) | Torch Sampler | TLLM C++ Sampler | KV Cache Reuse | Sliding Window Attention | Logits Post Processor | Guided Decoding | -| ------------------------------ | ----------------- | ---------- | -------------------------- | --------------------- | --------------- | --- | ---------------------------------- | ----------------------------------- | ------------------------- | ------------- | ---------------- | -------------- | ------------------------ | --------------------- | --------------- | -| `DeepseekV3ForCausalLM` | Yes | Yes | Yes | Yes | Yes [^1] | Yes | No | No | No | Yes | Yes | Yes [^2] | N/A | Yes | Yes | -| `DeepseekV32ForCausalLM` | Yes | Yes | Yes | Yes | Yes | Yes | No | No | No | Yes | Yes | Yes | N/A | Yes | Yes | -| `Glm4MoeForCausalLM` | Yes | Yes | Yes | Untested | Yes | Yes | No | No | No | Yes | Yes | Untested | N/A | Yes | Yes | -| `Qwen3MoeForCausalLM` | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | N/A | Yes | Yes | -| `Qwen3NextForCausalLM` [^3] | Yes | Yes | Yes | Untested | Yes | No | No | No | No | Yes | Yes | No | No | Untested | Untested | -| `Llama4ForConditionalGeneration` | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Untested | N/A | Yes | Yes | -| `GptOssForCausalLM` | Yes | Yes | Yes | Yes | Yes | No | Yes | No | Yes [^4] | Yes | Yes | Yes | N/A | Yes | Yes | -| `Qwen3_5MoeForCausalLM` [^5] | Yes | Yes | Untested | Untested | Yes | No | No | No | No | Yes | Untested | Yes | N/A | Untested | Untested | -| `Glm4MoeLiteForCausalLM` [^6] | Yes | Yes | Untested | Untested | Yes | No | No | No | No | Yes | Untested | Untested | N/A | Untested | Untested | -| `NemotronHForCausalLM` (Super) | Yes | Yes | Untested | Untested | Yes | Yes | No | No | No | Yes | Yes | Untested | N/A | Untested | Untested | +| Model Architecture/Feature | Overlap Scheduler | CUDA Graph | Attention Data Parallelism | Disaggregated Serving | Chunked Prefill | MTP | EAGLE-3(One Model Engine) | EAGLE-3(Two Model Engine) | Torch Sampler | TLLM C++ Sampler | KV Cache Reuse | Sliding Window Attention | Logits Post Processor | Guided Decoding | +| ------------------------------ | ----------------- | ---------- | -------------------------- | --------------------- | --------------- | --- | ------------------------- | ------------------------- | ------------- | ---------------- | -------------- | ------------------------ | --------------------- | --------------- | +| `DeepseekV3ForCausalLM` | Yes | Yes | Yes | Yes | Yes [^1] | Yes | No | No | Yes | Yes | Yes [^2] | N/A | Yes | Yes | +| `DeepseekV32ForCausalLM` | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | Yes | Yes | N/A | Yes | Yes | +| `Glm4MoeForCausalLM` | Yes | Yes | Yes | Untested | Yes | Yes | No | No | Yes | Yes | Untested | N/A | Yes | Yes | +| `Qwen3MoeForCausalLM` | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | N/A | Yes | Yes | +| `Qwen3NextForCausalLM` [^3] | Yes | Yes | Yes | Untested | Yes | No | No | No | Yes | Yes | No | No | Untested | Untested | +| `Llama4ForConditionalGeneration` | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Untested | N/A | Yes | Yes | +| `GptOssForCausalLM` | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes [^4] | Yes | Yes | Yes | N/A | Yes | Yes | +| `Qwen3_5MoeForCausalLM` [^5] | Yes | Yes | Untested | Untested | Yes | No | No | No | Yes | Untested | Yes | N/A | Untested | Untested | +| `Glm4MoeLiteForCausalLM` [^6] | Yes | Yes | Untested | Untested | Yes | No | No | No | Yes | Untested | Untested | N/A | Untested | Untested | +| `NemotronHForCausalLM` (Super) | Yes | Yes | Untested | Untested | Yes | Yes | No | No | Yes | Yes | Untested | N/A | Untested | Untested | [^1]: Chunked Prefill for MLA can only be enabled on SM100/SM103. [^2]: KV cache reuse for MLA can only be enabled on SM90/SM100/SM103 and in BF16/FP8 KV cache dtype. diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index ffc620feddce..076fec69717b 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -197,7 +197,6 @@ def add_llm_args(parser): default="llama3", choices=["llama3", "mistral_large3"], help="The model architecture of the eagle3 model.") - parser.add_argument('--max_total_draft_tokens', type=int, default=None) # Relaxed acceptance parser.add_argument('--use_relaxed_acceptance_for_thinking', @@ -232,29 +231,6 @@ def parse_arguments(): description="LLM models with the PyTorch workflow.") parser = add_llm_args(parser) parser.add_argument("--kv_cache_fraction", type=float, default=0.9) - parser.add_argument( - "--streaming", - action="store_true", - default=False, - help="Use streaming generate_async instead of generate.") - parser.add_argument( - "--concurrency", - type=int, - default=1, - help="Number of concurrent requests to submit simultaneously.") - parser.add_argument( - "--dataset", - type=str, - default=None, - help= - "Path to a JSONL dataset file. Each line: {\"question_id\": N, \"question\": [\"prompt\"]}." - ) - parser.add_argument( - "--num_samples", - type=int, - default=None, - required=False, - help="Limit number of samples from dataset (for fast debugging).") args = parser.parse_args() return args @@ -294,9 +270,7 @@ def setup_llm(args, **kwargs): use_dynamic_tree=args.use_dynamic_tree, dynamic_tree_max_topK=args.dynamic_tree_max_topK, allow_advanced_sampling=args.allow_advanced_sampling, - eagle3_model_arch=args.eagle3_model_arch, - max_total_draft_tokens=args.max_total_draft_tokens, - max_batch_size=args.max_batch_size) + eagle3_model_arch=args.eagle3_model_arch) elif spec_decode_algo == "DRAFT_TARGET": spec_config = DraftTargetDecodingConfig( max_draft_len=args.spec_decode_max_draft_len, @@ -389,17 +363,7 @@ def setup_llm(args, **kwargs): def main(): args = parse_arguments() - - if args.dataset: - prompts = [] - with open(args.dataset, 'r') as f: - for line in f: - entry = json.loads(line.strip()) - prompts.append(entry["question"][-1]) - if args.num_samples is not None: - prompts = prompts[:args.num_samples] - else: - prompts = args.prompt if args.prompt else example_prompts + prompts = args.prompt if args.prompt else example_prompts llm, sampling_params = setup_llm(args) new_prompts = [] @@ -411,132 +375,41 @@ def main(): tokenize=False, add_generation_prompt=True)) prompts = new_prompts - llm.generate(prompts, sampling_params) - - accept_rates = [] - total_tokens = 0 - total_iterations = 0 - - if args.streaming: - for i, prompt in enumerate(prompts): - num_tokens = 0 - num_iterations = 0 - for output in llm.generate_async(prompt, - sampling_params, - streaming=True): - new_tokens = output.outputs[0].token_ids - num_tokens = len(new_tokens) - num_iterations += 1 - if num_iterations > 0: - accept_rate = num_tokens / num_iterations - accept_rates.append(accept_rate) - total_tokens += num_tokens - total_iterations += num_iterations - print(f"[{i}] Accept rate: {accept_rate:.2f} " - f"(tokens={num_tokens}, iterations={num_iterations})") - generated_text = output.outputs[0].text - print( - f"[{i}] Prompt: {prompt[:80]!r}..., Generated text: {generated_text[:200]!r}..." - ) - - if accept_rates: - avg_accept_rate = sum(accept_rates) / len(accept_rates) - global_accept_rate = total_tokens / total_iterations - print("\n=== Accept Rate Summary ===") - print(f"Total prompts: {len(accept_rates)}") - print(f"Mean accept rate (per-prompt avg): {avg_accept_rate:.2f}") - print( - f"Global accept rate (total_tokens/total_iters): {global_accept_rate:.2f}" - ) + outputs = llm.generate(prompts, sampling_params) + + for i, output in enumerate(outputs): + prompt = output.prompt + for sequence_idx, sequence in enumerate(output.outputs): + generated_text = sequence.text + # Skip printing the beam_idx if no beam search was used + sequence_id_text = f"[{sequence_idx}]" if args.max_beam_width > 1 or args.n > 1 else "" print( - f"Total tokens: {total_tokens}, Total iterations: {total_iterations}" + f"[{i}]{sequence_id_text} Prompt: {prompt!r}, Generated text: {generated_text!r}" ) - return - - concurrency = args.concurrency - - if concurrency > 1: - # Submit all prompts concurrently in batches of `concurrency` - for batch_start in range(0, len(prompts), concurrency): - batch_end = min(batch_start + concurrency, len(prompts)) - batch_prompts = prompts[batch_start:batch_end] - - futures = [] - for prompt in batch_prompts: - future = llm.generate_async(prompt, - sampling_params, - streaming=True) - futures.append(future) - - for j, future in enumerate(futures): - i = batch_start + j - num_tokens = 0 - num_iterations = 0 - for output in future: - new_tokens = output.outputs[0].token_ids - num_tokens = len(new_tokens) - num_iterations += 1 - if num_iterations > 0: - accept_rate = num_tokens / num_iterations - accept_rates.append(accept_rate) - total_tokens += num_tokens - total_iterations += num_iterations - print(f"[{i}] Accept rate: {accept_rate:.2f} " - f"(tokens={num_tokens}, iterations={num_iterations})") - generated_text = output.outputs[0].text + if args.return_context_logits: print( - f"[{i}] Prompt: {prompt[:80]!r}..., Generated text: {generated_text[:200]!r}..." + f"[{i}]{sequence_id_text} Context logits: {output.context_logits}" ) - else: - # Use streaming generate_async to count decode iterations for accept rate - for i, prompt in enumerate(prompts): - num_tokens = 0 - num_iterations = 0 - for output in llm.generate_async(prompt, - sampling_params, - streaming=True): - new_tokens = output.outputs[0].token_ids - num_tokens = len(new_tokens) - num_iterations += 1 - if num_iterations > 0: - accept_rate = num_tokens / num_iterations - accept_rates.append(accept_rate) - total_tokens += num_tokens - total_iterations += num_iterations - print(f"[{i}] Accept rate: {accept_rate:.2f} " - f"(tokens={num_tokens}, iterations={num_iterations})") - for sequence_idx, sequence in enumerate(output.outputs): - generated_text = sequence.text - sequence_id_text = f"[{sequence_idx}]" if args.max_beam_width > 1 or args.n > 1 else "" + if args.return_generation_logits: print( - f"[{i}]{sequence_id_text} Prompt: {prompt[:80]!r}..., Generated text: {generated_text[:200]!r}..." + f"[{i}]{sequence_id_text} Generation logits: {sequence.generation_logits}" ) - if args.return_context_logits: - print( - f"[{i}]{sequence_id_text} Context logits: {output.context_logits}" - ) - if args.return_generation_logits: - print( - f"[{i}]{sequence_id_text} Generation logits: {sequence.generation_logits}" - ) - if args.prompt_logprobs is not None: - print( - f"[{i}]{sequence_id_text} Prompt logprobs: {sequence.prompt_logprobs}" - ) - if args.logprobs is not None: - print( - f"[{i}]{sequence_id_text} Logprobs: {sequence.logprobs}" - ) + if args.prompt_logprobs is not None: + print( + f"[{i}]{sequence_id_text} Prompt logprobs: {sequence.prompt_logprobs}" + ) + if args.logprobs is not None: + print(f"[{i}]{sequence_id_text} Logprobs: {sequence.logprobs}") - if args.additional_model_outputs: - for output_name in args.additional_model_outputs: - if sequence.additional_context_outputs: - print( - f"[{i}]{sequence_id_text} Context {output_name}: {sequence.additional_context_outputs[output_name]}" - ) + if args.additional_model_outputs: + for output_name in args.additional_model_outputs: + if sequence.additional_context_outputs: print( - f"[{i}]{sequence_id_text} Generation {output_name}: {sequence.additional_generation_outputs[output_name]}" + f"[{i}]{sequence_id_text} Context {output_name}: {sequence.additional_context_outputs[output_name]}" ) + print( + f"[{i}]{sequence_id_text} Generation {output_name}: {sequence.additional_generation_outputs[output_name]}" + ) if args.log_kv_cache_events: time.sleep(1) # Wait for events to be dispatched diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 64b6173a7800..600f655bc517 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -10,7 +10,7 @@ from typing_extensions import Self if TYPE_CHECKING: - from tensorrt_llm.llmapi.llm_args import SparseAttentionConfig + from ..speculative.utils import SpecDecodingTensor from ..speculative.interface import SpecMetadata from ..speculative.spec_tree_manager import SpecTreeManager @@ -383,7 +383,8 @@ def update_spec_dec_param( max_total_draft_tokens, model_is_wrapped: bool = False, spec_metadata: Optional['SpecMetadata'] = None, - spec_tree_manager: Optional['SpecTreeManager'] = None): + spec_tree_manager: Optional['SpecTreeManager'] = None, + spec_decoding_tensor: Optional['SpecDecodingTensor'] = None): """ Hook to be called when using TRTLLM attention backend in spec-dec mode. """ diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 3c6b8dd4cf4f..d9a410fe62f2 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -597,13 +597,15 @@ def update_spec_dec_param( model_is_wrapped: bool = False, spec_metadata: Optional['SpecMetadata'] = None, spec_tree_manager: Optional['SpecTreeManager'] = None, + spec_decoding_tensor: Optional['SpecDecodingTensor'] = None, ): """Update speculative decoding parameters and create expanded buffers.""" super().update_spec_dec_param(batch_size, is_spec_decoding_enabled, is_spec_dec_tree, is_spec_dec_dynamic_tree, max_draft_len, max_total_draft_tokens, model_is_wrapped, - spec_metadata, spec_tree_manager) + spec_metadata, spec_tree_manager, + spec_decoding_tensor) self.max_draft_tokens = max_draft_len init_shape = self.kv_lens_expanded_host.shape[0] if self.max_num_sequences * (1 + self.max_draft_tokens) != init_shape: diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 9f32039f145f..efcffad624ca 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -8,6 +8,7 @@ import torch if TYPE_CHECKING: + from ..speculative.utils import SpecDecodingTensor from ..speculative.interface import SpecMetadata from ..speculative.spec_tree_manager import SpecTreeManager @@ -508,18 +509,9 @@ def run( self.is_spec_decoding_enabled, self.use_spec_decoding, self.is_spec_dec_tree ] - # For dynamic tree, reshape 1D position_offsets to 2D for C++ kernel compatibility - position_offsets_for_cpp = self.spec_decoding_position_offsets - if (self.spec_decoding_position_offsets is not None - and self.spec_decoding_position_offsets.dim() == 1): - # Reshape 1D [max_num_requests * N] to 2D [max_num_requests, N] - # C++ kernel requires 2D to extract max_generation_length from sizes()[1] - position_offsets_for_cpp = self.spec_decoding_position_offsets.view( - self.max_num_requests, -1) - spec_decoding_tensor_params = [ - self.spec_decoding_generation_lengths, position_offsets_for_cpp, - self.spec_decoding_packed_mask + self.spec_decoding_generation_lengths, + self.spec_decoding_position_offsets, self.spec_decoding_packed_mask ] if self.is_sm_version_trtllm_gen_kernel(sm=get_sm_version()): spec_decoding_tensor_params.append( @@ -531,12 +523,6 @@ def run( self.helix_position_offsets, self.helix_is_inactive_rank ] - # Zero bl_tree_mask only at layer transitions (local layer_idx==0) - # to clear stale bits from atomicOr. Layers 1-31 reuse layer 0's mask; - # eagle head (also local_idx==0) gets fresh mask. - if self.spec_decoding_bl_tree_mask is not None and self.layer_idx == 0: - self.spec_decoding_bl_tree_mask.zero_() - if self.print_skip_softmax_stat: self.skip_softmax_stat.zero_() @@ -1468,59 +1454,48 @@ def compute_max_num_custom_mask_tiles_kv_upper_bound( def spec_decoding_param_prepare_for_blackwell(self) -> None: """ Prepare the blackwell parameters for the speculative decoding (Medusa and Eagle) generation-phase attention kernels. - Uses persistent buffers (only allocate if None) for CUDA graph compatibility. """ - if self.spec_decoding_bl_tree_mask_offset is None: - self.spec_decoding_bl_tree_mask_offset = torch.zeros( - [self.max_num_requests], - dtype=torch.int64, - device='cuda', - ) - - if self.spec_bl_tree_first_sparse_mask_offset_kv is None: - self.spec_bl_tree_first_sparse_mask_offset_kv = torch.zeros( - [self.max_num_requests], - dtype=torch.int32, - device='cuda', - ) - - if self.spec_decoding_bl_tree_mask is None: - # Custom mask covers the tree region (seqLenQ x seqLenQ). - # The mask buffer needs extra room for tile boundary crossing: - # when firstSparse is not aligned to stepKv (=tileSizeKv*numInstsKv), - # the mask can span one extra KV tile. Adding (stepKv - 1) to - # max_kv_len guarantees the upper-bound tile count is correct. - seqLenQ = self.spec_decoding_packed_mask.shape[ - 1] if self.spec_decoding_packed_mask is not None else 1 - tile_size_kv = 128 - tile_size_q = 128 - num_instances_q = 1 - num_instances_kv = 2 - tile_size_kv_per_cta = tile_size_kv * num_instances_kv - max_kv_len = seqLenQ + tile_size_kv_per_cta - 1 - tile_size_q_per_cta = tile_size_q * num_instances_q - max_num_custom_mask_tiles_kv = self.compute_max_num_custom_mask_tiles_kv_upper_bound( - max_kv_len, 0, tile_size_kv_per_cta) - max_num_tiles_q = math.ceil( - (seqLenQ * self.num_heads_per_kv) / tile_size_q_per_cta) - mask_size = int(self.max_num_requests * max_num_tiles_q * - max_num_custom_mask_tiles_kv * num_instances_q * - num_instances_kv * tile_size_q * tile_size_kv / 32) - self.spec_decoding_bl_tree_mask = torch.zeros( - mask_size, - dtype=torch.uint32, - device='cuda', - ) - - def update_blackwell_first_sparse_mask_offset(self) -> None: - """Fill gen slots [0:ng) with (kv_lens_cuda - seq_lens); in-place for CUDA graph.""" - nc = self.num_contexts - n = self.num_seqs - ng = n - nc - if ng > 0: - gen_offset = (self.kv_lens_cuda[nc:n] - - self._seq_lens_cuda[nc:n]).to(torch.int32) - self.spec_bl_tree_first_sparse_mask_offset_kv[:ng].copy_(gen_offset) + self.spec_decoding_bl_tree_mask_offset = torch.zeros( + [self.max_num_requests], + dtype=torch.int64, + device='cuda', + ) + max_kv_len = self.kv_lens[:self.num_seqs].max() + assert self.kv_lens_cuda[:self. + num_seqs] >= self._seq_lens_cuda[:self. + num_seqs], "kv_lens should be greater than seq_lens,please run prepare() first" + + # Only support seq_lens are equal in one batch + seq_lens_slice = self.seq_lens[:self.num_seqs] + assert seq_lens_slice.min() == seq_lens_slice.max(), \ + f"All elements in seq_lens must be equal in one batch, but got min={seq_lens_slice.min()}, max={seq_lens_slice.max()}" + + self.spec_bl_tree_first_sparse_mask_offset_kv = ( + self.kv_lens_cuda[:self.num_seqs] - + self._seq_lens_cuda[:self.num_seqs]).to(torch.int32) + min_first_sparse_mask_offset_kv = self.spec_bl_tree_first_sparse_mask_offset_kv.min( + ) + # tile_size_kv * tile_size_q * num_instances_q * num_instances_kv is the largest value that is used in the trtllm-gen kernels + tile_size_kv = 128 + tile_size_q = 128 + # num_instances_q * num_instances_kv <= 2 + num_instances_q = 1 + num_instances_kv = 2 + tile_size_kv_per_cta = tile_size_kv * num_instances_kv + tile_size_q_per_cta = tile_size_q * num_instances_q + max_num_custom_mask_tiles_kv = self.compute_max_num_custom_mask_tiles_kv_upper_bound( + max_kv_len, min_first_sparse_mask_offset_kv, tile_size_kv_per_cta) + max_num_tiles_q = math.ceil( + (self.seq_lens[:self.num_seqs].max() * self.num_heads_per_kv) / + tile_size_q_per_cta) + mask_size = int(self.max_num_requests * max_num_tiles_q * + max_num_custom_mask_tiles_kv * num_instances_q * + num_instances_kv * tile_size_q * tile_size_kv / 32) + self.spec_decoding_bl_tree_mask = torch.zeros( + mask_size, + dtype=torch.uint32, + device='cuda', + ) def update_spec_dec_param( self, @@ -1533,6 +1508,7 @@ def update_spec_dec_param( model_is_wrapped: bool = False, spec_metadata: Optional['SpecMetadata'] = None, spec_tree_manager: Optional['SpecTreeManager'] = None, + spec_decoding_tensor: Optional['SpecDecodingTensor'] = None, ) -> None: ''' Update the spec-dec parameters for the TRTLLM attention layer. @@ -1546,12 +1522,31 @@ def update_spec_dec_param( model_is_wrapped: Optional[bool] = False, whether the drafter model is wrapped (i.e, CDL). spec_metadata: Optional['SpecMetadata'] = None, the metadata of the spec-dec. spec_tree_manager: Optional['SpecTreeManager'] = None, the spec_tree_manager for draft token tree. + spec_decoding_tensor: Optional['SpecDecodingTensor'] = None, the spec_decoding_tensor for draft token tree. ''' + if spec_decoding_tensor is not None: + spec_decoding_position_offsets = spec_decoding_tensor.position_offsets + spec_decoding_packed_mask = spec_decoding_tensor.packed_mask + spec_decoding_generation_lengths = spec_decoding_tensor.generation_lengths + else: + spec_decoding_position_offsets = None + spec_decoding_packed_mask = None + spec_decoding_generation_lengths = None - self.is_spec_decoding_enabled = is_spec_decoding_enabled + # spec_dec mode should only be enabled for non-sm100 machines and when there's a spec-dec tree. + self.is_spec_decoding_enabled = is_spec_decoding_enabled and ( + not self.is_sm_version_trtllm_gen_kernel(sm=get_sm_version())) + + self.is_spec_dec_tree = spec_tree_manager is not None + self.is_spec_dec_dynamic_tree = spec_tree_manager is not None and spec_tree_manager.use_dynamic_tree + + if self.is_sm_version_trtllm_gen_kernel(sm=get_sm_version()): + if self.is_spec_dec_tree or self.is_spec_dec_dynamic_tree: + assert not self.is_spec_dec_tree, "Spec-dec tree is not supported on this machine. Please use a pre-Blackwell machine for a spec-dec tree." # use_spec_decoding is default to true by default, change in runtime by layers / requests self.use_spec_decoding = self.is_spec_decoding_enabled + self.is_spec_dec_tree = is_spec_dec_tree self.is_spec_dec_dynamic_tree = is_spec_dec_dynamic_tree @@ -1564,35 +1559,20 @@ def update_spec_dec_param( # These buffers are accessed more like removing input padding, # rather than using max_total_draft_tokens + 1 as the offset between different requests. if is_spec_dec_tree and self.spec_decoding_position_offsets is None: - if spec_tree_manager is not None and spec_tree_manager.use_dynamic_tree: - # Dynamic tree: use _internal_buf_dim which may be larger - # than max_total_draft_tokens+1 to accommodate K*max_draft_len - buf_dim = spec_tree_manager._internal_buf_dim - # Dynamic tree: 1D layout for flexible view() in drafting loop - self.spec_decoding_position_offsets = torch.empty( - (self.max_num_requests * buf_dim, ), - dtype=torch.int, - device='cuda', - ) - else: - # Static tree: keep 2D layout - self.spec_decoding_position_offsets = torch.empty( - [self.max_num_requests, max_total_draft_tokens + 1], - dtype=torch.int, - device='cuda', - ) + self.spec_decoding_position_offsets = torch.empty( + [self.max_num_requests, max_total_draft_tokens + 1], + dtype=torch.int, + device='cuda', + ) if is_spec_dec_tree and self.spec_decoding_packed_mask is None: - if spec_tree_manager is not None and spec_tree_manager.use_dynamic_tree: - buf_dim = spec_tree_manager._internal_buf_dim - else: - buf_dim = max_total_draft_tokens + 1 self.spec_decoding_packed_mask = torch.empty( - [self.max_num_requests, buf_dim, - math.ceil(buf_dim / 32)], + [ + self.max_num_requests, max_total_draft_tokens + 1, + math.ceil((max_total_draft_tokens + 1) / 32) + ], dtype=torch.int, device='cuda', ) - if self.spec_decoding_generation_lengths is None: self.spec_decoding_generation_lengths = torch.empty( [self.max_num_requests], @@ -1607,31 +1587,20 @@ def update_spec_dec_param( self.spec_decoding_bl_tree_mask = None self.spec_bl_tree_first_sparse_mask_offset_kv = None - # Case 1: dynamic tree — copy per-request params from spec_tree_manager. + # Case 1: dynamic tree if self.is_spec_dec_dynamic_tree: - assert spec_tree_manager is not None, "spec_tree_manager is required for dynamic tree" - n_dt = spec_tree_manager.max_total_draft_tokens + 1 - mask_width = spec_tree_manager.spec_dec_packed_mask.shape[-1] - - self.spec_decoding_position_offsets[:batch_size * n_dt].view( - batch_size, - n_dt).copy_(spec_tree_manager. - spec_dec_position_offsets[:batch_size], - non_blocking=True) - - if self.is_sm_version_trtllm_gen_kernel(sm=get_sm_version()): - self.spec_decoding_packed_mask[:batch_size].zero_() - self.spec_decoding_packed_mask[:batch_size, :n_dt, :].copy_( - spec_tree_manager.spec_dec_packed_mask[:batch_size], - non_blocking=True) + assert spec_decoding_position_offsets is not None, "spec_decoding_position_offsets is required for dynamic tree" + assert spec_decoding_packed_mask is not None, "spec_decoding_packed_mask is required for dynamic tree" + self.spec_decoding_position_offsets.copy_( + spec_decoding_position_offsets, non_blocking=True) + self.spec_decoding_packed_mask.copy_(spec_decoding_packed_mask, + non_blocking=True) + if spec_decoding_generation_lengths is not None: + self.spec_decoding_generation_lengths.copy_( + spec_decoding_generation_lengths, non_blocking=True) else: - total = batch_size * n_dt * mask_width - self.spec_decoding_packed_mask.view(-1)[:total].copy_( - spec_tree_manager.spec_dec_packed_mask[:batch_size]. - reshape(-1), - non_blocking=True) - - self.spec_decoding_generation_lengths[:batch_size].fill_(n_dt) + self.generate_spec_decoding_generation_length( + runtime_draft_len=max_total_draft_tokens) # Case 2/3: static tree elif self.is_spec_dec_tree and not self.is_spec_dec_dynamic_tree and spec_metadata is not None: @@ -1998,13 +1967,6 @@ def forward( self._compute_flash_mla_metadata(metadata) metadata._flash_mla_metadata_valid = True - # Blackwell first_sparse: refresh at layer 0 before plan (mask zero is in wrapper.run). - layer_idx = self.get_local_layer_idx(metadata) - if layer_idx == 0 and (metadata.spec_bl_tree_first_sparse_mask_offset_kv - is not None - and metadata._seq_lens_cuda is not None): - metadata.update_blackwell_first_sparse_mask_offset() - self.wrapper.plan( layer_idx=self.get_local_layer_idx(metadata), tokens_per_block=metadata.tokens_per_block, diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 3ead33963531..fd0714e25a2a 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -899,19 +899,6 @@ def forward( else: q, k, v = qkv, None, None - # For dynamic tree spec decoding with Python RoPE, adjust position_ids - # to use tree offsets (same as C++ kernel: past_seq_len + offset). - if (not self.rope_fusion - and getattr(attn_metadata, 'is_spec_dec_dynamic_tree', False) - and getattr(attn_metadata, 'use_spec_decoding', False) - and getattr(attn_metadata, 'spec_decoding_position_offsets', - None) is not None - and attn_metadata.spec_decoding_position_offsets.dim() == - 1 # 1D layout ⇒ dynamic tree - and position_ids is not None): - position_ids = self._adjust_position_ids_for_spec_dec( - position_ids, attn_metadata) - q, k, v = self.apply_rope(q, k, v, position_ids) q, k, v = self.convert_qkv(q, k, v) @@ -961,25 +948,6 @@ def apply_rope(self, q: torch.Tensor, k: Optional[torch.Tensor], q, k = self.rotary_emb(position_ids, [q, k]) return q, k, v - def _adjust_position_ids_for_spec_dec(self, position_ids, attn_metadata): - """Replicate C++ kernel's rotary_pos = past_seq_len + offset.""" - num_contexts = attn_metadata.num_contexts - num_gens = attn_metadata.num_seqs - num_contexts - if num_gens <= 0: - return position_ids - gen_len = int(attn_metadata.seq_lens[num_contexts]) - base_pos = attn_metadata.kv_lens_cuda[num_contexts:num_contexts + - num_gens] - gen_len - offsets = attn_metadata.spec_decoding_position_offsets[:num_gens * - gen_len].view( - num_gens, - gen_len) - adjusted = (base_pos.unsqueeze(1) + offsets).reshape(-1) - start = attn_metadata.num_ctx_tokens - end = start + num_gens * gen_len - position_ids[0, start:end] = adjusted - return position_ids - def apply_qk_norm(self, q, k): raise NotImplementedError( f"QK norm is not implemented for {self.__class__.__name__}. " diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index fc8c79a4b3e7..3074baf9a3fa 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -56,6 +56,7 @@ from ..speculative.drafting_loops import BaseDraftingLoopWrapper from ..speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata from ..speculative.spec_sampler_base import SampleStateTensorsSpec +from ..speculative.utils import SpecDecodingTensor from ..utils import (get_model_extra_attrs, set_per_request_piecewise_cuda_graph_flag, set_torch_compiling, with_model_extra_attrs) @@ -2361,10 +2362,13 @@ def _prepare_tp_inputs( assert spec_config.spec_dec_mode.support_overlap_scheduler( ), f"{spec_config.decoding_type} does not support overlap scheduler" - # For tree decoding, runtime_draft_len should match total tree - # tokens (not tree depth). py_executor resets it every iteration. - if spec_config is not None and not spec_config.is_linear_tree: - self.runtime_draft_len = self.max_total_draft_tokens + spec_resource_manager, spec_tree_manager = None, None + if spec_config is not None: + spec_resource_manager = resource_manager.get_resource_manager( + ResourceManagerType.SPEC_RESOURCE_MANAGER) + if spec_resource_manager is not None and hasattr( + spec_resource_manager, 'spec_tree_manager'): + spec_tree_manager = spec_resource_manager.spec_tree_manager # will contain previous batch indices of generation requests previous_batch_indices = [] @@ -2404,10 +2408,20 @@ def _prepare_tp_inputs( list( range(len(position_ids), len(position_ids) + 1 + num_draft_tokens))) - position_ids.extend( - list( - range(past_seen_token_num, - past_seen_token_num + 1 + num_draft_tokens))) + # For the target model + tree decoding + if not self.is_draft_model and not spec_config.is_linear_tree: + assert spec_tree_manager is not None + assert num_draft_tokens == spec_tree_manager.max_total_draft_tokens + position_ids.extend( + past_seen_token_num + + spec_tree_manager.spec_dec_position_offsets[ + 0] # [max_total_draft_tokens + 1] + ) + else: + position_ids.extend( + list( + range(past_seen_token_num, + past_seen_token_num + 1 + num_draft_tokens))) num_cached_tokens_per_seq.append(past_seen_token_num) request.cached_tokens = num_cached_tokens_per_seq[-1] # update batch index @@ -2427,10 +2441,20 @@ def _prepare_tp_inputs( list( range(len(position_ids), len(position_ids) + 1 + self.runtime_draft_len))) - position_ids.extend( - list( - range(past_seen_token_num, past_seen_token_num + 1 + - self.runtime_draft_len))) + # For the target model + tree decoding + if not self.is_draft_model and not spec_config.is_linear_tree: + assert spec_tree_manager is not None + position_ids.extend( + past_seen_token_num + + spec_tree_manager.spec_dec_position_offsets[ + 0] # [max_total_draft_tokens + 1] + ) + else: + position_ids.extend( + list( + range( + past_seen_token_num, past_seen_token_num + 1 + + self.runtime_draft_len))) # previous tensor previous_batch_indices.append(previous_batch_idx) previous_pos_indices.extend([previous_batch_idx] * @@ -3665,6 +3689,7 @@ def forward(self, new_tensors_device: Optional[SampleStateTensors] = None, gather_context_logits: bool = False, cache_indirection_buffer: Optional[torch.Tensor] = None, + spec_decoding_tensor: Optional[SpecDecodingTensor] = None, num_accepted_tokens_device: Optional[torch.Tensor] = None, req_id_to_old_request: Optional[Dict[int, LlmRequest]] = None): kv_cache_manager = resource_manager.get_resource_manager( @@ -3678,8 +3703,7 @@ def forward(self, spec_resource_manager = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) spec_tree_manager = None - if spec_resource_manager is not None and hasattr( - spec_resource_manager, 'spec_tree_manager'): + if isinstance(spec_resource_manager, Eagle3ResourceManager): spec_tree_manager = spec_resource_manager.spec_tree_manager spec_metadata = self._set_up_spec_metadata(spec_resource_manager, no_cache=kv_cache_manager @@ -3702,18 +3726,6 @@ def forward(self, sd_max_draft_len = self.original_max_draft_len sd_max_total = self._spec_dec_max_total_draft_tokens - # Gather dynamic tree data from stable slots before target forward - if (spec_tree_manager is not None - and spec_tree_manager.use_dynamic_tree - and not self.is_draft_model): - gen_requests = list(scheduled_requests.generation_requests) - num_gens = len(gen_requests) - if num_gens > 0: - gen_slot_ids, _ = spec_tree_manager.fill_gen_slot_ids( - gen_requests) - spec_tree_manager.gather_trees_from_slots( - gen_slot_ids, num_gens) - attn_metadata.update_spec_dec_param( batch_size=scheduled_requests.batch_size, is_spec_decoding_enabled=is_spec_dec_mode, @@ -3723,7 +3735,8 @@ def forward(self, max_total_draft_tokens=sd_max_total, model_is_wrapped=self.model_is_wrapped, spec_metadata=spec_metadata, - spec_tree_manager=spec_tree_manager) + spec_tree_manager=spec_tree_manager, + spec_decoding_tensor=spec_decoding_tensor) else: spec_resource_manager = None spec_metadata = None @@ -3769,25 +3782,6 @@ def forward(self, spec_metadata = self.spec_metadata else: spec_metadata = None - - # Fill slot-ID buffer for scatter inside draft loop - if (self.enable_spec_decode and spec_tree_manager is not None - and spec_tree_manager.use_dynamic_tree - and not self.is_draft_model): - dummy_slot = spec_tree_manager._dummy_slot_id - idx = 0 - for req in padded_requests.context_requests: - spec_tree_manager._all_slot_ids_buf[idx] = ( - req.py_seq_slot - if req.py_seq_slot is not None else dummy_slot) - idx += 1 - for req in padded_requests.generation_requests: - slot = req.py_seq_slot if ( - not getattr(req, 'is_cuda_graph_dummy', False) - and req.py_seq_slot is not None) else dummy_slot - spec_tree_manager._all_slot_ids_buf[idx] = slot - idx += 1 - inputs, gather_ids = self._prepare_inputs( padded_requests, kv_cache_manager, attn_metadata, spec_metadata, new_tensors_device, cache_indirection_buffer, @@ -3846,10 +3840,6 @@ def capture_postprocess_fn(inputs: Dict[str, Any]): return outputs - def _get_spec_worker(self): - """Access the spec_worker from DecoderModelForCausalLM (one-model spec dec).""" - return getattr(self.model, 'spec_worker', None) - def model_forward(self, **kwargs): attrs = get_model_extra_attrs() assert attrs is not None, "Model extra attrs is not set" diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 37e80f71f7a4..a4519b534688 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1771,7 +1771,7 @@ def _handle_dynamic_draft_len(self, self.model_engine.runtime_draft_len = runtime_draft_len else: - self.model_engine.runtime_draft_len = self.model_engine.max_total_draft_tokens + self.model_engine.runtime_draft_len = self.model_engine.max_draft_len def _can_queue(self, scheduled_batch): diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 9aff9d894f4b..202912750b6f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -377,10 +377,6 @@ def create_py_executor( has_draft_model_engine = spec_config.spec_dec_mode.has_draft_model() has_spec_drafter = spec_config.spec_dec_mode.has_spec_drafter() - if hasattr(spec_config, - 'max_batch_size') and spec_config.max_batch_size is None: - spec_config.max_batch_size = max_batch_size - # WAR for https://nvbugs/5807902 # Disable separate draft KV cache in disaggregated mode # Enable separate pool for None DI + Non-KVBM and Aggregated + KVBM @@ -465,16 +461,15 @@ def allocation_scope(current_stage: ExecutorMemoryType): def drafting_loop_wrapper(model): from tensorrt_llm._torch.speculative.drafting_loops import ( - LinearDraftingLoopWrapper, - StaticTreeDraftingLoopWrapper) + LinearDraftingLoopWrapper, TreeDraftingLoopWrapper) from tensorrt_llm.llmapi import EagleDecodingConfig - static_tree_drafter = isinstance( + use_tree_drafter = isinstance( draft_spec_config, EagleDecodingConfig - ) and draft_spec_config.eagle_choices is not None + ) and not draft_spec_config.is_linear_tree - if static_tree_drafter: - return StaticTreeDraftingLoopWrapper( + if use_tree_drafter: + return TreeDraftingLoopWrapper( spec_config.max_draft_len, spec_config.tokens_per_gen_step - 1, max_batch_size, model) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 3e661ad838cd..3c0fa9e60114 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -231,18 +231,13 @@ def _update_kv_cache_draft_token_location(cache_manager, ) == 1, "Currently, only one max attention window size is supported." if use_paged_kv_cache: - assert len(set(cache_manager.num_kv_heads_per_layer)) == 1, \ - "update_kv_cache_draft_token_location requires uniform num_kv_heads across all layers, " \ - f"but got {cache_manager.num_kv_heads_per_layer}" torch.ops.tensorrt_llm.update_kv_cache_draft_token_location( accepted_draft_token_offsets, packed_accepted_draft_tokens_indices, past_key_value_lengths, True, cache_manager.num_layers, - # Use TP-sharded num_kv_heads (per-rank) instead of the unsharded - # total so the C++ kernel computes correct strides and grid dims. - cache_manager.num_kv_heads_per_layer[0], + cache_manager.num_kv_heads, int(cache_manager.head_dim * kv_cache_dtype_byte_size), cache_manager.max_total_draft_tokens, cache_manager.max_attention_window_vec[0], @@ -806,6 +801,11 @@ def update_resources(self, scheduled_batch: ScheduledRequests, attn_metadata: "AttentionMetadata" = None, kv_cache_dtype_byte_size: float = None): + if not self.is_draft: + _update_kv_cache_draft_token_location(self, scheduled_batch, + attn_metadata, + kv_cache_dtype_byte_size) + # Rewind KV cache for requests with rejected draft tokens. # Skip: # - GENERATION_COMPLETE: finished requests diff --git a/tensorrt_llm/_torch/speculative/drafting_loops.py b/tensorrt_llm/_torch/speculative/drafting_loops.py index e3033fa6d303..8cf41bee5e86 100644 --- a/tensorrt_llm/_torch/speculative/drafting_loops.py +++ b/tensorrt_llm/_torch/speculative/drafting_loops.py @@ -162,9 +162,11 @@ def prepare_for_generation(self, attn_metadata: AttentionMetadata, batch_size = attn_metadata.num_seqs num_accepted_draft_tokens = spec_metadata.num_accepted_draft_tokens[: batch_size] + # Using attn_metadata.seq_lens_cuda[:batch_size] to get the max_draft_len + 1 seq_lens = attn_metadata.seq_lens_cuda[:batch_size] attn_metadata.kv_lens_cuda[: batch_size] -= seq_lens - num_accepted_draft_tokens - 1 + # Calculate last accepted token indices last_tokens_idx = torch.cumsum( seq_lens, dim=0, @@ -200,7 +202,7 @@ def prepare_for_generation(self, attn_metadata: AttentionMetadata, return new_position_ids -class StaticTreeDraftingLoopWrapper(BaseDraftingLoopWrapper): +class TreeDraftingLoopWrapper(BaseDraftingLoopWrapper): def __init__(self, max_draft_len: int, max_total_draft_tokens: int, max_batch_size: int, draft_model: torch.nn.Module): @@ -224,8 +226,11 @@ def __init__(self, max_draft_len: int, max_total_draft_tokens: int, def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor, attn_metadata: AttentionMetadata, spec_metadata: SpecMetadata, **kwargs) -> dict[str, torch.Tensor]: - assert isinstance(spec_metadata, Eagle3SpecMetadata) - spec_tree_manager = spec_metadata.eagle3_resource_manager.spec_tree_manager + spec_tree_manager = None + if isinstance(spec_metadata, Eagle3SpecMetadata): + spec_tree_manager = spec_metadata.eagle3_resource_manager.spec_tree_manager + + assert spec_tree_manager is not None logits = self.draft_model.forward(input_ids=input_ids, position_ids=position_ids, @@ -246,6 +251,7 @@ def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor, new_draft_tokens=new_draft_tokens, use_cuda_graph=attn_metadata.is_cuda_graph, spec_tree_manager=spec_tree_manager) + return_draft_logits = None with save_metadata_state(attn_metadata, spec_metadata): batch_size = attn_metadata.num_seqs @@ -260,12 +266,10 @@ def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor, # position_ids: [batch_size * (max_total_draft_tokens + 1)] # logits: [batch_size * (max_total_draft_tokens + 1), vocab_size] logits = self.draft_model.forward( - input_ids=self.draft_tokens_buffer[:batch_size, :self. - max_total_draft_tokens + - 1].reshape(-1), - position_ids=self. - position_ids_buffer[:batch_size, :self. - max_total_draft_tokens + 1].reshape(-1), + input_ids=self.draft_tokens_buffer[:batch_size, :].reshape( + -1), + position_ids=self.position_ids_buffer[:batch_size, :]. + reshape(-1), attn_metadata=attn_metadata, spec_metadata=spec_metadata, return_context_logits=True) @@ -440,6 +444,7 @@ def prepare_for_generation(self, attn_metadata: AttentionMetadata, self.position_ids_buffer[:batch_size, :-1] = position_start_idx.unsqueeze( 1) + spec_tree_manager.spec_dec_position_offsets[0, 1:].unsqueeze( 0) - 1 # exclude the root node + # 2) Prepare the attn_metadata ## 2.1) kv_lens_cuda attn_metadata.kv_lens_cuda[: @@ -467,9 +472,7 @@ def prepare_for_generation(self, attn_metadata: AttentionMetadata, ### attn_metadata.spec_decoding_position_offsets: [max_num_requests, max_total_draft_tokens + 1] attn_metadata.spec_decoding_position_offsets[:batch_size, :self. max_total_draft_tokens] = spec_tree_manager.spec_dec_position_offsets[ - 0, 1:self. - max_total_draft_tokens + - 1].unsqueeze( + 0, 1:].unsqueeze( 0 ) - 1 # exclude the root node attn_metadata.spec_decoding_position_offsets[:batch_size, self. @@ -497,7 +500,7 @@ def prepare_for_generation(self, attn_metadata: AttentionMetadata, last_tokens_idx] # [batch_size], already take the accepted tokens into account. ### shape: [batch_size, self.max_total_draft_tokens + 1] - hidden_states_read_indices_offset = spec_tree_manager.hidden_states_read_indices_offset_for_drafter_model[:self.max_total_draft_tokens + 1].repeat( + hidden_states_read_indices_offset = spec_tree_manager.hidden_states_read_indices_offset_for_drafter_model.repeat( batch_size).reshape(batch_size, self.max_total_draft_tokens + 1) hidden_states_read_indices_offset = hidden_states_read_indices_offset + start_idx.unsqueeze( 1) diff --git a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py deleted file mode 100644 index 2565e034aaf1..000000000000 --- a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py +++ /dev/null @@ -1,198 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -""" -Dynamic Tree Operations for EAGLE3 Speculative Decoding - -This module provides high-performance CUDA kernel wrappers for building and verifying -dynamic tree structures used in EAGLE3 speculative decoding. It integrates SGLang's -optimized CUDA kernels into TensorRT-LLM's PyTorch backend. - -Key Features: -- Efficient tree construction from layer-local parent indices -- Greedy tree verification with parallel traversal -- Buffer pre-allocation and reuse for minimal runtime overhead -""" - -import torch - - -class DynamicTreeOpsConverter: - """ - Converter for dynamic tree operations using CUDA kernels. - - This class handles data format conversion and CUDA kernel invocation for - building and verifying dynamic trees in EAGLE3 speculative decoding. - - Args: - dynamic_tree_max_topK: Maximum top-K tokens per node. - max_draft_len: Maximum draft length (tree depth). - max_total_draft_tokens: Total number of draft tokens. - max_batch_size: Maximum batch size. - device: CUDA device. - """ - - def __init__( - self, - dynamic_tree_max_topK: int, - max_draft_len: int, - max_total_draft_tokens: int, - max_batch_size: int, - device: torch.device, - ): - self.K = dynamic_tree_max_topK - self.depth = max_draft_len - - # Pre-allocated output buffers for verify_dynamic_tree_greedy_out_op - N = max_total_draft_tokens + 1 # tokens_per_gen_step (includes root) - max_path_len = max_draft_len + 1 - self._verify_predicts_buf = torch.zeros( - max_batch_size * N, dtype=torch.int64, device=device - ) - self._verify_accept_index_buf = torch.zeros( - max_batch_size, max_path_len, dtype=torch.int64, device=device - ) - self._verify_accept_token_num_buf = torch.zeros( - max_batch_size, dtype=torch.int64, device=device - ) - self._verify_accept_token_buf = torch.zeros( - max_batch_size, max_path_len, dtype=torch.int64, device=device - ) - - def build_dynamic_tree( - self, - history_draft_tokens_parent_buffer: torch.Tensor, - topk_score_indices: torch.Tensor, - tree_mask: torch.Tensor, - positions: torch.Tensor, - retrieve_index: torch.Tensor, - retrieve_next_token: torch.Tensor, - retrieve_next_sibling: torch.Tensor, - use_packed_mask: bool = False, - ) -> None: - """ - Build dynamic tree structure using CUDA kernel (in-place, writes to pre-allocated buffers). - - All output tensors are written in-place; nothing is returned. - - Args: - history_draft_tokens_parent_buffer: [bs, history_size] int64 - Parent indices (directly used as parentList). - topk_score_indices: [bs, max_total_draft_tokens] int64 - Selected token indices (directly used as selectedIndex). - tree_mask: [bs, N, packed_bits] int32 (packed) or [bs, N, N] bool - Pre-allocated output buffer for attention mask. - positions: [bs, num_draft_tokens] int32 - Pre-allocated output buffer for position IDs. - retrieve_index: [bs, num_draft_tokens] int32 - Pre-allocated output buffer for token retrieval indices. - retrieve_next_token: [bs, num_draft_tokens] int32 - Pre-allocated output buffer for first child indices. - retrieve_next_sibling: [bs, num_draft_tokens] int32 - Pre-allocated output buffer for next sibling indices. - use_packed_mask: bool - Use bit-packed mask for memory efficiency. - """ - bs = topk_score_indices.shape[0] - # +1 because num_draft_tokens includes root node in SGLang's convention - num_draft_tokens = topk_score_indices.shape[1] + 1 - tree_mask_mode = 2 if use_packed_mask else 1 # QLEN_ONLY_BITPACKING / QLEN_ONLY - # Packed layout: last dim is int32 row stride (may exceed ceil(N/32) if padded). - # Non-packed path ignores this (bool [bs,N,N] still has a last dim). - num_int32_per_row = tree_mask.shape[-1] - - # Call CUDA kernel in-place - try: - torch.ops.trtllm.build_dynamic_tree_op( - history_draft_tokens_parent_buffer[:bs], - topk_score_indices, - tree_mask, - positions, - retrieve_index, - retrieve_next_token, - retrieve_next_sibling, - self.K, - self.depth, - num_draft_tokens, - tree_mask_mode, - num_int32_per_row, - ) - except Exception as e: - raise RuntimeError( - f"build_dynamic_tree_op failed: {e}\n" - f"Inputs: bs={bs}, K={self.K}, depth={self.depth}, " - f"num_draft_tokens={num_draft_tokens}" - ) from e - - def verify_dynamic_tree_greedy_out( - self, - candidates: torch.Tensor, - retrieve_index: torch.Tensor, - retrieve_next_token: torch.Tensor, - retrieve_next_sibling: torch.Tensor, - target_predict: torch.Tensor, - num_gens: int, - num_spec_step: int, - tree_valid: torch.Tensor = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """ - In-place verify using pre-allocated output buffers (CUDA graph friendly). - - Args: - candidates: [num_gens, N] int64 candidate tokens. - retrieve_index: [num_gens, N] int32 retrieval indices. - retrieve_next_token: [num_gens, N] int32 next token indices. - retrieve_next_sibling: [num_gens, N] int32 next sibling indices. - target_predict: [num_gens, N] int64 target predictions. - num_gens: Number of generation requests. - num_spec_step: Number of speculative steps. - tree_valid: [num_gens] bool per-request flag. When False the - kernel early-returns with acceptTokenNum=0 (first-gen / - dummy requests). None means all trees are valid. - - Returns: - Tuple of (predicts, accept_index, accept_token_num, accept_token) - as slices of pre-allocated buffers. - """ - N = candidates.size(1) - predicts = self._verify_predicts_buf[: num_gens * N] - accept_index = self._verify_accept_index_buf[:num_gens] - accept_token_num = self._verify_accept_token_num_buf[:num_gens] - accept_token = self._verify_accept_token_buf[:num_gens] - - if tree_valid is None: - tree_valid = torch.ones(num_gens, dtype=torch.bool, device=candidates.device) - - try: - torch.ops.trtllm.verify_dynamic_tree_greedy_out_op( - candidates, - retrieve_index, - retrieve_next_token, - retrieve_next_sibling, - target_predict, - predicts, - accept_index, - accept_token_num, - accept_token, - tree_valid, - num_spec_step, - ) - except Exception as e: - raise RuntimeError( - f"verify_dynamic_tree_greedy_out_op failed: {e}\n" - f"Inputs: num_gens={num_gens}, N={N}, " - f"num_spec_step={num_spec_step}" - ) from e - - return predicts, accept_index, accept_token_num, accept_token diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 0191d59ad08e..0211f3e7c989 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -73,8 +73,7 @@ def __init__(self, self.spec_tree_manager = None if isinstance(config, - EagleDecodingConfig) and (config.eagle_choices is not None - or config.use_dynamic_tree): + EagleDecodingConfig) and config.eagle_choices is not None: self.spec_tree_manager = SpecTreeManager( max_num_requests=self.max_num_requests, use_dynamic_tree=config.use_dynamic_tree, @@ -123,43 +122,6 @@ def get_needed_resource_to_completion(self, request: LlmRequest): return 0 -class Eagle3OneModelDynamicTreeResourceManager(BaseResourceManager): - """ - Lightweight resource manager for one-model EAGLE3 dynamic tree mode. - Holds a SpecTreeManager so that model_engine and sampler can access it - for tree attention setup and verification routing. - """ - - def __init__(self, config: "EagleDecodingConfig", max_num_requests: int): - self.max_num_requests = max_num_requests - self.spec_tree_manager = SpecTreeManager( - max_num_requests=max_num_requests, - use_dynamic_tree=config.use_dynamic_tree, - max_draft_len=config.max_draft_len, - max_total_draft_tokens=config.tokens_per_gen_step - 1, - eagle_choices=config.eagle_choices, - dynamic_tree_max_topK=config.dynamic_tree_max_topK, - ) - - def free_resources(self, request: LlmRequest): - """Clear tree validity for the freed request slot.""" - if request.py_seq_slot is not None: - self.spec_tree_manager.mark_tree_invalid(request.py_seq_slot) - - def add_dummy_requests(self, request_ids: List[int]): - """Handle CUDA graph dummy request registration. - - Dummies use _dummy_slot_id on the spec_tree_manager, so no - per-request slot allocation is needed. - """ - - def get_max_resource_count(self) -> int: - return self.max_num_requests - - def get_needed_resource_to_completion(self, request: LlmRequest): - return 0 - - def _get_eagle3_default_capture_layers(num_layers: int): return (1, num_layers // 2 - 1, num_layers - 4) @@ -351,9 +313,6 @@ class Eagle3OneModelSpecMetadata(SpecMetadata): batch_indices_cuda: Optional[torch.Tensor] = None # Optional resource manager (used to access SA manager for EAGLE3+SA) spec_resource_manager: Optional[Eagle3ResourceManager] = None - # Dynamic tree flags - use_dynamic_tree: bool = False - eagle_choices: Optional[List[List[int]]] = None def __post_init__(self): if self.layers_to_capture is None: @@ -380,16 +339,11 @@ def __post_init__(self): device='cuda', ) - # Set tree flags based on config - if self.use_dynamic_tree: - self.is_spec_dec_tree = True - self.is_spec_dec_dynamic_tree = True - elif self.eagle_choices is not None: - self.is_spec_dec_tree = True - self.is_spec_dec_dynamic_tree = False - else: - self.is_spec_dec_tree = False - self.is_spec_dec_dynamic_tree = False + # currently Eagle3 only supports linear tree + self.is_spec_dec_tree = False + + # currently Eagle3 only supports static tree + self.is_spec_dec_dynamic_tree = False def is_layer_capture(self, layer_id: int): return layer_id in self.layers_to_capture @@ -404,11 +358,7 @@ def prepare(self): pin_memory=prefer_pinned()) self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) - if self.is_spec_dec_tree: - self.num_tokens -= ( - self.num_generations) * self.max_total_draft_tokens - else: - self.num_tokens -= (self.num_generations) * self.max_draft_len + self.num_tokens -= (self.num_generations) * self.runtime_draft_len sa_manager = getattr(self.spec_resource_manager, 'sa_manager', None) if sa_manager is not None: @@ -432,27 +382,12 @@ def maybe_capture_hidden_states( class Eagle3OneModelSampler(MTPSampler): - """Sampler for one-model EAGLE3 (linear and dynamic tree modes).""" - def __init__(self, args: TorchSampler.Args, spec_config=None): - self._spec_config = spec_config - super().__init__(args, nextn=args.max_total_draft_tokens) - - def _get_max_new_tokens(self, args: TorchSampler.Args, - draft_len: int) -> int: - """Dynamic tree: accepted path depth <= max_draft_len + 1.""" - if (self._spec_config is not None - and getattr(self._spec_config, 'use_dynamic_tree', False)): - return self._spec_config.max_draft_len + 1 - return self._get_max_tokens(args, draft_len) + def __init__(self, args: TorchSampler.Args): + super().__init__(args, nextn=args.max_draft_len) class Eagle3OneModelWorker(SpecWorkerBase): - """Eagle3 one-model worker for linear tree speculative decoding. - - For dynamic tree mode, use Eagle3OneModelDynamicTreeWorker from - eagle3_dynamic_tree.py instead. - """ def __init__(self, spec_config: "EagleDecodingConfig", @@ -464,8 +399,6 @@ def __init__(self, self.sa_enhancer: Optional[SADraftEnhancer] = None if getattr(spec_config, 'sa_config', None) is not None: self.sa_enhancer = SADraftEnhancer(spec_config.sa_config.threshold) - self.use_dynamic_tree = getattr(spec_config, 'use_dynamic_tree', False) - self.spec_tree_manager = None @property def max_draft_len(self) -> int: @@ -484,27 +417,6 @@ def _prepare_attn_metadata_for_spec_dec(self, attn_metadata): else: self._saved_kv_lens_cuda = None - # Save spec-dec params that the drafting loop will overwrite. - # Without this, CUDA graph warmup's second iteration would run - # the target model attention with stale draft-layer masks - # instead of the correct target-tree masks. - if attn_metadata.spec_decoding_packed_mask is not None: - self._saved_packed_mask = attn_metadata.spec_decoding_packed_mask[: - batch_size].clone( - ) - else: - self._saved_packed_mask = None - if attn_metadata.spec_decoding_position_offsets is not None: - self._saved_position_offsets = attn_metadata.spec_decoding_position_offsets.clone( - ) - else: - self._saved_position_offsets = None - if attn_metadata.spec_decoding_generation_lengths is not None: - self._saved_generation_lengths = attn_metadata.spec_decoding_generation_lengths[:batch_size].clone( - ) - else: - self._saved_generation_lengths = None - def _restore_attn_metadata_from_spec_dec(self, attn_metadata): super()._restore_attn_metadata_from_spec_dec(attn_metadata) if self._saved_kv_lens_cuda is not None: @@ -513,21 +425,6 @@ def _restore_attn_metadata_from_spec_dec(self, attn_metadata): self._saved_kv_lens_cuda) self._saved_kv_lens_cuda = None - if self._saved_packed_mask is not None: - batch_size = self._saved_packed_mask.shape[0] - attn_metadata.spec_decoding_packed_mask[:batch_size].copy_( - self._saved_packed_mask) - self._saved_packed_mask = None - if self._saved_position_offsets is not None: - attn_metadata.spec_decoding_position_offsets.copy_( - self._saved_position_offsets) - self._saved_position_offsets = None - if self._saved_generation_lengths is not None: - batch_size = self._saved_generation_lengths.shape[0] - attn_metadata.spec_decoding_generation_lengths[:batch_size].copy_( - self._saved_generation_lengths) - self._saved_generation_lengths = None - # Skip torch.compile for now since current Torch is not compatible with Triton 3.4 # @torch.compile(options={"max-autotune": True}) @@ -587,64 +484,16 @@ def forward(self, spec_metadata=spec_metadata, draft_model=draft_model) - # Predict draft tokens + next_draft_tokens = [] original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens # Get the draft KV cache manager if using separate layouts draft_kv_cache_manager = self.get_draft_kv_cache_manager( resource_manager) - next_draft_tokens = self._forward_draft_loop( - inputs, attn_metadata, spec_metadata, draft_model, - draft_kv_cache_manager, num_contexts, num_gens, batch_size, - num_accepted_tokens, original_all_rank_num_tokens, resource_manager) - # restore attn_metadata to support cuda graph - self._restore_attn_metadata_from_spec_dec(attn_metadata) - # restore all_rank_num_tokens for attention DP - if original_all_rank_num_tokens is not None: - attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens - - # prepare next new tokens to support overlap scheduler - next_new_tokens = self._prepare_next_new_tokens( - accepted_tokens, next_draft_tokens, - spec_metadata.batch_indices_cuda, batch_size, num_accepted_tokens) - - attn_metadata.use_spec_decoding = True - - return { - 'logits': raw_logits, - 'new_tokens': accepted_tokens, - 'new_tokens_lens': num_accepted_tokens, - 'next_draft_tokens': next_draft_tokens, - 'next_new_tokens': next_new_tokens, - } - - def _forward_draft_loop(self, inputs, attn_metadata, spec_metadata, - draft_model, draft_kv_cache_manager, num_contexts, - num_gens, batch_size, num_accepted_tokens, - original_all_rank_num_tokens, resource_manager): - """Dispatch to the appropriate draft loop. Subclasses can override.""" - return self._forward_linear_draft_loop(inputs, attn_metadata, - spec_metadata, draft_model, - draft_kv_cache_manager, - num_contexts, batch_size, - num_accepted_tokens, - original_all_rank_num_tokens) - - def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, - draft_model, draft_kv_cache_manager, - num_contexts, batch_size, - num_accepted_tokens, - original_all_rank_num_tokens): - """Original linear draft loop (1 token per layer).""" - runtime_draft_len = spec_metadata.runtime_draft_len - next_draft_tokens = [] - position_ids = inputs["position_ids"] - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): for i in range(runtime_draft_len): if i == 0: - num_gens = batch_size - num_contexts start_ids_gen = ( spec_metadata.batch_indices_cuda[:num_gens] * (runtime_draft_len + 1)).long() @@ -732,7 +581,26 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, gen_draft_tokens) next_draft_tokens[num_contexts:] = gen_draft_tokens - return next_draft_tokens + # restore attn_metadata to support cuda graph + self._restore_attn_metadata_from_spec_dec(attn_metadata) + # restore all_rank_num_tokens for attention DP + if original_all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + + # prepare next new tokens to support overlap scheduler + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, next_draft_tokens, + spec_metadata.batch_indices_cuda, batch_size, num_accepted_tokens) + + attn_metadata.use_spec_decoding = True + + return { + 'logits': raw_logits, + 'new_tokens': accepted_tokens, + 'new_tokens_lens': num_accepted_tokens, + 'next_draft_tokens': next_draft_tokens, + 'next_new_tokens': next_new_tokens, + } def sample_and_accept_draft_tokens( self, @@ -744,14 +612,8 @@ def sample_and_accept_draft_tokens( num_contexts = attn_metadata.num_contexts num_gens = batch_size - num_contexts - # Linear mode: reshape draft tokens for base implementation draft_tokens = spec_metadata.draft_tokens.reshape( - num_gens, - spec_metadata.runtime_draft_len) if num_gens > 0 else torch.empty( - 0, - spec_metadata.runtime_draft_len, - dtype=torch.int, - device=logits.device) + num_gens, spec_metadata.runtime_draft_len) return self._sample_and_accept_draft_tokens_base( logits, draft_tokens, num_contexts, batch_size, spec_metadata) diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py deleted file mode 100644 index f33a5824a455..000000000000 --- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py +++ /dev/null @@ -1,1049 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-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. -"""Eagle3 one-model dynamic tree speculative decoding.""" - -from typing import TYPE_CHECKING - -import torch -import triton -import triton.language as tl - -from tensorrt_llm._utils import get_sm_version, nvtx_range - -from ..attention_backend import AttentionMetadata -from .eagle3 import Eagle3OneModelWorker - -if TYPE_CHECKING: - from ...llmapi.llm_args import EagleDecodingConfig - - -@torch.compile(options={"max-autotune": True}) -def _sample_softmax_topk(logits: torch.Tensor, k: int) -> tuple[torch.Tensor, torch.Tensor]: - """Fused softmax+topk for draft token sampling.""" - last_p = torch.softmax(logits, dim=-1) - topk_values, topk_indices = torch.topk(last_p, k=k, dim=-1) - return topk_indices, topk_values - - -@torch.compile(options={"max-autotune": True}) -def _select_topk_draft_tokens( - new_draft_tokens: torch.Tensor, - new_draft_scores: torch.Tensor, - previous_draft_scores: torch.Tensor, - K: int, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Fused score accumulation → topk → gather for draft layer selection. - - Returns (real_tokens, topk_values, topk_indices, parents, - all_tokens_flat, all_scores_flat) where the last two are the full - K*K reshaped tensors needed for history buffer writes. - """ - all_tokens_flat = new_draft_tokens.reshape(-1, K * K) - all_scores_flat = (new_draft_scores * previous_draft_scores.unsqueeze(2)).reshape(-1, K * K) - topk_values, topk_indices = torch.topk(all_scores_flat, k=K, dim=-1) - real_tokens = torch.gather(all_tokens_flat, dim=1, index=topk_indices) - parents = topk_indices // K - return real_tokens, topk_values, topk_indices, parents, all_tokens_flat, all_scores_flat - - -@torch.compile(options={"max-autotune": True}) -def _resample_final_tokens( - history_scores: torch.Tensor, - history_tokens: torch.Tensor, - k: int, -) -> tuple[torch.Tensor, torch.Tensor]: - """Fused topk → sort → gather for final tree token resampling.""" - topk_score_indices = torch.topk(history_scores, k=k, dim=-1).indices - topk_score_indices = torch.sort(topk_score_indices).values - real_draft_tokens = torch.gather(history_tokens, dim=1, index=topk_score_indices) - return real_draft_tokens, topk_score_indices - - -@triton.jit -def _gather_repack_step0_kernel( - hidden_states_ptr, - accept_token_ptr, - position_ids_ptr, - accepted_indices_ptr, - num_accepted_ptr, - out_hs_ptr, - out_ids_ptr, - out_pos_ptr, - out_gather_ids_ptr, - num_ctx_tokens, - num_contexts, - tokens_per_gen_step, - max_path_len, - max_draft_len, - hidden_dim, - gather_id_offset, - BLOCK_H: tl.constexpr, -): - """Fused gather+repack for generation requests in step 0. - - Each program handles one (gen_request, path_position) pair: - - Gathers hidden_states from tree-topology to contiguous layout - - Copies input_id from accept_token - - Computes position_id from base_position + path_offset - - Computes gather_id for the last accepted position - """ - pid = tl.program_id(0) - gen_idx = pid // max_path_len - path_idx = pid % max_path_len - - # Source position in tree-topology layout - gen_start = num_ctx_tokens + gen_idx * tokens_per_gen_step - # tl.where evaluates both branches, so clamp index to avoid OOB when path_idx==0 - safe_idx = tl.maximum(path_idx - 1, 0) - raw_val = tl.load(accepted_indices_ptr + gen_idx * max_draft_len + safe_idx) - tree_pos = tl.where(path_idx == 0, 0, raw_val.to(tl.int64) + 1) - src_row = gen_start + tree_pos - - dst_row = gen_idx * max_path_len + path_idx - - # 1) Gather hidden_states: src[src_row, :] → out[dst_row, :] - for h in range(0, hidden_dim, BLOCK_H): - offsets = h + tl.arange(0, BLOCK_H) - mask = offsets < hidden_dim - vals = tl.load(hidden_states_ptr + src_row * hidden_dim + offsets, mask=mask) - tl.store(out_hs_ptr + dst_row * hidden_dim + offsets, vals, mask=mask) - - # 2) Copy input_id - token = tl.load(accept_token_ptr + gen_idx * max_path_len + path_idx) - tl.store(out_ids_ptr + dst_row, token) - - # 3) Compute position_id: base_pos + path_idx - base_pos = tl.load(position_ids_ptr + gen_start) - tl.store(out_pos_ptr + dst_row, base_pos + path_idx) - - # 4) Compute gather_id (only for last accepted token of each gen request) - # gather_id references the FINAL combined [ctx | gen] tensor - n_acc = tl.load(num_accepted_ptr + num_contexts + gen_idx).to(tl.int64) - if path_idx == 0: - gather_id = gather_id_offset + gen_idx * max_path_len + n_acc - 1 - tl.store(out_gather_ids_ptr + num_contexts + gen_idx, gather_id) - - -class Eagle3OneModelDynamicTreeWorker(Eagle3OneModelWorker): - """Eagle3 one-model worker with dynamic tree support. - - Inherits linear tree functionality from Eagle3OneModelWorker and adds - dynamic tree draft loop, verification, and tree construction. - """ - - def __init__( - self, spec_config: "EagleDecodingConfig", mapping, use_separate_draft_kv_cache: bool = False - ): - super().__init__(spec_config, mapping, use_separate_draft_kv_cache) - assert self.use_dynamic_tree, ( - "Eagle3OneModelDynamicTreeWorker requires use_dynamic_tree=True" - ) - - from .dynamic_tree_ops import DynamicTreeOpsConverter - - self.K = spec_config.dynamic_tree_max_topK - self.max_total_draft_tokens = spec_config.tokens_per_gen_step - 1 - self.tokens_per_gen_step = spec_config.tokens_per_gen_step - if spec_config.max_batch_size is None: - raise ValueError( - "Eagle3OneModelDynamicTreeWorker requires max_batch_size to be set " - "on Eagle3DecodingConfig when use_dynamic_tree=True." - ) - self._max_batch_size = spec_config.max_batch_size - - K = self.K - max_draft_len = spec_config.max_draft_len - max_batch_size = self._max_batch_size - loop_max_tokens = K * max_draft_len # draft loop working size - - # Pre-allocated 2D buffers - self.draft_tokens_buffer = torch.zeros( - max_batch_size, loop_max_tokens, dtype=torch.int64, device="cuda" - ) - self.position_ids_buffer = torch.zeros( - max_batch_size, loop_max_tokens, dtype=torch.int64, device="cuda" - ) - self.history_draft_tokens_buffer = torch.zeros( - (max_batch_size, (K + K * K * (max_draft_len - 1))), dtype=torch.int64, device="cuda" - ) - self.history_score_buffer = torch.zeros( - (max_batch_size, K + K * K * (max_draft_len - 1)), dtype=torch.float32, device="cuda" - ) - self.history_draft_tokens_parent_buffer = torch.zeros( - (max_batch_size, max(K * (max_draft_len - 1) + 1, K + 1)), - dtype=torch.int64, - device="cuda", - ) - self.tree_mask_buffer = torch.zeros( - (max_batch_size * loop_max_tokens * loop_max_tokens), - dtype=torch.int32, - device="cuda", - ) - self.tree_mask_init_buffer = ( - torch.eye(K, dtype=torch.int32, device="cuda").unsqueeze(0).repeat(max_batch_size, 1, 1) - ) - self.tree_ops_converter = DynamicTreeOpsConverter( - dynamic_tree_max_topK=K, - max_draft_len=max_draft_len, - max_total_draft_tokens=self.max_total_draft_tokens, - max_batch_size=max_batch_size, - device=torch.device("cuda"), - ) - - self._accepted_draft_indices_tensor = torch.full( - (max_batch_size, max_draft_len), -1, dtype=torch.int32, device="cuda" - ) - - self._accept_token = torch.zeros( - max_batch_size, max_draft_len + 1, dtype=torch.int64, device="cuda" - ) - self._last_selected_parents = None - self._kv_head_dim_bytes = None - self._max_path_len = max_draft_len + 1 - self._kv_correction = self.tokens_per_gen_step - self._max_path_len - - # Hidden state management buffers (lazily initialized in update_hidden_states) - self._hs_write_buffer = None - self._hs_read_map = torch.zeros( - max_batch_size, max_draft_len * K, dtype=torch.long, device="cuda" - ) - self._accumulated_hs = None - self._step0_hs = None - self._hs_dim = None - - # Step 0 buffers - self._step0_causal_mask = torch.tensor( - [(1 << (t + 1)) - 1 for t in range(self._max_path_len)], - dtype=torch.int32, - device="cuda", - ) - self._arange_max_batch = torch.arange(max_batch_size, device="cuda") - self._causal_offs = torch.arange(self._max_path_len, device="cuda", dtype=torch.int32) - self._parent_init_arange = torch.arange(-1, K, device="cuda", dtype=torch.int32) - - # Verification buffers - tokens_per_gen_step = self.tokens_per_gen_step - self._accepted_tokens_buf = torch.zeros( - max_batch_size, self._max_path_len, dtype=torch.int32, device="cuda" - ) - self._num_accepted_tokens_buf = torch.ones(max_batch_size, dtype=torch.int32, device="cuda") - self._target_tokens_buf = torch.zeros( - max_batch_size * tokens_per_gen_step, dtype=torch.int64, device="cuda" - ) - self._candidates_buf = torch.zeros( - max_batch_size, tokens_per_gen_step, dtype=torch.int64, device="cuda" - ) - self._target_predict_buf = torch.zeros( - max_batch_size, tokens_per_gen_step, dtype=torch.int64, device="cuda" - ) - - # Step 0 input buffers - max_total_tokens = max_batch_size * tokens_per_gen_step - self._step0_input_ids_buf = torch.zeros(max_total_tokens, dtype=torch.int64, device="cuda") - self._step0_position_ids_buf = torch.zeros( - max_total_tokens, dtype=torch.int64, device="cuda" - ) - self._step0_hidden_states_buf = None - self._gather_ids_buf = torch.zeros(max_total_tokens, dtype=torch.long, device="cuda") - self._current_mask_buf = torch.zeros( - max_batch_size, - loop_max_tokens, - loop_max_tokens, - dtype=torch.int32, - device="cuda", - ) - self._new_pos_offset_buf = torch.zeros( - max_batch_size, loop_max_tokens, dtype=torch.int32, device="cuda" - ) - - # Mask repack scratch (graph-safe; avoids .contiguous() in the draft loop). - buf_dim = max(self.max_total_draft_tokens + 1, K * max_draft_len) - mask_width = (buf_dim + 31) // 32 - self._mask_repack_buf = torch.zeros( - max_batch_size * buf_dim * mask_width, dtype=torch.int32, device="cuda" - ) - - # sm≥100 (except 120/121): prepareCustomMask keeps padded 3D; no 1D repack. - sm = get_sm_version() - self._needs_mask_repack = sm < 100 or sm in (120, 121) - - def _repack_mask_padded_to_packed(self, mask_buf, n_req, n_tok): - """XQA indexes mask flat via cuQSeqLens; padded [n_req, buf_dim, maskW] has - batch stride buf_dim*maskW, so when n_tok < buf_dim and n_req > 1 packed - reads are wrong. Copy [:n_req,:n_tok] through scratch into flat prefix.""" - buf_dim = mask_buf.shape[1] - if n_tok >= buf_dim or n_req <= 1: - return - mask_width = mask_buf.shape[2] - total_elems = n_req * n_tok * mask_width - scratch = self._mask_repack_buf[:total_elems].view(n_req, n_tok, mask_width) - scratch.copy_(mask_buf[:n_req, :n_tok, :]) - flat = mask_buf.view(-1) - flat[:total_elems] = scratch.view(-1) - - @nvtx_range("eagle3_dyn._ensure_spec_tree_manager") - def _ensure_spec_tree_manager(self, resource_manager): - """Lazily initialize spec_tree_manager and KV head metadata.""" - if self.spec_tree_manager is not None: - return - from ..pyexecutor.resource_manager import ResourceManagerType - - spec_rm = resource_manager.get_resource_manager(ResourceManagerType.SPEC_RESOURCE_MANAGER) - assert spec_rm is not None and hasattr(spec_rm, "spec_tree_manager"), ( - "Dynamic tree mode requires spec_tree_manager in resource_manager" - ) - self.spec_tree_manager = spec_rm.spec_tree_manager - - if self._kv_head_dim_bytes is None: - cache_mgr = resource_manager.get_resource_manager(ResourceManagerType.KV_CACHE_MANAGER) - if cache_mgr is not None: - from tensorrt_llm.bindings import DataType - - _dtype_bytes = { - DataType.HALF: 2, - DataType.BF16: 2, - DataType.FLOAT: 4, - DataType.FP8: 1, - DataType.INT8: 1, - DataType.NVFP4: 0.5, - } - self._kv_head_dim_bytes = int( - cache_mgr.head_dim * _dtype_bytes.get(cache_mgr.dtype, 0.5) - ) - - @nvtx_range("eagle3_dyn.forward") - def forward( - self, - input_ids, - position_ids, - hidden_states, - logits, - attn_metadata, - spec_metadata, - draft_model, - resource_manager=None, - ): - """Override to add accepted_draft_tokens_indices to output.""" - # Initialize spec_tree_manager before super().forward() which calls - # _forward_draft_loop needing spec_tree_manager. - if resource_manager is not None: - self._ensure_spec_tree_manager(resource_manager) - output = super().forward( - input_ids, - position_ids, - hidden_states, - logits, - attn_metadata, - spec_metadata, - draft_model, - resource_manager, - ) - batch_size = attn_metadata.num_seqs - output["accepted_draft_tokens_indices"] = self._accepted_draft_indices_tensor[:batch_size] - - return output - - @nvtx_range("eagle3_dyn._relocate_kv_eagerly") - def _relocate_kv_eagerly(self, attn_metadata, batch_size): - """Move accepted draft tokens' KV from tree to linear positions. - - Called inside sample_and_accept_draft_tokens (captured in CUDA graph). - All tensor arguments are pre-allocated so their device addresses are - stable across CUDA graph capture and replay. - - Must use attn_metadata.kv_cache_manager (the target model's live KV - cache), as it may differ from ResourceManagerType.KV_CACHE_MANAGER in - one-model spec decoding configurations. - """ - cache_mgr = getattr(attn_metadata, "kv_cache_manager", None) - if cache_mgr is None: - return - - assert len(set(cache_mgr.num_kv_heads_per_layer)) == 1, ( - "update_kv_cache_draft_token_location_2d requires uniform num_kv_heads across all layers, " - f"but got {cache_mgr.num_kv_heads_per_layer}" - ) - torch.ops.tensorrt_llm.update_kv_cache_draft_token_location_2d( - self._accepted_draft_indices_tensor[:batch_size], - self._num_accepted_tokens_buf[:batch_size], - attn_metadata.kv_lens_cuda[:batch_size], - True, - cache_mgr.num_layers, - # Use TP-sharded num_kv_heads (per-rank) instead of the unsharded - # total so the C++ kernel computes correct strides and grid dims. - cache_mgr.num_kv_heads_per_layer[0], - self._kv_head_dim_bytes, - cache_mgr.max_total_draft_tokens, - cache_mgr.max_attention_window_vec[0], - cache_mgr.kv_cache_pool_pointers, - attn_metadata.kv_cache_block_offsets, - cache_mgr.max_blocks_per_seq, - cache_mgr.tokens_per_block, - None, - ) - - @nvtx_range("eagle3_dyn.sample_and_accept_draft_tokens") - def sample_and_accept_draft_tokens(self, logits, attn_metadata, spec_metadata): - """Override to handle dynamic tree verification.""" - batch_size = attn_metadata.num_seqs - num_contexts = attn_metadata.num_contexts - num_gens = batch_size - num_contexts - - accepted_tokens, num_accepted_tokens = self._sample_and_accept_dynamic_tree( - logits, attn_metadata, spec_metadata, batch_size, num_contexts, num_gens - ) - if num_gens > 0: - self._relocate_kv_eagerly(attn_metadata, batch_size) - return accepted_tokens, num_accepted_tokens - - @nvtx_range("eagle3_dyn.prepare_1st_drafter_inputs") - def prepare_1st_drafter_inputs( - self, - input_ids, - position_ids, - hidden_states, - accepted_tokens, - attn_metadata, - spec_metadata, - draft_model, - ): - """Re-pack gen inputs from tree-topology to uniform-padded layout - (max_draft_len + 1 tokens per gen request) for the shared KV cache. - Uses a fused Triton kernel for gen requests to minimize kernel launches. - """ - num_contexts = attn_metadata.num_contexts - num_gens = attn_metadata.num_seqs - num_contexts - batch_size = attn_metadata.num_seqs - num_tokens = input_ids.shape[0] - num_ctx_tokens = attn_metadata.num_ctx_tokens - - hidden_size_up = spec_metadata.hidden_size * len(spec_metadata.layers_to_capture) - hidden_states = spec_metadata.hidden_states[:num_tokens, :hidden_size_up] - hidden_states = draft_model.apply_eagle3_fc(hidden_states) - - input_ids_ctx = self._prepare_context_input_ids( - input_ids, - num_ctx_tokens, - spec_metadata.gather_ids, - accepted_tokens, - num_contexts, - ) - - if num_gens > 0: - max_path_len = self._max_path_len - num_gen_tokens = num_gens * max_path_len - - hidden_dim = hidden_states.shape[-1] - if ( - self._step0_hidden_states_buf is None - or self._step0_hidden_states_buf.shape[-1] != hidden_dim - ): - self._step0_hidden_states_buf = torch.zeros( - self._step0_input_ids_buf.shape[0], - hidden_dim, - dtype=hidden_states.dtype, - device="cuda", - ) - - # Gen tokens: fused Triton kernel writes into gen-only buffers - # (buffers sized for max_batch_size * tokens_per_gen_step) - BLOCK_H = triton.next_power_of_2(hidden_dim) - _gather_repack_step0_kernel[(num_gens * max_path_len,)]( - hidden_states, - self._accept_token[:num_gens], - position_ids, - self._accepted_draft_indices_tensor[num_contexts:batch_size], - self._num_accepted_tokens_buf, - self._step0_hidden_states_buf, - self._step0_input_ids_buf, - self._step0_position_ids_buf, - self._gather_ids_buf, - num_ctx_tokens, - num_contexts, - self.tokens_per_gen_step, - max_path_len, - self.max_draft_len, - hidden_dim, - num_ctx_tokens, # gather_id references combined [ctx|gen] tensor - BLOCK_H=BLOCK_H, - ) - - # Concat context + gen - input_ids = torch.cat( - [input_ids_ctx, self._step0_input_ids_buf[:num_gen_tokens]], dim=0 - ) - position_ids = torch.cat( - [position_ids[:num_ctx_tokens], self._step0_position_ids_buf[:num_gen_tokens]], - dim=0, - ) - hidden_states = torch.cat( - [hidden_states[:num_ctx_tokens], self._step0_hidden_states_buf[:num_gen_tokens]], - dim=0, - ) - - attn_metadata._seq_lens[num_contexts:batch_size].fill_(max_path_len) - attn_metadata._seq_lens_cuda[num_contexts:batch_size].fill_(max_path_len) - attn_metadata.on_update() - else: - # Context-only (warmup). No gen tokens. - input_ids = input_ids_ctx - - return { - "input_ids": input_ids, - "position_ids": position_ids, - "hidden_states": hidden_states, - "attn_metadata": attn_metadata, - "spec_metadata": spec_metadata, - } - - def _forward_draft_loop( - self, - inputs, - attn_metadata, - spec_metadata, - draft_model, - draft_kv_cache_manager, - num_contexts, - num_gens, - batch_size, - num_accepted_tokens, - original_all_rank_num_tokens, - resource_manager, - ): - """Dynamic tree draft loop with growing context.""" - spec_tree_manager = self.spec_tree_manager - - assert batch_size <= self._max_batch_size, ( - f"batch_size {batch_size} exceeds pre-allocated max_batch_size {self._max_batch_size}" - ) - - # === Step 0: Initial forward === - num_step0_tokens = self._max_path_len - - # Triton kernel already wrote gen gather_ids into _gather_ids_buf; - # just prepend context gather_ids. - self._gather_ids_buf[:num_contexts].copy_(spec_metadata.gather_ids[:num_contexts]) - gather_ids = self._gather_ids_buf[:batch_size] - - if original_all_rank_num_tokens is not None: - attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens - - # Step-0 causal spec-dec (None in prefill-only warmup). - if attn_metadata.spec_decoding_generation_lengths is not None: - # Gen kernel uses base data_ptr (no num_contexts offset); fill rows [:num_gens]. - attn_metadata.spec_decoding_generation_lengths[:num_gens].fill_(num_step0_tokens) - - # Position stride num_step0_tokens matches C++ generation_input_length. - pos_2d = attn_metadata.spec_decoding_position_offsets[ - : num_gens * num_step0_tokens - ].view(num_gens, num_step0_tokens) - pos_2d[:] = self._causal_offs[:num_step0_tokens] - - attn_metadata.spec_decoding_packed_mask[:num_gens].fill_(0) - attn_metadata.spec_decoding_packed_mask[:num_gens, :num_step0_tokens, 0] = ( - self._step0_causal_mask[:num_step0_tokens] - ) - # Packed flat for cuQSeqLens when _needs_mask_repack; else keep padded layout. - if self._needs_mask_repack: - self._repack_mask_padded_to_packed( - attn_metadata.spec_decoding_packed_mask, num_gens, num_step0_tokens - ) - - attn_metadata.use_spec_decoding = num_gens > 0 - - # KV correction: target processed tokens_per_gen_step per request, - # but step 0 should only attend to max_draft_len + 1 accepted-path tokens. - if num_gens > 0 and hasattr(attn_metadata, "kv_lens_cuda"): - attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= self._kv_correction - - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): - hidden_states, hidden_states_to_save = draft_model.model(**inputs) - - step0_hs = hidden_states_to_save[gather_ids] - - logits = draft_model.logits_processor( - hidden_states[gather_ids], draft_model.lm_head, attn_metadata, True - ) - - new_draft_tokens, new_draft_scores = self.sample( - logits, self.K, draft_model=draft_model - ) - - previous_draft_scores = self.update_draft_tokens_and_scores( - cur_draft_idx=0, - new_draft_tokens=new_draft_tokens, - new_draft_scores=new_draft_scores, - previous_draft_scores=None, - batch_size=batch_size, - attn_metadata=attn_metadata, - ) - - self.update_hidden_states( - cur_draft_idx=0, - batch_size=batch_size, - step0_hs=step0_hs, - ) - - self.prepare_for_generation( - 0, - attn_metadata, - batch_size, - inputs=inputs, - gather_ids=gather_ids, - num_contexts=num_contexts, - num_gens=num_gens, - num_accepted_tokens=num_accepted_tokens, - ) - - for layer_idx in range(1, self.max_draft_len): - num_tokens_per_req = layer_idx * self.K - - if original_all_rank_num_tokens is not None: - if spec_metadata.all_rank_num_seqs is not None: - attn_metadata.all_rank_num_tokens = spec_metadata.all_rank_num_seqs - - # Growing context: process ALL accumulated tokens - num_infer_tokens = batch_size * num_tokens_per_req - - inp_hs = self._accumulated_hs[:batch_size, :num_tokens_per_req, :].reshape( - num_infer_tokens, -1 - ) - inp_ids = ( - self.draft_tokens_buffer[:batch_size, :num_tokens_per_req] - .reshape(-1) - .to(torch.int32) - ) - inp_pos = self.position_ids_buffer[:batch_size, :num_tokens_per_req].reshape(-1) - inputs = { - "input_ids": inp_ids, - "position_ids": inp_pos, - "hidden_states": inp_hs, - "attn_metadata": attn_metadata, - "spec_metadata": spec_metadata, - } - - hidden_states, hidden_states_to_save = draft_model.model(**inputs) - - # Take last K logits per request - hs_reshaped = hidden_states.reshape(batch_size, num_tokens_per_req, -1) - selected_hs = hs_reshaped[:, -self.K :, :].reshape(batch_size * self.K, -1) - logits = draft_model.logits_processor( - selected_hs, draft_model.lm_head, attn_metadata, True - ) - - new_draft_tokens, new_draft_scores = self.sample( - logits, self.K, draft_model=draft_model - ) - - # Reshape for update: [batch_size, K, K] - new_draft_tokens = new_draft_tokens.reshape(batch_size, self.K, self.K) - new_draft_scores = new_draft_scores.reshape(batch_size, self.K, self.K) - - previous_draft_scores = self.update_draft_tokens_and_scores( - cur_draft_idx=layer_idx, - new_draft_tokens=new_draft_tokens, - new_draft_scores=new_draft_scores, - previous_draft_scores=previous_draft_scores, - batch_size=batch_size, - attn_metadata=attn_metadata, - ) - - self.update_hidden_states( - cur_draft_idx=layer_idx, - batch_size=batch_size, - hidden_states_to_save=hidden_states_to_save, - selected_parents=self._last_selected_parents, - ) - - self.prepare_for_generation(layer_idx, attn_metadata, batch_size) - - # Resample final tokens and build tree - real_draft_tokens, topk_score_indices = self.resampling_final_draft_tokens(batch_size) - - if spec_tree_manager is not None: - # Gen-only; spec-dec batch is [:num_gens] (not num_contexts-offset). - self.tree_ops_converter.build_dynamic_tree( - history_draft_tokens_parent_buffer=self.history_draft_tokens_parent_buffer[ - num_contexts:batch_size - ], - topk_score_indices=topk_score_indices[num_contexts:], - tree_mask=spec_tree_manager.spec_dec_packed_mask[:num_gens], - positions=spec_tree_manager.spec_dec_position_offsets[:num_gens], - retrieve_index=spec_tree_manager.retrieve_index[:num_gens], - retrieve_next_token=spec_tree_manager.retrieve_next_token[:num_gens], - retrieve_next_sibling=spec_tree_manager.retrieve_next_sibling[:num_gens], - use_packed_mask=True, - ) - - # Scatter only gen trees to stable slot storage - gen_slots = spec_tree_manager._all_slot_ids_buf[num_contexts:batch_size] - spec_tree_manager.scatter_trees_to_slots(gen_slots, num_gens) - spec_tree_manager.mark_tree_valid(gen_slots, num_gens) - - return real_draft_tokens.to(torch.int32) - - def _sample_and_accept_dynamic_tree( - self, logits, attn_metadata, spec_metadata, batch_size, num_contexts, num_gens - ): - """Dynamic tree verification using CUDA kernel.""" - N = self.tokens_per_gen_step - max_path_len = self._max_path_len - - # Reset output buffers - self._accepted_tokens_buf[:batch_size].zero_() - accepted_tokens = self._accepted_tokens_buf[:batch_size, :max_path_len] - self._num_accepted_tokens_buf[:batch_size].fill_(1) - num_accepted_tokens = self._num_accepted_tokens_buf[:batch_size] - self._accepted_draft_indices_tensor[:batch_size].fill_(-1) - - num_flat_tokens = logits.shape[0] - torch.argmax(logits, dim=-1, out=self._target_tokens_buf[:num_flat_tokens]) - target_tokens = self._target_tokens_buf[:num_flat_tokens] - - # Context requests: accept sampled token - accepted_tokens[:num_contexts, 0] = target_tokens[:num_contexts].to(torch.int32) - - # Generation requests: tree verification - if num_gens > 0: - spec_tree_manager = self.spec_tree_manager - - target_predict = self._target_predict_buf[:num_gens] - target_predict[:] = target_tokens[num_contexts:].reshape(num_gens, N) - - if spec_tree_manager is None: - # CUDA graph warmup: accept only the first token per request - num_accepted_tokens[num_contexts:batch_size] = 1 - accepted_tokens[num_contexts:batch_size, 0] = target_predict[:, 0].to(torch.int32) - self._accepted_draft_indices_tensor[num_contexts:batch_size] = -1 - return accepted_tokens, num_accepted_tokens - - candidates = self._candidates_buf[:num_gens] - candidates[:, 1:] = spec_metadata.draft_tokens.reshape(num_gens, N - 1).to(torch.int64) - candidates[:, 0] = target_predict[:, 0] - - # Slots for gen rows: real py_seq_slot vs dummy; dummy -> slot_has_tree False. - gen_slot_ids = spec_tree_manager._all_slot_ids_buf[ - num_contexts : num_contexts + num_gens - ] - tree_valid = spec_tree_manager.slot_has_tree[gen_slot_ids] - - _, accept_index, accept_token_num, accept_token = ( - self.tree_ops_converter.verify_dynamic_tree_greedy_out( - candidates, - spec_tree_manager.retrieve_index[:num_gens], - spec_tree_manager.retrieve_next_token[:num_gens], - spec_tree_manager.retrieve_next_sibling[:num_gens], - target_predict, - num_gens, - self._max_path_len, - tree_valid=tree_valid, - ) - ) - - self._accept_token = accept_token - n_acc_draft = accept_token_num[:num_gens] - num_accepted_tokens[num_contexts:batch_size] = (n_acc_draft + 1).to(torch.int32) - accepted_tokens[num_contexts:batch_size] = accept_token[:num_gens].to(torch.int32) - # accept_index is 0-based from kernel; -1 converts padding (0) to sentinel (-1) - self._accepted_draft_indices_tensor[num_contexts:batch_size] = ( - accept_index[:num_gens, 1:max_path_len] - 1 - ).to(torch.int32) - - num_accepted_tokens = self._apply_force_accepted_tokens( - num_accepted_tokens, num_contexts, self.max_draft_len - ) - - return accepted_tokens, num_accepted_tokens - - @nvtx_range("eagle3_dyn.sample") - def sample( - self, logits: torch.Tensor, max_top_k: int, draft_model=None - ) -> tuple[torch.Tensor, torch.Tensor]: - """TopK sampling with softmax for dynamic tree.""" - topk_indices, topk_values = _sample_softmax_topk(logits, max_top_k) - # Apply draft-to-target vocab mapping if the draft model has it - if draft_model is not None and hasattr(draft_model.model, "d2t"): - d2t = draft_model.model.d2t.data - topk_indices = topk_indices + d2t[topk_indices] - return topk_indices, topk_values - - def update_draft_tokens_and_scores( - self, - cur_draft_idx: int, - new_draft_tokens: torch.Tensor, - new_draft_scores: torch.Tensor, - previous_draft_scores: torch.Tensor, - batch_size: int, - attn_metadata: AttentionMetadata = None, - ): - """Update draft tokens and scores, write contiguously to buffer.""" - if cur_draft_idx == 0: - new_draft_scores = new_draft_scores.reshape(batch_size, self.K) - - new_draft_tokens_2d = new_draft_tokens.reshape(batch_size, self.K) - self.draft_tokens_buffer[:batch_size, : self.K] = new_draft_tokens_2d - self.history_draft_tokens_buffer[:batch_size, : self.K] = new_draft_tokens_2d - self.history_score_buffer[:batch_size, : self.K] = new_draft_scores - - # Parent buffer: -1 for root, 0..K-1 for first layer - self.history_draft_tokens_parent_buffer[:batch_size, : self.K + 1] = ( - self._parent_init_arange - ) - - self.prepare_tree_mask_and_position_offset(cur_draft_idx, attn_metadata, None) - - return new_draft_scores - else: - # Fused: score accumulation → topk → gather → parent selection - ( - real_draft_tokens, - topk_values, - topk_indices, - selected_parents, - new_draft_tokens, - new_draft_scores, - ) = _select_topk_draft_tokens( - new_draft_tokens, new_draft_scores, previous_draft_scores, self.K - ) - - num_tokens_previous_layer = cur_draft_idx * self.K - num_tokens_current_layer = (cur_draft_idx + 1) * self.K - self.draft_tokens_buffer[ - :batch_size, num_tokens_previous_layer:num_tokens_current_layer - ] = real_draft_tokens - - write_history_start_offset = self.K + (cur_draft_idx - 1) * self.K * self.K - write_history_end_offset = write_history_start_offset + self.K * self.K - self.history_draft_tokens_buffer[ - :batch_size, write_history_start_offset:write_history_end_offset - ] = new_draft_tokens - self.history_score_buffer[ - :batch_size, write_history_start_offset:write_history_end_offset - ] = new_draft_scores - - self._last_selected_parents = selected_parents - self.prepare_tree_mask_and_position_offset( - cur_draft_idx, attn_metadata, selected_parents - ) - - # Update parent buffer for next layer - if cur_draft_idx < self.max_draft_len - 1: - next_layer_start = cur_draft_idx * self.K + 1 - next_layer_end = next_layer_start + self.K - parents_relative_indices = topk_indices + self.K**2 * (cur_draft_idx - 1) + self.K - self.history_draft_tokens_parent_buffer[ - :batch_size, next_layer_start:next_layer_end - ] = parents_relative_indices - - return topk_values - - def resampling_final_draft_tokens(self, batch_size: int): - """Reconstruct the tree based on history buffers.""" - return _resample_final_tokens( - self.history_score_buffer[:batch_size, :], - self.history_draft_tokens_buffer[:batch_size, :], - self.max_total_draft_tokens, - ) - - def prepare_tree_mask_and_position_offset( - self, - cur_draft_idx: int, - attn_metadata: AttentionMetadata, - selected_parents: torch.Tensor = None, - ): - """Prepare the mask and position offsets for the next layer.""" - if attn_metadata.spec_decoding_packed_mask is None: - return - spec_tree_manager = self.spec_tree_manager - batch_size = attn_metadata.num_seqs - num_tokens_current_layer = self.K * (cur_draft_idx + 1) - num_tokens_previous_layer = self.K * cur_draft_idx - if cur_draft_idx == 0: - attn_metadata.spec_decoding_packed_mask.fill_(0) - spec_tree_manager.compute_spec_dec_packed_mask( - self.tree_mask_init_buffer[:batch_size], - attn_metadata.spec_decoding_packed_mask[:batch_size], - ) - self.tree_mask_buffer[ - : batch_size * num_tokens_current_layer * num_tokens_current_layer - ].copy_(self.tree_mask_init_buffer[:batch_size].view(-1)) - attn_metadata.spec_decoding_position_offsets.fill_(0) - attn_metadata.spec_decoding_generation_lengths[:batch_size] = num_tokens_current_layer - else: - num_parent_mask = batch_size * cur_draft_idx * self.K * cur_draft_idx * self.K - parent_mask = self.tree_mask_buffer[:num_parent_mask].reshape( - batch_size, cur_draft_idx * self.K, cur_draft_idx * self.K - ) - - selected_parents_expanded = selected_parents.unsqueeze(-1).expand( - batch_size, self.K, parent_mask.size(-1) - ) - parent_mask_selected = torch.gather( - parent_mask[:, -self.K :, :], dim=1, index=selected_parents_expanded - ) - - # current_mask = cat([parent_mask_selected, tree_mask_init], dim=2) then - # current_mask = cat([mask_padding, current_mask], dim=1) - current_mask = self._current_mask_buf[ - :batch_size, :num_tokens_current_layer, :num_tokens_current_layer - ] - # Top rows: padding zeros - current_mask[:, :num_tokens_previous_layer, :].zero_() - # Bottom rows: [parent_mask_selected | tree_mask_init] - prev_cols = parent_mask.size(-1) # = num_tokens_previous_layer - current_mask[:, num_tokens_previous_layer:, :prev_cols].copy_(parent_mask_selected) - current_mask[:, num_tokens_previous_layer:, prev_cols:num_tokens_current_layer].copy_( - self.tree_mask_init_buffer[:batch_size] - ) - - spec_tree_manager.compute_spec_dec_packed_mask( - current_mask, attn_metadata.spec_decoding_packed_mask[:batch_size] - ) - self.tree_mask_buffer[ - : batch_size * num_tokens_current_layer * num_tokens_current_layer - ].copy_(current_mask.reshape(-1)) - - attn_metadata.spec_decoding_generation_lengths[:batch_size] = num_tokens_current_layer - - previous_position_offsets = attn_metadata.spec_decoding_position_offsets[ - : batch_size * num_tokens_previous_layer - ].view(batch_size, num_tokens_previous_layer) - - new_pos = self._new_pos_offset_buf[:batch_size, :num_tokens_current_layer] - new_pos[:, :num_tokens_previous_layer].copy_(previous_position_offsets) - new_pos[:, num_tokens_previous_layer:num_tokens_current_layer].copy_( - previous_position_offsets[:, -self.K :] + 1 - ) - attn_metadata.spec_decoding_position_offsets[ - : batch_size * num_tokens_current_layer - ] = new_pos.reshape(-1) - - # Hopper XQA needs packed 1D mask; Blackwell expects padded 3D. - if self._needs_mask_repack: - self._repack_mask_padded_to_packed( - attn_metadata.spec_decoding_packed_mask, batch_size, num_tokens_current_layer - ) - - def prepare_for_generation( - self, - cur_draft_idx: int, - attn_metadata: AttentionMetadata, - batch_size: int, - *, - inputs=None, - gather_ids=None, - num_contexts: int = 0, - num_gens: int = 0, - num_accepted_tokens=None, - ): - """Set up attn_metadata for the subsequent drafter layer.""" - num_tokens_current_layer = self.K * (cur_draft_idx + 1) - - if cur_draft_idx == 0: - base_pos = inputs["position_ids"][gather_ids] + 1 - self.position_ids_buffer[:batch_size, : self.K] = base_pos.unsqueeze(1).expand( - -1, self.K - ) - - attn_metadata._seq_lens[:batch_size].fill_(self.K) - attn_metadata._seq_lens_cuda[:batch_size].fill_(self.K) - attn_metadata.on_update() - - if inputs["attn_metadata"].kv_cache_manager is not None: - attn_metadata.host_request_types[: attn_metadata.num_contexts].fill_(1) - attn_metadata.num_contexts = 0 - - if hasattr(attn_metadata, "kv_lens_cuda"): - # KV rewind: remove unaccepted draft-path tokens - if num_gens > 0: - attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= ( - self._max_path_len - ) - num_accepted_tokens[num_contexts:batch_size] - attn_metadata.kv_lens_cuda[:batch_size] += self.K - - attn_metadata.use_spec_decoding = True - - else: - num_tokens_previous_layer = cur_draft_idx * self.K - prev_pos = self.position_ids_buffer[:batch_size, :num_tokens_previous_layer] - self.position_ids_buffer[ - :batch_size, num_tokens_previous_layer:num_tokens_current_layer - ] = prev_pos[:, -self.K :] + 1 - - attn_metadata._seq_lens[:batch_size].fill_(num_tokens_current_layer) - attn_metadata._seq_lens_cuda[:batch_size].fill_(num_tokens_current_layer) - attn_metadata.on_update() - attn_metadata.kv_lens_cuda[:batch_size] += self.K - - def update_hidden_states( - self, - cur_draft_idx: int, - batch_size: int, - step0_hs=None, - hidden_states_to_save=None, - selected_parents=None, - ): - """Manage hidden states for the growing context pattern.""" - if cur_draft_idx == 0: - hs_dim = step0_hs.shape[-1] - self._hs_dim = hs_dim - - if self._hs_write_buffer is None or self._hs_write_buffer.shape[2] != hs_dim: - self._hs_write_buffer = torch.zeros( - self._max_batch_size, - self.max_draft_len * self.K, - hs_dim, - device=step0_hs.device, - dtype=step0_hs.dtype, - ) - if self._accumulated_hs is None or self._accumulated_hs.shape[2] != hs_dim: - self._accumulated_hs = torch.zeros( - self._max_batch_size, - self.max_draft_len * self.K, - hs_dim, - device=step0_hs.device, - dtype=step0_hs.dtype, - ) - - # All K depth-0 tokens share step0_hs - self._accumulated_hs[:batch_size, : self.K] = step0_hs.unsqueeze(1).expand( - -1, self.K, -1 - ) - self._step0_hs = step0_hs - - else: - num_tokens_per_req = cur_draft_idx * self.K - - hs_to_save_reshaped = hidden_states_to_save.reshape(batch_size, num_tokens_per_req, -1) - - # 1) Write current forward's prenorm to write_buffer - self._hs_write_buffer[:batch_size, :num_tokens_per_req] = hs_to_save_reshaped - - # 2) Set read_map for new K tokens (depth cur_draft_idx) - parent_offset = (cur_draft_idx - 1) * self.K - self._hs_read_map[ - :batch_size, cur_draft_idx * self.K : (cur_draft_idx + 1) * self.K - ] = parent_offset + selected_parents - - # 3) Gather from write_buffer into accumulated_hs (positions K onwards). - # Positions 0:K retain step0_hs set in step 0 and are never modified. - num_tokens_next = (cur_draft_idx + 1) * self.K - read_idx = self._hs_read_map[:batch_size, self.K : num_tokens_next] - gathered = torch.gather( - self._hs_write_buffer[:batch_size], - 1, - read_idx.unsqueeze(-1).expand(-1, -1, self._hs_dim), - ) - self._accumulated_hs[:batch_size, self.K : num_tokens_next] = gathered diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index a87724930248..a795a19d6d87 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -100,13 +100,12 @@ def __init__(self, args: TorchSampler.Args, *, draft_len: int): seq_slots = args.max_num_sequences max_tokens = self._get_max_tokens(args, draft_len) - max_new_tokens = self._get_max_new_tokens(args, draft_len) draft_tokens_size = self._get_draft_tokens_storage_size(args, draft_len) self.max_beam_width = args.max_beam_width assert self.max_beam_width == 1, "beam width must be 1 for speculative decoding" self.store = self.Store( - new_tokens=int_tensor((max_new_tokens, seq_slots, self.max_beam_width)), + new_tokens=int_tensor((max_tokens, seq_slots, self.max_beam_width)), next_new_tokens=int_tensor((max_tokens, seq_slots, self.max_beam_width)), next_draft_tokens=int_tensor((seq_slots, draft_tokens_size)), new_tokens_lens=int_tensor((seq_slots,)), @@ -121,15 +120,6 @@ def _get_max_tokens(self, args: TorchSampler.Args, draft_len: int) -> int: """ return draft_len + 1 - def _get_max_new_tokens(self, args: TorchSampler.Args, draft_len: int) -> int: - """Max depth of accepted token path for new_tokens buffer. - - Defaults to _get_max_tokens (same size as next_new_tokens). - Override when accepted path depth differs from total draft tokens, - e.g. dynamic tree where max_draft_len < max_total_draft_tokens. - """ - return self._get_max_tokens(args, draft_len) - def _get_draft_tokens_storage_size(self, args: TorchSampler.Args, draft_len: int) -> int: """ Calculate storage size for next_draft_tokens tensor. @@ -248,30 +238,19 @@ def sample_async( o_next_new_tokens = outputs["next_new_tokens"][num_skip : num_skip + num_sampling_requests] runtime_draft_len = o_next_draft_tokens.shape[1] - # Pad or truncate to match fixed-size store buffers for index_copy_. - # Use actual store buffer dimensions (which may differ from draft_len - # when _get_max_new_tokens is overridden, e.g. dynamic tree mode). - new_tokens_width = self.store.new_tokens.shape[0] - next_new_tokens_width = self.store.next_new_tokens.shape[0] - draft_tokens_width = self.store.next_draft_tokens.shape[1] - if o_new_tokens.shape[1] < new_tokens_width: + # Pad to match fixed-size store buffers for index_copy_. + if o_new_tokens.shape[1] < (self.draft_len + 1): o_new_tokens = torch.nn.functional.pad( - o_new_tokens, (0, new_tokens_width - o_new_tokens.shape[1]) + o_new_tokens, (0, (self.draft_len + 1) - o_new_tokens.shape[1]) ) - elif o_new_tokens.shape[1] > new_tokens_width: - o_new_tokens = o_new_tokens[:, :new_tokens_width] - if o_next_draft_tokens.shape[1] < draft_tokens_width: + if o_next_draft_tokens.shape[1] < self.draft_len: o_next_draft_tokens = torch.nn.functional.pad( - o_next_draft_tokens, (0, draft_tokens_width - o_next_draft_tokens.shape[1]) + o_next_draft_tokens, (0, self.draft_len - o_next_draft_tokens.shape[1]) ) - elif o_next_draft_tokens.shape[1] > draft_tokens_width: - o_next_draft_tokens = o_next_draft_tokens[:, :draft_tokens_width] - if o_next_new_tokens.shape[1] < next_new_tokens_width: + if o_next_new_tokens.shape[1] < (self.draft_len + 1): o_next_new_tokens = torch.nn.functional.pad( - o_next_new_tokens, (0, next_new_tokens_width - o_next_new_tokens.shape[1]) + o_next_new_tokens, (0, (self.draft_len + 1) - o_next_new_tokens.shape[1]) ) - elif o_next_new_tokens.shape[1] > next_new_tokens_width: - o_next_new_tokens = o_next_new_tokens[:, :next_new_tokens_width] # Use index_copy_ for efficient copying (slots are unique) self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens) diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index d6e7cd826d06..bb1d4758798f 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -1,14 +1,10 @@ -import logging import math -from itertools import accumulate from typing import List import torch from tensorrt_llm._utils import prefer_pinned -logger = logging.getLogger(__name__) - class SpecTreeManager: use_dynamic_tree: bool # Whether using dynamic tree @@ -19,11 +15,11 @@ class SpecTreeManager: # Auxiliary buffers # The top k list for each draft layer. - top_k_list: list + top_k_list = [] # The user input eagle choices, only available when using static tree. - eagle_choices: List[List[int]] - # If dynamic tree, each request has their own tree. If static tree, all requests share the same tree. - num_trees: int + eagle_choices: List[List[int]] = None + # If dynamice tree, each request has their own tree. If static tree, all requests share the same tree. + num_trees: int = None # Convert the choice to a path. Each path is an array of indices from the root to other nodes in the tree. # shape: [num_trees, max_total_draft_tokens + 1, max_draft_len + 1] @@ -41,61 +37,46 @@ class SpecTreeManager: # shape: [num_trees, max_total_draft_tokens + 1], device tensor. spec_dec_position_offsets: torch.Tensor = None - ############################ Auxiliary buffers for the static tree. ############################ + # TODO: Optimized together with the subsequent dynamic tree. + # Auxiliary buffers for the static tree. # Considering that the static tree does not modify the tree structure during inference, we can calculate some buffers in advance. # NOTE: Most of these buffers are introduced due to limitations of XQA: # With tree attention, XQA cannot simply take the tokens to be processed in the next round as input. Instead, it needs to take ALL of their parent nodes as input. # This incurs additional computation, but it is unavoidable. # NOTE: The reason why most of these auxiliary buffers are with `len == max_draft_len - 1` is that: we do not need to prepare specific input data for the first draft layer. + # The top k value for each draft layer. Device tensor. top_k_list_cuda: list[torch.Tensor] = None # The max top k value for all draft layers. Which is used for torch.topk and cuda graph. max_top_k = -1 - # Gather the required draft tokens among the 'max_total_draft_tokens + 1' tokens. + # Gather the required draft tokens from all currently generated draft tokens as the input of the next draft layer. # Only the nodes has child(s) this layer and all their parents nodes will be gathered. + # Device tensor. len(tokens_gather_idx) == max_draft_len - 1. Each element is a tensor with shape [num_tokens_for_next_layer]. tokens_gather_idx_for_drafter_model: list[torch.Tensor] = None + # Gather the required logits from all currently generated logits. + # Device tensor. len(tokens_gather_idx) == max_draft_len - 1. + logits_gather_idx: list[torch.Tensor] = None + # The packed mask for the drafter model's attention (i.e., xqa). - # shape: [1, max_total_draft_tokens + 1, math.ceil((max_total_draft_tokens + 1) / 32)], device tensor. spec_dec_packed_mask_for_drafter_model: torch.Tensor = None # The read indices offset for the drafter model. - # shape: [max_total_draft_tokens + 1], device tensor. hidden_states_read_indices_offset_for_drafter_model: torch.Tensor = None # The write back start indices for the drafter tokens between different draft layers. - # shape: [max_draft_len + 1], device tensor. draft_tokens_indices_cumsum: torch.Tensor = None - ############################ Auxiliary buffers for the dynamic tree. ############################ - # CUDA kernel outputs for dynamic tree verification. - # These are produced by build_dynamic_tree CUDA kernel and used by verify_dynamic_tree_greedy. - # shape: [num_trees, max_total_draft_tokens + 1], int32, device tensor. - retrieve_index: torch.Tensor = None - retrieve_next_token: torch.Tensor = None - retrieve_next_sibling: torch.Tensor = None - def __init__(self, max_num_requests: int, use_dynamic_tree: bool, max_total_draft_tokens: int, max_draft_len: int, - eagle_choices: List[List[int]] | None, - dynamic_tree_max_topK: int): + eagle_choices: [List[List[int]]], dynamic_tree_max_topK: int): self.use_dynamic_tree = use_dynamic_tree self.max_total_draft_tokens = max_total_draft_tokens self.max_draft_len = max_draft_len - - # In dynamic tree mode the draft loop can produce up to - # K * max_draft_len tokens, which may exceed max_total_draft_tokens+1. - # Size the working buffers to the larger of the two so the masks and - # position-offset tensors never run out of columns/rows. - if use_dynamic_tree and dynamic_tree_max_topK > 0: - self._internal_buf_dim = max(max_total_draft_tokens + 1, - dynamic_tree_max_topK * max_draft_len) - else: - self._internal_buf_dim = max_total_draft_tokens + 1 self.eagle_choices = eagle_choices self.num_trees = max_num_requests if use_dynamic_tree else 1 self.dynamic_tree_max_topK = dynamic_tree_max_topK @@ -103,29 +84,23 @@ def __init__(self, max_num_requests: int, use_dynamic_tree: bool, self.top_k_list = [] # Initialize the buffers - # eagle_paths and spec_dec_mask_matrix are only used in static tree mode. - # Skip allocation in dynamic tree mode to save memory. - if not use_dynamic_tree: - self.eagle_paths = torch.ones( - (self.num_trees, self.max_total_draft_tokens + 1, - self.max_draft_len + 1), - dtype=torch.int32, - device='cpu', - pin_memory=prefer_pinned(), - ) * -1 - - self.spec_dec_mask_matrix = torch.eye( - self.max_total_draft_tokens + 1, - dtype=torch.int32, - device='cuda', - ).unsqueeze(0).repeat(self.num_trees, 1, 1) - - # CUDA kernel facing — rows = max_total_draft_tokens + 1, - # columns widened to match attn_metadata mask_width so that the - # Hopper flat copy in update_spec_dec_param needs no per-row padding. + self.eagle_paths = torch.ones( + (self.num_trees, self.max_total_draft_tokens + 1, + self.max_draft_len + 1), + dtype=torch.int32, + device='cpu', + pin_memory=prefer_pinned(), + ) * -1 + + self.spec_dec_mask_matrix = torch.eye( + self.max_total_draft_tokens + 1, + dtype=torch.int32, + device='cuda', + ).unsqueeze(0).repeat(self.num_trees, 1, 1) + self.spec_dec_packed_mask = torch.zeros( (self.num_trees, self.max_total_draft_tokens + 1, - math.ceil(self._internal_buf_dim / 32)), + math.ceil((self.max_total_draft_tokens + 1) / 32)), dtype=torch.int32, device='cuda', ) @@ -135,45 +110,14 @@ def __init__(self, max_num_requests: int, use_dynamic_tree: bool, device='cuda', ) - # Cached constants for compute_spec_dec_packed_mask (avoids per-call allocation) - self._pack_weights = ( - 1 << torch.arange(32, device='cuda', dtype=torch.int32)) - # Python-only internal buffers — enlarged to _internal_buf_dim - num_blocks = math.ceil(self._internal_buf_dim / 32) - total_bits = num_blocks * 32 - self._padded_mask_buf = torch.zeros(self.num_trees, - self._internal_buf_dim, - total_bits, - dtype=torch.int32, - device='cuda') - self._pack_result_buf = torch.zeros(self.num_trees, - self._internal_buf_dim, - num_blocks, - dtype=torch.int32, - device='cuda') - if self.use_dynamic_tree: self.init_tree_info_for_dynamic_tree() else: self.init_tree_info_for_static_tree() - def init_tree_info_for_dynamic_tree(self): - # Allocate retrieve buffers for CUDA kernel outputs - num_draft_with_root = self.max_total_draft_tokens + 1 - self.retrieve_index = torch.zeros((self.num_trees, num_draft_with_root), - dtype=torch.int32, - device='cuda') - self.retrieve_next_token = torch.full( - (self.num_trees, num_draft_with_root), - -1, - dtype=torch.int32, - device='cuda') - self.retrieve_next_sibling = torch.full( - (self.num_trees, num_draft_with_root), - -1, - dtype=torch.int32, - device='cuda') + # self.dump_tree_info() + def init_tree_info_for_dynamic_tree(self): # For the dynamic tree # To the internal layer, the number of nodes is the same as the dynamic_tree_max_topK. self.top_k_list = [ @@ -183,101 +127,7 @@ def init_tree_info_for_dynamic_tree(self): pin_memory=prefer_pinned()) * self.dynamic_tree_max_topK ] - # Per-py_seq_slot storage (+1 dummy row for graph); survives batch reorder. - S = self.num_trees + 1 - self._slot_packed_mask = torch.zeros( - (S, ) + self.spec_dec_packed_mask.shape[1:], - dtype=self.spec_dec_packed_mask.dtype, - device='cuda') - self._slot_position_offsets = torch.zeros( - (S, ) + self.spec_dec_position_offsets.shape[1:], - dtype=self.spec_dec_position_offsets.dtype, - device='cuda') - # [S, n_dt, 3]: index / next_token / next_sibling in one tensor (fewer gathers). - self._slot_retrieve_packed = torch.full((S, num_draft_with_root, 3), - -1, - dtype=torch.int32, - device='cuda') - self._slot_retrieve_packed[:, :, 0] = 0 # match retrieve_index default - self._all_slot_ids_buf = torch.zeros(self.num_trees, - dtype=torch.long, - device='cuda') - self._dummy_slot_id = self.num_trees - - # True after scatter; dummy row stays False. - self.slot_has_tree = torch.zeros(S, dtype=torch.bool, device='cuda') - - # Graph-safe: no per-forward host-to-device slot-id tensor alloc. - self._gather_gen_slot_ids_buf = torch.zeros(self.num_trees, - dtype=torch.long, - device='cuda') - - # Scatter staging (no per-call torch.stack). - self._scatter_retrieve_staging = torch.empty( - (self.num_trees, num_draft_with_root, 3), - dtype=torch.int32, - device='cuda') - - def fill_gen_slot_ids(self, gen_requests): - """Fill _gather_gen_slot_ids_buf; return (buf[:count], count). Gen LlmRequest only.""" - buf = self._gather_gen_slot_ids_buf - dummy = self._dummy_slot_id - count = 0 - for r in gen_requests: - buf[count] = r.py_seq_slot if r.py_seq_slot is not None else dummy - count += 1 - return buf[:count], count - - def mark_tree_valid(self, slot_ids, count): - """Set slot_has_tree True for scattered slots.""" - if count == 0: - return - # index_fill_: graph-capture-safe (no fancy indexing). - self.slot_has_tree.index_fill_(0, slot_ids[:count], True) - - def mark_tree_invalid(self, slot_id): - """Clear validity when a slot is freed.""" - self.slot_has_tree[slot_id] = False - - def scatter_trees_to_slots(self, slot_ids, count): - """Copy work buffers [:count] into per-slot storage (index_copy_ for graph capture).""" - if count == 0: - return - ids = slot_ids[:count] - self._slot_packed_mask.index_copy_(0, ids, - self.spec_dec_packed_mask[:count]) - self._slot_position_offsets.index_copy_( - 0, ids, self.spec_dec_position_offsets[:count]) - staging = self._scatter_retrieve_staging[:count] - staging[:, :, 0] = self.retrieve_index[:count] - staging[:, :, 1] = self.retrieve_next_token[:count] - staging[:, :, 2] = self.retrieve_next_sibling[:count] - self._slot_retrieve_packed.index_copy_(0, ids, staging) - - def gather_attn_params_from_slots(self, slot_ids, count): - """Copy mask and position offsets from slots into work buffers [:count].""" - if count == 0: - return - ids = slot_ids[:count] - self.spec_dec_packed_mask[:count] = self._slot_packed_mask[ids] - self.spec_dec_position_offsets[:count] = self._slot_position_offsets[ - ids] - - def gather_retrieve_from_slots(self, slot_ids, count): - """Copy retrieve tensors from slots into work buffers [:count].""" - if count == 0: - return - ids = slot_ids[:count] - packed = self._slot_retrieve_packed[ids] - self.retrieve_index[:count] = packed[..., 0] - self.retrieve_next_token[:count] = packed[..., 1] - self.retrieve_next_sibling[:count] = packed[..., 2] - - def gather_trees_from_slots(self, slot_ids, count): - """gather_attn_params_from_slots + gather_retrieve_from_slots.""" - self.gather_attn_params_from_slots(slot_ids, count) - self.gather_retrieve_from_slots(slot_ids, count) - + # For the static tree def init_tree_info_for_static_tree(self): self.index_mapping_set = {} self.nodes_list_per_layer = [[] for _ in range(self.max_draft_len + 1)] @@ -326,7 +176,9 @@ def init_tree_info_for_static_tree(self): pin_memory=prefer_pinned())) # 6) Compute the spec decoding according to the eagle_paths for the target model - self.compute_spec_dec_mask_matrix(0) + for i, path in enumerate(self.eagle_paths[0]): + indices = path[path > -1] + self.spec_dec_mask_matrix[0][i, indices] = 1 self.compute_spec_dec_packed_mask(self.spec_dec_mask_matrix, self.spec_dec_packed_mask) @@ -365,6 +217,7 @@ def init_tree_info_for_static_tree(self): num_nodes_per_layer = [0] num_nodes_per_layer.extend( [len(node_list) for node_list in self.nodes_list_per_layer[1:]]) + from itertools import accumulate self.draft_tokens_indices_cumsum = torch.tensor(list( accumulate(num_nodes_per_layer)), dtype=torch.int32, @@ -402,81 +255,89 @@ def init_tree_info_for_static_tree(self): device= 'cuda') + # Get the eagle_paths def get_eagle_paths(self, tree_idx=0): - if self.eagle_paths is None: - raise RuntimeError( - "get_eagle_paths() is not supported in dynamic tree mode; " - "use retrieve_index/retrieve_next_token/retrieve_next_sibling instead" - ) - return self.eagle_paths[0] + if self.use_dynamic_tree: + self.eagle_paths[tree_idx].fill_(-1) + # If dynamic tree, return the eagle_paths according to the mask. + for i in range(self.max_total_draft_tokens + 1): + self.eagle_paths[tree_idx][:, i, :] = self.spec_dec_mask_matrix[ + tree_idx][i, :].nonzero() + return self.eagle_paths[tree_idx] + else: + # If static tree, return the prepared eagle_paths. These paths are immutable. + return self.eagle_paths[0] # Get the topK list for the specific draft layer def get_top_k_list(self, draft_layer_id): assert draft_layer_id >= 0 return self.top_k_list[draft_layer_id] - def compute_spec_dec_mask_matrix(self, tree_idx=0): - if self.eagle_paths is None: - raise RuntimeError( - "compute_spec_dec_mask_matrix() is not supported in dynamic tree mode" - ) - for i, path in enumerate(self.eagle_paths[0]): - indices = path[path > -1] - self.spec_dec_mask_matrix[0][i, indices] = 1 - + # Compute the packed mask according to the mask matrix def compute_spec_dec_packed_mask(self, mask_matrix, packed_mask): - bs, num_tokens, num_tokens_attend = mask_matrix.shape - assert mask_matrix.ndim == 3, f"Expected 3D mask_matrix, got {mask_matrix.ndim}D" - assert packed_mask.ndim == 3, f"Expected 3D packed_mask, got {packed_mask.ndim}D" - assert bs <= self._padded_mask_buf.shape[0], \ - f"batch size {bs} exceeds pre-allocated buffer size {self._padded_mask_buf.shape[0]}" - num_blocks = packed_mask.shape[-1] - - # Use cached bit weights - weights = self._pack_weights - - # Pad into pre-allocated buffer - total_bits = num_blocks * 32 - padded_m = self._padded_mask_buf[:bs, :num_tokens, :total_bits] - padded_m.zero_() - src = mask_matrix if mask_matrix.dtype == torch.int32 else mask_matrix.to( - torch.int32) - padded_m[:, :, :num_tokens_attend].copy_(src) - - # Reshape last dim into [num_blocks, 32] for blocked packing - blocked_matrix = padded_m.view(bs, num_tokens, num_blocks, 32) - - # Vectorized dot product into pre-allocated result buffer - result = self._pack_result_buf[:bs, :num_tokens, :num_blocks] - torch.sum(blocked_matrix * weights, dim=-1, out=result) - - # Write results back to the output buffer - packed_mask[:, :num_tokens, :] = result - return packed_mask + # mask_matrix: shape: [num_trees, max_total_draft_tokens + 1, max_total_draft_tokens + 1] + # packed_mask: shape: [num_trees, max_total_draft_tokens + 1, math.ceil((max_total_draft_tokens + 1) / 32)] + assert mask_matrix.ndim == 3 + assert packed_mask.ndim == 3 + + num_trees = mask_matrix.shape[0] + assert mask_matrix.shape == (num_trees, self.max_total_draft_tokens + 1, + self.max_total_draft_tokens + 1) + assert packed_mask.shape == (num_trees, self.max_total_draft_tokens + 1, + math.ceil( + (self.max_total_draft_tokens + 1) / + 32)) + + num_blocks = math.ceil((self.max_total_draft_tokens + 1) / 32) + int_tensor = mask_matrix.reshape( + -1, self.max_total_draft_tokens + 1 + ) # shape: [num_trees * (self.max_total_draft_tokens + 1), self.max_total_draft_tokens + 1] + packed_mask = packed_mask.reshape( + -1, num_blocks + ) # shape: [num_trees * (self.max_total_draft_tokens + 1), num_blocks] + + for block_idx in range(num_blocks): + start_idx = block_idx * 32 + end_idx = min(start_idx + 32, self.max_total_draft_tokens + 1) + if end_idx < start_idx: + break + block_bits = int_tensor[:, start_idx:end_idx] + weight = torch.pow( + 2, + torch.arange(end_idx - start_idx, + dtype=torch.int32, + device=int_tensor.device)) + block_value = torch.sum(block_bits * weight, dim=-1) + packed_mask[:, block_idx] = block_value + + packed_mask = packed_mask.reshape(num_trees, + self.max_total_draft_tokens + 1, + num_blocks) + # Print the tree info def dump_tree_info(self): - logger.debug("TopK list: %s", self.top_k_list) - if self.use_dynamic_tree: - logger.debug("Dynamic max top k: %s", self.dynamic_tree_max_topK) - else: - logger.debug("Max top k list cuda: %s", self.max_top_k) - logger.debug("Eagle paths: %s", self.eagle_paths) - logger.debug("Index mapping set: %s", self.index_mapping_set) - logger.debug("Nodes list per layer: %s", self.nodes_list_per_layer) - logger.debug("Spec dec position offsets: %s", - self.spec_dec_position_offsets) - logger.debug("Spec dec mask matrix: %s", - self.spec_dec_mask_matrix.int()) - logger.debug("Spec dec pack mask: %s", self.spec_dec_packed_mask) - logger.debug("Auxiliary buffers for the static tree.") - logger.debug("TopK list cuda: %s", self.top_k_list_cuda) - logger.debug("Tokens gather idx for drafter model: %s", - self.tokens_gather_idx_for_drafter_model) - logger.debug("Draft tokens indices cumsum: %s", - self.draft_tokens_indices_cumsum) - logger.debug("Spec dec packed mask for drafter model: %s", - self.spec_dec_packed_mask_for_drafter_model) - logger.debug( - "Hidden states read indices offset for drafter model: %s", - self.hidden_states_read_indices_offset_for_drafter_model, + print(f"TopK list: {self.top_k_list}") + if not self.use_dynamic_tree: + print(f"Max top k list cuda: {self.max_top_k}") + print(f"Static tree: {self.eagle_paths}") + print(f"Index mapping set: {self.index_mapping_set}") + print(f"Nodes list per layer: {self.nodes_list_per_layer}") + print( + f"Spec dec position offsets: {self.spec_dec_position_offsets}") + print(f"Spec dec mask matrix: {self.spec_dec_mask_matrix.int()}") + print(f"Spec dec pack mask: {self.spec_dec_packed_mask}") + print("Auxiliary buffers for the static tree.") + print(f"TopK list cuda: {self.top_k_list_cuda}") + print( + f"Tokens gather idx for drafter model: {self.tokens_gather_idx_for_drafter_model}" + ) + print( + f"Draft tokens indices cumsum: {self.draft_tokens_indices_cumsum}" + ) + print(f"Logits gather idx: {self.logits_gather_idx}") + print( + f"Spec dec packed mask for drafter model: {self.spec_dec_packed_mask_for_drafter_model}" + ) + print( + f"Hidden states read indices offset for drafter model: {self.hidden_states_read_indices_offset_for_drafter_model}" ) diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index b74b61af126d..33f377063994 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -14,11 +14,9 @@ from .draft_target import (DraftTargetOneModelSampler, DraftTargetOneModelSpecMetadata, DraftTargetOneModelWorker) -from .eagle3 import (Eagle3OneModelDynamicTreeResourceManager, - Eagle3OneModelSampler, Eagle3OneModelSpecMetadata, +from .eagle3 import (Eagle3OneModelSampler, Eagle3OneModelSpecMetadata, Eagle3OneModelWorker, Eagle3ResourceManager, Eagle3SpecMetadata) -from .eagle3_dynamic_tree import Eagle3OneModelDynamicTreeWorker from .model_drafter import ModelDrafter from .mtp import (MTPEagleWorker, MTPHiddenStatesManager, MTPSampler, MTPSpecMetadata, MTPWorker) @@ -93,8 +91,6 @@ def get_spec_metadata(spec_config, layers_to_capture=spec_config.eagle3_layers_to_capture, allow_advanced_sampling=spec_config.allow_advanced_sampling, spec_resource_manager=spec_resource_manager, - use_dynamic_tree=spec_config.use_dynamic_tree, - eagle_choices=spec_config.eagle_choices, ) if spec_config.spec_dec_mode.is_pard(): return PARDSpecMetadata( @@ -186,10 +182,6 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): max_num_requests, sa_manager=sa_manager, ) - if spec_dec_mode.is_eagle3_one_model() and getattr( - spec_config, 'use_dynamic_tree', False): - return Eagle3OneModelDynamicTreeResourceManager(spec_config, - max_num_requests) if spec_dec_mode.is_eagle3_one_model(): sa_manager = None sa_cfg = getattr(spec_config, 'sa_config', None) @@ -252,7 +244,7 @@ def get_spec_decoder( # TorchSampler handles Eagle3 gracefully, by integrating d2t into the sampling process return TorchSampler(sampler_args) if spec_config.spec_dec_mode.is_eagle3_one_model(): - return Eagle3OneModelSampler(sampler_args, spec_config=spec_config) + return Eagle3OneModelSampler(sampler_args) if spec_config.spec_dec_mode.is_pard(): return MTPSampler(sampler_args, nextn=spec_config.tokens_per_gen_step - 1) @@ -315,9 +307,6 @@ def get_spec_worker(spec_config, return MTPEagleWorker(spec_config, model_config, use_separate_draft_kv_cache) if spec_dec_mode.is_eagle3_one_model(): - if getattr(spec_config, 'use_dynamic_tree', False): - return Eagle3OneModelDynamicTreeWorker(spec_config, mapping, - use_separate_draft_kv_cache) return Eagle3OneModelWorker(spec_config, mapping, use_separate_draft_kv_cache) if spec_dec_mode.is_pard(): diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 8e9dde25a490..3c59bff28be8 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1088,6 +1088,7 @@ def validate_eagle_config(self) -> 'EagleDecodingConfig': ) self.num_eagle_layers = self.max_draft_len + self.max_total_draft_tokens = self.max_draft_len # If using linear-tree, the max_total_draft_tokens is the same as max_draft_len if self.eagle3_model_arch == "mistral_large3" and self.eagle3_layers_to_capture is None: # FIXME find a better way to setup it. @@ -1116,8 +1117,7 @@ def validate_eagle_config(self) -> 'EagleDecodingConfig': self.max_total_draft_tokens = len(self.eagle_choices) # Dynamic tree logic - if self.use_dynamic_tree or self.dynamic_tree_max_topK is not None: - self.use_dynamic_tree = True + if self.use_dynamic_tree: if self.eagle_choices is not None: raise ValueError( "If use_dynamic_tree is True, eagle_choices should be None") @@ -1129,28 +1129,10 @@ def validate_eagle_config(self) -> 'EagleDecodingConfig': raise ValueError( "dynamic_tree_max_topK should be provided, which indicates the number of nodes to expand each time" ) - - default_max_total_draft_tokens = self.dynamic_tree_max_topK * self.max_draft_len - - if self.max_total_draft_tokens is None: - self.max_total_draft_tokens = default_max_total_draft_tokens - logger.warning( - f"max_total_draft_tokens is not provided, use the default value {default_max_total_draft_tokens} (default_max_total_draft_tokens = dynamic_tree_max_topK * max_draft_len)" + if self.max_total_draft_tokens is None or self.max_total_draft_tokens <= 0: + raise ValueError( + "max_total_draft_tokens should be provided, which indicates the total nodes of the final draft tree. (exclude the root node)" ) - else: - if self.max_total_draft_tokens < self.max_draft_len: - raise ValueError( - f"max_total_draft_tokens ({self.max_total_draft_tokens}) should be >= max_draft_len ({self.max_draft_len})" - ) - if self.max_total_draft_tokens > self.dynamic_tree_max_topK * self.max_draft_len: - raise ValueError( - f"max_total_draft_tokens ({self.max_total_draft_tokens}) should be <= " - f"dynamic_tree_max_topK * max_draft_len ({self.dynamic_tree_max_topK * self.max_draft_len})" - ) - - # Linear tree - if self.max_total_draft_tokens is None: - self.max_total_draft_tokens = self.max_draft_len return self @@ -1231,11 +1213,6 @@ class SAEnhancerConfig(StrictBaseModel): class Eagle3DecodingConfig(EagleDecodingConfig): decoding_type: Literal["Eagle3"] = "Eagle3" - max_batch_size: Optional[int] = Field( - default=None, - description="Max batch size for pre-allocating dynamic tree buffers. " - "Required when use_dynamic_tree=True.") - sa_config: Optional[SAEnhancerConfig] = Field( default=None, status="beta", @@ -2171,8 +2148,8 @@ def supports_backend(self, backend: str) -> bool: SpeculativeConfig: TypeAlias = Annotated[ Union[ DraftTargetDecodingConfig, - Eagle3DecodingConfig, # Must be before EagleDecodingConfig since it's a subclass EagleDecodingConfig, + Eagle3DecodingConfig, LookaheadDecodingConfig, MedusaDecodingConfig, MTPDecodingConfig, diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index a77c151c55f1..5a963d90f818 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -345,7 +345,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUT accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=True-v2_kv_cache=True] SKIP (https://nvbugs/6071081) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=True-v2_kv_cache=False] SKIP (https://nvbugs/6071081) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6071081) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[mtp_nextn=2-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False] SKIP (https://nvbugs/6071081) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=True] SKIP (https://nvbugs/6071081) accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[google_gemma-3-1b-it-False] SKIP (https://nvbugs/6059036) accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[meta-llama_Llama-3.1-8B-Instruct-False] SKIP (https://nvbugs/6059036) diff --git a/tests/unittest/_torch/modeling/test_modeling_llama.py b/tests/unittest/_torch/modeling/test_modeling_llama.py index 711f90b7a7dc..334e60a61e97 100644 --- a/tests/unittest/_torch/modeling/test_modeling_llama.py +++ b/tests/unittest/_torch/modeling/test_modeling_llama.py @@ -20,9 +20,7 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import ( KVCacheManager, _update_kv_cache_draft_token_location) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests -from tensorrt_llm._torch.speculative.interface import (SpecMetadata, - SpeculativeDecodingMode) -from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager +from tensorrt_llm._torch.speculative.utils import SpecDecodingTensor from tensorrt_llm._utils import get_sm_version from tensorrt_llm.bindings.executor import KvCacheConfig from tensorrt_llm.mapping import Mapping @@ -518,20 +516,6 @@ def run_forward(input_ids, position_ids, attn_metadata): is_spec_dec_dynamic_tree = True max_total_draft_tokens = gen_input_ids_0.size(-1) - 1 - spec_tree_mgr = SpecTreeManager( - max_num_requests=1, - use_dynamic_tree=is_spec_dec_dynamic_tree, - max_total_draft_tokens=max_total_draft_tokens, - max_draft_len=max_total_draft_tokens, - eagle_choices=None, - dynamic_tree_max_topK=10, - ) - # Populate with test data - spec_tree_mgr.spec_dec_position_offsets[:1, :].copy_( - spec_decoding_position_offsets.unsqueeze(0)) - spec_tree_mgr.spec_dec_packed_mask[:1, :, :].copy_( - spec_decoding_packed_mask) - attn_metadata_gen_phase_0 = metadata_cls( seq_lens=torch.tensor([gen_input_ids_0.size(-1)], dtype=torch.int), num_contexts=0, @@ -550,6 +534,10 @@ def run_forward(input_ids, position_ids, attn_metadata): is_spec_dec_dynamic_tree=is_spec_dec_dynamic_tree, num_heads_per_kv=num_heads_per_kv, ) + spec_decoding_tensor = SpecDecodingTensor( + position_offsets=spec_decoding_position_offsets, + packed_mask=spec_decoding_packed_mask) + attn_metadata_gen_phase_0.prepare() attn_metadata_gen_phase_0.update_spec_dec_param( batch_size=batch_size, @@ -559,7 +547,7 @@ def run_forward(input_ids, position_ids, attn_metadata): max_draft_len=max_total_draft_tokens, max_total_draft_tokens=max_total_draft_tokens, model_is_wrapped=False, - spec_tree_manager=spec_tree_mgr, + spec_decoding_tensor=spec_decoding_tensor, ) gen_position_ids_0 = [ @@ -606,35 +594,15 @@ def run_forward(input_ids, position_ids, attn_metadata): attn_metadata_gen_phase_0.spec_decoding_packed_mask = None attn_metadata_gen_phase_0.spec_decoding_generation_lengths = None attn_metadata_gen_phase_0.prepare() - is_tree_phase1 = is_spec_dec_tree if get_sm_version() < 100 else False - spec_tree_mgr_phase1 = None - spec_metadata_phase1 = None - if is_tree_phase1: - max_draft_1 = gen_input_ids_1.size(-1) - 1 - spec_tree_mgr_phase1 = SpecTreeManager( - max_num_requests=1, - use_dynamic_tree=True, - max_total_draft_tokens=max_draft_1, - max_draft_len=max_draft_1, - eagle_choices=None, - dynamic_tree_max_topK=10, - ) - spec_metadata_phase1 = SpecMetadata( - max_num_requests=1, - max_draft_len=max_draft_1, - max_total_draft_tokens=max_draft_1, - spec_dec_mode=SpeculativeDecodingMode.EAGLE3, - ) attn_metadata_gen_phase_0.update_spec_dec_param( batch_size=batch_size, is_spec_decoding_enabled=is_spec_decoding_enabled, - is_spec_dec_tree=is_tree_phase1, - is_spec_dec_dynamic_tree=is_tree_phase1, + is_spec_dec_tree=is_spec_dec_tree + if get_sm_version() < 100 else False, + is_spec_dec_dynamic_tree=False, max_draft_len=gen_input_ids_1.size(-1) - 1, max_total_draft_tokens=gen_input_ids_1.size(-1) - 1, - model_is_wrapped=False, - spec_metadata=spec_metadata_phase1, - spec_tree_manager=spec_tree_mgr_phase1) + model_is_wrapped=False) gen_position_ids_1 = [ torch.full( @@ -681,35 +649,15 @@ def run_forward(input_ids, position_ids, attn_metadata): attn_metadata_ref.spec_decoding_packed_mask = None attn_metadata_ref.spec_decoding_generation_lengths = None attn_metadata_ref.prepare() - is_tree_ref = is_spec_dec_tree if get_sm_version() < 100 else False - spec_tree_mgr_ref = None - spec_metadata_ref = None - if is_tree_ref: - max_draft_ref = gen_input_ids_ref.size(-1) - 1 - spec_tree_mgr_ref = SpecTreeManager( - max_num_requests=1, - use_dynamic_tree=True, - max_total_draft_tokens=max_draft_ref, - max_draft_len=max_draft_ref, - eagle_choices=None, - dynamic_tree_max_topK=10, - ) - spec_metadata_ref = SpecMetadata( - max_num_requests=1, - max_draft_len=max_draft_ref, - max_total_draft_tokens=max_draft_ref, - spec_dec_mode=SpeculativeDecodingMode.EAGLE3, - ) attn_metadata_ref.update_spec_dec_param( batch_size=batch_size, is_spec_decoding_enabled=is_spec_decoding_enabled, - is_spec_dec_tree=is_tree_ref, - is_spec_dec_dynamic_tree=is_tree_ref, + is_spec_dec_tree=is_spec_dec_tree + if get_sm_version() < 100 else False, + is_spec_dec_dynamic_tree=False, max_draft_len=gen_input_ids_ref.size(-1) - 1, max_total_draft_tokens=gen_input_ids_ref.size(-1) - 1, - model_is_wrapped=False, - spec_metadata=spec_metadata_ref, - spec_tree_manager=spec_tree_mgr_ref) + model_is_wrapped=False) gen_position_ids_ref = [ torch.full((gen_input_ids_ref.size(-1), ), @@ -727,12 +675,12 @@ def run_forward(input_ids, position_ids, attn_metadata): torch.cuda.synchronize() torch.testing.assert_close(gen_logits_1[0, :], gen_logits_ref[2, :], - atol=1.0, - rtol=1.0) + atol=0.4, + rtol=0.4) torch.testing.assert_close(gen_logits_1[1, :], gen_logits_ref[3, :], - atol=1.0, - rtol=1.0) + atol=0.4, + rtol=0.4) token_id_ref = torch.argmax(gen_logits_ref[3, :], dim=-1) token_id_gen = torch.argmax(gen_logits_1[1, :], dim=-1) diff --git a/tests/unittest/_torch/speculative/test_draft_token_prepare_for_generation.py b/tests/unittest/_torch/speculative/test_draft_token_prepare_for_generation.py index 78acbde58cb9..7293cf91582a 100644 --- a/tests/unittest/_torch/speculative/test_draft_token_prepare_for_generation.py +++ b/tests/unittest/_torch/speculative/test_draft_token_prepare_for_generation.py @@ -1,15 +1,16 @@ import math import os import sys +import unittest import torch from utils.llm_data import llm_models_root from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata -from tensorrt_llm._torch.speculative.drafting_loops import StaticTreeDraftingLoopWrapper +from tensorrt_llm._torch.speculative.drafting_loops import TreeDraftingLoopWrapper from tensorrt_llm._torch.speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager -from tensorrt_llm.llmapi.llm_args import Eagle3DecodingConfig +from tensorrt_llm.llmapi import Eagle3DecodingConfig sys.path.append(os.path.join(os.path.dirname(__file__), "..")) @@ -20,7 +21,6 @@ def __init__(self): self.model_config = None self.config = None self.model = {} - self.model_is_wrapped = True def forward(self, *args, **kwargs) -> torch.Tensor: pass @@ -138,8 +138,8 @@ def run_test( input_hidden_states_read_indices # set from input ) - # 3) Create StaticTreeDraftingLoopWrapper - static_tree_drafting_loop_wrapper = StaticTreeDraftingLoopWrapper( + # 3) Create TreeDraftingLoopWrapper + tree_drafting_loop_wrapper = TreeDraftingLoopWrapper( max_batch_size=max_batch_size, max_draft_len=max_draft_len, max_total_draft_tokens=max_total_draft_tokens, @@ -147,7 +147,7 @@ def run_test( ) # 3) Run the function - static_tree_drafting_loop_wrapper.prepare_for_generation( + tree_drafting_loop_wrapper.prepare_for_generation( attn_metadata=attn_metadata, spec_metadata=spec_metadata, spec_tree_manager=spec_tree_manager, @@ -156,8 +156,7 @@ def run_test( # Compare input_ids and position_ids print( - f"static_tree_drafting_loop_wrapper.position_ids_buffer: \ - {static_tree_drafting_loop_wrapper.position_ids_buffer}, \ + f"tree_drafting_loop_wrapper.position_ids_buffer: {tree_drafting_loop_wrapper.position_ids_buffer}, \ ref_output_position_ids: {ref_position_ids}" ) @@ -209,7 +208,7 @@ def run_test( ref_spec_metadata.hidden_states_write_indices: {ref_spec_metadata['hidden_states_write_indices']}" ) - assert torch.all(static_tree_drafting_loop_wrapper.position_ids_buffer == ref_position_ids) + assert torch.all(tree_drafting_loop_wrapper.position_ids_buffer == ref_position_ids) assert torch.all(attn_metadata.kv_lens_cuda == ref_attn_metadata["kv_lens_cuda"]) assert torch.all(attn_metadata._seq_lens == ref_attn_metadata["_seq_lens"]) assert torch.all(attn_metadata._seq_lens_cuda == ref_attn_metadata["_seq_lens_cuda"]) @@ -453,3 +452,7 @@ def run_test( ref_attn_metadata, ref_spec_metadata, ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/speculative/test_draft_token_tree_sampling.py b/tests/unittest/_torch/speculative/test_draft_token_tree_sampling.py index 7f75793cfa89..f62316758732 100644 --- a/tests/unittest/_torch/speculative/test_draft_token_tree_sampling.py +++ b/tests/unittest/_torch/speculative/test_draft_token_tree_sampling.py @@ -6,7 +6,7 @@ from utils.llm_data import llm_models_root from tensorrt_llm._torch.speculative.drafting_loops import \ - StaticTreeDraftingLoopWrapper + TreeDraftingLoopWrapper from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager from tensorrt_llm.llmapi import Eagle3DecodingConfig @@ -53,7 +53,7 @@ def run_test(max_batch_size, draft_layer_id, max_total_draft_tokens, ) # Create the chain drafter - tree_drafter = StaticTreeDraftingLoopWrapper( + tree_drafter = TreeDraftingLoopWrapper( max_batch_size=max_batch_size, max_draft_len=spec_config.max_draft_len, max_total_draft_tokens=spec_config.max_total_draft_tokens, diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index fb792149e5f7..c8d8880c31f4 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -10,7 +10,6 @@ import torch from test_common.llm_data import with_mocked_hf_download_for_single_gpu from utils.llm_data import llm_models_root -from utils.util import skip_blackwell from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata @@ -816,95 +815,5 @@ def test_eagle3_cdl_sampling(disable_overlap_scheduler: bool): llm_spec.shutdown() -@pytest.mark.parametrize("disable_overlap_scheduler", [False, True]) -@pytest.mark.parametrize("use_cuda_graph", [False, True]) -@pytest.mark.high_cuda_memory -@skip_blackwell -@with_mocked_hf_download_for_single_gpu -def test_llama_eagle3_dynamic_tree(use_cuda_graph: bool, - disable_overlap_scheduler: bool): - """Test EAGLE3 dynamic tree speculative decoding with one-model architecture.""" - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 35: - pytest.skip("Not enough memory to load target + draft model") - - models_path = llm_models_root() - target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - eagle_model = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" - - max_batch_size = 4 - max_draft_len = 6 - dynamic_tree_max_topK = 10 - max_total_draft_tokens = 60 - kv_cache_config = KvCacheConfig(enable_block_reuse=False, max_tokens=8192) - cuda_graph_config = CudaGraphConfig( - batch_sizes=[i for i in range(1, max_batch_size + - 1)]) if use_cuda_graph else None - - llm_common_config = dict( - model=target_model_dir, - attn_backend="TRTLLM", - disable_overlap_scheduler=disable_overlap_scheduler, - cuda_graph_config=cuda_graph_config, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - max_seq_len=8192, - ) - - spec_config = Eagle3DecodingConfig( - max_draft_len=max_draft_len, - speculative_model=eagle_model, - eagle3_one_model=True, - use_dynamic_tree=True, - dynamic_tree_max_topK=dynamic_tree_max_topK, - max_total_draft_tokens=max_total_draft_tokens, - max_batch_size=max_batch_size, - ) - - # Create the LLM instance - llm_spec = LLM(**llm_common_config, speculative_config=spec_config) - - # Acceptance rate tests - prompts = [ - "The capital of France is", - "The president of the United States is", - ] - tok_ids = [llm_spec.tokenizer.encode("The future of AI is")] - - sampling_params = SamplingParams(max_tokens=128, temperature=0) - - for i in range(len(tok_ids)): - num_tokens = 0 - num_drafted = 0 - num_accepted = 0 - - for output in llm_spec.generate_async(tok_ids[i], - sampling_params, - streaming=True): - new_tokens = output.outputs[0].token_ids - num_drafted += max_draft_len - num_accepted += len(new_tokens) - num_tokens - 1 - num_tokens = len(new_tokens) - - accept_rate = num_accepted / num_drafted - # Measured ~0.24 across all 4 configs (CG x overlap). - assert accept_rate > 0.20 - - # Output tests: verify spec decode matches reference - sampling_params = SamplingParams(max_tokens=10, temperature=0) - - results_spec = llm_spec.generate(prompts, sampling_params) - generated_text_spec = [result.outputs[0].text for result in results_spec] - llm_spec.shutdown() - - llm_ref = LLM(**llm_common_config) - results_ref = llm_ref.generate(prompts, sampling_params) - generated_text_ref = [result.outputs[0].text for result in results_ref] - llm_ref.shutdown() - - for text_spec, text_ref in zip(generated_text_spec, generated_text_ref): - assert text_spec == text_ref - - if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/_torch/thop/parallel/test_custom_ops.py b/tests/unittest/_torch/thop/parallel/test_custom_ops.py index db2920393a0a..ec6db275d8ff 100644 --- a/tests/unittest/_torch/thop/parallel/test_custom_ops.py +++ b/tests/unittest/_torch/thop/parallel/test_custom_ops.py @@ -100,9 +100,6 @@ def test_register_fake(custom_ops): "trtllm::mxe4m3_mxe2m1_block_scale_moe_runner", "trtllm::mxfp8_quantize", "trtllm::mamba2_mtp_ssm_cache_update", - "trtllm::build_dynamic_tree_op", - "trtllm::verify_dynamic_tree_greedy_op", - "trtllm::verify_dynamic_tree_greedy_out_op", } ops_missing_fake_impl = [] diff --git a/tests/unittest/others/test_kv_cache_update.py b/tests/unittest/others/test_kv_cache_update.py index b56a0f5aa794..5abb913edc4f 100644 --- a/tests/unittest/others/test_kv_cache_update.py +++ b/tests/unittest/others/test_kv_cache_update.py @@ -22,16 +22,6 @@ def is_float_type(torch_dtype): return torch_dtype in float_types -def _make_random_cache_data(shape, dtype, device='cuda'): - """Create random cache data of the given shape and dtype.""" - if is_integer_type(dtype): - return torch.randint(0, 100, shape, dtype=dtype, device=device) - elif is_float_type(dtype): - return torch.rand(shape, dtype=dtype, device=device) - else: - raise ValueError(f"Unsupported dtype: {dtype}") - - @pytest.mark.parametrize("num_kv_heads", [1, 4, 8]) @pytest.mark.parametrize("head_dim", [64, 67, 128]) @pytest.mark.parametrize("layer_count", [1, 32, 45]) @@ -162,194 +152,3 @@ def test_linear_kvcache_update(num_kv_heads: int, head_dim: int, ground_truth_layer_past_key_value = ground_truth_past_key_values[i] assert torch.allclose(layer_past_key_value, ground_truth_layer_past_key_value) - - -@pytest.mark.parametrize("num_kv_heads", [1, 4, 8]) -@pytest.mark.parametrize("head_dim", [64, 128]) -@pytest.mark.parametrize("layer_count", [1, 32]) -@pytest.mark.parametrize("batch_size", [1, 5]) -@pytest.mark.parametrize("kv_cache_type", [torch.int8, torch.float16]) -@pytest.mark.parametrize("rewind_draft_token_count", [5, 63]) -@pytest.mark.parametrize("max_kv_cache_length", [128, 256]) -def test_paged_kvcache_update_2d(num_kv_heads: int, head_dim: int, - layer_count: int, batch_size: int, - kv_cache_type: torch.dtype, - rewind_draft_token_count: int, - max_kv_cache_length: int): - """Test the 2D variant of update_kv_cache_draft_token_location (paged KV cache only). - - The 2D kernel takes a [batch_size, max_draft_len] indices tensor instead of - packed 1D indices with offsets. It reads KV entries from scattered draft - positions and compacts them into sequential positions. - """ - torch.cuda.set_device(0) - torch.cuda.manual_seed(1234) - - elt_size = torch.zeros(1, dtype=kv_cache_type).element_size() - head_size_in_bytes = head_dim * elt_size - bytes_per_token = num_kv_heads * head_size_in_bytes - - # tokens_per_block must be power of 2 - tokens_per_block = 64 - bytes_per_block = tokens_per_block * bytes_per_token - max_blocks_per_seq = (max_kv_cache_length + tokens_per_block - - 1) // tokens_per_block + 1 # +1 for oneMoreBlock - total_blocks = batch_size * max_blocks_per_seq - - # Allocate pool: each pool block stores all layers, each layer has K+V - # Layout per pool block: [layer_count * 2 * bytes_per_block] bytes - pool_bytes_per_block = layer_count * 2 * bytes_per_block - pool = torch.zeros(total_blocks * pool_bytes_per_block, - dtype=torch.uint8, - device='cuda') - - # Fill pool with random data via a view - # We'll use a structured view for ground truth computation - # Pool block layout per layer: [K: num_kv_heads, tokens_per_block, head_dim] - # [V: num_kv_heads, tokens_per_block, head_dim] - pool_int8 = pool.view(-1) - pool_int8[:] = torch.randint(0, - 256, (pool_int8.shape[0], ), - dtype=torch.uint8, - device='cuda') - - # pointer array: [primary_pool_ptr, secondary_pool_ptr] as int64 - pointer_array = torch.tensor( - [pool.data_ptr(), pool.data_ptr()], dtype=torch.int64, device='cuda') - - # offset array: [batch_size, 2, max_blocks_per_seq] as int32 - # Assign blocks sequentially: seq i gets blocks [i*max_blocks_per_seq, ...] - # Both K and V rows point to the same blocks (primary pool, high bit = 0) - offset_array = torch.zeros(batch_size * 2 * max_blocks_per_seq, - dtype=torch.int32, - device='cuda') - for seq_idx in range(batch_size): - for blk in range(max_blocks_per_seq): - block_id = seq_idx * max_blocks_per_seq + blk - # K row - offset_array[seq_idx * max_blocks_per_seq * 2 + blk] = block_id - # V row - offset_array[seq_idx * max_blocks_per_seq * 2 + max_blocks_per_seq + - blk] = block_id - - # Generate accepted draft token data - max_draft_len = rewind_draft_token_count - # num_accepted_tokens: the kernel uses max(num_accepted - 1, 0) as draft count - num_accepted_list = [ - random.randint(1, rewind_draft_token_count) for _ in range(batch_size) - ] - num_accepted_tokens = torch.tensor(num_accepted_list, - dtype=torch.int32, - device='cuda') - - # accepted_draft_tokens_indices_2d: [batch_size, max_draft_len] - # Each row has indices in [0, rewind_draft_token_count) for accepted tokens - accepted_indices_2d_cpu = torch.zeros(batch_size, - max_draft_len, - dtype=torch.int32) - for seq_idx in range(batch_size): - n_accepted = num_accepted_list[seq_idx] - draft_count = max(n_accepted - 1, 0) - if draft_count > 0: - perm = torch.randperm(rewind_draft_token_count, - dtype=torch.int32)[:draft_count] - accepted_indices_2d_cpu[seq_idx, :draft_count] = perm - accepted_indices_2d = accepted_indices_2d_cpu.cuda() - - past_key_value_lengths = torch.randint(rewind_draft_token_count, - max_kv_cache_length, (batch_size, ), - dtype=torch.int32, - device='cuda') - past_key_value_lengths_cpu = past_key_value_lengths.to('cpu') - - # Helper to read/write a KV element from the pool (CPU-side for ground truth) - pool_cpu = pool.clone().cpu() - - def _get_block_offset(seq_idx, token_pos, kv_idx): - """Return (pool_block_id, local_token_idx) for a given seq/token/kv.""" - blk_idx = token_pos // tokens_per_block - local_tok = token_pos % tokens_per_block - row_base = seq_idx * max_blocks_per_seq * 2 + kv_idx * max_blocks_per_seq - block_id = offset_array[row_base + blk_idx].item() - return block_id, local_tok - - def _pool_byte_offset(block_id, layer_idx, kv_idx, head_idx, local_tok, - chan): - """Compute byte offset into the pool for a specific element.""" - block_base = block_id * pool_bytes_per_block - layer_base = layer_idx * 2 * bytes_per_block - kv_base = kv_idx * bytes_per_block - # Within a KV block: [num_kv_heads, tokens_per_block, head_dim] in element units - elem_offset = (head_idx * tokens_per_block * head_dim + - local_tok * head_dim + chan) - return block_base + layer_base + kv_base + elem_offset * elt_size - - def _read_kv_token(pool_data, seq_idx, token_pos, layer_idx, kv_idx): - """Read [num_kv_heads, head_dim] for one token from pool.""" - block_id, local_tok = _get_block_offset(seq_idx, token_pos, kv_idx) - result = torch.zeros(num_kv_heads, head_dim, dtype=kv_cache_type) - for h in range(num_kv_heads): - for d in range(head_dim): - off = _pool_byte_offset(block_id, layer_idx, kv_idx, h, - local_tok, d) - raw = pool_data[off:off + elt_size] - result[h, d] = torch.frombuffer(raw.numpy().tobytes(), - dtype=kv_cache_type)[0] - return result - - def _write_kv_token(pool_data, seq_idx, token_pos, layer_idx, kv_idx, - values): - """Write [num_kv_heads, head_dim] for one token to pool.""" - block_id, local_tok = _get_block_offset(seq_idx, token_pos, kv_idx) - val_bytes = values.contiguous().numpy().view('uint8') - for h in range(num_kv_heads): - for d in range(head_dim): - off = _pool_byte_offset(block_id, layer_idx, kv_idx, h, - local_tok, d) - src_off = (h * head_dim + d) * elt_size - pool_data[off:off + elt_size] = torch.from_numpy( - val_bytes[src_off:src_off + elt_size].copy()) - - # Compute ground truth on CPU - gt_pool = pool_cpu.clone() - for layer_idx in range(layer_count): - for seq_idx in range(batch_size): - draft_count = max(num_accepted_list[seq_idx] - 1, 0) - if draft_count == 0: - continue - past_len = past_key_value_lengths_cpu[seq_idx].item() - token_start = past_len - rewind_draft_token_count - for target_idx in range(draft_count): - src_pos = accepted_indices_2d_cpu[seq_idx, target_idx].item() - for kv_idx in range(2): # K=0, V=1 - val = _read_kv_token(pool_cpu, seq_idx, - token_start + src_pos, layer_idx, - kv_idx) - _write_kv_token(gt_pool, seq_idx, token_start + target_idx, - layer_idx, kv_idx, val) - - torch.cuda.synchronize() - - # Run the 2D kernel - torch.ops.tensorrt_llm.update_kv_cache_draft_token_location_2d( - accepted_indices_2d, - num_accepted_tokens, - past_key_value_lengths, - True, - layer_count, - num_kv_heads, - head_size_in_bytes, - rewind_draft_token_count, - max_kv_cache_length, - pointer_array, - offset_array, - max_blocks_per_seq, - tokens_per_block, - None, - ) - torch.cuda.synchronize() - - # Compare - result_pool = pool.cpu() - assert torch.equal(result_pool, gt_pool), \ - "Paged KV cache 2D update mismatch" From 9d2d3e550fd6d152d34d3d31cd10901587078d54 Mon Sep 17 00:00:00 2001 From: Yiqing Yan Date: Tue, 14 Apr 2026 10:43:41 +0800 Subject: [PATCH 2/2] Unwaive tests marked as nvbug 6071081 Signed-off-by: Yiqing Yan --- tests/integration/test_lists/waives.txt | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 5a963d90f818..9479b247f95e 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -341,11 +341,6 @@ accuracy/test_llm_api_pytorch.py::TestKimiK2::test_nvfp4[4gpus] SKIP (https://nv accuracy/test_llm_api_pytorch.py::TestGLM4_5Air::test_nvfp4_2_model_mtp[2model_trtllm] SKIP (https://nvbugs/5981293) accuracy/test_llm_api_pytorch.py::TestGLM4_5Air::test_nvfp4_multi_gpus[throughput] SKIP (https://nvbugs/5981293) disaggregated/test_disaggregated.py::test_disaggregated_trtllm_sampler[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6069686) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=2-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False] SKIP (https://nvbugs/6071081) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=True-v2_kv_cache=True] SKIP (https://nvbugs/6071081) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=True-v2_kv_cache=False] SKIP (https://nvbugs/6071081) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6071081) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=True] SKIP (https://nvbugs/6071081) accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[google_gemma-3-1b-it-False] SKIP (https://nvbugs/6059036) accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[meta-llama_Llama-3.1-8B-Instruct-False] SKIP (https://nvbugs/6059036) accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[mistralai_Codestral-22B-v0.1-False] SKIP (https://nvbugs/6059036)