diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu index 1dcfdde71671..7a5e922d6ad1 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.cu @@ -15,338 +15,2243 @@ * limitations under the License. */ +#ifndef CUDART_VERSION +#error CUDART_VERSION Undefined! +#elif (CUDART_VERSION >= 11050) +#include +#else +#include "3rdparty/cub/cub.cuh" +#endif + #include "dynamicTreeKernels.h" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" - +#include "tensorrt_llm/common/reduceKernelUtils.cuh" +#include "tensorrt_llm/common/vec_dtypes.cuh" +#include "tensorrt_llm/kernels/decodingCommon.h" +#include +#include +#include +#include +#include +#include TRTLLM_NAMESPACE_BEGIN +using namespace tensorrt_llm::common; 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) +// --------------------------------------------------------------------------- +// Two-stage top-k / top-p masking kernels +// Mirrors the approach in invokeBatchTopKSampling (samplingTopKKernels.cu), +// but outputs a masked logits tensor instead of sampling a token. +// --------------------------------------------------------------------------- + +// Stage 1: Parallel top-k reduction across BLOCKS_PER_BEAM_ blocks per row. +// Each block handles (vocabSize / BLOCKS_PER_BEAM_) elements and finds its +// local top-k, writing (global_index, logit_value) pairs into the tmp buffers. +template +__global__ void topKProbStage1(T const* __restrict__ logits, T* tmpLogProbs, int32_t* topKTmpIdBuf, T* topKTmpValBuf, + int32_t maxTopK, int32_t const* topKs, int32_t vocabSize) { - int32_t bid = blockIdx.x; - int32_t tid = threadIdx.x; + typedef cub::BlockReduce, BLOCK_SIZE_> BlockReduce; + __shared__ typename BlockReduce::TempStorage tempStorage; - if (tid >= draftTokenNum) - { - return; - } + auto const tid = static_cast(threadIdx.x); + auto const bid = static_cast(blockIdx.x); + auto const rowId = bid / BLOCKS_PER_BEAM_; + auto const blockLane = bid % BLOCKS_PER_BEAM_; // chunk index within the row - // treeMask layout: [batchSize, draftTokenNum, draftTokenNum] (QLEN_ONLY mode) - int32_t tokenTreeIdx = draftTokenNum * draftTokenNum * bid + draftTokenNum * tid + 1; + auto const k = (topKs != nullptr) ? topKs[rowId] : maxTopK; - treeMask[tokenTreeIdx - 1] = 1; // self-attention diagonal - for (int32_t i = 0; i < draftTokenNum - 1; i++) - { - treeMask[tokenTreeIdx + i] = 0; - } + bool const IS_FP16 = std::is_same::value; + T const MAX_T_VAL = IS_FP16 ? HALF_FLT_MAX : FLT_MAX; - int32_t position = 0; + // Base offset into the flat (nRows * vocabSize) logits array for this row. + auto const rowOffset = rowId * vocabSize; + // Base offset into the tmp buffers for this (row, blockLane). + auto const tmpIdxBase = rowId * BLOCKS_PER_BEAM_ * maxTopK + blockLane * k; - if (tid == 0) + // Copy this block's chunk of logits into tmpLogProbs scratch space. + for (auto elemId = tid + blockLane * BLOCK_SIZE_; elemId < vocabSize; elemId += BLOCK_SIZE_ * BLOCKS_PER_BEAM_) { - 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; + tmpLogProbs[rowOffset + elemId] = logits[rowOffset + elemId]; } - else + __syncthreads(); + + // Iteratively find the top-k values via max-reduction, zeroing each found max. + TopK_2 partial; + for (int32_t ite = 0; ite < k; ite++) { - // Walk up to root, setting treeMask ancestor bits and counting depth - int32_t curPosition = tid - 1; - while (position < depth + 1) + partial.init(); + for (auto elemId = tid + blockLane * BLOCK_SIZE_; elemId < vocabSize; elemId += BLOCK_SIZE_ * BLOCKS_PER_BEAM_) { - position += 1; - treeMask[tokenTreeIdx + curPosition] = 1; + partial.insert(tmpLogProbs[rowOffset + elemId], rowOffset + elemId); + } - int64_t parentTbIdx = selectedIndex[bid * (draftTokenNum - 1) + curPosition] / topK; - if (parentTbIdx == 0) - { - break; - } + TopK_2 total = BlockReduce(tempStorage).Reduce(partial, reduce_topk_op_2); - 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) + if (tid == 0) + { + topKTmpIdBuf[tmpIdxBase + ite] = total.p; // global index (rowOffset + vocabIdx) + topKTmpValBuf[tmpIdxBase + ite] = total.u; // logit value + if (total.p >= 0) { - break; + tmpLogProbs[total.p] = -MAX_T_VAL; // zero out so next iteration finds next-best } } - positions[bid * draftTokenNum + tid] = position; + __syncthreads(); } } -//! 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) +// Stage 2: Merge BLOCKS_PER_BEAM_ * k candidates per row, apply optional top-p, +// then scatter selected logit values back to an output logits tensor (all other +// positions are set to -inf so that a subsequent softmax produces 0 probability). +template +__global__ void topKProbStage2ForLogits(int32_t const* __restrict__ topKTmpIdBuf, T* topKTmpValBuf, float* outputLogits, + int32_t maxTopK, int32_t const* topKs, float const* topPs, int32_t vocabSize) { - int32_t bid = blockIdx.x; - int32_t tid = threadIdx.x; + bool const IS_FP16 = std::is_same::value; + T const MAX_T_VAL = IS_FP16 ? HALF_FLT_MAX : FLT_MAX; - if (tid >= draftTokenNum) - { - return; - } + auto const tid = static_cast(threadIdx.x); + auto const rowId = static_cast(blockIdx.x); - int32_t rowBaseIdx = (bid * draftTokenNum + tid) * numInt32PerRow; + auto const k = (topKs != nullptr) ? topKs[rowId] : maxTopK; + // size: number of valid candidates written by Stage 1 for this row. + // stride: row pitch in the tmp buffers (same as in invokeBatchTopKSampling). + auto const size = k * BLOCKS_PER_BEAM_; + auto const stride = maxTopK * BLOCKS_PER_BEAM_; - treeMask[rowBaseIdx] = 1; // bit 0 = root, always visible + typedef cub::BlockReduce, BLOCK_SIZE_> BlockReduce; + __shared__ typename BlockReduce::TempStorage tempStorage; + extern __shared__ char sharedArray[]; + // Shared layout: sId[maxTopK] | sVal2[maxTopK] + auto* sId = reinterpret_cast(sharedArray); + auto* sVal2 = reinterpret_cast(sId + maxTopK); - int32_t position = 0; + // Pointer to this row's candidates in the tmp value buffer (modified in-place during reduction). + T* sVal = topKTmpValBuf + rowId * stride; - if (tid == 0) + // Step 1: Initialize output row to -inf (all threads cooperate for bandwidth). + float* outRow = outputLogits + rowId * vocabSize; + float const negInf = -std::numeric_limits::infinity(); + for (int32_t i = tid; i < vocabSize; i += BLOCK_SIZE_) { - positions[bid * draftTokenNum] = 0; + outRow[i] = negInf; + } + __syncthreads(); - for (int32_t i = draftTokenNum - 1; i > 0; --i) + // Step 2: k-round block-reduction over the k * BLOCKS_PER_BEAM_ valid candidates. + // (Only the first 'size' entries of the row's tmp buffer were written by Stage 1.) + TopK_2 partial; + __shared__ float sMaxLogit; + for (int32_t ite = 0; ite < k; ite++) + { + partial.init(); + for (int32_t i = tid; i < size; i += BLOCK_SIZE_) { - 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; - } - } - } + partial.insert(static_cast(sVal[i]), i); + } - if (parentPosition == draftTokenNum) - { - printf("WARNING: Invalid dynamic tree! Detected a token with no parent token selected.\n"); - continue; - } + TopK_2 total = BlockReduce(tempStorage).Reduce(partial, reduce_topk_op_2); - if (retrieveNextToken[bid * draftTokenNum + parentPosition] == -1) - { - retrieveNextToken[bid * draftTokenNum + parentPosition] = i; - } - else + if (tid == 0) + { + if (ite == 0) { - int32_t originNextToken = retrieveNextToken[bid * draftTokenNum + parentPosition]; - retrieveNextToken[bid * draftTokenNum + parentPosition] = i; - retrieveNextSibling[bid * draftTokenNum + i] = originNextToken; + sMaxLogit = total.u; } + sId[ite] = total.p; + sVal[total.p] = -MAX_T_VAL; // zero out so next iteration finds next-best + sVal2[ite] = total.u; // store raw logit value (not exponentiated) } - retrieveIndex[bid * draftTokenNum] = 0; + __syncthreads(); } - else + + // Step 3: Determine top-p cutoff (tid=0 only). + // sVal2 contains logit values in descending order; we exponentiate to get unnormalized probs. + if (tid == 0) { - int32_t curPosition = tid - 1; - while (position < depth + 1) + int32_t cutoff = k; + if (topPs != nullptr) { - 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) + float const topP = topPs[rowId]; + if (topP < 1.0f) { - break; + // Compute unnormalized probabilities and their sum. + float sSum = 0.0f; + for (int32_t ki = 0; ki < k; ki++) + { + sVal2[ki] = __expf(sVal2[ki] - sMaxLogit); // reuse sVal2 to hold exp probs + sSum += sVal2[ki]; + } + // Walk in descending-probability order; stop as soon as cumulative prob >= topP. + float cumProb = 0.0f; + for (int32_t ki = 0; ki < k; ki++) + { + cumProb += sVal2[ki] / sSum; + if (cumProb >= topP) + { + cutoff = ki + 1; // always keep at least this token + break; + } + } } + } - int64_t tokenIdx = parentList[bid * (topK * (depth - 1) + 1) + parentTbIdx]; - for (curPosition = 0; curPosition < draftTokenNum; ++curPosition) + // Step 4: Scatter selected logit values back to output. + // topKTmpIdBuf stores (rowOffset + vocabIdx); recover vocabIdx with % vocabSize. + auto const rowStride = rowId * stride; + for (int32_t ki = 0; ki < cutoff; ki++) + { + auto const candidateIdx = sId[ki]; + auto const globalIdx = topKTmpIdBuf[rowStride + candidateIdx]; + if (globalIdx >= 0) { - if (selectedIndex[bid * (draftTokenNum - 1) + curPosition] == tokenIdx) + auto const vocabIdx = globalIdx % vocabSize; + // sVal2 was overwritten with exp probs when topP < 1; we need the original logit. + // Re-read from the original tmp buffer — the stored value IS the logit (set in Stage 1). + // However sVal[candidateIdx] was zeroed during Stage 2 reduction; but + // topKTmpValBuf still holds the original value at that index (sVal points there). + // We stored the logit as sVal2[ite] = total.u BEFORE any exp, so if topP was not + // applied we can use sVal2[ki] directly. If topP was applied, sVal2[ki] now holds + // the exp prob — we cannot recover the logit. To handle both cases cleanly we use + // log(sVal2[ki]) + sMaxLogit when topPs was applied, otherwise sVal2[ki] directly. + float logitVal; + if (topPs != nullptr && topPs[rowId] < 1.0f) { - break; + // sVal2[ki] = exp(logit - sMaxLogit), so logit = log(sVal2[ki]) + sMaxLogit + logitVal = __logf(sVal2[ki]) + sMaxLogit; } - } - if (curPosition == draftTokenNum) - { - break; + else + { + logitVal = sVal2[ki]; // still the raw logit + } + outRow[vocabIdx] = logitVal; } } - 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) +#define CASE_K_PROB(K_MAX, BLOCK_SIZE_1_, BLOCK_SIZE_2_, BLOCKS_PER_BEAM_) \ + do \ + { \ + topKProbStage1 \ + <<>>( \ + logits, tmpLogProbs, topKTmpIdBuf, topKTmpValBuf, maxTopK, topKs, vocabSize); \ + topKProbStage2ForLogits \ + <<>>( \ + topKTmpIdBuf, topKTmpValBuf, outputLogits, maxTopK, topKs, topPs, vocabSize); \ + } while (0) + +// Host launcher: allocates workspace tensors internally and dispatches the two-stage kernels. +// logits [nRows, vocabSize] – temperature-scaled input (float or half) +// outputLogits [nRows, vocabSize] – output: -inf everywhere except selected top-k-p positions +// topKs [nRows] – per-row k values (int32, on device) +// topPs [nRows] or nullptr – per-row p values (float, on device) +// maxTopK – maximum k across all rows (CPU scalar, 1–1024) +template +void invokeTopKTopPMaskingForProbs(T const* logits, float* outputLogits, int32_t const* topKs, float const* topPs, + int32_t maxTopK, int32_t nRows, int32_t vocabSize, cudaStream_t stream) { - dim3 grid(batchSize); - dim3 block(numDraftTokens); + constexpr int32_t BLOCKS_PER_BEAM = 8; - if (treeMaskMode == TreeMaskMode::QLEN_ONLY_BITPACKING) + // Workspace buffers (allocated as CUDA device tensors via ATen). + auto opts = at::TensorOptions().dtype(torch::kFloat32).device(at::kCUDA); + auto tmpLogProbsTensor = torch::empty({nRows * vocabSize}, opts); + auto topKTmpIdBufTensor + = torch::empty({nRows * BLOCKS_PER_BEAM * maxTopK}, at::TensorOptions().dtype(torch::kInt32).device(at::kCUDA)); + // topKTmpValBuf uses the same dtype as T; we allocate as float and reinterpret for half if needed. + auto topKTmpValBufTensor = torch::empty({nRows * BLOCKS_PER_BEAM * maxTopK}, opts); + + T* tmpLogProbs = reinterpret_cast(tmpLogProbsTensor.data_ptr()); + int32_t* topKTmpIdBuf = topKTmpIdBufTensor.data_ptr(); + T* topKTmpValBuf = reinterpret_cast(topKTmpValBufTensor.data_ptr()); + + int32_t logMaxTopK = 0; + int32_t recursor = maxTopK - 1; + while (recursor >>= 1) { - 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); + ++logMaxTopK; } - else + + switch (logMaxTopK) { - buildDynamicTreeKernel<<>>(parentList, selectedIndex, static_cast(treeMask), - positions, retrieveIndex, retrieveNextToken, retrieveNextSibling, topK, depth, numDraftTokens); + case 0: + case 1: + case 2: + case 3: // 0 < maxTopK <= 16 + CASE_K_PROB(16, 128, 128, 8); + break; + case 4: // 16 < maxTopK <= 32 + CASE_K_PROB(32, 256, 128, 8); + break; + case 5: // 32 < maxTopK <= 64 + CASE_K_PROB(64, 256, 256, 8); + break; + case 6: + case 7: + case 8: + case 9: // 64 < maxTopK <= 1024 + CASE_K_PROB(1024, 256, 256, 8); + break; + default: TLLM_CHECK_WITH_INFO(false, "topKProbMasking supports 1 <= k <= 1024 but got k=%d", maxTopK); } +} - sync_check_cuda_error(stream); +#undef CASE_K_PROB + +namespace +{ +constexpr double kGreedyTempThreshold = 1e-4; + +bool isTopPEnabled(torch::optional const& topP) +{ + return topP.has_value() && topP->defined() && topP->lt(1.0).any().item(); } -//! \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) +torch::Tensor computeSoftmaxForProbOp(torch::Tensor logits) { - uint32_t bx = blockIdx.x; - uint32_t batchOffset = bx * numDraftTokens; + TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor"); + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); - // First-gen or dummy request: no valid tree, accept only the bonus token - if (treeValid != nullptr && !treeValid[bx]) + auto probs = logits.contiguous().to(torch::kFloat32); + auto stream = at::cuda::getCurrentCUDAStream(probs.device().index()); + + BiasSoftmaxParams biasSoftmaxParams; + biasSoftmaxParams.logits = probs.data_ptr(); + biasSoftmaxParams.probs = probs.data_ptr(); + biasSoftmaxParams.batchSize = static_cast(probs.size(0)); + biasSoftmaxParams.maxBatchSize = static_cast(probs.size(0)); + biasSoftmaxParams.maxBeamWidth = 1; + biasSoftmaxParams.vocabSize = static_cast(probs.size(1)); + biasSoftmaxParams.vocabSizePadded = static_cast(probs.size(1)); + biasSoftmaxParams.skipSoftMax = false; + biasSoftmaxParams.batchSlotsLogits = false; + biasSoftmaxParams.checkParams(); + + invokeAddBiasSoftMax(biasSoftmaxParams, stream); + return probs; +} + +struct DraftProbMaxFloatOp +{ + __device__ __forceinline__ float operator()(float a, float b) const + { + return a > b ? a : b; + } +}; + +template +__global__ void computeDraftProbsSkipAllKernel(float const* draftLogits, int32_t const* d2t, float* draftProbs, + int32_t nRows, int32_t draftVocabSize, int32_t targetVocabSize) +{ + int32_t const rowId = static_cast(blockIdx.x); + int32_t const tid = static_cast(threadIdx.x); + if (rowId >= nRows) { - 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; + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tempStorage; + __shared__ float sMaxLogit; + __shared__ float sExpSum; - // Root token: target prediction at root position - acceptToken[bx * numSpeculativeTokens] = targetPredict[batchOffset + lastAcceptedLocalIdx]; + float const* rowLogits = draftLogits + static_cast(rowId) * draftVocabSize; + float* rowProbs = draftProbs + static_cast(rowId) * targetVocabSize; - for (uint32_t j = 1; j < numSpeculativeTokens; ++j) + for (int32_t v = tid; v < targetVocabSize; v += BLOCK_SIZE) { - curIndex = retrieveNextToken[batchOffset + curIndex]; + rowProbs[v] = 0.0f; + } - while (curIndex != -1) + float localMax = -FLT_MAX; + for (int32_t v = tid; v < draftVocabSize; v += BLOCK_SIZE) + { + localMax = fmaxf(localMax, rowLogits[v]); + } + + float const blockMax = BlockReduce(tempStorage).Reduce(localMax, DraftProbMaxFloatOp{}); + if (tid == 0) + { + sMaxLogit = blockMax; + } + __syncthreads(); + + float localSum = 0.0f; + for (int32_t v = tid; v < draftVocabSize; v += BLOCK_SIZE) + { + localSum += __expf(rowLogits[v] - sMaxLogit); + } + + float const blockSum = BlockReduce(tempStorage).Sum(localSum); + if (tid == 0) + { + constexpr float kFloatSoftmaxEpsilon = 1e-6f; + sExpSum = blockSum + kFloatSoftmaxEpsilon; + } + __syncthreads(); + + for (int32_t v = tid; v < draftVocabSize; v += BLOCK_SIZE) + { + int64_t const targetIdx + = d2t != nullptr ? static_cast(v) + static_cast(d2t[v]) : static_cast(v); + if (targetIdx >= 0 && targetIdx < targetVocabSize) + { + rowProbs[targetIdx] = __expf(rowLogits[v] - sMaxLogit) / sExpSum; + } + } +} + +torch::Tensor computeDraftProbsSkipAllForDynamicTreeRejection(torch::Tensor const& draftLogits, int64_t batchSize, + SizeType32 const numDraftProbRows, SizeType32 const targetVocabSize, torch::optional const& d2t) +{ + auto const draftVocabSize = draftLogits.size(1); + bool const hasD2T = d2t.has_value() && d2t->defined(); + + auto draftLogitsFloat = draftLogits.contiguous().to(torch::kFloat32); + if (!hasD2T && draftVocabSize == targetVocabSize) + { + return computeSoftmaxForProbOp(draftLogitsFloat).reshape({batchSize, numDraftProbRows, targetVocabSize}); + } + + auto fullDraftProbs = torch::empty({draftLogitsFloat.size(0), targetVocabSize}, + torch::TensorOptions().dtype(torch::kFloat32).device(draftLogitsFloat.device())); + torch::Tensor d2tInt; + int32_t const* d2tPtr = nullptr; + if (hasD2T) + { + d2tInt = d2t->contiguous().to(torch::kInt32); + d2tPtr = d2tInt.data_ptr(); + } + + constexpr int32_t kBlockSize = 1024; + dim3 grid(draftLogitsFloat.size(0)); + dim3 block(kBlockSize); + auto stream = at::cuda::getCurrentCUDAStream(draftLogitsFloat.device().index()); + computeDraftProbsSkipAllKernel<<>>(draftLogitsFloat.data_ptr(), d2tPtr, + fullDraftProbs.data_ptr(), static_cast(draftLogitsFloat.size(0)), + static_cast(draftVocabSize), static_cast(targetVocabSize)); + sync_check_cuda_error(stream); + + return fullDraftProbs.reshape({batchSize, numDraftProbRows, targetVocabSize}); +} + +// Fast path for top-K (and optional top-P) filtering using torch::topk instead of a +// full vocab-size sort. kMax must be provided as a CPU integer (the caller computes it +// via topK.max().item() on the Python side). When kMax == 0 or kMax >= vocabSize the +// function falls back to the original sort-based path. +// +// Key advantages over the full-sort path: +// 1. torch::topk with small kMax is O(V * log kMax) vs O(V * log V) for full sort. +// 2. The topk index tensor is [nRows, kMax] instead of [nRows, V] — much smaller. +// 3. No scatter-back of sorted indices needed; masking is done directly on logits. +// 4. For combined top-K + top-P, softmax/cumsum are computed on kMax values (not V). +torch::Tensor applyTopKTopPForProbOp(torch::Tensor logits, torch::optional const& topK, + torch::optional const& topP, int32_t kMax) +{ + int64_t const vocabSize = logits.size(1); + // Host-only checks: the caller is expected to pass nullopt when filtering is fully + // disabled (see SpecMetadata.skip_top_k / skip_top_p). Probing the tensor contents + // via `.item()` here would force a host-device sync and break CUDA graph + // capture; the per-row `effectiveTopK` formula below already handles disabled rows. + bool const hasTopK = topK.has_value() && topK->defined(); + bool const hasTopP = topP.has_value() && topP->defined(); + + if (!hasTopK && !hasTopP) + { + return logits; + } + + torch::Tensor effectiveTopK; + if (hasTopK) + { + auto topKLong = topK->to(torch::kLong); + effectiveTopK + = torch::where(topKLong > 0, topKLong, torch::full_like(topKLong, vocabSize)).clamp_max(vocabSize); + } + + // Fast path uses `topk(kMax)` which is unsafe when any row has effective top-k > kMax + // (i.e. disabled rows expand to the full vocab). Detecting this requires a tensor + // reduction + `.item()`, which is incompatible with CUDA graph capture. Only + // probe when the caller explicitly opted into the fast path via kMax > 0 (today only + // the dynamic-tree caller, which is not graph-captured). + bool hasDisabledTopKRows = false; + if (hasTopK && kMax > 0 && kMax < vocabSize) + { + auto topKLong = topK->to(torch::kLong); + hasDisabledTopKRows = topKLong.le(0).any().item(); + } + + if (hasTopK && !hasDisabledTopKRows && kMax > 0 && kMax < vocabSize) + { + // Fast topk path ───────────────────────────────────────────────────────────── + // topKValues/topKIdx: [nRows, kMax], values in descending order + auto [topKValues, topKIdx] = logits.topk(kMax, /*dim=*/-1, /*largest=*/true, /*sorted=*/true); + + // validTopK[i, j]: True when position j falls within top-K[i] for row i + auto kArange = torch::arange(kMax, torch::TensorOptions().dtype(torch::kInt64).device(logits.device())) + .unsqueeze(0); // [1, kMax] + auto kVals = effectiveTopK.to(torch::kInt64).unsqueeze(1); // [nRows, 1] + auto validTopK = kArange < kVals; // [nRows, kMax] + + // Start with everything masked; scatter will unmark the kept positions. + auto mask = torch::ones( + {logits.size(0), vocabSize}, torch::TensorOptions().dtype(torch::kBool).device(logits.device())); + + if (hasTopP) + { + // Compute top-P on the kMax descending-sorted values only (much cheaper). + // Positions beyond K[i] are treated as -inf so their probability ≈ 0. + auto validTopKValues = topKValues.masked_fill(~validTopK, -std::numeric_limits::infinity()); + auto sortedProbs = validTopKValues.softmax(/*dim=*/-1); // [nRows, kMax] + auto cumsum = sortedProbs.cumsum(/*dim=*/-1); // [nRows, kMax] + // Mask positions where the cumulative probability *before* this token + // already reaches topP — i.e. we have enough probability mass already. + auto topPMask = (cumsum - sortedProbs) >= topP->unsqueeze(1); // [nRows, kMax] + topPMask.select(/*dim=*/1, /*index=*/0).fill_(false); // always keep the top-1 token + // combinedMask: True → mask this vocab position + // False → keep this vocab position + auto combinedMask = topPMask | (~validTopK); // [nRows, kMax] + mask.scatter_(/*dim=*/1, /*index=*/topKIdx, /*src=*/combinedMask); + } + else + { + // Top-K only: unmark the first K[i] positions (those within validTopK). + // ~validTopK is True for positions j >= K[i] → they should stay masked. + mask.scatter_(/*dim=*/1, /*index=*/topKIdx, /*src=*/(~validTopK)); + } + + return logits.masked_fill(mask, -std::numeric_limits::infinity()); + } + + // Fallback: full-sort path (used for top-P only, or when kMax == 0) ──────────── + auto sortResult = logits.sort(/*dim=*/-1, /*descending=*/false); + auto logitsSort = std::get<0>(sortResult); + auto logitsIdx = std::get<1>(sortResult); + + if (hasTopK) + { + auto topKMask = logitsSort.size(1) - effectiveTopK; + topKMask = topKMask.clamp_min(0); + auto topKThreshold = logitsSort.gather(1, topKMask.unsqueeze(1)); + auto mask = logitsSort < topKThreshold; + logitsSort.masked_fill_(mask, -std::numeric_limits::infinity()); + } + + if (hasTopP) + { + auto probsSort = logitsSort.softmax(/*dim=*/-1); + auto probsSum = probsSort.cumsum(/*dim=*/-1, /*dtype=*/probsSort.scalar_type()); + auto topPMask = probsSum <= (1.0 - topP->unsqueeze(1)); + topPMask.select(/*dim=*/1, /*index=*/logitsSort.size(1) - 1).fill_(false); + logitsSort.masked_fill_(topPMask, -std::numeric_limits::infinity()); + } + + return logitsSort.scatter(/*dim=*/-1, /*index=*/logitsIdx, /*src=*/logitsSort); +} + +} // namespace + +torch::Tensor computeProbsFromLogits(torch::Tensor const& logits, torch::Tensor const& temperatures, + torch::optional const& topK, torch::optional const& topP, bool skipTemperature, + int32_t kMax) +{ + TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor"); + TORCH_CHECK(temperatures.is_cuda(), "temperatures must be a CUDA tensor"); + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); + TORCH_CHECK(temperatures.dim() == 1, "temperatures must be a 1D tensor"); + TORCH_CHECK(logits.size(0) == temperatures.size(0), "logits and temperatures size mismatch"); + if (topK.has_value() && topK->defined()) + { + TORCH_CHECK(topK->is_cuda(), "top_k must be a CUDA tensor"); + TORCH_CHECK(topK->dim() == 1, "top_k must be a 1D tensor"); + TORCH_CHECK(topK->size(0) == logits.size(0), "top_k and logits size mismatch"); + } + if (topP.has_value() && topP->defined()) + { + TORCH_CHECK(topP->is_cuda(), "top_p must be a CUDA tensor"); + TORCH_CHECK(topP->dim() == 1, "top_p must be a 1D tensor"); + TORCH_CHECK(topP->size(0) == logits.size(0), "top_p and logits size mismatch"); + } + + auto const isGreedy = temperatures <= kGreedyTempThreshold; + auto const safeTemperatures = torch::where(isGreedy, torch::ones_like(temperatures), temperatures); + auto scaledLogits + = (skipTemperature ? logits : logits.div(safeTemperatures.unsqueeze(1))).contiguous().to(torch::kFloat32); + + int64_t const vocabSize = scaledLogits.size(1); + int64_t const nRows = scaledLogits.size(0); + // Host-only presence checks; see comment in applyTopKTopPForProbOp() for why we + // avoid probing tensor contents (would sync and break CUDA graph capture). + bool const hasTopKPresence = topK.has_value() && topK->defined(); + bool const hasTopPPresence = topP.has_value() && topP->defined(); + + // The kernel path produces -inf for rows whose top_k value is 0, so it is only + // safe when every row has an active top_k filter. Determining that requires a + // host-device sync, so only probe when the caller has opted into the kernel + // path (kMax > 0). The kMax > 0 callers (dynamic-tree) are not graph-captured. + bool useKernelPath = false; + if (hasTopKPresence && kMax > 0 && kMax < vocabSize) + { + useKernelPath = torch::logical_and(topK->gt(0), topK->lt(vocabSize)).any().item(); + } + + torch::Tensor maskedLogits; + if (useKernelPath) + { + // Two-stage CUDA top-k/top-p masking (mirrors invokeBatchTopKSampling). + maskedLogits = torch::empty_like(scaledLogits); + auto topKForKernel = topK->to(torch::kInt32).contiguous(); + auto topPForKernel = hasTopPPresence ? topP->to(torch::kFloat32).contiguous() : torch::Tensor(); + auto stream = at::cuda::getCurrentCUDAStream(scaledLogits.device().index()); + invokeTopKTopPMaskingForProbs(scaledLogits.data_ptr(), maskedLogits.data_ptr(), + topKForKernel.data_ptr(), hasTopPPresence ? topPForKernel.data_ptr() : nullptr, kMax, + static_cast(nRows), static_cast(vocabSize), stream); + } + else + { + // Fallback: PyTorch-based sort path (top-P only or kMax == 0). + maskedLogits = applyTopKTopPForProbOp(scaledLogits, topK, topP, kMax); + } + + auto probs = computeSoftmaxForProbOp(maskedLogits); + + auto argmaxIds = maskedLogits.argmax(/*dim=*/-1, /*keepdim=*/true); + auto oneHot = torch::zeros_like(probs).scatter_(1, argmaxIds, 1.0); + return torch::where(isGreedy.unsqueeze(1), oneHot, probs); +} + +torch::Tensor computeDraftProbsForDynamicTreeRejection(torch::Tensor const& draftLogits, + torch::Tensor const& temperatures, SizeType32 const numDraftProbRows, torch::optional const& topK, + torch::optional const& topP, SizeType32 const targetVocabSize, bool skipTemperature, + torch::optional const& d2t, SizeType32 const kMax, bool skipAllSamplingParams) +{ + TORCH_CHECK(draftLogits.is_cuda(), "draftLogits must be a CUDA tensor"); + TORCH_CHECK(temperatures.is_cuda(), "temperatures must be a CUDA tensor"); + TORCH_CHECK(draftLogits.dim() == 2, "draftLogits must be a 2D tensor"); + TORCH_CHECK(temperatures.dim() == 1, "temperatures must be a 1D tensor"); + TORCH_CHECK(numDraftProbRows > 0, "numDraftProbRows must be positive"); + + auto const batchSize = temperatures.size(0); + auto const draftVocabSize = draftLogits.size(1); + + TORCH_CHECK(batchSize > 0, "batchSize must be positive"); + TORCH_CHECK( + draftLogits.size(0) == batchSize * numDraftProbRows, "draftLogits row count does not match numDraftProbRows"); + TORCH_CHECK(targetVocabSize >= draftVocabSize, "targetVocabSize must be >= draft vocab size"); + + if (topK.has_value() && topK->defined()) + { + TORCH_CHECK(topK->is_cuda(), "top_k must be a CUDA tensor"); + TORCH_CHECK(topK->dim() == 1, "top_k must be a 1D tensor"); + TORCH_CHECK(topK->size(0) == batchSize, "top_k size mismatch"); + } + if (topP.has_value() && topP->defined()) + { + TORCH_CHECK(topP->is_cuda(), "top_p must be a CUDA tensor"); + TORCH_CHECK(topP->dim() == 1, "top_p must be a 1D tensor"); + TORCH_CHECK(topP->size(0) == batchSize, "top_p size mismatch"); + } + if (d2t.has_value() && d2t->defined()) + { + TORCH_CHECK(d2t->is_cuda(), "d2t must be a CUDA tensor"); + TORCH_CHECK(d2t->dim() == 1, "d2t must be a 1D tensor"); + TORCH_CHECK(d2t->size(0) >= draftVocabSize, "d2t size mismatch"); + } + + if (skipAllSamplingParams) + { + return computeDraftProbsSkipAllForDynamicTreeRejection( + draftLogits, batchSize, numDraftProbRows, targetVocabSize, d2t); + } + + auto draftTemps = temperatures.repeat_interleave(numDraftProbRows); + auto draftTopK = topK.has_value() && topK->defined() + ? torch::optional(topK->repeat_interleave(numDraftProbRows)) + : torch::optional(); + auto draftTopP = isTopPEnabled(topP) ? torch::optional(topP->repeat_interleave(numDraftProbRows)) + : torch::optional(); + + auto draftProbs = computeProbsFromLogits(draftLogits, draftTemps, draftTopK, draftTopP, skipTemperature, kMax) + .reshape({batchSize, numDraftProbRows, draftVocabSize}); + + if (draftVocabSize == targetVocabSize) + { + return draftProbs; + } + + auto fullDraftProbs = torch::zeros({batchSize, numDraftProbRows, targetVocabSize}, + torch::TensorOptions().dtype(torch::kFloat32).device(draftProbs.device())); + if (d2t.has_value() && d2t->defined()) + { + auto srcIdx + = torch::arange(draftVocabSize, torch::TensorOptions().dtype(torch::kInt64).device(draftProbs.device())); + auto targetIdx = srcIdx + d2t->slice(0, 0, draftVocabSize).to(torch::kInt64); + auto expandedTargetIdx + = targetIdx.view({1, 1, draftVocabSize}).expand({batchSize, numDraftProbRows, draftVocabSize}); + fullDraftProbs.scatter_(2, expandedTargetIdx, draftProbs); + } + else + { + fullDraftProbs.slice(/*dim=*/2, /*start=*/0, /*end=*/draftVocabSize).copy_(draftProbs); + } + + return fullDraftProbs; +} + +std::tuple computeTargetProbsForDynamicTreeRejection( + torch::Tensor const& targetLogits, torch::Tensor const& temperatures, SizeType32 const numDraftTokens, + torch::optional const& topK, torch::optional const& topP, bool skipTemperature, + SizeType32 const kMax, bool skipAllSamplingParams) +{ + TORCH_CHECK(targetLogits.is_cuda(), "targetLogits must be a CUDA tensor"); + TORCH_CHECK(temperatures.is_cuda(), "temperatures must be a CUDA tensor"); + TORCH_CHECK(targetLogits.dim() == 2, "targetLogits must be a 2D tensor"); + TORCH_CHECK(temperatures.dim() == 1, "temperatures must be a 1D tensor"); + TORCH_CHECK(numDraftTokens > 1, "numDraftTokens must be greater than 1"); + + auto const batchSize = temperatures.size(0); + auto const targetVocabSize = targetLogits.size(1); + auto const nRows = batchSize * numDraftTokens; + + TORCH_CHECK(batchSize > 0, "batchSize must be positive"); + TORCH_CHECK( + targetLogits.size(0) == batchSize * numDraftTokens, "targetLogits row count does not match numDraftTokens"); + + if (topK.has_value() && topK->defined()) + { + TORCH_CHECK(topK->is_cuda(), "top_k must be a CUDA tensor"); + TORCH_CHECK(topK->dim() == 1, "top_k must be a 1D tensor"); + TORCH_CHECK(topK->size(0) == batchSize, "top_k size mismatch"); + } + if (topP.has_value() && topP->defined()) + { + TORCH_CHECK(topP->is_cuda(), "top_p must be a CUDA tensor"); + TORCH_CHECK(topP->dim() == 1, "top_p must be a 1D tensor"); + TORCH_CHECK(topP->size(0) == batchSize, "top_p size mismatch"); + } + + if (skipAllSamplingParams) + { + auto targetSupportIndices + = torch::empty({0}, torch::TensorOptions().dtype(torch::kInt32).device(targetLogits.device())); + auto targetSupportLengths + = torch::empty({0}, torch::TensorOptions().dtype(torch::kInt32).device(targetLogits.device())); + auto targetProbs = computeSoftmaxForProbOp(targetLogits); + return std::make_tuple(targetProbs.reshape({batchSize, numDraftTokens, targetVocabSize}), targetSupportIndices, + targetSupportLengths); + } + + auto targetTemps = temperatures.repeat_interleave(numDraftTokens); + auto targetTopK = topK.has_value() && topK->defined() + ? torch::optional(topK->repeat_interleave(numDraftTokens)) + : torch::optional(); + auto targetTopP = isTopPEnabled(topP) ? torch::optional(topP->repeat_interleave(numDraftTokens)) + : torch::optional(); + + bool const hasTopK = targetTopK.has_value() && targetTopK->defined(); + bool const hasTopP = isTopPEnabled(targetTopP); + bool const hasFiltering = hasTopK || hasTopP; + torch::Tensor effectiveTargetTopK; + bool hasDisabledTopKRows = false; + + auto const isGreedy = targetTemps <= kGreedyTempThreshold; + auto const safeTargetTemps = torch::where(isGreedy, torch::ones_like(targetTemps), targetTemps); + auto scaledTargetLogits = (skipTemperature ? targetLogits : targetLogits.div(safeTargetTemps.unsqueeze(1))) + .contiguous() + .to(torch::kFloat32); + + if (hasTopK) + { + auto targetTopKLong = targetTopK->to(torch::kLong); + effectiveTargetTopK + = torch::where(targetTopKLong > 0, targetTopKLong, torch::full_like(targetTopKLong, targetVocabSize)) + .clamp_max(targetVocabSize); + hasDisabledTopKRows = targetTopKLong.le(0).any().item(); + } + + torch::Tensor maskedTargetLogits; + torch::Tensor targetSupportIndices; + torch::Tensor targetSupportLengths; + + if (!hasFiltering) + { + // No filtering: use full-vocab probs; sparse support is not applicable. + maskedTargetLogits = scaledTargetLogits; + targetSupportIndices + = torch::empty({0}, torch::TensorOptions().dtype(torch::kInt32).device(targetLogits.device())); + targetSupportLengths + = torch::empty({0}, torch::TensorOptions().dtype(torch::kInt32).device(targetLogits.device())); + } + else if (hasTopK && !hasDisabledTopKRows && kMax > 0 && static_cast(kMax) < targetVocabSize) + { + // Fast two-stage CUDA path for masked logits. + maskedTargetLogits = torch::empty_like(scaledTargetLogits); + auto topKForKernel = effectiveTargetTopK.to(torch::kInt32).contiguous(); + auto topPForKernel = hasTopP ? targetTopP->to(torch::kFloat32).contiguous() : torch::Tensor(); + auto stream = at::cuda::getCurrentCUDAStream(scaledTargetLogits.device().index()); + invokeTopKTopPMaskingForProbs(scaledTargetLogits.data_ptr(), maskedTargetLogits.data_ptr(), + topKForKernel.data_ptr(), hasTopP ? topPForKernel.data_ptr() : nullptr, kMax, + static_cast(nRows), static_cast(targetVocabSize), stream); + + // Extract support indices: the finite positions after masking (at most kMax per row). + auto [topKVals, topKIdx] = maskedTargetLogits.topk(kMax, /*dim=*/-1, /*largest=*/true, /*sorted=*/true); + auto validMask = topKVals.isfinite(); // [nRows, kMax] + auto supportLengthsLong = validMask.sum(/*dim=*/-1, /*keepdim=*/false, torch::kLong); // [nRows] + auto supportIndicesRaw + = torch::where(validMask, topKIdx.to(torch::kInt32), torch::full_like(topKIdx.to(torch::kInt32), -1)); + + targetSupportIndices = supportIndicesRaw.reshape({batchSize, numDraftTokens, kMax}); + targetSupportLengths = supportLengthsLong.to(torch::kInt32).reshape({batchSize, numDraftTokens}); + } + else + { + // Sort-based fallback: top-P only, or kMax == 0 / kMax >= vocabSize. + auto sortResult = scaledTargetLogits.sort(/*dim=*/-1, /*descending=*/false); + auto logitsSort = std::get<0>(sortResult); + auto logitsIdx = std::get<1>(sortResult); + + if (hasTopK) + { + auto topKMask = logitsSort.size(1) - effectiveTargetTopK; + topKMask = topKMask.clamp_min(0); + auto topKThreshold = logitsSort.gather(1, topKMask.unsqueeze(1)); + auto mask = logitsSort < topKThreshold; + logitsSort.masked_fill_(mask, -std::numeric_limits::infinity()); + } + + if (hasTopP) + { + auto probsSort = logitsSort.softmax(/*dim=*/-1); + auto probsSum = probsSort.cumsum(/*dim=*/-1, /*dtype=*/probsSort.scalar_type()); + auto topPMask = probsSum <= (1.0 - targetTopP->unsqueeze(1)); + topPMask.select(/*dim=*/1, /*index=*/logitsSort.size(1) - 1).fill_(false); + logitsSort.masked_fill_(topPMask, -std::numeric_limits::infinity()); + } + + maskedTargetLogits = logitsSort.scatter(/*dim=*/-1, /*index=*/logitsIdx, /*src=*/logitsSort); + + // Compact support indices: finite values are at the END of the ascending-sorted logitsSort. + auto supportLengthsLong + = logitsSort.isfinite().sum(/*dim=*/-1, /*keepdim=*/false, /*dtype=*/torch::kLong); // [nRows] + auto supportLengths1D = supportLengthsLong.to(torch::kInt32); + + int64_t maxSupportSize = targetVocabSize; + if (hasTopK && effectiveTargetTopK.defined()) + { + maxSupportSize = std::min(targetVocabSize, effectiveTargetTopK.max().item()); + } + + auto compactPositions + = torch::arange(maxSupportSize, torch::TensorOptions().dtype(torch::kLong).device(targetLogits.device())) + .unsqueeze(0) + .expand({nRows, maxSupportSize}); + auto supportStart = targetVocabSize - supportLengthsLong.unsqueeze(1); + auto gatherPositions = (supportStart + compactPositions).clamp_max(targetVocabSize - 1); + auto gatheredSupportIndices = logitsIdx.gather(1, gatherPositions); + auto validMask = compactPositions < supportLengthsLong.unsqueeze(1); + auto invalidFill = torch::full_like(gatheredSupportIndices, -1L); + targetSupportIndices = torch::where(validMask, gatheredSupportIndices, invalidFill) + .to(torch::kInt32) + .reshape({batchSize, numDraftTokens, maxSupportSize}); + targetSupportLengths = supportLengths1D.reshape({batchSize, numDraftTokens}); + } + + auto targetProbs = computeSoftmaxForProbOp(maskedTargetLogits); + + auto argmaxIds = maskedTargetLogits.argmax(/*dim=*/-1, /*keepdim=*/true); + auto oneHot = torch::zeros_like(targetProbs).scatter_(1, argmaxIds, 1.0); + targetProbs = torch::where(isGreedy.unsqueeze(1), oneHot, targetProbs); + + return std::make_tuple( + targetProbs.reshape({batchSize, numDraftTokens, targetVocabSize}), targetSupportIndices, targetSupportLengths); +} + +//! \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); +} + +__global__ void buildDraftProbIndicesKernel( + int64_t const* topkScoreIndices, int32_t* draftProbIndices, SizeType32 topK, SizeType32 numDraftTokens) +{ + int32_t const batchIdx = blockIdx.x; + int32_t const tokenIdx = threadIdx.x; + + if (tokenIdx > numDraftTokens) + { + return; + } + + int32_t* draftProbIndicesRow = draftProbIndices + batchIdx * (numDraftTokens + 1); + + if (tokenIdx == 0) + { + draftProbIndicesRow[0] = 0; + return; + } + + int64_t const histIdx = topkScoreIndices[batchIdx * numDraftTokens + (tokenIdx - 1)]; + int32_t draftProbRow = 0; + + if (histIdx >= topK) + { + int64_t const relative = histIdx - topK; + int64_t const depthBucket = relative / (topK * topK); + int64_t const parentK = (relative % (topK * topK)) / topK; + draftProbRow = static_cast(1 + depthBucket * topK + parentK); + } + + draftProbIndicesRow[tokenIdx] = draftProbRow; +} + +void invokeBuildDraftProbIndices(int64_t const* topkScoreIndices, int32_t* draftProbIndices, SizeType32 batchSize, + SizeType32 topK, SizeType32 numDraftTokens, cudaStream_t stream) +{ + dim3 const grid(batchSize); + dim3 const block(numDraftTokens + 1); + + buildDraftProbIndicesKernel<<>>(topkScoreIndices, draftProbIndices, topK, 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/acceptToken +//! (= numSpecStep = max_path_len, >= 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); +} + +// ------------------------------------------------------------ +// Background: Speculative Sampling Theory +// ------------------------------------------------------------ +// +// Goal: reuse draft model samples to speed up generation while keeping the +// final output distribution strictly equal to the target distribution q. +// +// For a given token x: +// p(x) = draft_probs[x] (draft model probability) +// q(x) = target_probs[x] (target model probability) +// +// Step 1 - The draft model proposes token x sampled from p. +// Step 2 - Accept x with probability min(1, q(x)/p(x)). +// Equivalently: accept when u * p(x) < q(x), where u ~ Uniform(0,1). +// +// Why does this work? +// x is proposed with probability p(x) and then accepted with probability +// min(1, q(x)/p(x)), so its total probability mass reaching the output is: +// p(x) * min(1, q(x)/p(x)) = min(p(x), q(x)) +// +// This covers only the min(p, q) portion of the target mass. +// The remaining portion q - min(p, q) = relu(q - p) is not yet covered. +// +// Therefore, if the draft token is rejected, we must resample from the +// residual distribution relu(q - p) (normalised) to fill the gap and +// restore the full target distribution. +// +// Example: +// p = [0.6, 0.3, 0.1] tokens [A, B, C] +// q = [0.2, 0.5, 0.3] +// +// Accept probabilities: +// A: min(1, 0.2/0.6) = 1/3 B: min(1, 0.5/0.3) = 1 C: min(1, 0.3/0.1) = 1 +// +// Case 1 - draft proposes A (prob 0.6): +// Accept (1/3): contributes 0.6 * 1/3 = 0.2 to output A. +// Reject (2/3): total rejected mass = 0.6 * 2/3 = 0.4. +// relu(q-p) = [0, 0.2, 0.2] -> normalised [0, 0.5, 0.5] +// contributes 0.4*0.5 = 0.2 to B and 0.4*0.5 = 0.2 to C. +// Case 2 - draft proposes B (prob 0.3): always accepted -> 0.3 to B. +// Case 3 - draft proposes C (prob 0.1): always accepted -> 0.1 to C. +// +// Final output distribution: +// A = 0.2, B = 0.3 + 0.2 = 0.5, C = 0.1 + 0.2 = 0.3 -> exactly q. +// +// Tree extension: +// The same logic applies depth-by-depth along the draft tree. At each +// depth the kernel tries siblings in score order; the first accepted +// sibling extends the current path. If every sibling at a depth is +// rejected the kernel samples a correction token from relu(q-p) and +// terminates traversal for that request. +// ------------------------------------------------------------ + +#include + +/// Map curand_uniform (0, 1] to [0, 1) so that cumulative-sum sampling +/// never falls off the end of a probability distribution due to float32 +/// rounding. 1.0 is mapped to 0.0 (probability mass epsilon). +__device__ __forceinline__ float curand_uniform_open_right(curandStatePhilox4_32_10_t& state) +{ + float u = curand_uniform(&state); // (0, 1] + return u < 1.0f ? u : 0.0f; // [0, 1) +} + +__device__ int64_t sampleFromDistribution(curandStatePhilox4_32_10_t& state, float const* probs, uint32_t vocabSize) +{ + float r = curand_uniform_open_right(state); // [0, 1) + float cumsum = 0.0f; + int64_t sampledTok = 0; + + for (uint32_t v = 0; v < vocabSize; ++v) + { + cumsum += probs[v]; + if (r < cumsum) + { + sampledTok = static_cast(v); + return sampledTok; + } + } + + // Float32 cumsum may not reach 1.0 for large vocabs. + // Fall back to the last token with positive probability. + for (int64_t v = static_cast(vocabSize) - 1; v >= 0; --v) + { + if (probs[v] > 0.0f) + { + return v; + } + } + return static_cast(vocabSize) - 1; +} + +__device__ int64_t sampleFromIndexedDistribution(curandStatePhilox4_32_10_t& state, float const* probs, + int32_t const* supportIndices, uint32_t supportSize, uint32_t vocabSize) +{ + float r = curand_uniform_open_right(state); // [0, 1) + float cumsum = 0.0f; + int64_t sampledTok = static_cast(vocabSize) - 1; + + for (uint32_t i = 0; i < supportSize; ++i) + { + int32_t const tok = supportIndices[i]; + cumsum += probs[tok]; + if (r < cumsum) + { + return static_cast(tok); + } + } + + // Fallback: last support token with positive probability. + for (int64_t i = static_cast(supportSize) - 1; i >= 0; --i) + { + if (probs[supportIndices[i]] > 0.0f) + { + return static_cast(supportIndices[i]); + } + } + return sampledTok; +} + +struct MinInt32Op +{ + __device__ __forceinline__ int32_t operator()(int32_t a, int32_t b) const + { + return a < b ? a : b; + } +}; + +struct MaxInt32Op +{ + __device__ __forceinline__ int32_t operator()(int32_t a, int32_t b) const + { + return a > b ? a : b; + } +}; + +struct MaxFloatOp +{ + __device__ __forceinline__ float operator()(float a, float b) const + { + return a > b ? a : b; + } +}; + +struct SoftmaxStats +{ + float maxVal; + float sumVal; + int32_t argmax; +}; + +//! \param acceptIndex [out] accepted path as tree positions [bs, numSpecStep]. int64. +//! \param acceptTokenNum [out] number of accepted draft tokens (excl. root) [bs]. int64. +//! \param acceptToken [out] emitted token ids [bs, numSpecStep]. int64. +//! \param candidates [in] candidate token ids [bs, numDraftTokens]; col 0 = root. int64. +//! \param draftProbs [in] unique draft probs [bs, numDraftProbRows, vocabSize]. float32. +//! \param targetProbs [in] target probs [bs, numDraftTokens, vocabSize]; index 0 = root. float32. +//! \param targetSupportIndices [in] compact target support per tree position +//! [bs, numDraftTokens, maxTargetSupportSize]. int32, or nullptr. +//! \param targetSupportLengths [in] support length per tree position [bs, numDraftTokens]. int32, or nullptr. +//! \param draftProbIndices [in] tree position -> draftProbs row [bs, numDraftTokens], root unused. int32. +//! \param retrieveNextToken [in] first-child pointer [bs, numDraftTokens], -1=none. int32. +//! \param retrieveNextSibling [in] next-sibling pointer [bs, numDraftTokens], -1=none. int32. +//! \param treeValid [in] per-request tree validity flag [bs]. bool. +//! \param batchSize batch size. +//! \param numDraftProbRows unique draft-prob rows per request. +//! \param maxTargetSupportSize support-array width. Zero when targetSupportIndices is null. +//! \param numSpecStep second dim of acceptIndex/acceptToken +//! (= max_path_len = max_draft_len + 1). +//! \param numDraftTokens total tree nodes per batch (including root). +//! \param vocabSize vocabulary size. +//! \param seed [1] int64 on GPU. Philox RNG seed. +//! \param offset [1] int64 on GPU. Philox RNG offset. +template +__global__ void verifyDynamicTreeRejectionKernel(int64_t* acceptIndex, int64_t* acceptTokenNum, int64_t* acceptToken, + int64_t const* candidates, float const* draftInputs, float const* targetInputs, int32_t const* targetSupportIndices, + int32_t const* targetSupportLengths, int32_t const* draftProbIndices, int32_t const* retrieveNextToken, + int32_t const* retrieveNextSibling, bool const* treeValid, uint32_t batchSize, uint32_t numDraftProbRows, + uint32_t maxTargetSupportSize, uint32_t numSpeculativeTokens, uint32_t numDraftTokens, uint32_t vocabSize, + uint32_t draftVocabSize, int32_t const* targetToDraft, int64_t const* seed, int64_t const* offset, + float const* temperatures) +{ + uint32_t bx = blockIdx.x; + int32_t const tid = static_cast(threadIdx.x); + constexpr uint32_t kVecSize = 4; + if (bx >= batchSize) + { + return; + } + + using BlockReduce = cub::BlockReduce; + using BlockReduceInt = cub::BlockReduce; + using BlockScan = cub::BlockScan; + + __shared__ union + { + typename BlockReduce::TempStorage reduce; + typename BlockReduceInt::TempStorage reduceInt; + typename BlockScan::TempStorage scan; + } tempStorage; + + __shared__ int32_t sLastAcceptedLocalIdx; + __shared__ uint32_t sNumAcceptedTokens; + __shared__ int32_t sFirstChild; + __shared__ bool sHasTerminalToken; + __shared__ float sDiffSum; + __shared__ float sTargetMass; + __shared__ float sPrefixBase; + __shared__ int32_t sWinnerIndex; + __shared__ int32_t sLastValidIndex; + __shared__ int64_t sSampledToken; + __shared__ float sLogitsMax; + __shared__ float sLogitsSum; + __shared__ int32_t sLogitsArgmax; + + // The first sibling that passes the rejection test at the current depth. + __shared__ int32_t sAccSibIdx; + __shared__ int64_t sAccSibTok; + __shared__ int32_t sNumAccSiblings; + + uint32_t batchOffset = bx * numDraftTokens; + + curandStatePhilox4_32_10_t state; + if (tid == 0) + { + curand_init( + static_cast(seed[0]), static_cast(bx), static_cast(offset[0]), &state); + } + __syncthreads(); + bool const hasCompactTargetSupport = targetSupportIndices != nullptr && targetSupportLengths != nullptr; + bool const isGreedyRequest + = USE_LOGITS && temperatures != nullptr && temperatures[bx] <= static_cast(kGreedyTempThreshold); + float const* draftProbs = draftInputs; + float const* targetProbs = targetInputs; + uint32_t const draftRowStride = USE_LOGITS ? draftVocabSize : vocabSize; + + auto canVectorizeLoad = [&](float const* probs, uint32_t rowSize) -> bool + { + constexpr uint32_t kLoadAlignmentBytes = kVecSize * sizeof(float); + return (rowSize % kVecSize == 0) && (reinterpret_cast(probs) % kLoadAlignmentBytes == 0); + }; + + auto loadProbVec = [&](float const* probs, uint32_t base, bool useVectorizedLoads, uint32_t rowSize) + { + flashinfer::vec_t probVec; + probVec.fill(0.0f); + if (useVectorizedLoads && base + kVecSize <= rowSize) + { + probVec.cast_load(probs + base); + } + else + { +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + uint32_t const v = base + j; + if (v < rowSize) + { + probVec[j] = probs[v]; + } + } + } + return probVec; + }; + + auto computeLogitsStats = [&](float const* logitsRow, uint32_t rowSize) -> SoftmaxStats + { + float threadMax = -FLT_MAX; + for (uint32_t v = static_cast(tid); v < rowSize; v += BLOCK_SIZE) + { + threadMax = fmaxf(threadMax, logitsRow[v]); + } + + float const blockMax = BlockReduce(tempStorage.reduce).Reduce(threadMax, MaxFloatOp{}); + if (tid == 0) + { + sLogitsMax = blockMax; + } + __syncthreads(); + + float threadSum = 0.0f; + int32_t localArgmax = static_cast(rowSize); + for (uint32_t v = static_cast(tid); v < rowSize; v += BLOCK_SIZE) + { + float const logit = logitsRow[v]; + threadSum += __expf(logit - sLogitsMax); + if (logit == sLogitsMax && localArgmax == static_cast(rowSize)) + { + localArgmax = static_cast(v); + } + } + + float const blockSum = BlockReduce(tempStorage.reduce).Sum(threadSum); + __syncthreads(); + int32_t const blockArgmax = BlockReduceInt(tempStorage.reduceInt).Reduce(localArgmax, MinInt32Op{}); + if (tid == 0) + { + sLogitsSum = blockSum; + sLogitsArgmax = blockArgmax; + } + __syncthreads(); + + return SoftmaxStats{sLogitsMax, sLogitsSum, sLogitsArgmax}; + }; + + auto probFromLogits = [&](float const* logitsRow, uint32_t tokenId, uint32_t rowSize, SoftmaxStats stats) -> float + { + if (tokenId >= rowSize) + { + return 0.0f; + } + if (isGreedyRequest) + { + return tokenId == static_cast(stats.argmax) ? 1.0f : 0.0f; + } + constexpr float kFloatSoftmaxEpsilon = 1e-6f; + return __expf(logitsRow[tokenId] - stats.maxVal) / (stats.sumVal + kFloatSoftmaxEpsilon); + }; + + auto targetTokenToDraftToken = [&](uint32_t targetTokenId) -> int32_t + { + if (targetTokenId >= vocabSize) + { + return -1; + } + if (targetToDraft != nullptr) + { + return targetToDraft[targetTokenId]; + } + return targetTokenId < draftVocabSize ? static_cast(targetTokenId) : -1; + }; + + auto draftProbFromTargetToken + = [&](float const* draftLogitsRow, uint32_t targetTokenId, SoftmaxStats stats) -> float + { + int32_t const draftTokenId = targetTokenToDraftToken(targetTokenId); + if (draftTokenId < 0) { - int32_t draftLocalIdx = retrieveIndex[batchOffset + curIndex]; - int64_t draftTokenId = candidates[batchOffset + curIndex]; - int64_t targetTokenId = targetPredict[batchOffset + lastAcceptedLocalIdx]; + return 0.0f; + } + return probFromLogits(draftLogitsRow, static_cast(draftTokenId), draftVocabSize, stats); + }; - if (draftTokenId == targetTokenId) + auto sampleProbTile = [&](float(&value)[kVecSize], uint32_t base) -> bool + { + float const tileSum = BlockReduce(tempStorage.reduce).template Sum(value); + if (tid == 0) + { + sDiffSum = tileSum; + } + __syncthreads(); + + int32_t localLastValid = -1; +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + uint32_t const v = base + j; + if (v < vocabSize && value[j] > 0.0f) + { + localLastValid = static_cast(v); + } + } + int32_t const blockLastValid = BlockReduceInt(tempStorage.reduceInt).Reduce(localLastValid, MaxInt32Op{}); + if (tid == 0 && blockLastValid >= 0) + { + sLastValidIndex = blockLastValid; + } + __syncthreads(); + + if (sPrefixBase + sDiffSum > sTargetMass) + { + float inclusive[kVecSize]; + BlockScan(tempStorage.scan).template InclusiveSum(value, inclusive); + __syncthreads(); + + int32_t localWinner = static_cast(vocabSize); +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + uint32_t const v = base + j; + if (v < vocabSize && value[j] > 0.0f && sPrefixBase + inclusive[j] > sTargetMass) + { + localWinner = static_cast(v); + break; + } + } + + int32_t const blockWinner = BlockReduceInt(tempStorage.reduceInt).Reduce(localWinner, MinInt32Op{}); + if (tid == 0) + { + sWinnerIndex = blockWinner; + sSampledToken = blockWinner < static_cast(vocabSize) ? static_cast(blockWinner) + : static_cast(vocabSize) - 1; + } + __syncthreads(); + return true; + } + + if (tid == 0) + { + sPrefixBase += sDiffSum; + } + __syncthreads(); + return false; + }; + + auto sampleTargetFullVocab = [&](float const* tProbs) + { + bool const useVectorizedLoads = canVectorizeLoad(tProbs, vocabSize); + uint32_t const numIters = (vocabSize + BLOCK_SIZE * kVecSize - 1) / (BLOCK_SIZE * kVecSize); + if (tid == 0) + { + sPrefixBase = 0.0f; + sWinnerIndex = static_cast(vocabSize); + sLastValidIndex = -1; + sSampledToken = static_cast(vocabSize) - 1; + sTargetMass = curand_uniform_open_right(state); + } + __syncthreads(); + +#pragma unroll 2 + for (uint32_t i = 0; i < numIters; ++i) + { + uint32_t const base = (i * BLOCK_SIZE + static_cast(tid)) * kVecSize; + auto const qVec = loadProbVec(tProbs, base, useVectorizedLoads, vocabSize); + float value[kVecSize]; +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + value[j] = qVec[j]; + } + + if (sampleProbTile(value, base)) + { + break; + } + } + + if (tid == 0 && sWinnerIndex >= static_cast(vocabSize) && sLastValidIndex >= 0) + { + sSampledToken = static_cast(sLastValidIndex); + } + __syncthreads(); + }; + + auto sampleResidualFullVocab = [&](float const* tProbs, float const* dProbs) + { + bool const useVectorizedTargetLoads = canVectorizeLoad(tProbs, vocabSize); + bool const useVectorizedDraftLoads = canVectorizeLoad(dProbs, vocabSize); + uint32_t const numIters = (vocabSize + BLOCK_SIZE * kVecSize - 1) / (BLOCK_SIZE * kVecSize); + float totalSum = 0.0f; +#pragma unroll 2 + for (uint32_t i = 0; i < numIters; ++i) + { + uint32_t const base = (i * BLOCK_SIZE + static_cast(tid)) * kVecSize; + auto const qVec = loadProbVec(tProbs, base, useVectorizedTargetLoads, vocabSize); + auto const pVec = loadProbVec(dProbs, base, useVectorizedDraftLoads, vocabSize); + float value[kVecSize]; +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + value[j] = fmaxf(qVec[j] - pVec[j], 0.0f); + } + totalSum += BlockReduce(tempStorage.reduce).template Sum(value); + __syncthreads(); + } + + if (tid == 0) + { + sDiffSum = totalSum; + sPrefixBase = 0.0f; + sWinnerIndex = static_cast(vocabSize); + sLastValidIndex = -1; + sSampledToken = static_cast(vocabSize) - 1; + sTargetMass = totalSum > 1e-10f ? curand_uniform_open_right(state) * totalSum : 0.0f; + if (totalSum <= 1e-10f) + { + sSampledToken = sampleFromDistribution(state, tProbs, vocabSize); + } + } + __syncthreads(); + + if (sDiffSum <= 1e-10f) + { + return; + } + +#pragma unroll 2 + for (uint32_t i = 0; i < numIters; ++i) + { + uint32_t const base = (i * BLOCK_SIZE + static_cast(tid)) * kVecSize; + auto const qVec = loadProbVec(tProbs, base, useVectorizedTargetLoads, vocabSize); + auto const pVec = loadProbVec(dProbs, base, useVectorizedDraftLoads, vocabSize); + float value[kVecSize]; +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + value[j] = fmaxf(qVec[j] - pVec[j], 0.0f); + } + + if (sampleProbTile(value, base)) + { + break; + } + } + + if (tid == 0 && sWinnerIndex >= static_cast(vocabSize) && sLastValidIndex >= 0) + { + sSampledToken = static_cast(sLastValidIndex); + } + __syncthreads(); + }; + + auto sampleTargetLogitsFullVocabWithStats = [&](float const* tLogits, SoftmaxStats targetStats) + { + if (tid == 0) + { + constexpr float kFloatSoftmaxEpsilon = 1e-6f; + sPrefixBase = 0.0f; + sWinnerIndex = static_cast(vocabSize); + sLastValidIndex = -1; + sSampledToken = static_cast(vocabSize) - 1; + sTargetMass = isGreedyRequest + ? 0.0f + : curand_uniform_open_right(state) * (targetStats.sumVal + kFloatSoftmaxEpsilon); + if (isGreedyRequest) + { + sSampledToken = static_cast(targetStats.argmax); + } + } + __syncthreads(); + + if (isGreedyRequest) + { + return; + } + + bool const useVectorizedLoads = canVectorizeLoad(tLogits, vocabSize); + uint32_t const numIters = (vocabSize + BLOCK_SIZE * kVecSize - 1) / (BLOCK_SIZE * kVecSize); +#pragma unroll 2 + for (uint32_t i = 0; i < numIters; ++i) + { + uint32_t const base = (i * BLOCK_SIZE + static_cast(tid)) * kVecSize; + auto const qVec = loadProbVec(tLogits, base, useVectorizedLoads, vocabSize); + float value[kVecSize]; +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + uint32_t const v = base + j; + value[j] = v < vocabSize ? __expf(qVec[j] - targetStats.maxVal) : 0.0f; + } + + if (sampleProbTile(value, base)) + { + break; + } + } + + if (tid == 0 && sWinnerIndex >= static_cast(vocabSize) && sLastValidIndex >= 0) + { + sSampledToken = static_cast(sLastValidIndex); + } + __syncthreads(); + }; + + auto sampleTargetLogitsFullVocab = [&](float const* tLogits) + { + auto const targetStats = computeLogitsStats(tLogits, vocabSize); + sampleTargetLogitsFullVocabWithStats(tLogits, targetStats); + }; + + auto sampleResidualLogitsFullVocab + = [&](float const* tLogits, float const* dLogits, SoftmaxStats targetStats, SoftmaxStats draftStats) + { + if (isGreedyRequest) + { + if (tid == 0) + { + sSampledToken = static_cast(targetStats.argmax); + } + __syncthreads(); + return; + } + + bool const useVectorizedTargetLoads = canVectorizeLoad(tLogits, vocabSize); + bool const useVectorizedDraftLoads = targetToDraft == nullptr && canVectorizeLoad(dLogits, draftVocabSize); + uint32_t const numIters = (vocabSize + BLOCK_SIZE * kVecSize - 1) / (BLOCK_SIZE * kVecSize); + constexpr float kFloatSoftmaxEpsilon = 1e-6f; + float totalSum = 0.0f; +#pragma unroll 2 + for (uint32_t i = 0; i < numIters; ++i) + { + uint32_t const base = (i * BLOCK_SIZE + static_cast(tid)) * kVecSize; + auto const qVec = loadProbVec(tLogits, base, useVectorizedTargetLoads, vocabSize); + flashinfer::vec_t pVec; + pVec.fill(0.0f); + if (targetToDraft == nullptr) + { + pVec = loadProbVec(dLogits, base, useVectorizedDraftLoads, draftVocabSize); + } + float value[kVecSize]; +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + uint32_t const v = base + j; + if (v < vocabSize) + { + float const q = __expf(qVec[j] - targetStats.maxVal) / (targetStats.sumVal + kFloatSoftmaxEpsilon); + int32_t const draftTokenId = targetTokenToDraftToken(v); + float p = 0.0f; + if (draftTokenId >= 0 && draftTokenId < static_cast(draftVocabSize)) + { + float const draftLogit = targetToDraft == nullptr ? pVec[j] : dLogits[draftTokenId]; + p = __expf(draftLogit - draftStats.maxVal) / (draftStats.sumVal + kFloatSoftmaxEpsilon); + } + value[j] = fmaxf(q - p, 0.0f); + } + else + { + value[j] = 0.0f; + } + } + totalSum += BlockReduce(tempStorage.reduce).template Sum(value); + __syncthreads(); + } + + if (tid == 0) + { + sDiffSum = totalSum; + sPrefixBase = 0.0f; + sWinnerIndex = static_cast(vocabSize); + sLastValidIndex = -1; + sSampledToken = static_cast(vocabSize) - 1; + sTargetMass = totalSum > 1e-10f ? curand_uniform_open_right(state) * totalSum : 0.0f; + } + __syncthreads(); + + if (sDiffSum <= 1e-10f) + { + sampleTargetLogitsFullVocabWithStats(tLogits, targetStats); + return; + } + +#pragma unroll 2 + for (uint32_t i = 0; i < numIters; ++i) + { + uint32_t const base = (i * BLOCK_SIZE + static_cast(tid)) * kVecSize; + auto const qVec = loadProbVec(tLogits, base, useVectorizedTargetLoads, vocabSize); + flashinfer::vec_t pVec; + pVec.fill(0.0f); + if (targetToDraft == nullptr) + { + pVec = loadProbVec(dLogits, base, useVectorizedDraftLoads, draftVocabSize); + } + float value[kVecSize]; +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) + { + uint32_t const v = base + j; + if (v < vocabSize) + { + float const q = __expf(qVec[j] - targetStats.maxVal) / (targetStats.sumVal + kFloatSoftmaxEpsilon); + int32_t const draftTokenId = targetTokenToDraftToken(v); + float p = 0.0f; + if (draftTokenId >= 0 && draftTokenId < static_cast(draftVocabSize)) + { + float const draftLogit = targetToDraft == nullptr ? pVec[j] : dLogits[draftTokenId]; + p = __expf(draftLogit - draftStats.maxVal) / (draftStats.sumVal + kFloatSoftmaxEpsilon); + } + value[j] = fmaxf(q - p, 0.0f); + } + else + { + value[j] = 0.0f; + } + } + + if (sampleProbTile(value, base)) { - 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; } + } + + if (tid == 0 && sWinnerIndex >= static_cast(vocabSize) && sLastValidIndex >= 0) + { + sSampledToken = static_cast(sLastValidIndex); + } + __syncthreads(); + }; + + // First-gen or dummy request: no valid tree exists yet. Sample directly + // from the target distribution at the root and skip tree traversal. + if (treeValid != nullptr && !treeValid[bx]) + { + float const* tProbs = targetProbs + static_cast(bx) * numDraftTokens * vocabSize; + if (hasCompactTargetSupport) + { + if (tid == 0) + { + uint32_t const supportOffset = static_cast(bx) * numDraftTokens * maxTargetSupportSize; + uint32_t const supportSize = static_cast(targetSupportLengths[batchOffset]); + sSampledToken = sampleFromIndexedDistribution( + state, tProbs, targetSupportIndices + supportOffset, supportSize, vocabSize); + } + } + else + { + if constexpr (USE_LOGITS) + { + sampleTargetLogitsFullVocab(tProbs); + } else { - curIndex = retrieveNextSibling[batchOffset + curIndex]; + sampleTargetFullVocab(tProbs); + } + } + if (tid == 0) + { + acceptIndex[bx * numSpeculativeTokens] = 0; + acceptTokenNum[bx] = 0; + acceptToken[bx * numSpeculativeTokens] = sSampledToken; + } + return; + } + + // Root (depth 0): initialize path state at tree position 0. + // + // Example tree used in code review discussions: + // root: E + // children of E: F1, F2, F3 + // children of F1: G1, G2 + // children of F2: G3 + // + // In that example the per-request inputs are conceptually: + // candidates = [E, F1, F2, F3, G1, G2, G3] + // draftProbs = [p(.|E), p(.|F1), p(.|F2)] + // draftProbIndices = [0, 0, 0, 0, 1, 1, 2] + // targetProbs = [q(.|E), q(.|F1), q(.|F2), q(.|F3), q(.|G1), q(.|G2), q(.|G3)] + // + // draftProbs stores one row per unique parent context, so siblings that share + // the same parent also share the same draftProbs row via draftProbIndices. + // targetProbs remains aligned to all tree positions, including the root at slot 0. + // + // Output convention: + // - acceptIndex stores the accepted draft path as tree positions, with slot 0 + // reserved for the root position. + // - acceptToken stores the emitted token sequence, matching the greedy kernel: + // slot 0 = first emitted token + // slot numAcceptedTokens = final bonus/correction token + // - acceptTokenNum stores the number of accepted draft tokens only. The caller + // adds 1 to obtain the total number of emitted tokens. + if (tid == 0) + { + sLastAcceptedLocalIdx = 0; + acceptIndex[bx * numSpeculativeTokens] = sLastAcceptedLocalIdx; + sNumAcceptedTokens = 0; + sHasTerminalToken = false; + } + __syncthreads(); + for (uint32_t j = 1; j < numSpeculativeTokens; ++j) + { + // Get first child of the last accepted node. + if (tid == 0) + { + sFirstChild = retrieveNextToken[batchOffset + sLastAcceptedLocalIdx]; + } + __syncthreads(); + + // Leaf node: no children at this depth. + // Emit bonus token from the target distribution at the last accepted position. + if (sFirstChild == -1) + { + float const* tProbs + = targetProbs + (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize; + if (hasCompactTargetSupport) + { + if (tid == 0) + { + uint32_t const supportOffset + = (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * maxTargetSupportSize; + uint32_t const supportSize + = static_cast(targetSupportLengths[batchOffset + sLastAcceptedLocalIdx]); + sSampledToken = sampleFromIndexedDistribution( + state, tProbs, targetSupportIndices + supportOffset, supportSize, vocabSize); + } } + else + { + if constexpr (USE_LOGITS) + { + sampleTargetLogitsFullVocab(tProbs); + } + else + { + sampleTargetFullVocab(tProbs); + } + } + if (tid == 0) + { + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sSampledToken; + sHasTerminalToken = true; + } + __syncthreads(); + break; } - if (curIndex == -1) + // Test siblings in linked-list order. Once a sibling passes the + // Bernoulli rejection test, accept it immediately and skip the rest. + int32_t const firstDraftProbRow = draftProbIndices[batchOffset + sFirstChild]; + float const* siblingTargetRow + = targetProbs + (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize; + float const* siblingDraftRow + = draftProbs + (static_cast(bx) * numDraftProbRows + firstDraftProbRow) * draftRowStride; + SoftmaxStats siblingTargetStats{}; + SoftmaxStats siblingDraftStats{}; + if constexpr (USE_LOGITS) + { + siblingTargetStats = computeLogitsStats(siblingTargetRow, vocabSize); + siblingDraftStats = computeLogitsStats(siblingDraftRow, draftVocabSize); + } + + if (tid == 0) + { + sNumAccSiblings = 0; + int32_t childIdx = sFirstChild; + while (childIdx != -1) + { + int64_t const draftTokenId = candidates[batchOffset + childIdx]; + int32_t const draftProbRow = draftProbIndices[batchOffset + childIdx]; + uint32_t const tokenId = static_cast(draftTokenId); + float const pDraft = USE_LOGITS + ? draftProbFromTargetToken(siblingDraftRow, tokenId, siblingDraftStats) + : draftProbs[(static_cast(bx) * numDraftProbRows + draftProbRow) * vocabSize + + draftTokenId]; + float const pTarget = USE_LOGITS + ? probFromLogits(siblingTargetRow, tokenId, vocabSize, siblingTargetStats) + : targetProbs[(static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize + + draftTokenId]; + + float const acceptProb = fminf(1.0f, pTarget / (pDraft + 1e-10f)); + float const u = curand_uniform_open_right(state); + + if (u < acceptProb) + { + sAccSibIdx = childIdx; + sAccSibTok = draftTokenId; + sNumAccSiblings = 1; + break; + } + childIdx = retrieveNextSibling[batchOffset + childIdx]; + } + } + __syncthreads(); + + // Select the first accepted sibling or emit correction when all siblings reject. + if (sNumAccSiblings > 0) + { + if (tid == 0) + { + int32_t const childIdx = sAccSibIdx; + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sAccSibTok; + ++sNumAcceptedTokens; + acceptIndex[bx * numSpeculativeTokens + sNumAcceptedTokens] = childIdx; + sLastAcceptedLocalIdx = childIdx; + } + __syncthreads(); + } + else + { + // All siblings rejected -> sample correction token from relu(q - p). + { + int32_t const draftProbRow = draftProbIndices[batchOffset + sFirstChild]; + float const* tProbs + = targetProbs + (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize; + float const* dProbs + = draftProbs + (static_cast(bx) * numDraftProbRows + draftProbRow) * draftRowStride; + int32_t const* tProbIndices = nullptr; + uint32_t targetSupportSize = vocabSize; + if (hasCompactTargetSupport) + { + uint32_t const supportOffset + = (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * maxTargetSupportSize; + tProbIndices = targetSupportIndices + supportOffset; + targetSupportSize + = static_cast(targetSupportLengths[batchOffset + sLastAcceptedLocalIdx]); + } + + if (hasCompactTargetSupport) + { + if (tid == 0) + { + float diffSum = 0.0f; + for (uint32_t i = 0; i < targetSupportSize; ++i) + { + uint32_t const v = static_cast(tProbIndices[i]); + float const diff = tProbs[v] - dProbs[v]; + if (diff > 0.0f) + { + diffSum += diff; + } + } + + int64_t corrTok = static_cast(vocabSize) - 1; + bool const useDiff = (diffSum > 1e-10f); + + if (useDiff) + { + float const r = curand_uniform_open_right(state); + float cumsum = 0.0f; + for (uint32_t i = 0; i < targetSupportSize; ++i) + { + uint32_t const v = static_cast(tProbIndices[i]); + float const diff = tProbs[v] - dProbs[v]; + float const prob = (diff > 0.0f) ? diff / diffSum : 0.0f; + cumsum += prob; + if (r <= cumsum) + { + corrTok = static_cast(v); + break; + } + } + } + else + { + corrTok = sampleFromIndexedDistribution( + state, tProbs, tProbIndices, targetSupportSize, vocabSize); + } + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = corrTok; + sHasTerminalToken = true; + } + } + else + { + if constexpr (USE_LOGITS) + { + sampleResidualLogitsFullVocab(tProbs, dProbs, siblingTargetStats, siblingDraftStats); + } + else + { + sampleResidualFullVocab(tProbs, dProbs); + } + + if (tid == 0) + { + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sSampledToken; + sHasTerminalToken = true; + } + } + } + __syncthreads(); break; + } } - acceptTokenNum[bx] = numAcceptedTokens; - // Bonus token from target model at the last accepted position - predicts[batchOffset + lastAcceptedLocalIdx] = targetPredict[batchOffset + lastAcceptedLocalIdx]; + if (!sHasTerminalToken) + { + // Reached max speculative depth while continuing to accept the draft path. + // Emit the final bonus token from the last accepted position. + float const* tProbs + = targetProbs + (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * vocabSize; + if (hasCompactTargetSupport) + { + if (tid == 0) + { + uint32_t const supportOffset + = (static_cast(bx) * numDraftTokens + sLastAcceptedLocalIdx) * maxTargetSupportSize; + uint32_t const supportSize + = static_cast(targetSupportLengths[batchOffset + sLastAcceptedLocalIdx]); + sSampledToken = sampleFromIndexedDistribution( + state, tProbs, targetSupportIndices + supportOffset, supportSize, vocabSize); + } + } + else + { + if constexpr (USE_LOGITS) + { + sampleTargetLogitsFullVocab(tProbs); + } + else + { + sampleTargetFullVocab(tProbs); + } + } + if (tid == 0) + { + acceptToken[bx * numSpeculativeTokens + sNumAcceptedTokens] = sSampledToken; + } + } + + if (tid == 0) + { + acceptTokenNum[bx] = sNumAcceptedTokens; + } } -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) +void invokeVerifyDynamicTreeRejection(int64_t* acceptIndex, int64_t* acceptTokenNum, int64_t* acceptToken, + int64_t const* candidates, float const* draftProbs, float const* targetProbs, int32_t const* targetSupportIndices, + int32_t const* targetSupportLengths, int32_t const* draftProbIndices, int32_t const* retrieveNextToken, + int32_t const* retrieveNextSibling, bool const* treeValid, SizeType32 batchSize, SizeType32 numDraftProbRows, + SizeType32 maxTargetSupportSize, SizeType32 numDraftTokens, SizeType32 numSpecStep, SizeType32 vocabSize, + int64_t const* seed, int64_t const* offset, cudaStream_t stream) { + constexpr int32_t kVerifyDynamicTreeRejectionBlockSize = 1024; dim3 grid(batchSize); - dim3 block(1); + dim3 block(kVerifyDynamicTreeRejectionBlockSize); - verifyDynamicTreeGreedyKernel<<>>(predicts, acceptIndex, acceptTokenNum, acceptToken, - candidates, retrieveIndex, retrieveNextToken, retrieveNextSibling, targetPredict, treeValid, batchSize, - numSpecStep, numDraftTokens); + verifyDynamicTreeRejectionKernel<<>>( + acceptIndex, acceptTokenNum, acceptToken, candidates, draftProbs, targetProbs, targetSupportIndices, + targetSupportLengths, draftProbIndices, retrieveNextToken, retrieveNextSibling, treeValid, batchSize, + numDraftProbRows, maxTargetSupportSize, numSpecStep, numDraftTokens, vocabSize, vocabSize, + /*targetToDraft=*/nullptr, seed, offset, nullptr); sync_check_cuda_error(stream); } diff --git a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h index c48b0f25005a..b7adbdc2f573 100644 --- a/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h +++ b/cpp/tensorrt_llm/kernels/speculativeDecoding/dynamicTreeKernels.h @@ -19,6 +19,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/runtime/common.h" #include +#include TRTLLM_NAMESPACE_BEGIN @@ -60,6 +61,18 @@ void invokeBuildDynamicTree(int64_t const* parentList, int64_t const* selectedIn runtime::SizeType32 topK, runtime::SizeType32 depth, runtime::SizeType32 numDraftTokens, TreeMaskMode treeMaskMode, cudaStream_t stream, runtime::SizeType32 numInt32PerRow); +//! \brief Build tree-position -> unique draft-prob row mapping for rejection sampling. +//! \param topkScoreIndices [batchSize, numDraftTokens], on GPU. int64. +//! History-buffer indices selected for each final tree position. +//! \param draftProbIndices output [batchSize, numDraftTokens + 1], on GPU. int32. +//! Column 0 is reserved for the root and set to 0. +//! \param batchSize runtime::SizeType32. Batch size. +//! \param topK runtime::SizeType32. Tree top-K branching factor. +//! \param numDraftTokens runtime::SizeType32. Total number of non-root draft positions. +//! \param stream cuda stream. +void invokeBuildDraftProbIndices(int64_t const* topkScoreIndices, int32_t* draftProbIndices, + runtime::SizeType32 batchSize, runtime::SizeType32 topK, runtime::SizeType32 numDraftTokens, cudaStream_t stream); + //! \brief Verify dynamic tree using greedy strategy //! Verifies draft tokens against target model predictions using tree traversal. //! All index/token pointer parameters use int64. @@ -93,6 +106,77 @@ void invokeVerifyDynamicTreeGreedy(int64_t* predicts, int64_t* acceptIndex, int6 runtime::SizeType32 batchSize, runtime::SizeType32 numDraftTokens, runtime::SizeType32 numSpecStep, cudaStream_t stream); +//! \brief Verify dynamic tree using rejection sampling. +//! For each request, traverses the tree depth-by-depth. At each depth, siblings are tried +//! in order; the first sibling accepted by rejection sampling (p_target/p_draft) continues +//! the path. If all siblings are rejected, a correction token is sampled from (target-draft)_+. +//! \param acceptIndex output [batchSize, numSpecStep] int64 — tree positions of accepted tokens. +//! \param acceptTokenNum output [batchSize] int64 — # accepted draft tokens (excl. root). +//! \param acceptToken output [batchSize, numSpecStep] int64 — accepted/correction token ids. +//! \param candidates [batchSize, numDraftTokens] int64; col 0 = root (target sample). +//! \param draftProbs [batchSize, numDraftProbRows, vocabSize] float32. +//! Unique draft probability rows per request; tree positions map into this +//! tensor via draftProbIndices. +//! \param targetProbs [batchSize, numDraftTokens, vocabSize] float32; index 0 = root. +//! \param targetSupportIndices [batchSize, numDraftTokens, maxTargetSupportSize] int32; compact token ids that +//! survive top-k/top-p filtering for each row, padded with -1. May be empty when no filtering is active. +//! \param targetSupportLengths [batchSize, numDraftTokens] int32; valid support length per row. May be empty when +//! no filtering is active. +//! \param draftProbIndices [batchSize, numDraftTokens] int32; maps each tree position to the +//! corresponding row in draftProbs. Root is unused. +//! \param retrieveNextToken [batchSize, numDraftTokens] int32 first-child pointer, -1=none. +//! \param retrieveNextSibling [batchSize, numDraftTokens] int32 next-sibling pointer, -1=none. +//! \param treeValid [batchSize] bool; false means no valid tree exists for this request. +//! \param batchSize runtime::SizeType32. +//! \param numDraftProbRows runtime::SizeType32. Number of unique draft-prob rows per request. +//! \param maxTargetSupportSize runtime::SizeType32. Third dim of targetSupportIndices. Can be zero. +//! \param numDraftTokens runtime::SizeType32. Total tree nodes per request (including root). +//! \param numSpecStep runtime::SizeType32. Second dim of acceptIndex/acceptToken. +//! \param vocabSize runtime::SizeType32. Vocabulary size. +//! \param seed [1] int64 on GPU. Philox RNG seed. +//! \param offset [1] int64 on GPU. Philox RNG offset. +//! \param stream cudaStream_t. +void invokeVerifyDynamicTreeRejection(int64_t* acceptIndex, int64_t* acceptTokenNum, int64_t* acceptToken, + int64_t const* candidates, float const* draftProbs, float const* targetProbs, int32_t const* targetSupportIndices, + int32_t const* targetSupportLengths, int32_t const* draftProbIndices, int32_t const* retrieveNextToken, + int32_t const* retrieveNextSibling, bool const* treeValid, runtime::SizeType32 batchSize, + runtime::SizeType32 numDraftProbRows, runtime::SizeType32 maxTargetSupportSize, runtime::SizeType32 numDraftTokens, + runtime::SizeType32 numSpecStep, runtime::SizeType32 vocabSize, int64_t const* seed, int64_t const* offset, + cudaStream_t stream); + +//! \brief Compute draft probabilities for dynamic-tree rejection sampling from logits. +//! \param draftLogits [batchSize * numDraftProbRows, draftVocabSize], on GPU. +//! \param temperatures [batchSize], on GPU. +//! \param numDraftProbRows runtime::SizeType32. Unique draft-prob rows per request. +//! \param topK Optional [batchSize], on GPU. +//! \param topP Optional [batchSize], on GPU. +//! \param targetVocabSize runtime::SizeType32. Output vocabulary size after optional d2t expansion. +//! \param d2t Optional [draftVocabSize], on GPU. +//! \return [batchSize, numDraftProbRows, targetVocabSize] float32 probabilities. +torch::Tensor computeDraftProbsForDynamicTreeRejection(torch::Tensor const& draftLogits, + torch::Tensor const& temperatures, runtime::SizeType32 numDraftProbRows, torch::optional const& topK, + torch::optional const& topP, runtime::SizeType32 targetVocabSize, bool skipTemperature, + torch::optional const& d2t, runtime::SizeType32 kMax = 0, bool skipAllSamplingParams = false); + +//! \brief Compute target probabilities for dynamic-tree rejection sampling from logits. +//! \param targetLogits [batchSize * numDraftTokens, targetVocabSize], on GPU. +//! \param temperatures [batchSize], on GPU. +//! \param numDraftTokens runtime::SizeType32. Total tree nodes per request (including root). +//! \param topK Optional [batchSize], on GPU. +//! \param topP Optional [batchSize], on GPU. +//! \param kMax runtime::SizeType32. Max top-K value across the batch; enables the fast topk +//! path when > 0. Must be computed on CPU (e.g. topK.max().item()). Default 0 = fallback +//! to full sort. +//! \return Tuple of: +//! 1. [batchSize, numDraftTokens, targetVocabSize] float32 probabilities. +//! 2. [batchSize, numDraftTokens, maxTargetSupportSize] int32 compact token ids that survive +//! top-k/top-p filtering for each row, padded with -1. Empty when no filtering is active. +//! 3. [batchSize, numDraftTokens] int32 support lengths. Empty when no filtering is active. +std::tuple computeTargetProbsForDynamicTreeRejection( + torch::Tensor const& targetLogits, torch::Tensor const& temperatures, runtime::SizeType32 numDraftTokens, + torch::optional const& topK, torch::optional const& topP, bool skipTemperature, + runtime::SizeType32 kMax = 0, bool skipAllSamplingParams = false); + } // namespace kernels::speculative_decoding TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp b/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp index 2c70fcfc5cb1..7f0bd16b93da 100644 --- a/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp +++ b/cpp/tensorrt_llm/thop/dynamicTreeOp.cpp @@ -25,9 +25,45 @@ namespace tk = tensorrt_llm::kernels::speculative_decoding; TRTLLM_NAMESPACE_BEGIN +namespace kernels::speculative_decoding +{ +th::Tensor computeProbsFromLogits(th::Tensor const& logits, th::Tensor const& temperatures, + th::optional const& topK, th::optional const& topP, bool skipTemperature, + runtime::SizeType32 kMax); +void invokeBuildDraftProbIndices(int64_t const* topkScoreIndices, int32_t* draftProbIndices, + runtime::SizeType32 batchSize, runtime::SizeType32 topK, runtime::SizeType32 numDraftTokens, cudaStream_t stream); +th::Tensor computeDraftProbsForDynamicTreeRejection(th::Tensor const& draftLogits, th::Tensor const& temperatures, + runtime::SizeType32 numDraftProbRows, th::optional const& topK, th::optional const& topP, + runtime::SizeType32 targetVocabSize, bool skipTemperature, th::optional const& d2t, + runtime::SizeType32 kMax, bool skipAllSamplingParams); +std::tuple computeTargetProbsForDynamicTreeRejection(th::Tensor const& targetLogits, + th::Tensor const& temperatures, runtime::SizeType32 numDraftTokens, th::optional const& topK, + th::optional const& topP, bool skipTemperature, runtime::SizeType32 kMax, bool skipAllSamplingParams); +} // namespace kernels::speculative_decoding + namespace torch_ext { +void build_draft_prob_indices_out_op( + th::Tensor& topkScoreIndices, th::Tensor& draftProbIndices, int64_t topK, int64_t numDraftTokens) +{ + TORCH_CHECK(topkScoreIndices.is_cuda(), "topkScoreIndices must be a CUDA tensor"); + TORCH_CHECK(draftProbIndices.is_cuda(), "draftProbIndices must be a CUDA tensor"); + TORCH_CHECK(topkScoreIndices.dim() == 2, "topkScoreIndices must be a 2D tensor"); + TORCH_CHECK(draftProbIndices.dim() == 2, "draftProbIndices must be a 2D tensor"); + TORCH_CHECK(topkScoreIndices.scalar_type() == torch::kInt64, "topkScoreIndices must be int64 tensor"); + TORCH_CHECK(draftProbIndices.scalar_type() == torch::kInt32, "draftProbIndices must be int32 tensor"); + TORCH_CHECK(topkScoreIndices.size(1) == numDraftTokens, "topkScoreIndices size mismatch"); + TORCH_CHECK(draftProbIndices.size(0) == topkScoreIndices.size(0), "Batch size mismatch"); + TORCH_CHECK(draftProbIndices.size(1) == numDraftTokens + 1, "draftProbIndices size mismatch"); + TORCH_CHECK(topK > 0, "topK must be positive"); + TORCH_CHECK(numDraftTokens + 1 <= 1024, "numDraftTokens + 1 exceeds CUDA block size limit of 1024"); + + auto stream = at::cuda::getCurrentCUDAStream(topkScoreIndices.device().index()); + tk::invokeBuildDraftProbIndices(topkScoreIndices.data_ptr(), draftProbIndices.data_ptr(), + topkScoreIndices.size(0), topK, numDraftTokens, stream); +} + //! \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, @@ -175,12 +211,178 @@ void verify_dynamic_tree_greedy_out_op(th::Tensor& candidates, th::Tensor& retri batchSize, numDraftTokens, numSpecStep, stream); } +//! \brief In-place tree rejection sampling verify op. +//! Accepts draft tokens by rejection sampling at each depth using pre-computed probabilities. +void verify_dynamic_tree_rejection_out_op(th::Tensor& candidates, th::Tensor& draftProbs, th::Tensor& targetProbs, + th::Tensor& targetSupportIndices, th::Tensor& targetSupportLengths, th::Tensor& draftProbIndices, + th::Tensor& retrieveNextToken, th::Tensor& retrieveNextSibling, th::Tensor& treeValid, th::Tensor& acceptIndex, + th::Tensor& acceptTokenNum, th::Tensor& acceptToken, int64_t numSpecStep, th::Tensor& seed, th::Tensor& offset) +{ + TORCH_CHECK(candidates.dim() == 2, "candidates must be 2D tensor"); + TORCH_CHECK(draftProbs.dim() == 3, "draftProbs must be 3D tensor"); + TORCH_CHECK(targetProbs.dim() == 3, "targetProbs must be 3D tensor"); + TORCH_CHECK(targetSupportIndices.dim() == 1 || targetSupportIndices.dim() == 3, + "targetSupportIndices must be 1D or 3D tensor"); + TORCH_CHECK(targetSupportLengths.dim() == 1 || targetSupportLengths.dim() == 2, + "targetSupportLengths must be 1D or 2D tensor"); + TORCH_CHECK(draftProbIndices.dim() == 2, "draftProbIndices 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(treeValid.dim() == 1, "treeValid must be 1D tensor"); + TORCH_CHECK(candidates.scalar_type() == torch::kInt64, "candidates must be int64 tensor"); + TORCH_CHECK(draftProbs.scalar_type() == torch::kFloat32, "draftProbs must be float32 tensor"); + TORCH_CHECK(targetProbs.scalar_type() == torch::kFloat32, "targetProbs must be float32 tensor"); + TORCH_CHECK(targetSupportIndices.scalar_type() == torch::kInt32, "targetSupportIndices must be int32 tensor"); + TORCH_CHECK(targetSupportLengths.scalar_type() == torch::kInt32, "targetSupportLengths must be int32 tensor"); + TORCH_CHECK(draftProbIndices.scalar_type() == torch::kInt32, "draftProbIndices must be int32 tensor"); + TORCH_CHECK(treeValid.scalar_type() == torch::kBool, "treeValid must be bool tensor"); + TORCH_CHECK(candidates.is_cuda(), "candidates must be a CUDA tensor"); + TORCH_CHECK(draftProbs.is_cuda(), "draftProbs must be a CUDA tensor"); + TORCH_CHECK(targetProbs.is_cuda(), "targetProbs must be a CUDA tensor"); + TORCH_CHECK(draftProbIndices.is_cuda(), "draftProbIndices must be a CUDA tensor"); + TORCH_CHECK(retrieveNextToken.is_cuda(), "retrieveNextToken must be a CUDA tensor"); + TORCH_CHECK(retrieveNextSibling.is_cuda(), "retrieveNextSibling must be a CUDA tensor"); + TORCH_CHECK(treeValid.is_cuda(), "treeValid must be a CUDA tensor"); + TORCH_CHECK(acceptIndex.is_cuda(), "acceptIndex must be a CUDA tensor"); + TORCH_CHECK(acceptTokenNum.is_cuda(), "acceptTokenNum must be a CUDA tensor"); + TORCH_CHECK(acceptToken.is_cuda(), "acceptToken must be a CUDA tensor"); + + int64_t batchSize = candidates.size(0); + int64_t numDraftProbRows = draftProbs.size(1); + int64_t numDraftTokens = candidates.size(1); + int64_t vocabSize = targetProbs.size(2); + int64_t maxTargetSupportSize = targetSupportIndices.dim() == 3 ? targetSupportIndices.size(2) : 0; + + TORCH_CHECK(draftProbs.size(0) == batchSize, "draftProbs batch size mismatch"); + TORCH_CHECK(draftProbs.size(2) == vocabSize, "draftProbs vocabSize mismatch"); + TORCH_CHECK(targetProbs.size(0) == batchSize, "targetProbs batch size mismatch"); + TORCH_CHECK(targetProbs.size(1) == numDraftTokens, "targetProbs numDraftTokens mismatch"); + if (targetSupportIndices.numel() > 0) + { + TORCH_CHECK(targetSupportIndices.dim() == 3, "targetSupportIndices must be 3D when non-empty"); + TORCH_CHECK(targetSupportIndices.size(0) == batchSize, "targetSupportIndices batch size mismatch"); + TORCH_CHECK(targetSupportIndices.size(1) == numDraftTokens, "targetSupportIndices numDraftTokens mismatch"); + TORCH_CHECK(targetSupportIndices.is_cuda(), "targetSupportIndices must be a CUDA tensor when non-empty"); + TORCH_CHECK(targetSupportIndices.device() == candidates.device(), + "targetSupportIndices must be on the same device as candidates"); + } + if (targetSupportLengths.numel() > 0) + { + TORCH_CHECK(targetSupportLengths.dim() == 2, "targetSupportLengths must be 2D when non-empty"); + TORCH_CHECK(targetSupportLengths.size(0) == batchSize, "targetSupportLengths batch size mismatch"); + TORCH_CHECK(targetSupportLengths.size(1) == numDraftTokens, "targetSupportLengths numDraftTokens mismatch"); + TORCH_CHECK(targetSupportLengths.is_cuda(), "targetSupportLengths must be a CUDA tensor when non-empty"); + TORCH_CHECK(targetSupportLengths.device() == candidates.device(), + "targetSupportLengths must be on the same device as candidates"); + } + TORCH_CHECK((targetSupportIndices.numel() == 0) == (targetSupportLengths.numel() == 0), + "targetSupportIndices and targetSupportLengths must both be empty or both be non-empty"); + TORCH_CHECK(draftProbIndices.size(0) == batchSize, "draftProbIndices batch size mismatch"); + TORCH_CHECK(draftProbIndices.size(1) == numDraftTokens, "draftProbIndices size mismatch"); + TORCH_CHECK(retrieveNextToken.size(0) == batchSize, "retrieveNextToken batch size mismatch"); + TORCH_CHECK(retrieveNextToken.size(1) == numDraftTokens, "retrieveNextToken size mismatch"); + TORCH_CHECK(retrieveNextSibling.size(0) == batchSize, "retrieveNextSibling batch size mismatch"); + TORCH_CHECK(retrieveNextSibling.size(1) == numDraftTokens, "retrieveNextSibling size mismatch"); + TORCH_CHECK(treeValid.size(0) >= batchSize, "treeValid buffer too small"); + TORCH_CHECK(draftProbs.device() == candidates.device(), "draftProbs must be on the same device as candidates"); + TORCH_CHECK(targetProbs.device() == candidates.device(), "targetProbs must be on the same device as candidates"); + TORCH_CHECK( + draftProbIndices.device() == candidates.device(), "draftProbIndices must be on the same device as candidates"); + TORCH_CHECK(retrieveNextToken.device() == candidates.device(), + "retrieveNextToken must be on the same device as candidates"); + TORCH_CHECK(retrieveNextSibling.device() == candidates.device(), + "retrieveNextSibling must be on the same device as candidates"); + TORCH_CHECK(treeValid.device() == candidates.device(), "treeValid must be on the same device as candidates"); + 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(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"); + TORCH_CHECK(acceptIndex.device() == candidates.device(), "acceptIndex must be on the same device as candidates"); + TORCH_CHECK( + acceptTokenNum.device() == candidates.device(), "acceptTokenNum must be on the same device as candidates"); + TORCH_CHECK(acceptToken.device() == candidates.device(), "acceptToken must be on the same device as candidates"); + TORCH_CHECK(seed.scalar_type() == torch::kInt64, "seed must be int64 tensor"); + TORCH_CHECK(offset.scalar_type() == torch::kInt64, "offset must be int64 tensor"); + TORCH_CHECK(seed.numel() >= 1, "seed tensor must have at least one element"); + TORCH_CHECK(offset.numel() >= 1, "offset tensor must have at least one element"); + TORCH_CHECK(seed.is_cuda(), "seed must be CUDA tensor"); + TORCH_CHECK(offset.is_cuda(), "offset must be CUDA tensor"); + TORCH_CHECK(seed.device() == candidates.device(), "seed must be on the same device as candidates"); + TORCH_CHECK(offset.device() == candidates.device(), "offset must be on the same device as candidates"); + + auto stream = at::cuda::getCurrentCUDAStream(candidates.device().index()); + + acceptIndex.zero_(); + acceptTokenNum.zero_(); + acceptToken.zero_(); + + tk::invokeVerifyDynamicTreeRejection(acceptIndex.data_ptr(), acceptTokenNum.data_ptr(), + acceptToken.data_ptr(), candidates.data_ptr(), draftProbs.data_ptr(), + targetProbs.data_ptr(), + targetSupportIndices.numel() > 0 ? targetSupportIndices.data_ptr() : nullptr, + targetSupportLengths.numel() > 0 ? targetSupportLengths.data_ptr() : nullptr, + draftProbIndices.data_ptr(), retrieveNextToken.data_ptr(), + retrieveNextSibling.data_ptr(), treeValid.data_ptr(), batchSize, numDraftProbRows, + maxTargetSupportSize, numDraftTokens, numSpecStep, vocabSize, seed.data_ptr(), + offset.data_ptr(), stream); +} + +th::Tensor compute_draft_probs_for_dynamic_tree_rejection_op(th::Tensor draftLogits, th::Tensor temperatures, + int64_t numDraftProbRows, int64_t targetVocabSize, th::optional topK, th::optional topP, + bool skipTemperature, th::optional d2t, int64_t topKMax, bool skipAllSamplingParams) +{ + return tk::computeDraftProbsForDynamicTreeRejection(draftLogits, temperatures, numDraftProbRows, topK, topP, + targetVocabSize, skipTemperature, d2t, topKMax, skipAllSamplingParams); +} + +std::tuple compute_target_probs_for_dynamic_tree_rejection_op( + th::Tensor targetLogits, th::Tensor temperatures, int64_t numDraftTokens, th::optional topK, + th::optional topP, bool skipTemperature, int64_t topKMax, bool skipAllSamplingParams) +{ + return tk::computeTargetProbsForDynamicTreeRejection( + targetLogits, temperatures, numDraftTokens, topK, topP, skipTemperature, topKMax, skipAllSamplingParams); +} + +th::Tensor compute_probs_from_logits_op(th::Tensor logits, th::Tensor temperatures, th::optional topK, + th::optional topP, bool skipTemperature) +{ + TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor"); + TORCH_CHECK(temperatures.is_cuda(), "temperatures must be a CUDA tensor"); + TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor"); + TORCH_CHECK(temperatures.dim() == 1, "temperatures must be a 1D tensor"); + TORCH_CHECK(logits.size(0) == temperatures.size(0), "logits and temperatures size mismatch"); + if (topK.has_value() && topK->defined()) + { + TORCH_CHECK(topK->is_cuda(), "top_k must be a CUDA tensor"); + } + if (topP.has_value() && topP->defined()) + { + TORCH_CHECK(topP->is_cuda(), "top_p must be a CUDA tensor"); + } + + return tk::computeProbsFromLogits(logits, temperatures, topK, topP, skipTemperature, /*kMax=*/0); +} + } // namespace torch_ext TRTLLM_NAMESPACE_END //////////////////////////////////////////////////////////////////////////////////////////////////////////// +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "build_draft_prob_indices_out_op(Tensor topkScoreIndices, Tensor(a!) draftProbIndices, " + "int topK, int numDraftTokens) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("build_draft_prob_indices_out_op", &tensorrt_llm::torch_ext::build_draft_prob_indices_out_op); +} + TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( @@ -227,3 +429,62 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("verify_dynamic_tree_greedy_out_op", &tensorrt_llm::torch_ext::verify_dynamic_tree_greedy_out_op); } + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "verify_dynamic_tree_rejection_out_op(" + "Tensor candidates, Tensor draftProbs, Tensor targetProbs, Tensor targetSupportIndices, " + "Tensor targetSupportLengths, Tensor draftProbIndices, " + "Tensor retrieveNextToken, Tensor retrieveNextSibling, Tensor treeValid, " + "Tensor(a!) acceptIndex, Tensor(b!) acceptTokenNum, Tensor(c!) acceptToken, " + "int numSpecStep, Tensor seed, Tensor offset) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("verify_dynamic_tree_rejection_out_op", &tensorrt_llm::torch_ext::verify_dynamic_tree_rejection_out_op); +} + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "compute_draft_probs_for_dynamic_tree_rejection_op(" + "Tensor draftLogits, Tensor temperatures, int numDraftProbRows, int targetVocabSize, " + "Tensor? top_k=None, Tensor? top_p=None, bool skip_temperature=False, Tensor? d2t=None, " + "int top_k_max=0, bool skip_all_sampling_params=False) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("compute_draft_probs_for_dynamic_tree_rejection_op", + &tensorrt_llm::torch_ext::compute_draft_probs_for_dynamic_tree_rejection_op); +} + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "compute_target_probs_for_dynamic_tree_rejection_op(" + "Tensor targetLogits, Tensor temperatures, int numDraftTokens, " + "Tensor? top_k=None, Tensor? top_p=None, bool skip_temperature=False, int top_k_max=0, " + "bool skip_all_sampling_params=False) -> (Tensor, Tensor, Tensor)"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("compute_target_probs_for_dynamic_tree_rejection_op", + &tensorrt_llm::torch_ext::compute_target_probs_for_dynamic_tree_rejection_op); +} + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "compute_probs_from_logits_op(" + "Tensor logits, Tensor temperatures, Tensor? top_k=None, Tensor? top_p=None, " + "bool skip_temperature=False) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("compute_probs_from_logits_op", &tensorrt_llm::torch_ext::compute_probs_from_logits_op); +} diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index ce0465dfa294..097edc9e2e6e 100644 --- a/docs/source/features/speculative-decoding.md +++ b/docs/source/features/speculative-decoding.md @@ -59,7 +59,8 @@ To enable dynamic tree mode, set `use_dynamic_tree=True` on the `Eagle3DecodingC * `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. + +When `use_dynamic_tree=True`, the dynamic tree CUDA buffers are pre-allocated based on the `LLM`'s `max_batch_size`, which is propagated internally and must not be passed directly to `Eagle3DecodingConfig`. ```python from tensorrt_llm.llmapi import Eagle3DecodingConfig @@ -70,7 +71,6 @@ speculative_config = Eagle3DecodingConfig( 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) diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index 8d66de1c11ba..10901d87c520 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -296,8 +296,7 @@ def setup_llm(args, **kwargs): dynamic_tree_max_topK=args.dynamic_tree_max_topK, allow_advanced_sampling=args.allow_advanced_sampling, eagle3_model_arch=args.eagle3_model_arch, - max_total_draft_tokens=args.max_total_draft_tokens, - max_batch_size=args.max_batch_size) + max_total_draft_tokens=args.max_total_draft_tokens) elif spec_decode_algo == "DFLASH": spec_config = DFlashDecodingConfig( max_draft_len=args.spec_decode_max_draft_len, diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 75f4f9d84465..a43d6e6f530b 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1190,3 +1190,72 @@ def _(like: torch.Tensor, out_shape = shape if shape is not None else list(like.shape) dtype = out_dtype if out_dtype is not None else like.dtype return like.new_empty(out_shape, dtype=dtype), output_buffer_kind + + @torch.library.register_fake("trtllm::compute_probs_from_logits_op") + def _(logits: torch.Tensor, + temperatures: torch.Tensor, + top_k: Optional[torch.Tensor] = None, + top_p: Optional[torch.Tensor] = None, + skip_temperature: bool = False) -> torch.Tensor: + return logits.new_empty(list(logits.shape), dtype=torch.float32) + + @torch.library.register_fake("trtllm::build_draft_prob_indices_out_op") + def _(topkScoreIndices: torch.Tensor, draftProbIndices: torch.Tensor, + topK: int, numDraftTokens: int) -> None: + return None + + @torch.library.register_fake("trtllm::verify_dynamic_tree_rejection_out_op") + def _(candidates: torch.Tensor, draftProbs: torch.Tensor, + targetProbs: torch.Tensor, targetSupportIndices: torch.Tensor, + targetSupportLengths: torch.Tensor, draftProbIndices: torch.Tensor, + retrieveNextToken: torch.Tensor, retrieveNextSibling: torch.Tensor, + treeValid: torch.Tensor, acceptIndex: torch.Tensor, + acceptTokenNum: torch.Tensor, acceptToken: torch.Tensor, + numSpecStep: int, seed: torch.Tensor, offset: torch.Tensor) -> None: + return None + + @torch.library.register_fake( + "trtllm::compute_draft_probs_for_dynamic_tree_rejection_op") + def _(draftLogits: torch.Tensor, + temperatures: torch.Tensor, + numDraftProbRows: int, + targetVocabSize: int, + top_k: Optional[torch.Tensor] = None, + top_p: Optional[torch.Tensor] = None, + skip_temperature: bool = False, + d2t: Optional[torch.Tensor] = None, + top_k_max: int = 0, + skip_all_sampling_params: bool = False) -> torch.Tensor: + batch_size = temperatures.shape[0] + return draftLogits.new_empty( + (batch_size, numDraftProbRows, targetVocabSize), + dtype=torch.float32) + + @torch.library.register_fake( + "trtllm::compute_target_probs_for_dynamic_tree_rejection_op") + def _( + targetLogits: torch.Tensor, + temperatures: torch.Tensor, + numDraftTokens: int, + top_k: Optional[torch.Tensor] = None, + top_p: Optional[torch.Tensor] = None, + skip_temperature: bool = False, + top_k_max: int = 0, + skip_all_sampling_params: bool = False + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size = temperatures.shape[0] + target_vocab_size = targetLogits.shape[-1] + target_probs = targetLogits.new_empty( + (batch_size, numDraftTokens, target_vocab_size), + dtype=torch.float32) + has_filtering = (top_k is not None) or (top_p is not None) + if skip_all_sampling_params or not has_filtering: + support_indices = targetLogits.new_empty((0, ), dtype=torch.int32) + support_lengths = targetLogits.new_empty((0, ), dtype=torch.int32) + else: + support_dim = top_k_max if top_k_max > 0 else target_vocab_size + support_indices = targetLogits.new_empty( + (batch_size, numDraftTokens, support_dim), dtype=torch.int32) + support_lengths = targetLogits.new_empty( + (batch_size, numDraftTokens), dtype=torch.int32) + return target_probs, support_indices, support_lengths diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 522157840f55..fa33e0ce8643 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -446,9 +446,13 @@ 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 + # Eagle3DecodingConfig._max_batch_size is internally managed: the + # dynamic-tree worker pre-allocates batch-indexed CUDA buffers sized + # by this value, and runtime indexes them with no bounds check. It + # MUST equal the global max_batch_size to avoid OOB; we populate it + # here as the single source of truth. + if hasattr(spec_config, '_max_batch_size'): + spec_config._max_batch_size = max_batch_size # WAR for https://nvbugs/5807902 # Disable separate draft KV cache in disaggregated mode diff --git a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py index 2565e034aaf1..77edcb7c1cb1 100644 --- a/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py +++ b/tensorrt_llm/_torch/speculative/dynamic_tree_ops.py @@ -51,6 +51,7 @@ def __init__( max_batch_size: int, device: torch.device, ): + """Allocate reusable CUDA buffers for dynamic-tree build/verify ops.""" self.K = dynamic_tree_max_topK self.depth = max_draft_len @@ -70,6 +71,39 @@ def __init__( max_batch_size, max_path_len, dtype=torch.int64, device=device ) + # Pre-allocated output buffers for verify_dynamic_tree_rejection_out + self._rej_accept_index_buf = torch.zeros( + max_batch_size, max_path_len, dtype=torch.int64, device=device + ) + self._rej_accept_token_num_buf = torch.zeros( + max_batch_size, dtype=torch.int64, device=device + ) + self._rej_accept_token_buf = torch.zeros( + max_batch_size, max_path_len, dtype=torch.int64, device=device + ) + self._rej_seed_buf = torch.zeros(1, dtype=torch.int64, device=device) + self._rej_offset_buf = torch.zeros(1, dtype=torch.int64, device=device) + + def _get_rejection_rng_tensor( + self, + value: int | torch.Tensor, + buffer: torch.Tensor, + name: str, + ) -> torch.Tensor: + """Normalize a rejection RNG input to a one-element int64 CUDA tensor.""" + if isinstance(value, int): + buffer.fill_(value) + return buffer + if not isinstance(value, torch.Tensor): + raise TypeError(f"{name} must be an int or torch.Tensor, got {type(value)!r}") + if value.dtype != torch.int64: + raise TypeError(f"{name} must be int64 tensor, got {value.dtype}") + if value.numel() != 1: + raise ValueError(f"{name} tensor must have exactly one element, got {value.numel()}") + if value.device != buffer.device: + raise ValueError(f"{name} tensor must be on device {buffer.device}, got {value.device}") + return value.reshape(-1) + def build_dynamic_tree( self, history_draft_tokens_parent_buffer: torch.Tensor, @@ -112,6 +146,14 @@ def build_dynamic_tree( # Non-packed path ignores this (bool [bs,N,N] still has a last dim). num_int32_per_row = tree_mask.shape[-1] + # The CUDA builder writes only active tree links/bits. Clear the reused + # work buffers first so stale slot/tree state cannot leak into this tree. + tree_mask[:bs].zero_() + positions[:bs].zero_() + retrieve_index[:bs].zero_() + retrieve_next_token[:bs].fill_(-1) + retrieve_next_sibling[:bs].fill_(-1) + # Call CUDA kernel in-place try: torch.ops.trtllm.build_dynamic_tree_op( @@ -196,3 +238,113 @@ def verify_dynamic_tree_greedy_out( ) from e return predicts, accept_index, accept_token_num, accept_token + + def verify_dynamic_tree_rejection_from_logits_out( + self, + candidates: torch.Tensor, + draft_logits_tree: torch.Tensor, + target_logits_tree: torch.Tensor, + draft_prob_indices: torch.Tensor, + retrieve_next_token: torch.Tensor, + retrieve_next_sibling: torch.Tensor, + tree_valid: torch.Tensor, + temperatures: torch.Tensor, + top_k: torch.Tensor | None, + top_p: torch.Tensor | None, + skip_temperature: bool, + num_gens: int, + num_spec_step: int, + seed: int | torch.Tensor = 0, + offset: int | torch.Tensor = 0, + d2t: torch.Tensor | None = None, + skip_all_sampling_params: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Tree-aware rejection sampling from logits (three CUDA ops). + + This path keeps draft/target logits as inputs, computes unique draft + and target probabilities with separate CUDA ops, then runs the tree + rejection kernel as a third CUDA op. `draft_prob_indices` maps each + tree position to its shared draft-prob row. `tree_valid` guards + first-gen and dummy requests that do not have a usable tree yet. + """ + accept_index = self._rej_accept_index_buf[:num_gens] + accept_token = self._rej_accept_token_buf[:num_gens] + accept_tok_num = self._rej_accept_token_num_buf[:num_gens] + seed_tensor = self._get_rejection_rng_tensor(seed, self._rej_seed_buf, "seed") + offset_tensor = self._get_rejection_rng_tensor(offset, self._rej_offset_buf, "offset") + num_draft_tokens = candidates.shape[1] + if num_gens <= 0: + raise ValueError(f"num_gens must be positive, got {num_gens}") + if draft_logits_tree.shape[0] % num_gens != 0: + raise ValueError( + f"draft_logits_tree rows ({draft_logits_tree.shape[0]}) must be divisible by " + f"num_gens ({num_gens})" + ) + num_draft_prob_rows = draft_logits_tree.shape[0] // num_gens + target_vocab_size = target_logits_tree.shape[-1] + + if tree_valid is None: + tree_valid = torch.ones(num_gens, dtype=torch.bool, device=candidates.device) + tree_valid = tree_valid.contiguous() + + if top_k is None: + top_k_max = 0 + else: + enabled_top_k = top_k[(top_k > 0) & (top_k < target_vocab_size)] + top_k_max = int(enabled_top_k.max().item()) if enabled_top_k.numel() > 0 else 0 + + try: + draft_probs_tree = torch.ops.trtllm.compute_draft_probs_for_dynamic_tree_rejection_op( + draft_logits_tree, + temperatures, + num_draft_prob_rows, + target_vocab_size, + top_k, + top_p, + skip_temperature, + d2t=d2t, + top_k_max=top_k_max, + skip_all_sampling_params=skip_all_sampling_params, + ) + + ( + target_probs_tree, + target_support_indices, + target_support_lengths, + ) = torch.ops.trtllm.compute_target_probs_for_dynamic_tree_rejection_op( + target_logits_tree, + temperatures, + num_draft_tokens, + top_k, + top_p, + skip_temperature, + top_k_max=top_k_max, + skip_all_sampling_params=skip_all_sampling_params, + ) + + torch.ops.trtllm.verify_dynamic_tree_rejection_out_op( + candidates, + draft_probs_tree, + target_probs_tree, + target_support_indices, + target_support_lengths, + draft_prob_indices, + retrieve_next_token, + retrieve_next_sibling, + tree_valid, + accept_index, + accept_tok_num, + accept_token, + num_spec_step, + seed_tensor, + offset_tensor, + ) + except Exception as e: + raise RuntimeError( + f"dynamic tree rejection op chain failed: {e}\n" + f"Inputs: num_gens={num_gens}, N={candidates.shape[1]}, " + f"draft_vocab={draft_logits_tree.shape[-1]}, " + f"target_vocab={target_logits_tree.shape[-1]}, num_spec_step={num_spec_step}" + ) from e + + return target_support_indices, accept_index, accept_tok_num, accept_token diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index b7f95df29c81..828bb4b868cc 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -221,6 +221,7 @@ def __post_init__(self): self.is_spec_dec_dynamic_tree = False def prepare(self): + super().prepare() is_first_draft = self.eagle3_resource_manager.is_first_draft spec_tree_manager = self.eagle3_resource_manager.spec_tree_manager # Update start indices @@ -408,6 +409,7 @@ def is_layer_capture(self, layer_id: int): return layer_id in self.layers_to_capture def prepare(self): + super().prepare() assert self.request_ids is not None # update batch indices num_seqs = len(self.request_ids) @@ -652,6 +654,7 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, """Original linear draft loop (1 token per layer).""" runtime_draft_len = spec_metadata.runtime_draft_len next_draft_tokens = [] + draft_logits_list = [] position_ids = inputs["position_ids"] with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): @@ -704,6 +707,9 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, d2t, draft_step=i) + if spec_metadata.use_rejection_sampling: + draft_logits_list.append(logits.clone()) + new_draft_token = self.draft_decoder(logits, draft_model) next_draft_tokens.append(new_draft_token) # update inputs @@ -745,6 +751,12 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, gen_draft_tokens) next_draft_tokens[num_contexts:] = gen_draft_tokens + if spec_metadata.use_rejection_sampling and draft_logits_list: + d2t_param = getattr(draft_model.model, "d2t", None) + spec_metadata.d2t = d2t_param.data if d2t_param is not None else None + self._compute_and_store_draft_probs(draft_logits_list, + spec_metadata, batch_size) + return next_draft_tokens def sample_and_accept_draft_tokens( @@ -765,8 +777,8 @@ def sample_and_accept_draft_tokens( 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) + return self._accept_draft_tokens(logits, draft_tokens, num_contexts, + batch_size, spec_metadata) def draft_decoder( self, diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py index f33a5824a455..4fad34a74131 100644 --- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py @@ -14,7 +14,7 @@ # limitations under the License. """Eagle3 one-model dynamic tree speculative decoding.""" -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import torch import triton @@ -146,6 +146,7 @@ class Eagle3OneModelDynamicTreeWorker(Eagle3OneModelWorker): def __init__( self, spec_config: "EagleDecodingConfig", mapping, use_separate_draft_kv_cache: bool = False ): + """Initialize dynamic-tree specific buffers and helper ops.""" super().__init__(spec_config, mapping, use_separate_draft_kv_cache) assert self.use_dynamic_tree, ( "Eagle3OneModelDynamicTreeWorker requires use_dynamic_tree=True" @@ -156,12 +157,14 @@ def __init__( 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 + # Eagle3DecodingConfig._max_batch_size is auto-populated by + # py_executor_creator from the global max_batch_size, so it should + # always be set by the time we get here. + assert spec_config._max_batch_size is not None, ( + "Eagle3DecodingConfig._max_batch_size was not populated; " + "py_executor_creator should have set it from the global max_batch_size." + ) + self._max_batch_size = spec_config._max_batch_size K = self.K max_draft_len = spec_config.max_draft_len @@ -279,6 +282,26 @@ def __init__( sm = get_sm_version() self._needs_mask_repack = sm < 100 or sm in (120, 121) + # Rejection sampling buffers + # Unique draft logits per request, one row per distinct parent context. + # Shape: [max_batch_size, 1 + (max_draft_len - 1) * K, vocab_size] + # row 0 -> p(.|root) + # rows [1 : 1 + K] -> p(.|depth-0 parent_k) + # rows [1 + K : 1 + 2K] -> p(.|depth-1 parent_k) + # This avoids materializing repeated per-tree-position logits such as + # [p(.|E), p(.|E), p(.|E), p(.|F1), p(.|F1), p(.|F2)]. + self._draft_depth_logits_cat: Optional[torch.Tensor] = None + # topk_score_indices from resampling_final_draft_tokens for path tracing. + # Shape: [max_batch_size, max_total_draft_tokens] + self._topk_score_indices_buf = torch.zeros( + max_batch_size, self.max_total_draft_tokens, dtype=torch.int64, device="cuda" + ) + # Tree position -> unique draft-prob row mapping for rejection sampling. + # Shape: [max_batch_size, max_total_draft_tokens + 1], root row is unused and kept at 0. + self._draft_prob_indices_buf = torch.zeros( + max_batch_size, self.max_total_draft_tokens + 1, dtype=torch.int32, device="cuda" + ) + 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 @@ -524,6 +547,7 @@ def _forward_draft_loop( ): """Dynamic tree draft loop with growing context.""" spec_tree_manager = self.spec_tree_manager + self._d2t = getattr(draft_model.model, "d2t", None) assert batch_size <= self._max_batch_size, ( f"batch_size {batch_size} exceeds pre-allocated max_batch_size {self._max_batch_size}" @@ -577,6 +601,11 @@ def _forward_draft_loop( hidden_states[gather_ids], draft_model.lm_head, attn_metadata, True ) + # Capture unique draft logits for the root parent context. + if spec_metadata.use_rejection_sampling and num_gens > 0: + self._lazy_alloc_draft_logits_buf(logits.shape[-1], logits.dtype, logits.device) + self._draft_depth_logits_cat[:num_gens, 0].copy_(logits[num_contexts:]) + new_draft_tokens, new_draft_scores = self.sample( logits, self.K, draft_model=draft_model ) @@ -643,6 +672,13 @@ def _forward_draft_loop( selected_hs, draft_model.lm_head, attn_metadata, True ) + # Capture unique draft logits for the K parent contexts at this depth. + if spec_metadata.use_rejection_sampling and num_gens > 0: + row_start = 1 + (layer_idx - 1) * self.K + self._draft_depth_logits_cat[:num_gens, row_start : row_start + self.K].copy_( + logits[num_contexts * self.K :].reshape(num_gens, self.K, -1) + ) + new_draft_tokens, new_draft_scores = self.sample( logits, self.K, draft_model=draft_model ) @@ -672,6 +708,15 @@ def _forward_draft_loop( # Resample final tokens and build tree real_draft_tokens, topk_score_indices = self.resampling_final_draft_tokens(batch_size) + # Save topk_score_indices for rejection sampling path tracing. + # Rejection sampling needs to map each final draft token back to its history + # buffer index to retrieve the corresponding draft logits. Greedy verification + # doesn't need this mapping since it only compares token IDs. + # Only save the gen rows (skip context rows) so that + # _build_draft_prob_indices can read them back at [:num_gens]. + if spec_metadata.use_rejection_sampling: + self._topk_score_indices_buf[:num_gens].copy_(topk_score_indices[num_contexts:]) + 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( @@ -698,6 +743,11 @@ def _sample_and_accept_dynamic_tree( self, logits, attn_metadata, spec_metadata, batch_size, num_contexts, num_gens ): """Dynamic tree verification using CUDA kernel.""" + if num_gens > self._max_batch_size: + raise RuntimeError( + f"Dynamic tree batch size {num_gens} exceeds configured " + f"max_batch_size {self._max_batch_size}" + ) N = self.tokens_per_gen_step max_path_len = self._max_path_len @@ -733,12 +783,91 @@ def _sample_and_accept_dynamic_tree( 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. + # 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] + if self._can_use_rejection_sampling(spec_metadata): + vocab_size = logits.shape[-1] + num_ctx_tokens = logits.shape[0] - num_gens * N + device = logits.device + + draft_logits_tree = self._get_unique_draft_logits(num_gens) + draft_prob_indices = self._build_draft_prob_indices(num_gens) + target_logits_tree = logits[num_ctx_tokens:].reshape(-1, vocab_size) + gen_slice = slice(num_contexts, num_contexts + num_gens) + skip_top_k = getattr(spec_metadata, "skip_top_k", False) + skip_top_p = getattr(spec_metadata, "skip_top_p", False) + skip_temperature = getattr(spec_metadata, "skip_temperature", False) + + if spec_metadata.request_temperatures is None: + temps = torch.ones(num_gens, dtype=torch.float32, device=device) + skip_temperature = True + else: + temps = spec_metadata.request_temperatures[gen_slice] + + top_ks = None + if not skip_top_k and spec_metadata.request_top_ks is not None: + top_ks = spec_metadata.request_top_ks[gen_slice] + + top_ps = None + if not skip_top_p and spec_metadata.request_top_ps is not None: + top_ps = spec_metadata.request_top_ps[gen_slice] + + skip_all_sampling_params = ( + skip_temperature + and skip_top_k + and skip_top_p + and not getattr(spec_metadata, "has_greedy_requests", False) + ) + + # Lazily initialize seed/offset tensors on correct device + if self.seed is None: + self.seed = torch.tensor([0], dtype=torch.int64, device=device) + self.offset = torch.tensor([0], dtype=torch.int64, device=device) + # Use in-place operations for CUDA graph compatibility + self.seed.add_(1).remainder_(2**31) + + _, accept_index, accept_token_num, accept_token = ( + self.tree_ops_converter.verify_dynamic_tree_rejection_from_logits_out( + candidates, + draft_logits_tree, + target_logits_tree, + draft_prob_indices, + spec_tree_manager.retrieve_next_token[:num_gens], + spec_tree_manager.retrieve_next_sibling[:num_gens], + tree_valid, + temps, + top_ks, + top_ps, + skip_temperature, + num_gens, + self._max_path_len, + seed=self.seed, + offset=self.offset, + d2t=self._d2t, + skip_all_sampling_params=skip_all_sampling_params, + ) + ) + + self._finalize_dynamic_tree_verify_outputs( + accept_index=accept_index, + accept_token_num=accept_token_num, + accept_token=accept_token, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + num_gens=num_gens, + max_path_len=max_path_len, + ) + num_accepted_tokens = self._apply_force_accepted_tokens( + num_accepted_tokens, num_contexts, self.max_draft_len + ) + return accepted_tokens, num_accepted_tokens + _, accept_index, accept_token_num, accept_token = ( self.tree_ops_converter.verify_dynamic_tree_greedy_out( candidates, @@ -752,14 +881,17 @@ def _sample_and_accept_dynamic_tree( ) ) - 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) + self._finalize_dynamic_tree_verify_outputs( + accept_index=accept_index, + accept_token_num=accept_token_num, + accept_token=accept_token, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + num_gens=num_gens, + max_path_len=max_path_len, + ) num_accepted_tokens = self._apply_force_accepted_tokens( num_accepted_tokens, num_contexts, self.max_draft_len @@ -767,6 +899,91 @@ def _sample_and_accept_dynamic_tree( return accepted_tokens, num_accepted_tokens + def _can_use_rejection_sampling(self, spec_metadata) -> bool: + """Check if rejection sampling can be used for dynamic tree verification. + + Dynamic tree uses its own compact unique-logit buffer + (_draft_depth_logits_cat) instead of spec_metadata.draft_logits, so we + check that buffer's allocation status rather than draft_logits_valid. + The buffer is lazily allocated and populated during the first forward + pass with generation requests. + + Args: + spec_metadata: Speculative decoding metadata + + Returns: + True if rejection sampling is enabled and the draft logit buffer is allocated + """ + return spec_metadata.use_rejection_sampling and self._draft_depth_logits_cat is not None + + def _finalize_dynamic_tree_verify_outputs( + self, + accept_index: torch.Tensor, + accept_token_num: torch.Tensor, + accept_token: torch.Tensor, + accepted_tokens: torch.Tensor, + num_accepted_tokens: torch.Tensor, + num_contexts: int, + batch_size: int, + num_gens: int, + max_path_len: int, + ) -> None: + """Write dynamic-tree verify outputs back to the shared step buffers.""" + self._accept_token = accept_token + accepted_draft_count = accept_token_num[:num_gens] + num_accepted_tokens[num_contexts:batch_size] = (accepted_draft_count + 1).to(torch.int32) + accepted_tokens[num_contexts:batch_size] = accept_token[:num_gens].to(torch.int32) + # accept_index stores root at slot 0. Subtract 1 so padding/root value 0 becomes sentinel -1. + self._accepted_draft_indices_tensor[num_contexts:batch_size] = ( + accept_index[:num_gens, 1:max_path_len] - 1 + ).to(torch.int32) + + def _lazy_alloc_draft_logits_buf(self, vocab_size: int, dtype, device): + """Lazily allocate unique draft-logit capture buffer.""" + if self._draft_depth_logits_cat is None: + rows = 1 + (self.max_draft_len - 1) * self.K + self._draft_depth_logits_cat = torch.empty( + self._max_batch_size, rows, vocab_size, dtype=dtype, device=device + ) + + def _get_unique_draft_logits( + self, + num_gens: int, + ) -> torch.Tensor: + """Return compact unique draft logits for rejection sampling. + + The returned rows are unique per parent context, not duplicated per + final tree position. + + Args: + num_gens: Number of generation requests. + + Returns: + draft_logits: [num_gens * (1 + (max_draft_len - 1) * K), draft_vocab_size] + """ + return self._draft_depth_logits_cat[:num_gens].reshape( + -1, self._draft_depth_logits_cat.shape[-1] + ) + + def _build_draft_prob_indices( + self, + num_gens: int, + ) -> torch.Tensor: + """Build tree-position -> unique draft-prob row mapping. + + For a final tree position: + - depth 0 children map to row 0, which stores p(.|root) + - deeper nodes map to row 1 + depth_bucket * K + parent_k + """ + draft_prob_indices = self._draft_prob_indices_buf[:num_gens] + torch.ops.trtllm.build_draft_prob_indices_out_op( + self._topk_score_indices_buf[:num_gens], + draft_prob_indices, + self.K, + self.max_total_draft_tokens, + ) + return draft_prob_indices + @nvtx_range("eagle3_dyn.sample") def sample( self, logits: torch.Tensor, max_top_k: int, draft_model=None diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index d39ace8988cd..c7a8358124cf 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -25,6 +25,10 @@ if IS_FLASHINFER_AVAILABLE: import flashinfer +from .one_model_sampler import (compute_probs_from_logits, + rejection_sampling_one_model, + sampling_batch_spec_dec_one_model) + # Environment variable name for forcing the number of accepted tokens in speculative decoding FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR = "TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS" @@ -434,10 +438,34 @@ class SpecMetadata: # For non-greedy sampling on 1-model. allow_advanced_sampling: bool = False + # Whether to use rejection sampling for one-model speculative decoding. + use_rejection_sampling: bool = False # Sampling parameters for non-greedy sampling (per-request) temperatures: Optional[torch.Tensor] = None top_ks: Optional[torch.Tensor] = None top_ps: Optional[torch.Tensor] = None + # Whether top-k/top-p/temperature are globally disabled for the current batch. + skip_temperature: bool = False + skip_top_k: bool = False + skip_top_p: bool = False + has_greedy_requests: bool = False + # Sampling parameters indexed per request. + request_temperatures: Optional[torch.Tensor] = None + request_top_ks: Optional[torch.Tensor] = None + request_top_ps: Optional[torch.Tensor] = None + # Whether to use sampling parameters when sampling draft tokens. + use_sampling_params_for_draft_tokens: bool = False + # Vocab size used for draft_probs buffer allocation. + vocab_size: int = 0 + # Draft probabilities buffer for rejection sampling, stored flat. + draft_probs: Optional[torch.Tensor] = None + draft_probs_vocab_size: int = 0 + # Whether draft_probs contains valid data. + draft_probs_valid: bool = False + # Last dimension size of the draft logits/probs stored in draft_probs. + draft_probs_last_dim: int = 0 + # Draft-to-target vocab offset tensor. + d2t: Optional[torch.Tensor] = None def __post_init__(self): pass @@ -446,6 +474,14 @@ def prepare(self): """ Hook to be called before the forward step of the model. """ + if (self.use_rejection_sampling and self.draft_probs is None + and self.vocab_size > 0): + buffer_size = (self.max_num_requests * self.max_draft_len * + self.vocab_size) + self.draft_probs = torch.empty(buffer_size, + dtype=torch.float32, + device='cuda') + self.draft_probs_vocab_size = self.vocab_size def create_cuda_graph_metadata(self, max_batch_size: int): """ @@ -494,6 +530,13 @@ def populate_sampling_params_for_one_model( temperatures = [] top_ks = [] top_ps = [] + request_temperatures = [] + request_top_ks = [] + request_top_ps = [] + top_k_enabled = False + top_p_enabled = False + has_greedy_requests = False + temperature_enabled = False # Need to use a very small value for temperature when disabled to avoid division by 0 DISABLE_TEMP_VAL = 1e-5 @@ -501,47 +544,104 @@ def populate_sampling_params_for_one_model( DISABLE_TOPK_VAL = torch.iinfo(torch.int32).max DISABLE_TOPP_VAL = 1.0 - for request in requests: - sampling_config = request.sampling_config - temp = sampling_config.temperature - temp_val = temp[0] if temp is not None and len(temp) > 0 else None + def _first_or_none(values): + """Return the first sampling parameter value when present.""" + return values[0] if values is not None and len(values) > 0 else None + + def _normalize_request_sampling_params( + *, + temperature: Optional[float], + top_k: Optional[int], + top_p: Optional[float], + ) -> tuple[float, int, float, bool, bool, bool, bool]: + """Convert request sampling params into normalized per-request scalars.""" + is_greedy = SamplingParams.params_imply_greedy_decoding( + temperature=temperature, + top_k=top_k, + top_p=top_p, + use_beam_search=False) - tk = sampling_config.top_k - tk_val = tk[0] if tk is not None and len(tk) > 0 else None + use_temperature = (not is_greedy + and temperature not in (None, 0, 1)) + use_top_k = not is_greedy and top_k is not None and top_k > 0 + use_top_p = not is_greedy and top_p is not None and top_p < 1.0 + + normalized_temperature = (DISABLE_TEMP_VAL + if is_greedy or temperature is None + or temperature == 0 else temperature) + normalized_top_k = DISABLE_TOPK_VAL if not use_top_k else top_k + normalized_top_p = (DISABLE_TOPP_VAL + if is_greedy or top_p is None else top_p) + + return ( + normalized_temperature, + normalized_top_k, + normalized_top_p, + use_temperature, + use_top_k, + use_top_p, + is_greedy, + ) - tp = sampling_config.top_p - tp_val = tp[0] if tp is not None and len(tp) > 0 else None + for request in requests: + sampling_config = request.sampling_config + temp_val = _first_or_none(sampling_config.temperature) + tk_val = _first_or_none(sampling_config.top_k) + tp_val = _first_or_none(sampling_config.top_p) # Context requests have no draft tokens yet. num_tokens = 1 + self.runtime_draft_len if request.state == LlmRequestState.GENERATION_IN_PROGRESS else 1 - is_greedy = SamplingParams.params_imply_greedy_decoding( + ( + temp_val, + tk_val, + tp_val, + use_temperature, + use_top_k, + use_top_p, + is_greedy, + ) = _normalize_request_sampling_params( temperature=temp_val, top_k=tk_val, top_p=tp_val, - use_beam_search=False) + ) - temp_val = DISABLE_TEMP_VAL if is_greedy or temp_val is None or temp_val == 0 else temp_val - tk_val = DISABLE_TOPK_VAL if is_greedy or tk_val is None or tk_val <= 0 else tk_val - tp_val = DISABLE_TOPP_VAL if is_greedy or tp_val is None else tp_val + temperature_enabled |= use_temperature + top_k_enabled |= use_top_k + top_p_enabled |= use_top_p + has_greedy_requests |= is_greedy + request_temperatures.append(temp_val) + request_top_ks.append(tk_val) + request_top_ps.append(tp_val) temperatures.extend(temp_val for _ in range(num_tokens)) top_ks.extend(tk_val for _ in range(num_tokens)) top_ps.extend(tp_val for _ in range(num_tokens)) - if self.temperatures is None: - self.temperatures = torch.ones( - (self.max_draft_len + 1) * self.max_num_requests, - dtype=torch.float32, - device='cuda') - self.top_ks = torch.zeros( - (self.max_draft_len + 1) * self.max_num_requests, - dtype=torch.int32, - device='cuda') - self.top_ps = torch.ones( - (self.max_draft_len + 1) * self.max_num_requests, - dtype=torch.float32, - device='cuda') + tokens_per_request = (self.max_total_draft_tokens + 1 if + self.is_spec_dec_tree else self.max_draft_len + 1) + required_flat_size = tokens_per_request * self.max_num_requests + + if self.temperatures is None or self.temperatures.numel( + ) < required_flat_size: + self.temperatures = torch.ones(required_flat_size, + dtype=torch.float32, + device='cuda') + self.top_ks = torch.zeros(required_flat_size, + dtype=torch.int32, + device='cuda') + self.top_ps = torch.ones(required_flat_size, + dtype=torch.float32, + device='cuda') + self.request_temperatures = torch.ones(self.max_num_requests, + dtype=torch.float32, + device='cuda') + self.request_top_ks = torch.zeros(self.max_num_requests, + dtype=torch.int32, + device='cuda') + self.request_top_ps = torch.ones(self.max_num_requests, + dtype=torch.float32, + device='cuda') self.temperatures[:len(temperatures)].copy_(torch.tensor( temperatures, dtype=torch.float32, pin_memory=prefer_pinned()), @@ -552,6 +652,27 @@ def populate_sampling_params_for_one_model( self.top_ps[:len(top_ps)].copy_(torch.tensor( top_ps, dtype=torch.float32, pin_memory=prefer_pinned()), non_blocking=True) + self.request_temperatures[:len(request_temperatures)].copy_( + torch.tensor(request_temperatures, + dtype=torch.float32, + pin_memory=prefer_pinned()), + non_blocking=True) + self.request_top_ks[:len(request_top_ks)].copy_( + torch.tensor(request_top_ks, + dtype=torch.int32, + pin_memory=prefer_pinned()), + non_blocking=True, + ) + self.request_top_ps[:len(request_top_ps)].copy_( + torch.tensor(request_top_ps, + dtype=torch.float32, + pin_memory=prefer_pinned()), + non_blocking=True, + ) + self.skip_temperature = not temperature_enabled + self.skip_top_k = not top_k_enabled + self.skip_top_p = not top_p_enabled + self.has_greedy_requests = has_greedy_requests class SpecWorkerBase(nn.Module, ABC): @@ -862,6 +983,115 @@ def _sample_and_accept_draft_tokens_base( return accepted_tokens, num_accepted_tokens + def _accept_draft_tokens(self, logits, draft_tokens, num_contexts, + batch_size, spec_metadata): + """ + Accept draft tokens with optional rejection sampling support. + """ + if self._can_use_rejection_sampling(spec_metadata, num_contexts): + draft_len = draft_tokens.shape[1] + stored_vocab = (spec_metadata.draft_probs_last_dim + if spec_metadata.draft_probs_last_dim > 0 else + spec_metadata.draft_probs_vocab_size) + draft_probs = spec_metadata.draft_probs[:batch_size * draft_len * + stored_vocab].reshape( + batch_size, draft_len, + stored_vocab) + return self._sample_and_accept_draft_tokens_rejection( + logits, draft_tokens, draft_probs, batch_size, spec_metadata) + return self._sample_and_accept_draft_tokens_base( + logits, draft_tokens, num_contexts, batch_size, spec_metadata) + + def _can_use_rejection_sampling(self, spec_metadata: SpecMetadata, + num_contexts: int) -> bool: + return (spec_metadata.use_rejection_sampling + and spec_metadata.draft_probs_valid and num_contexts == 0) + + def _sample_and_accept_draft_tokens_rejection( + self, + logits: torch.Tensor, + draft_tokens: torch.Tensor, + draft_probs: torch.Tensor, + batch_size: int, + spec_metadata, + ): + """ + Rejection-sampling acceptance for one-model speculative decoding. + """ + device = logits.device + vocab_size = logits.shape[-1] + + if logits.dim() == 1: + logits = logits.unsqueeze(0) + + runtime_draft_len = draft_tokens.shape[1] + draft_vocab_size = draft_probs.shape[-1] + num_target_tokens = batch_size * (runtime_draft_len + 1) + + temperatures = spec_metadata.temperatures[:num_target_tokens] + # Pass None instead of an all-disabled tensor so the C++ op can short-circuit + # on a host-side check rather than a `.item()` sync, which would break + # CUDA graph capture. + top_ks = None if spec_metadata.skip_top_k else spec_metadata.top_ks[: + num_target_tokens] + top_ps = None if spec_metadata.skip_top_p else spec_metadata.top_ps[: + num_target_tokens] + + target_probs_flat = compute_probs_from_logits(logits.clone(), + temperatures, top_ks, + top_ps) + target_probs = target_probs_flat.reshape(batch_size, + runtime_draft_len + 1, + vocab_size) + + assert draft_probs.shape[1] == runtime_draft_len, ( + f"draft_probs draft length mismatch: {draft_probs.shape[1]} != " + f"{runtime_draft_len}") + d2t = getattr(spec_metadata, "d2t", None) + if draft_vocab_size != vocab_size: + full_draft_probs = torch.zeros( + (batch_size, runtime_draft_len, vocab_size), + dtype=torch.float32, + device=device) + if d2t is not None: + assert d2t.numel() == draft_vocab_size, ( + f"d2t size mismatch: {d2t.numel()} != {draft_vocab_size}") + d2t = d2t.to(device=device) + source_indices = torch.arange(draft_vocab_size, + device=device, + dtype=torch.long) + target_indices = (source_indices + d2t) % vocab_size + full_draft_probs[:, :runtime_draft_len, + target_indices] = draft_probs + else: + assert draft_vocab_size < vocab_size + full_draft_probs[:, :runtime_draft_len, :draft_vocab_size] = ( + draft_probs) + else: + full_draft_probs = draft_probs + + full_draft_tokens = draft_tokens.to(torch.int32).contiguous() + + if self.seed is None: + self.seed = torch.tensor([0], dtype=torch.int64, device=device) + if self.offset is None: + self.offset = torch.tensor([0], dtype=torch.int64, device=device) + self.seed += 1 + self.seed %= 2**31 + + accepted_tokens, num_accepted_tokens = rejection_sampling_one_model( + draft_probs=full_draft_probs, + draft_token_ids=full_draft_tokens, + target_probs=target_probs, + deterministic=True, + seed=self.seed, + offset=self.offset, + ) + + num_accepted_tokens = self._apply_force_accepted_tokens( + num_accepted_tokens, 0, draft_tokens.shape[1]) + return accepted_tokens, num_accepted_tokens + def _draft_sampler_greedy(self, logits: torch.Tensor, d2t=None): """ Simple greedy draft token sampling using argmax. @@ -881,6 +1111,48 @@ def _draft_sampler_greedy(self, logits: torch.Tensor, d2t=None): return draft_tokens.type(torch.int32) + def _compute_and_store_draft_probs( + self, + draft_logits_list: List[torch.Tensor], + spec_metadata: SpecMetadata, + batch_size: int, + ): + """ + Compute draft probabilities and store them for next-step rejection sampling. + """ + draft_tokens_per_request = len(draft_logits_list) + vocab_size = draft_logits_list[0].shape[-1] + device = draft_logits_list[0].device + + draft_logits = torch.stack(draft_logits_list, dim=0) + draft_logits_flat = draft_logits.transpose(0, 1).reshape(-1, vocab_size) + + num_draft_tokens = batch_size * draft_tokens_per_request + if spec_metadata.request_temperatures is not None: + draft_temps = spec_metadata.request_temperatures[:batch_size].repeat_interleave( + draft_tokens_per_request) + draft_top_ks = ( + spec_metadata.request_top_ks[:batch_size].repeat_interleave( + draft_tokens_per_request) if not spec_metadata.skip_top_k + and spec_metadata.request_top_ks is not None else None) + draft_top_ps = ( + spec_metadata.request_top_ps[:batch_size].repeat_interleave( + draft_tokens_per_request) if not spec_metadata.skip_top_p + and spec_metadata.request_top_ps is not None else None) + else: + draft_temps = torch.ones(num_draft_tokens, device=device) + draft_top_ks = None + draft_top_ps = None + + draft_probs_flat = compute_probs_from_logits(draft_logits_flat, + draft_temps, draft_top_ks, + draft_top_ps) + num_elements = batch_size * draft_tokens_per_request * vocab_size + spec_metadata.draft_probs[:num_elements].copy_( + draft_probs_flat.flatten()) + spec_metadata.draft_probs_last_dim = vocab_size + spec_metadata.draft_probs_valid = True + def _execute_guided_decoder_if_present(self, logits): """Execute guided decoder on target model logits if available.""" if self.guided_decoder is not None: @@ -1011,8 +1283,6 @@ def _sample_tokens_for_batch( sampled_tokens: [num_tokens] - Sampled token ids """ if spec_metadata.allow_advanced_sampling: - from .one_model_sampler import sampling_batch_spec_dec_one_model - num_gens = batch_size - num_contexts num_tokens = num_contexts + num_gens * ( spec_metadata.runtime_draft_len + 1) diff --git a/tensorrt_llm/_torch/speculative/one_model_sampler.py b/tensorrt_llm/_torch/speculative/one_model_sampler.py index 7d49aa85dd17..7e3b06b383cc 100644 --- a/tensorrt_llm/_torch/speculative/one_model_sampler.py +++ b/tensorrt_llm/_torch/speculative/one_model_sampler.py @@ -1,7 +1,14 @@ -from typing import Optional +from typing import Optional, Tuple import torch -from flashinfer.sampling import top_k_top_p_sampling_from_logits + +from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE + +if IS_FLASHINFER_AVAILABLE: + from flashinfer.sampling import chain_speculative_sampling, top_k_top_p_sampling_from_logits +else: + chain_speculative_sampling = None + top_k_top_p_sampling_from_logits = None def forward_native( @@ -92,6 +99,74 @@ def sampling_batch_spec_dec_one_model( """ logits = apply_temperature(logits, temperatures) if use_flashinfer: + if top_k_top_p_sampling_from_logits is None: + raise RuntimeError("FlashInfer sampling requested but flashinfer is unavailable") return top_k_top_p_sampling_from_logits(logits, top_k, top_p, seed=seed, offset=offset) random_sampled = forward_native(logits, top_k, top_p) return random_sampled + + +def compute_probs_from_logits( + logits: torch.Tensor, + temperatures: torch.Tensor, + top_k: Optional[torch.Tensor], + top_p: Optional[torch.Tensor], + skip_temperature: bool = False, +) -> torch.Tensor: + """ + Compute probabilities from logits with temperature, top-k, and top-p applied. + """ + if logits.is_cuda: + return torch.ops.trtllm.compute_probs_from_logits_op( + logits, temperatures, top_k, top_p, skip_temperature + ) + + if not skip_temperature: + logits = apply_temperature(logits, temperatures) + logits = apply_top_k_top_p(logits, top_k, top_p) + probs = logits.softmax(dim=-1, dtype=torch.float32) + + # Greedy rows should remain exactly one-hot so rejection sampling does not + # spuriously reject numerically-near argmax tokens. + greedy_temp_threshold = 1e-4 + is_greedy = temperatures <= greedy_temp_threshold + argmax_ids = logits.argmax(dim=-1, keepdim=True) + one_hot = torch.zeros_like(probs).scatter_(1, argmax_ids, 1.0) + return torch.where(is_greedy.unsqueeze(1), one_hot, probs) + + +def rejection_sampling_one_model( + draft_probs: torch.Tensor, + draft_token_ids: torch.Tensor, + target_probs: torch.Tensor, + deterministic: bool = True, + seed: Optional[int] = None, + offset: Optional[int] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + CUDA-graph compatible rejection sampling for one-model speculative decoding. + """ + batch_size = draft_token_ids.shape[0] + device = draft_token_ids.device + + output_accepted_token_num = torch.zeros(batch_size, dtype=torch.int32, device=device) + output_emitted_draft_token_num = torch.zeros(batch_size, dtype=torch.int32, device=device) + + if chain_speculative_sampling is None: + raise RuntimeError( + "Rejection sampling for one-model speculative decoding requires flashinfer" + ) + + accepted_tokens, _, output_emitted_draft_token_num = chain_speculative_sampling( + draft_probs, + draft_token_ids, + target_probs, + maybe_output_accepted_token_num=output_accepted_token_num, + maybe_output_emitted_draft_token_num=output_emitted_draft_token_num, + deterministic=deterministic, + generator=None, + seed=seed, + offset=offset, + ) + + return accepted_tokens, output_emitted_draft_token_num + 1 diff --git a/tensorrt_llm/_torch/speculative/spec_tree_manager.py b/tensorrt_llm/_torch/speculative/spec_tree_manager.py index d6e7cd826d06..80d9228e6582 100644 --- a/tensorrt_llm/_torch/speculative/spec_tree_manager.py +++ b/tensorrt_llm/_torch/speculative/spec_tree_manager.py @@ -234,6 +234,8 @@ def mark_tree_valid(self, slot_ids, count): return # index_fill_: graph-capture-safe (no fancy indexing). self.slot_has_tree.index_fill_(0, slot_ids[:count], True) + # CUDA graph padding may use the dummy slot; it must never carry a real tree. + self.slot_has_tree.narrow(0, self._dummy_slot_id, 1).fill_(False) def mark_tree_invalid(self, slot_id): """Clear validity when a slot is freed.""" diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 4243f27bfe7f..68ef347b0533 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -40,6 +40,9 @@ def get_spec_metadata(spec_config, spec_resource_manager=None, is_draft_model=False, max_seq_len=262144): + use_rejection_sampling = getattr(spec_config, "use_rejection_sampling", + False) + vocab_size = getattr(model_config, "vocab_size", 0) if spec_config.spec_dec_mode.is_mtp_one_model(): return MTPSpecMetadata( max_draft_len=spec_config.max_draft_len, @@ -95,6 +98,8 @@ def get_spec_metadata(spec_config, max_num_tokens=max_num_tokens, layers_to_capture=spec_config.eagle3_layers_to_capture, allow_advanced_sampling=spec_config.allow_advanced_sampling, + use_rejection_sampling=use_rejection_sampling, + vocab_size=vocab_size, spec_resource_manager=spec_resource_manager, use_dynamic_tree=spec_config.use_dynamic_tree, eagle_choices=spec_config.eagle_choices, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 3db16e5cabfb..151db668ad80 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -914,6 +914,14 @@ class DecodingBaseConfig(StrictBaseModel): "to 1-model code paths; non-greedy sampling is always enabled on 2-model paths." ) + use_rejection_sampling: bool = Field( + default=False, + status="prototype", + description= + "If true, enables rejection sampling for one-model speculative decoding paths. " + "This is intended for non-greedy sampling configurations on the PyTorch backend. " + "The non-dynamic-tree one-model path requires FlashInfer.") + # If set, drafting is allowed to use chain drafter. _allow_chain_drafter: bool = PrivateAttr(True) # If set, drafting uses greedy sampling, irrespective of sampling parameters. @@ -966,6 +974,17 @@ def validate_draft_len_schedule_and_sort(cls, v, info): return dict(sorted(v.items(), key=lambda x: x[0])) return v + @model_validator(mode='after') + def validate_rejection_sampling_config(self): + """Reject SA-enhanced configurations that invalidate rejection sampling.""" + if self.use_rejection_sampling and getattr(self, 'sa_config', + None) is not None: + raise ValueError( + "use_rejection_sampling is incompatible with sa_config " + "because SA enhancement may override the proposed draft tokens." + ) + return self + @model_validator(mode='after') # 1. Validate that max_concurrency and draft_len_schedule are mutually exclusive. # 2. If max_concurrency is set, translate it to the corresponding draft_len_schedule. @@ -1348,10 +1367,16 @@ 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.") + # Backs the dynamic-tree worker's pre-allocated, batch-indexed CUDA buffers + # (draft_tokens_buffer, history_*_buffer, tree_mask_buffer, etc. in + # Eagle3OneModelDynamicTreeWorker.__init__). This MUST equal the global + # max_batch_size: the worker indexes those buffers with batch_idx in + # [0, global_max_batch_size) at runtime with no bounds check, so any value + # smaller than the global will OOB during warmup or generation as soon as + # batch_idx exceeds this capacity (illegal memory access). It is therefore + # exposed as a PrivateAttr -- not a user-tunable knob -- and is + # auto-populated by py_executor_creator from the global max_batch_size. + _max_batch_size: Optional[int] = PrivateAttr(default=None) sa_config: Optional[SAEnhancerConfig] = Field( default=None, @@ -4120,6 +4145,13 @@ def validate_speculative_config(self): exclude={"decoding_type"}) self.speculative_config = Eagle3DecodingConfig(**eagle_data) + if self.speculative_config.use_rejection_sampling: + if not isinstance(self.speculative_config, + Eagle3DecodingConfig): + raise ValueError( + "use_rejection_sampling is only supported for " + "PyTorch Eagle3 one-model speculative decoding paths.") + if isinstance(self.speculative_config, PARDDecodingConfig): assert self.speculative_config.max_draft_len > 0, "PARD max_draft_len must be > 0" diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 51504b530938..13bf1bb2a9be 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -187,6 +187,54 @@ def test_dummy_load_format(self): task = MMLU(self.MODEL_NAME) task.evaluate(llm, is_integration_test=True) + @pytest.mark.parametrize("use_dynamic_tree", [False, True], + ids=["no_dynamic_tree", "dynamic_tree"]) + def test_eagle3_rejection_dynamic_tree_smoke(self, use_dynamic_tree, + mocker): + """Smoke-test one-model Eagle3 rejection sampling with both tree modes.""" + mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", 128) + + eagle_model_dir = f"{llm_models_root()}/EAGLE3-LLaMA3.1-Instruct-8B" + spec_config_kwargs = dict( + max_draft_len=4, + speculative_model=eagle_model_dir, + eagle3_one_model=True, + allow_advanced_sampling=True, + use_rejection_sampling=True, + ) + max_batch_size = 1 + if use_dynamic_tree: + spec_config_kwargs.update( + use_dynamic_tree=True, + dynamic_tree_max_topK=4, + max_total_draft_tokens=16, + ) + + llm = LLM( + self.MODEL_PATH, + tensor_parallel_size=1, + pipeline_parallel_size=1, + attn_backend="TRTLLM", + disable_overlap_scheduler=True, + cuda_graph_config=None, + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4, + dtype="auto"), + max_seq_len=4096, + max_batch_size=max_batch_size, + speculative_config=Eagle3DecodingConfig(**spec_config_kwargs), + ) + + with llm: + task = GSM8K(self.MODEL_NAME) + sampling_params = SamplingParams(temperature=1.0, + top_p=1.0, + max_tokens=128, + truncate_prompt_tokens=2048) + task.evaluate(llm, + sampling_params=sampling_params, + extra_evaluator_kwargs=dict(apply_chat_template=True), + is_integration_test=True) + @pytest.mark.skip_less_device_memory(32000) @parametrize_with_ids("torch_compile", [False, True]) @parametrize_with_ids("attn_backend", ["TRTLLM", "FLASHINFER"]) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 0044105fb351..4d4a6a6daa73 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -585,6 +585,8 @@ accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_dummy_load_forma accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=False-eagle3_one_model=False-overlap_scheduler=False] accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=False-eagle3_one_model=True-overlap_scheduler=True] accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3[sampler_async_worker=True-eagle3_one_model=True-overlap_scheduler=True] +accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_rejection_dynamic_tree_smoke[no_dynamic_tree] +accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_rejection_dynamic_tree_smoke[dynamic_tree] accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_sa accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_sa_global_pool accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8[fp8kv=False-attn_backend=FLASHINFER-torch_compile=False] diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index d7d8f0a6e813..20c188578d41 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -825,6 +825,69 @@ def test_eagle3_cdl_sampling(disable_overlap_scheduler: bool): llm_spec.shutdown() +@pytest.mark.parametrize("use_dynamic_tree", [False, True], + ids=["no_dynamic_tree", "dynamic_tree"]) +@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_rejection_sampling_modes(use_dynamic_tree: bool, + use_cuda_graph: bool): + """Test one-model rejection sampling with and without dynamic tree.""" + 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 = 1 + max_draft_len = 6 + dynamic_tree_max_top_k = 10 + max_total_draft_tokens = 60 + kv_cache_config = KvCacheConfig(enable_block_reuse=False, max_tokens=8192) + cuda_graph_config = CudaGraphConfig( + batch_sizes=[1]) if use_cuda_graph else None + + llm_common_config = dict( + model=target_model_dir, + attn_backend="TRTLLM", + disable_overlap_scheduler=True, + cuda_graph_config=cuda_graph_config, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + max_seq_len=8192, + ) + + spec_config_kwargs = dict( + max_draft_len=max_draft_len, + speculative_model=eagle_model, + eagle3_one_model=True, + allow_advanced_sampling=True, + use_rejection_sampling=True, + ) + if use_dynamic_tree: + spec_config_kwargs.update( + use_dynamic_tree=True, + dynamic_tree_max_topK=dynamic_tree_max_top_k, + max_total_draft_tokens=max_total_draft_tokens, + ) + + llm_spec = LLM(**llm_common_config, + speculative_config=Eagle3DecodingConfig( + **spec_config_kwargs)) + + prompts = ["The president of the United States is"] + sampling_params = SamplingParams(max_tokens=20, temperature=1.0, top_p=1.0) + + results = llm_spec.generate(prompts, sampling_params) + llm_spec.shutdown() + + assert len(results) == len(prompts) + assert len(results[0].outputs[0].token_ids) > 0 + + @pytest.mark.parametrize("use_cuda_graph", [True, False]) def test_eagle3_lora(use_cuda_graph: bool): """Test LoRA with 3 requests and max_batch_size=4. @@ -930,7 +993,6 @@ def test_llama_eagle3_dynamic_tree(use_cuda_graph: bool, 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