diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 36f92aa7a0aa..1c838bbcf12c 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -3157,6 +3157,7 @@ int AttentionOp::initialize() noexcept fixedParams.isSpecDecoding = mIsSpecDecodingEnabled; fixedParams.hasAlibi = isALiBi(); fixedParams.useTllmGenSparseAttention = useTllmGenSparseAttention(); + fixedParams.specDecodingTargetMaxGenLen = mSpecDecodingTargetMaxGenLen; mXqaDispatcher.reset(new XqaDispatcher(fixedParams)); diff --git a/cpp/tensorrt_llm/common/attentionOp.h b/cpp/tensorrt_llm/common/attentionOp.h index f7822947b1ea..c8f0d33e4705 100644 --- a/cpp/tensorrt_llm/common/attentionOp.h +++ b/cpp/tensorrt_llm/common/attentionOp.h @@ -491,6 +491,8 @@ class AttentionOp bool mIsSpecDecTree = true; bool mSpecDecodingIsGenerationLengthVariable = false; int32_t mSpecDecodingMaxGenerationLength = 1; + // Static spec-dec tree length used by FMHA autotuning. + int32_t mSpecDecodingTargetMaxGenLen = 0; bool mIsMLAEnabled = false; bool mIsGenerationMLA = false; bool mUseGenFlashMLA = false; @@ -559,13 +561,14 @@ class AttentionOp mCrossAttention, mMaxDistance, mPosShiftEnabled, mPagedContextFMHA, mFP8ContextFMHA, mFP8AttenOutput, mFP8ContextMLA, mFP8GenerationMLA, mChunkPrefillBufferBatchSize, mDenseContextFMHA, mHasFullAttentionMask, mIsSpecDecodingEnabled, mUseSpecDecoding, mIsSpecDecTree, mSpecDecodingIsGenerationLengthVariable, - mSpecDecodingMaxGenerationLength, mIsMLAEnabled, mIsGenerationMLA, mUseGenFlashMLA, mUseSparseAttention, - mUseTllmGenSparseAttentionPaged, mUseTllmGenSparseAttention, mMLAParams.data(), mCpSize, mCpRank, mCpGroup, - mNumAttnHeads, mNumAttnKVHeads, mNumKVHeadsOrigin, mAttnTpSize, mAttnTpRank, mAttnCpSize, mAttnCpRank, - mUlyssesMQABroadcast, mEnableContextFMHA, mFMHAForceFP32Acc, mMultiBlockMode, mEnableXQA, mUseKVCache, - mSkipAttn, mFuseFp4Quant, mNbMultiBlockSemaphores, mAttentionChunkSize.value_or(-1), - mSkipSoftmaxThresholdScaleFactorPrefill, mSkipSoftmaxThresholdScaleFactorDecode, mSageAttnNumEltsPerBlkQ, - mSageAttnNumEltsPerBlkK, mSageAttnNumEltsPerBlkV, mSageAttnQkInt8); + mSpecDecodingMaxGenerationLength, mSpecDecodingTargetMaxGenLen, mIsMLAEnabled, mIsGenerationMLA, + mUseGenFlashMLA, mUseSparseAttention, mUseTllmGenSparseAttentionPaged, mUseTllmGenSparseAttention, + mMLAParams.data(), mCpSize, mCpRank, mCpGroup, mNumAttnHeads, mNumAttnKVHeads, mNumKVHeadsOrigin, + mAttnTpSize, mAttnTpRank, mAttnCpSize, mAttnCpRank, mUlyssesMQABroadcast, mEnableContextFMHA, + mFMHAForceFP32Acc, mMultiBlockMode, mEnableXQA, mUseKVCache, mSkipAttn, mFuseFp4Quant, + mNbMultiBlockSemaphores, mAttentionChunkSize.value_or(-1), mSkipSoftmaxThresholdScaleFactorPrefill, + mSkipSoftmaxThresholdScaleFactorDecode, mSageAttnNumEltsPerBlkQ, mSageAttnNumEltsPerBlkK, + mSageAttnNumEltsPerBlkV, mSageAttnQkInt8); }; private: diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index aa90a39146e6..355934d6fe28 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -990,6 +990,8 @@ class TllmGenFmhaKernel options.mIsCustomSpecDecodingGen = !isContext && params.mMaxSeqLenQ > 1 && params.mIsSpecDecTree; options.mIsCausalSpecDecodingGen = !isContext && params.mMaxSeqLenQ > 1 && !params.mIsSpecDecTree; options.mNumSpecDecodingTokens = !isContext && params.mMaxSeqLenQ > 1 ? params.mMaxSeqLenQ : 0; + // Carry static tree length into FMHA kernel selection. + options.mSpecDecodingTargetMaxGenLen = params.mSpecDecodingTargetMaxGenLen; options.mIsTrtllmLayout = true; } @@ -1013,9 +1015,24 @@ class TllmGenFmhaKernel // loop. And the number of loops are not the same in different tasks. sstream << "\"checksTaskSchedules\": false,\n"; + bool hasCompileDefs = false; + auto writeCompileDef = [&](char const* compileDef) + { + if (!hasCompileDefs) + { + sstream << "\"compileDefs\": ["; + hasCompileDefs = true; + } + else + { + sstream << ", "; + } + sstream << "\"" << compileDef << "\""; + }; + if (options.mIsExportingCubin) { - sstream << "\"compileDefs\": [\"-DTLLM_EXPORT_CUBIN\"],\n"; + writeCompileDef("-DTLLM_EXPORT_CUBIN"); } // Set compile flags for E2M1 KV kernel benchmark. @@ -1023,7 +1040,18 @@ class TllmGenFmhaKernel if (options.mChecksResults == 0 && options.mDtypeKv == tg::Dtype::E2m1) { TLLM_LOG_INFO("Forcing -DTLLM_BENCHMARK_E2M1_KV_CACHE for E2m1 Kv. The results are not correct."); - sstream << "\"compileDefs\": [\"-DTLLM_BENCHMARK_E2M1_KV_CACHE\"],\n"; + writeCompileDef("-DTLLM_BENCHMARK_E2M1_KV_CACHE"); + } + + // SwapsMmaAb NVRTC kernels already emit __launch_bounds__; avoid a CUDA 13 .reqntid/.maxntid conflict. + if (shouldUseNvrtc(options) && options.mFmhaKernelType == FmhaKernelType::SwapsMmaAbForGeneration) + { + writeCompileDef("-DTLLM_DISABLE_BLOCK_SIZE"); + } + + if (hasCompileDefs) + { + sstream << "],\n"; } // Enable programmatic dependent launch. diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h index 19486cd6cf50..369e634fb61c 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h @@ -342,6 +342,7 @@ struct TllmGenFmhaRunnerParams // When seqlensQPtr[i] < mPackedMaskMaxSeqLenQ, the packed mask tensor has // row stride ceilDiv(mPackedMaskMaxSeqLenQ, 32) rather than ceilDiv(seqLenQ, 32). int32_t mPackedMaskMaxSeqLenQ = 0; + int32_t mSpecDecodingTargetMaxGenLen = 0; // set the attention mask type TllmGenFmhaRunnerParams& setAttentionMaskType(std::int8_t maskType) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/prepareCustomMask.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/prepareCustomMask.cu index 799fb6d9ae85..32a9a213087f 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/prepareCustomMask.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/prepareCustomMask.cu @@ -19,7 +19,6 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/common/logger.h" #include -#include #include TRTLLM_NAMESPACE_BEGIN @@ -28,68 +27,103 @@ namespace kernels { //////////////////////////////////////////////////////////////////////////////////////////////////// +constexpr int32_t kCustomMaskOffsetScanBlockSize = 256; + __device__ __host__ inline int32_t ceilDiv(int32_t a, int32_t b) { return (a + b - 1) / b; } -__global__ void computeCustomMaskOffsetsKernel( - TllmGenFmhaRunnerParams runnerParams, int32_t stepQ, int32_t stepKv, unsigned long long* globalCounter) +__device__ inline int64_t computeKeepsMmaAbCustomMaskSize(int32_t seqLenQ, int32_t seqLenKv, + int32_t firstSparseMaskOffsetKv, int32_t numHeadsQPerKv, int32_t stepQ, int32_t stepKv) { - int32_t batchSize = runnerParams.mBatchSize; - int32_t numHeadsQPerKv = runnerParams.mNumHeadsQPerKv; - int32_t tileSizeQPerCta = stepQ; - int32_t tileSizeKvPerCta = stepKv; - int32_t const* seqLensKvPtr = runnerParams.seqLensKvPtr; - int32_t const* firstSparseMaskOffsetsKvPtr = runnerParams.firstSparseMaskOffsetsKvPtr; + int32_t const numTilesQ = ceilDiv(seqLenQ * numHeadsQPerKv, stepQ); + int32_t const firstSparseTile = firstSparseMaskOffsetKv / stepKv; + int32_t const numCustomMaskTilesKv = ceilDiv(seqLenKv, stepKv) - firstSparseTile; + + return static_cast(numTilesQ) * numCustomMaskTilesKv * stepQ * stepKv / 32; +} - typedef cub::BlockScan BlockScan; - __shared__ typename BlockScan::TempStorage temp_storage; +__device__ inline int64_t computeSwapsMmaAbCustomMaskSize(int32_t seqLenQ, int32_t seqLenKv, + int32_t firstSparseMaskOffsetKv, int32_t numHeadsQPerKv, int32_t stepQ, int32_t stepKv, int32_t tileSizeQRaw, + int32_t tileSizeKv) +{ + int32_t const numInstsQ = stepQ / tileSizeQRaw; + int32_t const numInstsKv = stepKv / tileSizeKv; + int32_t const tileSizeKvPerCta = stepKv; + int32_t const tileSizeQ = ceilDiv(tileSizeQRaw, 32) * 32; + int32_t const tileSizeQPerCta = tileSizeQ * numInstsQ; + int32_t const numTilesQPerToken = ceilDiv(numHeadsQPerKv, tileSizeQPerCta); + int32_t const numTilesQ = seqLenQ * numTilesQPerToken; + int32_t const firstSparseTile = firstSparseMaskOffsetKv / tileSizeKvPerCta; + int32_t const numCustomMaskTilesKv = ceilDiv(seqLenKv, tileSizeKvPerCta) - firstSparseTile; + int32_t const perTileSize = numInstsQ * numInstsKv * (tileSizeQ * tileSizeKv) / 32; + + return static_cast(numTilesQ) * numCustomMaskTilesKv * perTileSize; +} - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int64_t maskSize = 0; +__global__ void computeCustomMaskOffsetsParallelKernel( + TllmGenFmhaRunnerParams runnerParams, int32_t stepQ, int32_t stepKv) +{ + // One CTA computes deterministic batch-prefix offsets. + __shared__ int64_t threadSums[kCustomMaskOffsetScanBlockSize]; - if (idx < batchSize) + int32_t const batchSize = runnerParams.mBatchSize; + int32_t const numHeadsQPerKv = runnerParams.mNumHeadsQPerKv; + int32_t const* seqLensKvPtr = runnerParams.seqLensKvPtr; + int32_t const* firstSparseMaskOffsetsKvPtr = runnerParams.firstSparseMaskOffsetsKvPtr; + int32_t const itemsPerThread = ceilDiv(batchSize, static_cast(blockDim.x)); + int32_t const threadIdxX = static_cast(threadIdx.x); + int32_t const startIdx = threadIdxX * itemsPerThread; + int32_t endIdx = startIdx + itemsPerThread; + if (endIdx > batchSize) { + endIdx = batchSize; + } - int32_t seqLenQ = runnerParams.seqLensQPtr[idx]; - int32_t seqLenKv = seqLensKvPtr[idx]; - int32_t firstSparseMaskOffsetKv = firstSparseMaskOffsetsKvPtr[idx]; + int64_t threadSum = 0; + for (int32_t idx = startIdx; idx < endIdx; ++idx) + { + int32_t const seqLenQ = runnerParams.seqLensQPtr[idx]; + int32_t const seqLenKv = seqLensKvPtr[idx]; + int32_t const firstSparseMaskOffsetKv = firstSparseMaskOffsetsKvPtr[idx]; + threadSum += computeKeepsMmaAbCustomMaskSize( + seqLenQ, seqLenKv, firstSparseMaskOffsetKv, numHeadsQPerKv, stepQ, stepKv); + } - int32_t numTilesQ = (seqLenQ * numHeadsQPerKv + tileSizeQPerCta - 1) / tileSizeQPerCta; - int32_t firstSparseTile = firstSparseMaskOffsetKv / tileSizeKvPerCta; - int32_t numCustomMaskTilesKv = (seqLenKv + tileSizeKvPerCta - 1) / tileSizeKvPerCta - firstSparseTile; + threadSums[threadIdxX] = threadSum; + __syncthreads(); - maskSize = static_cast(numTilesQ * numCustomMaskTilesKv * stepQ * stepKv / 32); + for (int32_t stride = 1; stride < blockDim.x; stride <<= 1) + { + int64_t partialSum = 0; + if (threadIdxX >= stride) + { + partialSum = threadSums[threadIdxX - stride]; + } + __syncthreads(); + threadSums[threadIdxX] += partialSum; + __syncthreads(); } - int64_t prefixOffset; - int64_t blockSum; - BlockScan(temp_storage).ExclusiveSum(maskSize, prefixOffset, blockSum); - - __shared__ unsigned long long blockBase; - if (threadIdx.x == 0) - blockBase = atomicAdd(globalCounter, (unsigned long long) blockSum); - __syncthreads(); + int64_t localOffset = threadSums[threadIdxX] - threadSum; + for (int32_t idx = startIdx; idx < endIdx; ++idx) + { + runnerParams.customMaskOffsetsPtr[idx] = localOffset; - if (idx < batchSize) - runnerParams.customMaskOffsetsPtr[idx] = static_cast(blockBase) + prefixOffset; + int32_t const seqLenQ = runnerParams.seqLensQPtr[idx]; + int32_t const seqLenKv = seqLensKvPtr[idx]; + int32_t const firstSparseMaskOffsetKv = firstSparseMaskOffsetsKvPtr[idx]; + localOffset += computeKeepsMmaAbCustomMaskSize( + seqLenQ, seqLenKv, firstSparseMaskOffsetKv, numHeadsQPerKv, stepQ, stepKv); + } } void launchComputeCustomMaskOffsetsKernel( TllmGenFmhaRunnerParams const& runnerParams, int32_t stepQ, int32_t stepKv, cudaStream_t stream) { - int32_t batchSize = runnerParams.mBatchSize; - - unsigned long long* d_globalCounter; - cudaMallocAsync(&d_globalCounter, sizeof(unsigned long long), stream); - cudaMemsetAsync(d_globalCounter, 0, sizeof(unsigned long long), stream); - - int blockSize = 128; - int gridSize = (batchSize + blockSize - 1) / blockSize; - computeCustomMaskOffsetsKernel<<>>(runnerParams, stepQ, stepKv, d_globalCounter); - - cudaFreeAsync(d_globalCounter, stream); + computeCustomMaskOffsetsParallelKernel<<<1, kCustomMaskOffsetScanBlockSize, 0, stream>>>( + runnerParams, stepQ, stepKv); } // Input: customMaskInput (generalPackedCustoMaskPtr) shape: [batch_size, seqLenQ, ceilDiv(seqLenKv-firstSparse, 32)] @@ -132,12 +166,10 @@ __global__ void prepareCustomMaskBuffersKernelForKeepsMmaAb( // The sequence length of tensor KV. int32_t const seqLenKv = seqLensKvPtr[batchIdx]; - // The packed mask tensor row stride: use mPackedMaskMaxSeqLenQ when available, - // because the packed mask tensor has shape [bs, maxSeqLenQ, ceilDiv(maxSeqLenQ, 32)] - // even when the actual seqLenQ is smaller (e.g., drafter layers in dynamic tree). - int32_t const packedMaskSeqLenQ + // Use the padded mask row stride from the Python buffer. + int32_t const packedMaskMaxSeqLenQ = runnerParams.mPackedMaskMaxSeqLenQ > 0 ? runnerParams.mPackedMaskMaxSeqLenQ : seqLenQ; - int32_t const packedMaskNumBlocks = ceilDiv(packedMaskSeqLenQ, 32); + int32_t const packedMaskNumBlocks = ceilDiv(packedMaskMaxSeqLenQ, 32); // Calculate global Q token index (flattened across heads) int32_t const qTokensPerBlock = static_cast(blockDim.x); @@ -173,12 +205,13 @@ __global__ void prepareCustomMaskBuffersKernelForKeepsMmaAb( else { // Sparse region: check the input mask - // Input mask shape: [bs, packedMaskSeqLenQ, packedMaskNumBlocks] // The KV dimension in the mask corresponds to Q positions (tree mask) int32_t const qPosInTree = tokenIdxKv - firstSparseMaskOffsetKv; if (qPosInTree < seqLenQ) { - int32_t const qMaskBaseIdx = (batchIdx * packedMaskSeqLenQ + tokenIdxQ) * packedMaskNumBlocks; + // Use padded mask stride; FMHA cumSeqLensQ does not describe mask storage. + int32_t const rowOffset = batchIdx * packedMaskMaxSeqLenQ + tokenIdxQ; + int32_t const qMaskBaseIdx = rowOffset * packedMaskNumBlocks; int32_t const packedMaskIdx = qMaskBaseIdx + (qPosInTree >> 5); int32_t const bitPos = qPosInTree & 0x1F; randomMask = (customMaskInputPtr[packedMaskIdx] >> bitPos) & 1; @@ -267,6 +300,220 @@ void launchPrepareCustomMaskBuffersKernelForKeepsMmaAb(TllmGenFmhaRunnerParams c } } +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// SwapsMmaAb custom mask for groupsTokensHeadsQ=false. +__global__ void prepareCustomMaskBuffersKernelForSwapsMmaAb( + TllmGenFmhaRunnerParams runnerParams, int32_t stepQ, int32_t stepKv, int32_t tileSizeQRaw, int32_t tileSizeKv) +{ + int32_t const batchSize = runnerParams.mBatchSize; + int32_t const numHeadsQPerKv = runnerParams.mNumHeadsQPerKv; + int32_t const numInstsQ = stepQ / tileSizeQRaw; + int32_t const numInstsKv = stepKv / tileSizeKv; + int32_t const tileSizeKvPerCta = stepKv; + // Pad tileSizeQ for uint32 packing. + int32_t const tileSizeQ = ((tileSizeQRaw + 31) / 32) * 32; + int32_t const tileSizeQPerCta = tileSizeQ * numInstsQ; + int32_t const* seqLensKvPtr = runnerParams.seqLensKvPtr; + int64_t* customMaskOffsetsPtr = runnerParams.customMaskOffsetsPtr; + uint32_t* customMaskPtr = runnerParams.customMaskPtr; + int32_t const* customMaskInputPtr = runnerParams.generalPackedCustoMaskPtr; + int32_t* firstSparseMaskOffsetsKvPtr = runnerParams.firstSparseMaskOffsetsKvPtr; + + int32_t const batchIdx = static_cast(blockIdx.x); + int32_t const tokenThreadIdx = static_cast(threadIdx.x); + int32_t const tokenGroupIdx = static_cast(blockIdx.y); + int32_t const kvThreadIdx = static_cast(threadIdx.y); + int32_t const kvGroupIdx = static_cast(blockIdx.z); + + if (batchIdx >= batchSize) + { + return; + } + + int32_t const firstSparseMaskOffsetKv = firstSparseMaskOffsetsKvPtr[batchIdx]; + int32_t const firstSparseMaskTileOffsetKv = firstSparseMaskOffsetKv / tileSizeKvPerCta; + int32_t const adjustedFirstSparseMaskOffsetKv = firstSparseMaskTileOffsetKv * tileSizeKvPerCta; + + int32_t const seqLenQ = runnerParams.seqLensQPtr[batchIdx]; + int32_t const seqLenKv = seqLensKvPtr[batchIdx]; + + // Use padded mask row stride; see KeepsMmaAb path. + int32_t const packedMaskMaxSeqLenQ + = runnerParams.mPackedMaskMaxSeqLenQ > 0 ? runnerParams.mPackedMaskMaxSeqLenQ : seqLenQ; + int32_t const packedMaskNumBlocks = ceilDiv(packedMaskMaxSeqLenQ, 32); + + int32_t const tokensPerBlock = static_cast(blockDim.x); + int32_t const tokenIdxQ = tokenGroupIdx * tokensPerBlock + tokenThreadIdx; + if (tokenIdxQ >= seqLenQ) + { + return; + } + + int32_t const kvTokensPerBlock = static_cast(blockDim.y); + int32_t const globalKvIdx = kvGroupIdx * kvTokensPerBlock + kvThreadIdx; + int32_t const tokenIdxKv = adjustedFirstSparseMaskOffsetKv + globalKvIdx; + if (tokenIdxKv >= seqLenKv) + { + return; + } + + int32_t randomMask = 0; + if (tokenIdxKv < firstSparseMaskOffsetKv) + { + randomMask = 1; + } + else + { + int32_t const qPosInTree = tokenIdxKv - firstSparseMaskOffsetKv; + if (qPosInTree < seqLenQ) + { + // Use padded mask row stride. + int32_t const rowOffset = batchIdx * packedMaskMaxSeqLenQ + tokenIdxQ; + int32_t const qMaskBaseIdx = rowOffset * packedMaskNumBlocks; + int32_t const packedMaskIdx = qMaskBaseIdx + (qPosInTree >> 5); + int32_t const bitPos = qPosInTree & 0x1F; + randomMask = (customMaskInputPtr[packedMaskIdx] >> bitPos) & 1; + } + } + + if (randomMask) + { + int32_t const numCustomMaskTilesKv = ceilDiv(seqLenKv, tileSizeKvPerCta) - firstSparseMaskTileOffsetKv; + int64_t const customMaskOffset = customMaskOffsetsPtr[batchIdx]; + uint32_t* localCustomMaskPtr = customMaskPtr + customMaskOffset; + + // One Q tile per token when heads fit in the padded tile. + int32_t const numTilesQPerToken = ceilDiv(numHeadsQPerKv, tileSizeQPerCta); + + int32_t const customMaskKvIdx = tokenIdxKv - adjustedFirstSparseMaskOffsetKv; + int32_t const tileIdxKv = customMaskKvIdx / tileSizeKvPerCta; + int32_t const instIdxKv = (customMaskKvIdx % tileSizeKvPerCta) / tileSizeKv; + int32_t const kvInTile = customMaskKvIdx % tileSizeKv; + + // Match trtllm-gen SwapsMmaAb LDTM bit layout. + for (int32_t headIdxInGrp = 0; headIdxInGrp < numHeadsQPerKv; ++headIdxInGrp) + { + int32_t const customMaskTokenIdxQ = headIdxInGrp; + int32_t tileIdxQ = customMaskTokenIdxQ / tileSizeQPerCta; + tileIdxQ += tokenIdxQ * numTilesQPerToken; + int32_t const instIdxQ = (customMaskTokenIdxQ % tileSizeQPerCta) / tileSizeQ; + int32_t const tokenIdxInTileQ = (customMaskTokenIdxQ % tileSizeQPerCta) % tileSizeQ; + + int64_t const tileOffset = static_cast(tileIdxQ) * numCustomMaskTilesKv + tileIdxKv; + int64_t const instOffset = tileOffset * numInstsQ * numInstsKv + (instIdxQ * numInstsKv + instIdxKv); + int64_t maskOffset = instOffset * tileSizeQ * tileSizeKv; + + int32_t const tokenIdxInTileKv = kvInTile; + int32_t const threadIdxQ = (tokenIdxInTileQ % 8) / 2; + int32_t const threadIdxKv = (tokenIdxInTileKv % 8) + (tokenIdxInTileKv / 32) * 8; + int32_t const tokenIdxInWarpTileKv = tokenIdxInTileKv % 32; + int32_t const eltIdxInThread = (tokenIdxInTileQ % 2) + ((tokenIdxInWarpTileKv / 8) % 2) * 2 + + (tokenIdxInTileQ / 8) * 4 + (tokenIdxInWarpTileKv / 16) * 4 * (tileSizeQRaw / 8); + maskOffset += (threadIdxKv * 4 + threadIdxQ) * 32 + eltIdxInThread; + + int64_t const offsetAsUInt32 = maskOffset / 32; + int32_t const bitPosInUInt32 = maskOffset % 32; + atomicOr(&localCustomMaskPtr[offsetAsUInt32], (1U << bitPosInUInt32)); + } + } +} + +void launchPrepareCustomMaskBuffersKernelForSwapsMmaAb(TllmGenFmhaRunnerParams const& runnerParams, int32_t stepQ, + int32_t stepKv, int32_t tileSizeQ, int32_t tileSizeKv, cudaStream_t stream) +{ + int32_t const batchSize = runnerParams.mBatchSize; + int32_t const maxSeqLenQ = runnerParams.mMaxSeqLenQ; + int32_t const tileSizeKvPerCta = stepKv; + + int32_t const maxKvRangeLength = maxSeqLenQ + (tileSizeKvPerCta - 1); + + int32_t const tokensPerBlock = 64; + int32_t const kvTokensPerBlock = 4; + + int32_t const numBlocksY = ceilDiv(maxSeqLenQ, tokensPerBlock); + int32_t const numBlocksZ = ceilDiv(maxKvRangeLength, kvTokensPerBlock); + + dim3 gridDim(batchSize, numBlocksY, numBlocksZ); + dim3 blockDim(tokensPerBlock, kvTokensPerBlock, 1); + + prepareCustomMaskBuffersKernelForSwapsMmaAb<<>>( + runnerParams, stepQ, stepKv, tileSizeQ, tileSizeKv); + { + int const blockSize = 128; + int const gridSize = (batchSize + blockSize - 1) / blockSize; + adjustFirstSparseMaskOffsetsKernel<<>>(runnerParams, stepKv); + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +__global__ void computeCustomMaskOffsetsParallelKernelForSwapsMmaAb( + TllmGenFmhaRunnerParams runnerParams, int32_t stepQ, int32_t stepKv, int32_t tileSizeQRaw, int32_t tileSizeKv) +{ + // One CTA computes deterministic batch-prefix offsets. + __shared__ int64_t threadSums[kCustomMaskOffsetScanBlockSize]; + + int32_t const batchSize = runnerParams.mBatchSize; + int32_t const numHeadsQPerKv = runnerParams.mNumHeadsQPerKv; + int32_t const* seqLensKvPtr = runnerParams.seqLensKvPtr; + int32_t const* firstSparseMaskOffsetsKvPtr = runnerParams.firstSparseMaskOffsetsKvPtr; + int32_t const itemsPerThread = ceilDiv(batchSize, static_cast(blockDim.x)); + int32_t const threadIdxX = static_cast(threadIdx.x); + int32_t const startIdx = threadIdxX * itemsPerThread; + int32_t endIdx = startIdx + itemsPerThread; + if (endIdx > batchSize) + { + endIdx = batchSize; + } + + int64_t threadSum = 0; + for (int32_t idx = startIdx; idx < endIdx; ++idx) + { + int32_t const seqLenQ = runnerParams.seqLensQPtr[idx]; + int32_t const seqLenKv = seqLensKvPtr[idx]; + int32_t const firstSparseMaskOffsetKv = firstSparseMaskOffsetsKvPtr[idx]; + threadSum += computeSwapsMmaAbCustomMaskSize( + seqLenQ, seqLenKv, firstSparseMaskOffsetKv, numHeadsQPerKv, stepQ, stepKv, tileSizeQRaw, tileSizeKv); + } + + threadSums[threadIdxX] = threadSum; + __syncthreads(); + + for (int32_t stride = 1; stride < blockDim.x; stride <<= 1) + { + int64_t partialSum = 0; + if (threadIdxX >= stride) + { + partialSum = threadSums[threadIdxX - stride]; + } + __syncthreads(); + threadSums[threadIdxX] += partialSum; + __syncthreads(); + } + + int64_t localOffset = threadSums[threadIdxX] - threadSum; + for (int32_t idx = startIdx; idx < endIdx; ++idx) + { + runnerParams.customMaskOffsetsPtr[idx] = localOffset; + + int32_t const seqLenQ = runnerParams.seqLensQPtr[idx]; + int32_t const seqLenKv = seqLensKvPtr[idx]; + int32_t const firstSparseMaskOffsetKv = firstSparseMaskOffsetsKvPtr[idx]; + localOffset += computeSwapsMmaAbCustomMaskSize( + seqLenQ, seqLenKv, firstSparseMaskOffsetKv, numHeadsQPerKv, stepQ, stepKv, tileSizeQRaw, tileSizeKv); + } +} + +void launchComputeCustomMaskOffsetsKernelForSwapsMmaAb(TllmGenFmhaRunnerParams const& runnerParams, int32_t stepQ, + int32_t stepKv, int32_t tileSizeQ, int32_t tileSizeKv, cudaStream_t stream) +{ + computeCustomMaskOffsetsParallelKernelForSwapsMmaAb<<<1, kCustomMaskOffsetScanBlockSize, 0, stream>>>( + runnerParams, stepQ, stepKv, tileSizeQ, tileSizeKv); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + void runPrepareCustomMask(TllmGenFmhaRunnerParams const& runnerParams, FmhaKernelType kernelType, int32_t stepQ, int32_t stepKv, int32_t tileSizeQ, int32_t tileSizeKv, cudaStream_t stream) { @@ -280,12 +527,16 @@ void runPrepareCustomMask(TllmGenFmhaRunnerParams const& runnerParams, FmhaKerne "formula: tile_size_q * tile_size_k * num_instances_q * num_instances_k"); return; } - // Step 1: Compute offsets on GPU using prefix sum launchComputeCustomMaskOffsetsKernel(runnerParams, stepQ, stepKv, stream); - // Step 2: Compute custom mask buffers launchPrepareCustomMaskBuffersKernelForKeepsMmaAb(runnerParams, stepQ, stepKv, tileSizeQ, tileSizeKv, stream); TLLM_CUDA_CHECK(cudaGetLastError()); } + else if (isSwapsMmaAbForGenerationKernel(static_cast(kernelType))) + { + launchComputeCustomMaskOffsetsKernelForSwapsMmaAb(runnerParams, stepQ, stepKv, tileSizeQ, tileSizeKv, stream); + launchPrepareCustomMaskBuffersKernelForSwapsMmaAb(runnerParams, stepQ, stepKv, tileSizeQ, tileSizeKv, stream); + TLLM_CUDA_CHECK(cudaGetLastError()); + } else { TLLM_LOG_ERROR("TRTLLM-GEN does not support kernel type: %d for custom mask preparation", kernelType); diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/trtllm/dev/Utils.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/trtllm/dev/Utils.h index b11b5c8a9874..c0edb65b628d 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/trtllm/dev/Utils.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/trtllm/dev/Utils.h @@ -32,7 +32,7 @@ namespace dev { // __block_size__ is only supported in CUDA 13 and later. // We can always emit the macro, and it will simply be ignored in CUDA 12. -#if defined(__CUDACC_VER_MAJOR__) && __CUDACC_VER_MAJOR__ >= 13 +#if !defined(TLLM_DISABLE_BLOCK_SIZE) && defined(__CUDACC_VER_MAJOR__) && __CUDACC_VER_MAJOR__ >= 13 #define TLLM_BLOCK_SIZE(bx, by, bz) __block_size__((bx, by, bz)) #else #define TLLM_BLOCK_SIZE(bx, by, bz) diff --git a/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp b/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp index 37f248687e91..4864af1d0187 100644 --- a/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp +++ b/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp @@ -282,6 +282,7 @@ bool XqaDispatcher::isSupported() tllmRunnerParams.mMaskType = mFixedParams.isSpecDecoding ? TrtllmGenAttentionMaskType::Custom : TrtllmGenAttentionMaskType::Causal; tllmRunnerParams.mIsSpecDecTree = mFixedParams.isSpecDecoding; + tllmRunnerParams.mSpecDecodingTargetMaxGenLen = mFixedParams.specDecodingTargetMaxGenLen; tllmRunnerParams.mKernelType = FmhaKernelType::Generation; tllmRunnerParams.mTileScheduler = TileScheduler::Static; tllmRunnerParams.mMultiCtasKvMode = true; @@ -494,7 +495,19 @@ void XqaDispatcher::runImpl( // It is used to construct contiguous kv cache TMA descriptors. tllmRunnerParams.mMaxSeqLenCacheKv = params.max_attention_window_size; tllmRunnerParams.mMaxSeqLenQ = params.generation_input_length; - tllmRunnerParams.mMaxSeqLenKv = params.max_past_kv_length; + bool const isSpecDecTree = params.is_spec_dec_tree && params.multi_query_tokens; + if (isSpecDecTree) + { + TLLM_CHECK_WITH_INFO(params.spec_decoding_max_generation_length > 0, + "spec_decoding_max_generation_length must be positive for spec-dec tree."); + tllmRunnerParams.mMaxSeqLenQ + = std::min(tllmRunnerParams.mMaxSeqLenQ, params.spec_decoding_max_generation_length); + tllmRunnerParams.mMaxSeqLenKv = params.max_past_kv_length + tllmRunnerParams.mMaxSeqLenQ; + } + else + { + tllmRunnerParams.mMaxSeqLenKv = params.max_past_kv_length; + } tllmRunnerParams.mJITWarmup = params.trtllm_gen_jit_warmup; tllmRunnerParams.mJITWarmupMaxNumRequests = params.trtllm_gen_jit_warmup_max_num_requests; tllmRunnerParams.mJITWarmupMaxSeqLenQ = params.trtllm_gen_jit_warmup_max_seq_len_q; @@ -520,6 +533,7 @@ void XqaDispatcher::runImpl( tllmRunnerParams.seqLensQPtr = params.spec_decoding_generation_lengths; tllmRunnerParams.generalPackedCustoMaskPtr = params.spec_decoding_packed_mask; tllmRunnerParams.mPackedMaskMaxSeqLenQ = params.spec_decoding_max_generation_length; + tllmRunnerParams.mSpecDecodingTargetMaxGenLen = mFixedParams.specDecodingTargetMaxGenLen; tllmRunnerParams.customMaskPtr = params.spec_decoding_bl_tree_mask; tllmRunnerParams.customMaskOffsetsPtr = params.spec_decoding_bl_tree_mask_offset; tllmRunnerParams.firstSparseMaskOffsetsKvPtr = params.spec_bl_tree_first_sparse_mask_offset_kv; diff --git a/cpp/tensorrt_llm/kernels/xqaDispatcher.h b/cpp/tensorrt_llm/kernels/xqaDispatcher.h index 9441c5a76859..83eb3febe46b 100644 --- a/cpp/tensorrt_llm/kernels/xqaDispatcher.h +++ b/cpp/tensorrt_llm/kernels/xqaDispatcher.h @@ -61,6 +61,8 @@ struct XqaFixedParams bool isPagedKv; // Is speculative decoding enabled. bool isSpecDecoding; + // Spec-dec tree bound for FMHA kernel selection. + int specDecodingTargetMaxGenLen = 0; // Do we apply alibi ? bool hasAlibi; // Use trtllm-gen sparse attention kernel. diff --git a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp index 993de4743193..659064f4a5a7 100644 --- a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp @@ -173,7 +173,8 @@ void initBindings(nb::module_& m) nb::arg("sage_attn_num_elts_per_blk_k") = 0, nb::arg("sage_attn_num_elts_per_blk_v") = 0, nb::arg("sage_attn_qk_int8") = false, nb::arg("num_contexts") = 0, nb::arg("num_ctx_tokens") = 0, nb::arg("trtllm_gen_jit_warmup") = false, nb::arg("compressed_kv_cache_pool_ptr") = std::nullopt, - "Multi-head attention operation", nb::call_guard()); + nb::arg("spec_decoding_target_max_draft_tokens") = std::nullopt, "Multi-head attention operation", + nb::call_guard()); m.def( "get_helix_workspace_size_per_rank", diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 2cc7adb7fa1b..07d10c28efa1 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -859,7 +859,17 @@ class Runner : public RunnerBase enqueue_params.spec_decoding_packed_mask = spec_decoding_packed_mask->data_ptr(); enqueue_params.spec_decoding_is_generation_length_variable = true; TLLM_CHECK(spec_decoding_position_offsets_for_cpp->dim() == 2); // [batch_size, max_draft_len + 1] - enqueue_params.spec_decoding_max_generation_length = spec_decoding_position_offsets_for_cpp->sizes()[1]; + if (useTllmGen) + { + // Blackwell uses the padded packed-mask row dim as the mask stride. + TLLM_CHECK(spec_decoding_packed_mask->dim() == 3); + enqueue_params.spec_decoding_max_generation_length = spec_decoding_packed_mask->sizes()[1]; + } + else + { + enqueue_params.spec_decoding_max_generation_length + = spec_decoding_position_offsets_for_cpp->sizes()[1]; + } } // Current mlaGeneration will using fmha to do attention, so we don't go into enqueueGeneration @@ -965,7 +975,7 @@ void attention(torch::Tensor q, std::optional k, std::optional flash_mla_tile_scheduler_metadata, std::optional flash_mla_num_splits, int64_t sage_attn_num_elts_per_blk_q, int64_t sage_attn_num_elts_per_blk_k, int64_t sage_attn_num_elts_per_blk_v, bool sage_attn_qk_int8, int64_t num_contexts, int64_t num_ctx_tokens, bool trtllm_gen_jit_warmup, - std::optional compressed_kv_cache_pool_ptr) + std::optional compressed_kv_cache_pool_ptr, std::optional spec_decoding_target_max_draft_tokens) { TLLM_LOG_TRACE("Attention op starts at layer %d", local_layer_idx); // Use these tensors to infer if the attention is using KV cache @@ -1096,6 +1106,11 @@ void attention(torch::Tensor q, std::optional k, std::optionalmIsSpecDecodingEnabled = is_spec_decoding_enabled; op->mUseSpecDecoding = use_spec_decoding; op->mIsSpecDecTree = is_spec_dec_tree; + // Include static tree length in the AttentionOp cache key. + if (spec_decoding_target_max_draft_tokens.has_value() && op->mSpecDecodingTargetMaxGenLen == 0) + { + op->mSpecDecodingTargetMaxGenLen = static_cast(spec_decoding_target_max_draft_tokens.value()) + 1; + } op->mUseSparseAttention = false; op->mUseTllmGenSparseAttentionPaged = false; diff --git a/cpp/tensorrt_llm/thop/attentionOp.h b/cpp/tensorrt_llm/thop/attentionOp.h index 0d4ea46cdfd9..f55e30925177 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.h +++ b/cpp/tensorrt_llm/thop/attentionOp.h @@ -90,7 +90,8 @@ void attention(torch::Tensor q, std::optional k, std::optional flash_mla_num_splits = std::nullopt, int64_t sage_attn_num_elts_per_blk_q = 0, int64_t sage_attn_num_elts_per_blk_k = 0, int64_t sage_attn_num_elts_per_blk_v = 0, bool sage_attn_qk_int8 = false, int64_t num_contexts = 0, int64_t num_ctx_tokens = 0, bool trtllm_gen_jit_warmup = false, - std::optional compressed_kv_cache_pool_ptr = std::nullopt); + std::optional compressed_kv_cache_pool_ptr = std::nullopt, + std::optional spec_decoding_target_max_draft_tokens = std::nullopt); struct KvCachePoolPointers { diff --git a/cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp b/cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp index 61722e9953fc..61617934f236 100644 --- a/cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp +++ b/cpp/tests/unit_tests/kernels/prepareCustomMaskTest.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -374,4 +374,12 @@ TEST_F(PrepareCustomMaskTest, MediumBatch) /* numHeadsQPerKv */ 8); } +TEST_F(PrepareCustomMaskTest, LargeBatchParallelOffsetScan) +{ + testPrepareCustomMask(/* batchSize */ 128, + /* maxSeqLenQ */ 16, + /* maxSeqLenKv */ 128, + /* numHeadsQPerKv */ 4); +} + } // namespace diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 216e2a95df5c..933c66b9d0f8 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -118,6 +118,8 @@ class TrtllmAttentionMetadata(AttentionMetadata): # C++ attention op requires a 2-D position_offsets tensor and reads # sizes()[1] as the generation length / packed-mask row stride. spec_decoding_position_offsets_cpp: Optional[torch.Tensor] = None + # Compact Hopper C++ row stride for 1D dynamic-tree offsets. + position_offsets_stride: int = 0 spec_decoding_packed_mask: Optional[torch.Tensor] = None spec_decoding_generation_lengths: Optional[torch.Tensor] = None spec_decoding_bl_tree_mask_offset: Optional[torch.Tensor] = None @@ -170,10 +172,13 @@ def effective_workspace(self) -> Optional[torch.Tensor]: @property def spec_decoding_position_offsets_for_cpp(self) -> Optional[torch.Tensor]: """``spec_decoding_position_offsets`` reshaped to the 2D layout the C++ - kernel expects. 1D inputs (dynamic-tree shorthand) are viewed as - ``(max_num_requests, -1)``.""" + kernel expects.""" offsets = self.spec_decoding_position_offsets if offsets is not None and offsets.dim() == 1: + if (self.spec_decoding_position_offsets_cpp is not None + and not self.is_sm_version_trtllm_gen_kernel( + sm=get_sm_version())): + return self.spec_decoding_position_offsets_cpp return offsets.view(self.max_num_requests, -1) return offsets @@ -240,14 +245,17 @@ def update_position_offsets_for_cpp(self, query_len: int) -> None: offsets = self.spec_decoding_position_offsets if offsets is None or offsets.dim() != 1: self.spec_decoding_position_offsets_cpp = offsets + self.position_offsets_stride = 0 return if self.max_num_requests > 0 and query_len > 0: + self.position_offsets_stride = query_len total = self.max_num_requests * query_len self.spec_decoding_position_offsets_cpp = offsets[:total].view( self.max_num_requests, query_len) else: self.spec_decoding_position_offsets_cpp = offsets + self.position_offsets_stride = 0 def _post_init_with_buffers(self, buffers) -> None: @@ -908,15 +916,17 @@ def update_spec_dec_param( ``[num_contexts:batch_size]`` rather than ``[:batch_size]``. ''' - # Disable spec decoding on Blackwell (sm100+). The trtllmGen FMHA - # kernels do not yet support speculative decoding mode. + # Blackwell trtllm-gen spec-dec is enabled only for dynamic-tree masks. self.is_spec_decoding_enabled = is_spec_decoding_enabled and ( - not self.is_sm_version_trtllm_gen_kernel(sm=get_sm_version())) + not self.is_sm_version_trtllm_gen_kernel(sm=get_sm_version()) + or is_spec_dec_dynamic_tree) # use_spec_decoding is default to true by default, change in runtime by layers / requests self.use_spec_decoding = self.is_spec_decoding_enabled self.is_spec_dec_tree = is_spec_dec_tree self.is_spec_dec_dynamic_tree = is_spec_dec_dynamic_tree + # Forward static tree length to FMHA kernel selection. + self.max_total_draft_tokens = max_total_draft_tokens # Parameters can be fixed and not changed during runtime if the if self.is_spec_decoding_enabled: @@ -1004,9 +1014,18 @@ def update_spec_dec_param( mask_src = torch.index_select( slot_storage.packed_mask, 0, slot_ids)[:, :, :actual_mask_width] - total = num_gens * n_dt * actual_mask_width - self.spec_decoding_packed_mask.view(-1)[:total].copy_( - mask_src.reshape(-1), non_blocking=True) + if self.is_sm_version_trtllm_gen_kernel( + sm=get_sm_version()): + # Blackwell reads the padded 3D mask layout. + self.spec_decoding_packed_mask[:num_gens, :n_dt, : + actual_mask_width].copy_( + mask_src, + non_blocking=True) + else: + # Hopper XQA reads a compact flat prefix. + total = num_gens * n_dt * actual_mask_width + self.spec_decoding_packed_mask.view(-1)[:total].copy_( + mask_src.reshape(-1), non_blocking=True) self.spec_decoding_generation_lengths[:batch_size].fill_(n_dt) cpp_query_len = n_dt @@ -1673,6 +1692,8 @@ def _run( # stay as literal ``None`` until DeepSeek V4 sparse-MLA lands. sparse_mla_topk_lens=None, compressed_kv_cache_pool_ptr=None, + spec_decoding_target_max_draft_tokens=getattr( + metadata, 'max_total_draft_tokens', None), ) if self.print_skip_softmax_stat: diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 3d463d5052bc..8a393f21ad35 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -954,7 +954,6 @@ def test_eagle3_lora(use_cuda_graph: bool): @pytest.mark.parametrize("disable_overlap_scheduler", [False]) @pytest.mark.parametrize("use_cuda_graph", [True]) @pytest.mark.high_cuda_memory -@skip_blackwell @with_mocked_hf_download_for_single_gpu def test_llama_eagle3_dynamic_tree(use_cuda_graph: bool, disable_overlap_scheduler: bool):