diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu new file mode 100644 index 000000000000..1dcfdde71671 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu @@ -0,0 +1,356 @@ +/* + * 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 new file mode 100644 index 000000000000..c48b0f25005a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h @@ -0,0 +1,98 @@ +/* + * 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 8d1ca4530d87..40fe23b8f8a8 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.cu +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.cu @@ -135,8 +135,164 @@ __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, @@ -241,6 +397,55 @@ 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 f8551db9b75c..fd3cbd53512f 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.h +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.h @@ -208,6 +208,21 @@ 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 0d5d386bef00..8e9b85e58ae4 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h @@ -295,6 +295,10 @@ 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 af267c59013b..adfed54696ed 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/prepareCustomMask.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/prepareCustomMask.cu @@ -73,6 +73,13 @@ __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; @@ -107,12 +114,12 @@ __global__ void prepareCustomMaskBuffersKernelForKeepsMmaAb( else { // Sparse region: check the input mask - // Input mask shape: [bs, seqLenQ, ceilDiv(seqLenQ, 32)] + // Input mask shape: [bs, packedMaskSeqLenQ, packedMaskNumBlocks] // 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 * seqLenQ + tokenIdxQ) * ceilDiv(seqLenQ, 32); + int32_t const qMaskBaseIdx = (batchIdx * packedMaskSeqLenQ + tokenIdxQ) * packedMaskNumBlocks; 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 30c9e2b703ad..a55a17c69302 100644 --- a/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp +++ b/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp @@ -491,6 +491,7 @@ 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 0d077c4b030e..30ee674d8534 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -108,6 +108,7 @@ 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 new file mode 100644 index 000000000000..2c70fcfc5cb1 --- /dev/null +++ b/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp @@ -0,0 +1,229 @@ +/* + * 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 400cf81033ba..3a25d903bf32 100644 --- a/cpp/tensorrt_llm/thop/parallelDecodeKVCacheUpdateOp.cpp +++ b/cpp/tensorrt_llm/thop/parallelDecodeKVCacheUpdateOp.cpp @@ -113,9 +113,73 @@ 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 7dbfdea2b07c..2d74fadc5598 100644 --- a/docs/source/features/feature-combination-matrix.md +++ b/docs/source/features/feature-combination-matrix.md @@ -1,23 +1,24 @@ # 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) | 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 | --- | +| 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 | --- | diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index d1673deaffea..5364d9b6fcd7 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). -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`. +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. 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,6 +50,36 @@ 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). @@ -199,6 +229,18 @@ 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 0492533a7986..ea688acebe48 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -41,18 +41,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) | 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 | +| 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 | [^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 55ffaf6f52b2..d382ab9f7830 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -195,6 +195,7 @@ 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', @@ -229,6 +230,29 @@ 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 @@ -268,7 +292,9 @@ 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) + eagle3_model_arch=args.eagle3_model_arch, + max_total_draft_tokens=args.max_total_draft_tokens, + max_batch_size=args.max_batch_size) elif spec_decode_algo == "DRAFT_TARGET": spec_config = DraftTargetDecodingConfig( max_draft_len=args.spec_decode_max_draft_len, @@ -361,7 +387,17 @@ def setup_llm(args, **kwargs): def main(): args = parse_arguments() - prompts = args.prompt if args.prompt else example_prompts + + 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 llm, sampling_params = setup_llm(args) new_prompts = [] @@ -373,41 +409,132 @@ def main(): tokenize=False, add_generation_prompt=True)) prompts = new_prompts - 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 "" + 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}]{sequence_id_text} Prompt: {prompt!r}, Generated text: {generated_text!r}" + f"[{i}] Prompt: {prompt[:80]!r}..., Generated text: {generated_text[:200]!r}..." ) - if args.return_context_logits: - print( - f"[{i}]{sequence_id_text} Context logits: {output.context_logits}" - ) - if args.return_generation_logits: + + 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}" + ) + print( + f"Total tokens: {total_tokens}, Total iterations: {total_iterations}" + ) + 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 print( - f"[{i}]{sequence_id_text} Generation logits: {sequence.generation_logits}" + f"[{i}] Prompt: {prompt[:80]!r}..., Generated text: {generated_text[:200]!r}..." ) - if args.prompt_logprobs is not None: + 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 "" print( - f"[{i}]{sequence_id_text} Prompt logprobs: {sequence.prompt_logprobs}" + f"[{i}]{sequence_id_text} Prompt: {prompt[:80]!r}..., Generated text: {generated_text[:200]!r}..." ) - if args.logprobs is not None: - print(f"[{i}]{sequence_id_text} Logprobs: {sequence.logprobs}") + 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.additional_model_outputs: - for output_name in args.additional_model_outputs: - if sequence.additional_context_outputs: + 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]}" + ) print( - f"[{i}]{sequence_id_text} Context {output_name}: {sequence.additional_context_outputs[output_name]}" + f"[{i}]{sequence_id_text} Generation {output_name}: {sequence.additional_generation_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 600f655bc517..64b6173a7800 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 ..speculative.utils import SpecDecodingTensor + from tensorrt_llm.llmapi.llm_args import SparseAttentionConfig from ..speculative.interface import SpecMetadata from ..speculative.spec_tree_manager import SpecTreeManager @@ -383,8 +383,7 @@ 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_decoding_tensor: Optional['SpecDecodingTensor'] = None): + spec_tree_manager: Optional['SpecTreeManager'] = 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 6d237cf01feb..2e3428838725 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -597,15 +597,13 @@ 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_decoding_tensor) + spec_metadata, spec_tree_manager) 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 7f1b7c0c4dfa..541c53b59bee 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -8,7 +8,6 @@ import torch if TYPE_CHECKING: - from ..speculative.utils import SpecDecodingTensor from ..speculative.interface import SpecMetadata from ..speculative.spec_tree_manager import SpecTreeManager @@ -507,9 +506,18 @@ 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, - self.spec_decoding_position_offsets, self.spec_decoding_packed_mask + self.spec_decoding_generation_lengths, position_offsets_for_cpp, + self.spec_decoding_packed_mask ] if self.is_sm_version_trtllm_gen_kernel(sm=get_sm_version()): spec_decoding_tensor_params.append( @@ -521,6 +529,12 @@ 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_() @@ -1448,48 +1462,59 @@ 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. """ - 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', - ) + 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) def update_spec_dec_param( self, @@ -1502,7 +1527,6 @@ 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. @@ -1516,31 +1540,12 @@ 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 - - # 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." + self.is_spec_decoding_enabled = is_spec_decoding_enabled # 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 @@ -1553,20 +1558,35 @@ 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: - self.spec_decoding_position_offsets = torch.empty( - [self.max_num_requests, max_total_draft_tokens + 1], - dtype=torch.int, - device='cuda', - ) + 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', + ) 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, max_total_draft_tokens + 1, - math.ceil((max_total_draft_tokens + 1) / 32) - ], + [self.max_num_requests, buf_dim, + math.ceil(buf_dim / 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], @@ -1581,20 +1601,31 @@ 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 + # Case 1: dynamic tree — copy per-request params from spec_tree_manager. if self.is_spec_dec_dynamic_tree: - 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) + 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) else: - self.generate_spec_decoding_generation_length( - runtime_draft_len=max_total_draft_tokens) + 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) # 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: @@ -1961,6 +1992,13 @@ 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 6946294259fe..e93b7eef3bc7 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -890,6 +890,19 @@ 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) @@ -939,6 +952,25 @@ 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 7c2ce12a7f4b..65e09660c9ac 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -56,7 +56,6 @@ 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) @@ -2357,13 +2356,10 @@ def _prepare_tp_inputs( assert spec_config.spec_dec_mode.support_overlap_scheduler( ), f"{spec_config.decoding_type} does not support overlap scheduler" - 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 + # 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 # will contain previous batch indices of generation requests previous_batch_indices = [] @@ -2403,20 +2399,10 @@ def _prepare_tp_inputs( list( range(len(position_ids), len(position_ids) + 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))) + 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 @@ -2436,20 +2422,10 @@ def _prepare_tp_inputs( list( range(len(position_ids), len(position_ids) + 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))) + 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] * @@ -3653,7 +3629,6 @@ 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( @@ -3667,7 +3642,8 @@ def forward(self, spec_resource_manager = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) spec_tree_manager = None - if isinstance(spec_resource_manager, Eagle3ResourceManager): + if spec_resource_manager is not None and hasattr( + spec_resource_manager, 'spec_tree_manager'): spec_tree_manager = spec_resource_manager.spec_tree_manager spec_metadata = self._set_up_spec_metadata(spec_resource_manager, no_cache=kv_cache_manager @@ -3690,6 +3666,18 @@ 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, @@ -3699,8 +3687,7 @@ 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_decoding_tensor=spec_decoding_tensor) + spec_tree_manager=spec_tree_manager) else: spec_resource_manager = None spec_metadata = None @@ -3746,6 +3733,25 @@ 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, @@ -3804,6 +3810,10 @@ 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 783aa17b9c5e..2893ebf13d38 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1740,7 +1740,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_draft_len + self.model_engine.runtime_draft_len = self.model_engine.max_total_draft_tokens 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 e0aa739d8697..1aa00944744d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -370,6 +370,10 @@ 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 @@ -454,15 +458,16 @@ def allocation_scope(current_stage: ExecutorMemoryType): def drafting_loop_wrapper(model): from tensorrt_llm._torch.speculative.drafting_loops import ( - LinearDraftingLoopWrapper, TreeDraftingLoopWrapper) + LinearDraftingLoopWrapper, + StaticTreeDraftingLoopWrapper) from tensorrt_llm.llmapi import EagleDecodingConfig - use_tree_drafter = isinstance( + static_tree_drafter = isinstance( draft_spec_config, EagleDecodingConfig - ) and not draft_spec_config.is_linear_tree + ) and draft_spec_config.eagle_choices is not None - if use_tree_drafter: - return TreeDraftingLoopWrapper( + if static_tree_drafter: + return StaticTreeDraftingLoopWrapper( 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 971ff8c5402d..53047dd7262d 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -231,13 +231,18 @@ 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, - cache_manager.num_kv_heads, + # 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], int(cache_manager.head_dim * kv_cache_dtype_byte_size), cache_manager.max_total_draft_tokens, cache_manager.max_attention_window_vec[0], @@ -801,11 +806,6 @@ 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 8cf41bee5e86..e3033fa6d303 100644 --- a/tensorrt_llm/_torch/speculative/drafting_loops.py +++ b/tensorrt_llm/_torch/speculative/drafting_loops.py @@ -162,11 +162,9 @@ 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, @@ -202,7 +200,7 @@ def prepare_for_generation(self, attn_metadata: AttentionMetadata, return new_position_ids -class TreeDraftingLoopWrapper(BaseDraftingLoopWrapper): +class StaticTreeDraftingLoopWrapper(BaseDraftingLoopWrapper): def __init__(self, max_draft_len: int, max_total_draft_tokens: int, max_batch_size: int, draft_model: torch.nn.Module): @@ -226,11 +224,8 @@ 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]: - 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 + assert isinstance(spec_metadata, Eagle3SpecMetadata) + spec_tree_manager = spec_metadata.eagle3_resource_manager.spec_tree_manager logits = self.draft_model.forward(input_ids=input_ids, position_ids=position_ids, @@ -251,7 +246,6 @@ 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 @@ -266,10 +260,12 @@ 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, :].reshape( - -1), - position_ids=self.position_ids_buffer[:batch_size, :]. - reshape(-1), + 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), attn_metadata=attn_metadata, spec_metadata=spec_metadata, return_context_logits=True) @@ -444,7 +440,6 @@ 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[: @@ -472,7 +467,9 @@ 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:].unsqueeze( + 0, 1:self. + max_total_draft_tokens + + 1].unsqueeze( 0 ) - 1 # exclude the root node attn_metadata.spec_decoding_position_offsets[:batch_size, self. @@ -500,7 +497,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.repeat( + hidden_states_read_indices_offset = spec_tree_manager.hidden_states_read_indices_offset_for_drafter_model[:self.max_total_draft_tokens + 1].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 new file mode 100644 index 000000000000..2565e034aaf1 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py @@ -0,0 +1,198 @@ +# 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 0211f3e7c989..0191d59ad08e 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -73,7 +73,8 @@ def __init__(self, self.spec_tree_manager = None if isinstance(config, - EagleDecodingConfig) and config.eagle_choices is not None: + EagleDecodingConfig) and (config.eagle_choices is not None + or config.use_dynamic_tree): self.spec_tree_manager = SpecTreeManager( max_num_requests=self.max_num_requests, use_dynamic_tree=config.use_dynamic_tree, @@ -122,6 +123,43 @@ 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) @@ -313,6 +351,9 @@ 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: @@ -339,11 +380,16 @@ def __post_init__(self): device='cuda', ) - # 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 + # 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 def is_layer_capture(self, layer_id: int): return layer_id in self.layers_to_capture @@ -358,7 +404,11 @@ def prepare(self): pin_memory=prefer_pinned()) self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) - self.num_tokens -= (self.num_generations) * self.runtime_draft_len + 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 sa_manager = getattr(self.spec_resource_manager, 'sa_manager', None) if sa_manager is not None: @@ -382,12 +432,27 @@ def maybe_capture_hidden_states( class Eagle3OneModelSampler(MTPSampler): + """Sampler for one-model EAGLE3 (linear and dynamic tree modes).""" - def __init__(self, args: TorchSampler.Args): - super().__init__(args, nextn=args.max_draft_len) + 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) 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", @@ -399,6 +464,8 @@ 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: @@ -417,6 +484,27 @@ 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: @@ -425,6 +513,21 @@ 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}) @@ -484,16 +587,64 @@ def forward(self, spec_metadata=spec_metadata, draft_model=draft_model) - next_draft_tokens = [] + # Predict 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() @@ -581,26 +732,7 @@ def forward(self, gen_draft_tokens) next_draft_tokens[num_contexts:] = gen_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, - } + return next_draft_tokens def sample_and_accept_draft_tokens( self, @@ -612,8 +744,14 @@ 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) + 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) 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 new file mode 100644 index 000000000000..f33a5824a455 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py @@ -0,0 +1,1049 @@ +# 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 a795a19d6d87..a87724930248 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -100,12 +100,13 @@ 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_tokens, seq_slots, self.max_beam_width)), + new_tokens=int_tensor((max_new_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,)), @@ -120,6 +121,15 @@ 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. @@ -238,19 +248,30 @@ 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 to match fixed-size store buffers for index_copy_. - if o_new_tokens.shape[1] < (self.draft_len + 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: o_new_tokens = torch.nn.functional.pad( - o_new_tokens, (0, (self.draft_len + 1) - o_new_tokens.shape[1]) + o_new_tokens, (0, new_tokens_width - o_new_tokens.shape[1]) ) - if o_next_draft_tokens.shape[1] < self.draft_len: + 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: o_next_draft_tokens = torch.nn.functional.pad( - o_next_draft_tokens, (0, self.draft_len - o_next_draft_tokens.shape[1]) + o_next_draft_tokens, (0, draft_tokens_width - o_next_draft_tokens.shape[1]) ) - if o_next_new_tokens.shape[1] < (self.draft_len + 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: o_next_new_tokens = torch.nn.functional.pad( - o_next_new_tokens, (0, (self.draft_len + 1) - o_next_new_tokens.shape[1]) + o_next_new_tokens, (0, next_new_tokens_width - 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 bb1d4758798f..d6e7cd826d06 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -1,10 +1,14 @@ +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 @@ -15,11 +19,11 @@ class SpecTreeManager: # Auxiliary buffers # The top k list for each draft layer. - top_k_list = [] + top_k_list: list # The user input eagle choices, only available when using static tree. - 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 + 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 # 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] @@ -37,46 +41,61 @@ class SpecTreeManager: # shape: [num_trees, max_total_draft_tokens + 1], device tensor. spec_dec_position_offsets: torch.Tensor = None - # TODO: Optimized together with the subsequent dynamic tree. - # Auxiliary buffers for the static 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 from all currently generated draft tokens as the input of the next draft layer. + # Gather the required draft tokens among the 'max_total_draft_tokens + 1' tokens. # 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]]], dynamic_tree_max_topK: int): + eagle_choices: List[List[int]] | None, + 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 @@ -84,23 +103,29 @@ def __init__(self, max_num_requests: int, use_dynamic_tree: bool, self.top_k_list = [] # Initialize the buffers - 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) - + # 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.spec_dec_packed_mask = torch.zeros( (self.num_trees, self.max_total_draft_tokens + 1, - math.ceil((self.max_total_draft_tokens + 1) / 32)), + math.ceil(self._internal_buf_dim / 32)), dtype=torch.int32, device='cuda', ) @@ -110,14 +135,45 @@ 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() - # self.dump_tree_info() - 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') + # 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 = [ @@ -127,7 +183,101 @@ def init_tree_info_for_dynamic_tree(self): pin_memory=prefer_pinned()) * self.dynamic_tree_max_topK ] - # For the static tree + # 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) + 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)] @@ -176,9 +326,7 @@ 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 - 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_mask_matrix(0) self.compute_spec_dec_packed_mask(self.spec_dec_mask_matrix, self.spec_dec_packed_mask) @@ -217,7 +365,6 @@ 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, @@ -255,89 +402,81 @@ def init_tree_info_for_static_tree(self): device= 'cuda') - # Get the eagle_paths def get_eagle_paths(self, tree_idx=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] + 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] # 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] - # Compute the packed mask according to the mask matrix - def compute_spec_dec_packed_mask(self, mask_matrix, 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 + 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 - packed_mask = packed_mask.reshape(num_trees, - self.max_total_draft_tokens + 1, - num_blocks) + 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 - # Print the tree info def dump_tree_info(self): - 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}" + 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, ) diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 33f377063994..b74b61af126d 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -14,9 +14,11 @@ from .draft_target import (DraftTargetOneModelSampler, DraftTargetOneModelSpecMetadata, DraftTargetOneModelWorker) -from .eagle3 import (Eagle3OneModelSampler, Eagle3OneModelSpecMetadata, +from .eagle3 import (Eagle3OneModelDynamicTreeResourceManager, + Eagle3OneModelSampler, Eagle3OneModelSpecMetadata, Eagle3OneModelWorker, Eagle3ResourceManager, Eagle3SpecMetadata) +from .eagle3_dynamic_tree import Eagle3OneModelDynamicTreeWorker from .model_drafter import ModelDrafter from .mtp import (MTPEagleWorker, MTPHiddenStatesManager, MTPSampler, MTPSpecMetadata, MTPWorker) @@ -91,6 +93,8 @@ 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( @@ -182,6 +186,10 @@ 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) @@ -244,7 +252,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) + return Eagle3OneModelSampler(sampler_args, spec_config=spec_config) if spec_config.spec_dec_mode.is_pard(): return MTPSampler(sampler_args, nextn=spec_config.tokens_per_gen_step - 1) @@ -307,6 +315,9 @@ 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 9f001b4e5aeb..f72263ec83ad 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1087,7 +1087,6 @@ 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,7 +1115,8 @@ def validate_eagle_config(self) -> 'EagleDecodingConfig': self.max_total_draft_tokens = len(self.eagle_choices) # Dynamic tree logic - if self.use_dynamic_tree: + if self.use_dynamic_tree or self.dynamic_tree_max_topK is not None: + self.use_dynamic_tree = True if self.eagle_choices is not None: raise ValueError( "If use_dynamic_tree is True, eagle_choices should be None") @@ -1128,10 +1128,28 @@ 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" ) - 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)" + + 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)" ) + 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 @@ -1212,6 +1230,11 @@ 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", @@ -2147,8 +2170,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/unittest/_torch/modeling/test_modeling_llama.py b/tests/unittest/_torch/modeling/test_modeling_llama.py index 334e60a61e97..711f90b7a7dc 100644 --- a/tests/unittest/_torch/modeling/test_modeling_llama.py +++ b/tests/unittest/_torch/modeling/test_modeling_llama.py @@ -20,7 +20,9 @@ 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.utils import SpecDecodingTensor +from tensorrt_llm._torch.speculative.interface import (SpecMetadata, + SpeculativeDecodingMode) +from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager from tensorrt_llm._utils import get_sm_version from tensorrt_llm.bindings.executor import KvCacheConfig from tensorrt_llm.mapping import Mapping @@ -516,6 +518,20 @@ 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, @@ -534,10 +550,6 @@ 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, @@ -547,7 +559,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_decoding_tensor=spec_decoding_tensor, + spec_tree_manager=spec_tree_mgr, ) gen_position_ids_0 = [ @@ -594,15 +606,35 @@ 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_spec_dec_tree - if get_sm_version() < 100 else False, - is_spec_dec_dynamic_tree=False, + is_spec_dec_tree=is_tree_phase1, + is_spec_dec_dynamic_tree=is_tree_phase1, 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) + model_is_wrapped=False, + spec_metadata=spec_metadata_phase1, + spec_tree_manager=spec_tree_mgr_phase1) gen_position_ids_1 = [ torch.full( @@ -649,15 +681,35 @@ 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_spec_dec_tree - if get_sm_version() < 100 else False, - is_spec_dec_dynamic_tree=False, + is_spec_dec_tree=is_tree_ref, + is_spec_dec_dynamic_tree=is_tree_ref, 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) + model_is_wrapped=False, + spec_metadata=spec_metadata_ref, + spec_tree_manager=spec_tree_mgr_ref) gen_position_ids_ref = [ torch.full((gen_input_ids_ref.size(-1), ), @@ -675,12 +727,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=0.4, - rtol=0.4) + atol=1.0, + rtol=1.0) torch.testing.assert_close(gen_logits_1[1, :], gen_logits_ref[3, :], - atol=0.4, - rtol=0.4) + atol=1.0, + rtol=1.0) 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 7293cf91582a..78acbde58cb9 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,16 +1,15 @@ 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 TreeDraftingLoopWrapper +from tensorrt_llm._torch.speculative.drafting_loops import StaticTreeDraftingLoopWrapper from tensorrt_llm._torch.speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager -from tensorrt_llm.llmapi import Eagle3DecodingConfig +from tensorrt_llm.llmapi.llm_args import Eagle3DecodingConfig sys.path.append(os.path.join(os.path.dirname(__file__), "..")) @@ -21,6 +20,7 @@ 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 TreeDraftingLoopWrapper - tree_drafting_loop_wrapper = TreeDraftingLoopWrapper( + # 3) Create StaticTreeDraftingLoopWrapper + static_tree_drafting_loop_wrapper = StaticTreeDraftingLoopWrapper( 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 - tree_drafting_loop_wrapper.prepare_for_generation( + static_tree_drafting_loop_wrapper.prepare_for_generation( attn_metadata=attn_metadata, spec_metadata=spec_metadata, spec_tree_manager=spec_tree_manager, @@ -156,7 +156,8 @@ def run_test( # Compare input_ids and position_ids print( - f"tree_drafting_loop_wrapper.position_ids_buffer: {tree_drafting_loop_wrapper.position_ids_buffer}, \ + f"static_tree_drafting_loop_wrapper.position_ids_buffer: \ + {static_tree_drafting_loop_wrapper.position_ids_buffer}, \ ref_output_position_ids: {ref_position_ids}" ) @@ -208,7 +209,7 @@ def run_test( ref_spec_metadata.hidden_states_write_indices: {ref_spec_metadata['hidden_states_write_indices']}" ) - assert torch.all(tree_drafting_loop_wrapper.position_ids_buffer == ref_position_ids) + assert torch.all(static_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"]) @@ -452,7 +453,3 @@ 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 f62316758732..7f75793cfa89 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 \ - TreeDraftingLoopWrapper + StaticTreeDraftingLoopWrapper 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 = TreeDraftingLoopWrapper( + tree_drafter = StaticTreeDraftingLoopWrapper( 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 c8d8880c31f4..fb792149e5f7 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -10,6 +10,7 @@ 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 @@ -815,5 +816,95 @@ 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 ec6db275d8ff..db2920393a0a 100644 --- a/tests/unittest/_torch/thop/parallel/test_custom_ops.py +++ b/tests/unittest/_torch/thop/parallel/test_custom_ops.py @@ -100,6 +100,9 @@ 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 5abb913edc4f..b56a0f5aa794 100644 --- a/tests/unittest/others/test_kv_cache_update.py +++ b/tests/unittest/others/test_kv_cache_update.py @@ -22,6 +22,16 @@ 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]) @@ -152,3 +162,194 @@ 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"