diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 339da7c527a1..0c940a092181 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -295,7 +295,7 @@ bool AttentionOp::convertMMHAParamsToXQAParams(tensorrt_llm::kernels::XQAParams& xqaParams.start_token_idx_sf = generationsParams.start_token_idx_sf; // Parameters for sparse attention xqaParams.sparse_params = mRuntimeSparseAttentionParams; - xqaParams.use_sparse_attention = useTllmGenSparseAttention(); + xqaParams.use_sparse_attention_gen_paged = useTllmGenSparseAttentionPaged(); // Skip softmax threshold. xqaParams.skip_softmax_threshold_scale_factor = mSkipSoftmaxThresholdScaleFactorDecode; // Cross attention parameters. @@ -939,7 +939,7 @@ size_t AttentionOp::getWorkspaceSizeForGeneration(nvinfer1::DataType type, int32 size_t const cu_kv_seqlens_size = sizeof(int) * (batch_beam + 1); size_t const rotary_inv_freq_size = sizeof(float) * batch_beam * mRotaryEmbeddingDim / 2; // Two workspaces for sparse attention. One for the sequence lengths, and one for kv block offsets. - size_t const sparse_attn_cache_size = useTllmGenSparseAttention() + size_t const sparse_attn_cache_size = useTllmGenSparseAttentionPaged() ? sizeof(int) * (batch_beam + batch_beam * 2 * max_blocks_per_sequence) * mNumKVHeads : 0; xqa_workspaces[0] = cu_seqlens_size; @@ -1111,14 +1111,14 @@ int AttentionOp::mlaGeneration( = reinterpret_cast(params.bmm1_scale) + bmm1_scale_offset; } - // Set the following parameters if sparseMLA is used. + // Set the following parameters if sparseAttention is used. if (useSparseMLA()) { - tllmRunnerParams.mSparseMla = true; - tllmRunnerParams.mSparseMlaTopK = mRuntimeSparseAttentionParams.sparse_mla_topk; + tllmRunnerParams.mSparseAttention = true; + tllmRunnerParams.mSparseTopK = mRuntimeSparseAttentionParams.sparse_topk; tllmRunnerParams.kvPageIdxPtr = reinterpret_cast( mRuntimeSparseAttentionParams.sparse_attn_indices); - tllmRunnerParams.kvPtr = mRuntimeSparseAttentionParams.sparse_mla_kv_cache_pool; + tllmRunnerParams.kvPtr = mRuntimeSparseAttentionParams.sparse_kv_cache_pool; } mTllmGenFMHARunner->run(tllmRunnerParams); @@ -1883,7 +1883,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea fmhaParams.softmaxStatsPtr = params.softmax_stats; // Sparse attention parameters - if (useSparseMLA()) + if (useTllmGenSparseAttention()) { fmhaParams.sparse_params = mRuntimeSparseAttentionParams; } @@ -2751,6 +2751,10 @@ int AttentionOp::initialize() noexcept { fmhaParams.attentionInputLayout = AttentionInputLayout::PACKED_QKV; } + else if (useTllmGenSparseAttention()) + { + fmhaParams.attentionInputLayout = AttentionInputLayout::Q_PAGED_KV; + } else { fmhaParams.attentionInputLayout = (mPagedKVCache && mPagedContextFMHA) ? AttentionInputLayout::Q_PAGED_KV @@ -2799,6 +2803,7 @@ int AttentionOp::initialize() noexcept fmhaParams.hasAlibi = isALiBi(); fmhaParams.scaleAlibi = isAliBiWithScale(); fmhaParams.useSparseMLA = useSparseMLA(); + fmhaParams.useTllmGenSparseAttention = useTllmGenSparseAttention(); // Load kernels from the pre-compiled cubins. mFmhaDispatcher.reset(new FmhaDispatcher(fmhaParams)); diff --git a/cpp/tensorrt_llm/common/attentionOp.h b/cpp/tensorrt_llm/common/attentionOp.h index 4776b3d92dac..e7b2a5e7c5cd 100644 --- a/cpp/tensorrt_llm/common/attentionOp.h +++ b/cpp/tensorrt_llm/common/attentionOp.h @@ -367,9 +367,9 @@ class AttentionOp return mUseSparseAttention && mPagedKVCache && mEnableXQA; } - [[nodiscard]] bool useTllmGenSparseAttention() const + [[nodiscard]] bool useTllmGenSparseAttentionPaged() const { - return mUseTllmGenSparseAttention && useSparseAttention(); + return mUseTllmGenSparseAttentionPaged && useSparseAttention(); } [[nodiscard]] bool useSparseMLA() const @@ -377,6 +377,11 @@ class AttentionOp return mUseSparseAttention && mUseTllmGen && mIsMLAEnabled; } + [[nodiscard]] bool useTllmGenSparseAttention() const + { + return useSparseMLA() || (mUseSparseAttention && mUseTllmGen && mUseTllmGenSparseAttention); + } + [[nodiscard]] int smVersion() const { return mSM; @@ -457,6 +462,7 @@ class AttentionOp bool mIsGenerationMLA = false; bool mUseGenFlashMLA = false; bool mUseSparseAttention = false; + bool mUseTllmGenSparseAttentionPaged = false; bool mUseTllmGenSparseAttention = false; tensorrt_llm::kernels::MlaMetaParams mMLAParams; int mCpSize = 1; @@ -514,10 +520,10 @@ class AttentionOp mPosShiftEnabled, mPagedContextFMHA, mFP8ContextFMHA, mFP8AttenOutput, mFP8ContextMLA, mFP8GenerationMLA, mChunkPrefillBufferBatchSize, mDenseContextFMHA, mHasFullAttentionMask, mIsSpecDecodingEnabled, mUseSpecDecoding, mIsSpecDecTree, mSpecDecodingIsGenerationLengthVariable, mSpecDecodingMaxGenerationLength, - mIsMLAEnabled, mIsGenerationMLA, mUseGenFlashMLA, mUseSparseAttention, mUseTllmGenSparseAttention, - mMLAParams.data(), mCpSize, mCpRank, mCpGroup, mNumAttnHeads, mNumAttnKVHeads, mNumKVHeadsOrigin, - mAttnTpSize, mAttnTpRank, mAttnCpSize, mAttnCpRank, mUlyssesMQABroadcast, mEnableContextFMHA, - mFMHAForceFP32Acc, mMultiBlockMode, mEnableXQA, mUseKVCache, mSkipAttn, mFuseFp4Quant, + 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); }; diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h index 9679be86fcc6..b4c3b79f7503 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h @@ -143,6 +143,8 @@ struct MHARunnerFixedParams int sageBlockSizeV = 0; // Use sparse MLA ? bool useSparseMLA = false; + // Use sparse attention in trtllm-gen ? + bool useTllmGenSparseAttention = false; // Convert to string for debug. std::string convertToStrOutput() @@ -193,6 +195,8 @@ struct MHARunnerFixedParams output += ", sageBlockSizeQ = " + std::to_string(sageBlockSizeQ); output += ", sageBlockSizeK = " + std::to_string(sageBlockSizeK); output += ", sageBlockSizeV = " + std::to_string(sageBlockSizeV); + output += ", useSparseMLA = " + std::string(useSparseMLA ? "true" : "false"); + output += ", useTllmGenSparseAttention = " + std::string(useTllmGenSparseAttention ? "true" : "false"); return output; } diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.h b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.h index eb907edff1d5..60a5524de09e 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.h +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.h @@ -288,7 +288,7 @@ void buildXQALaunchParams(XQALaunchParam& launchParams, void*& in launchParams.bmm2_scale_ptr = reinterpret_cast(workspace); workspace = tensorrt_llm::common::nextWorkspacePtrWithAlignment(workspace, bmm2_scale_size); // Used for block sparse attention - if (params.use_sparse_attention) + if (params.use_sparse_attention_gen_paged) { launchParams.sparse_kv_block_offsets = reinterpret_cast(workspace); workspace = tensorrt_llm::common::nextWorkspacePtrWithAlignment(workspace, kv_block_offsets_size); diff --git a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h index 406bf54b1ffd..3fb7a753d0c6 100644 --- a/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h +++ b/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h @@ -116,7 +116,7 @@ struct XQAParams // sparse attention parameters SparseAttentionParams sparse_params; - bool use_sparse_attention = false; + bool use_sparse_attention_gen_paged = false; // Skip softmax threshold. float skip_softmax_threshold_scale_factor = 0.0f; @@ -197,7 +197,7 @@ struct XQAParams << "fp8_out_scale :" << fp8_out_scale << std ::endl << "encoder_input_lengths: " << encoder_input_lengths << std::endl << "sparse_params: " << sparse_params.toString() << std::endl - << "use_sparse_attention :" << (use_sparse_attention ? "true" : "false") << std ::endl + << "use_sparse_attention_gen_paged :" << (use_sparse_attention_gen_paged ? "true" : "false") << std ::endl << "skip_softmax_threshold_scale_factor :" << skip_softmax_threshold_scale_factor << std ::endl << "stream :" << stream; diff --git a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp index 58bfdb9ac048..682f30caa65a 100644 --- a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp +++ b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp @@ -120,10 +120,10 @@ bool FmhaDispatcher::isSupported() // the kernel is supported. tllmRunnerParams.mChunkedAttentionSize = INT_MAX; tllmRunnerParams.mAttentionWindowSize = INT_MAX; - // Set the kernel type and mask type if sparseMLA is used. - if (mFixedParams.useSparseMLA) + // Set the kernel type and mask type if sparse attention is used. + if (mFixedParams.useTllmGenSparseAttention) { - tllmRunnerParams.mSparseMla = true; + tllmRunnerParams.mSparseAttention = true; tllmRunnerParams.mKernelType = FmhaKernelType::Generation; tllmRunnerParams.mMaskType = TrtllmGenAttentionMaskType::Dense; } @@ -231,16 +231,24 @@ void FmhaDispatcher::run(MHARunnerParams runnerParams) // For skip softmax tllmRunnerParams.mSkipSoftmaxThresholdScaleFactor = runnerParams.skipSoftmaxThresholdScaleFactor; tllmRunnerParams.stream = runnerParams.stream; - // Set the sparse attention parameters if sparseMLA is used. - if (mFixedParams.useSparseMLA) + // Set the sparse attention parameters if trtllm-gen sparse attention is used. + if (mFixedParams.useTllmGenSparseAttention) { - tllmRunnerParams.mSparseMla = true; - tllmRunnerParams.mSparseMlaTopK = runnerParams.sparse_params.sparse_mla_topk; + tllmRunnerParams.mSparseAttention = true; + tllmRunnerParams.mSparseTopK = runnerParams.sparse_params.sparse_topk; tllmRunnerParams.mKernelType = FmhaKernelType::Generation; tllmRunnerParams.mMaskType = TrtllmGenAttentionMaskType::Dense; - tllmRunnerParams.kvPageIdxPtr - = reinterpret_cast(runnerParams.sparse_params.sparse_attn_indices); - tllmRunnerParams.kvPtr = runnerParams.sparse_params.sparse_mla_kv_cache_pool; + if (mFixedParams.useSparseMLA) + { + tllmRunnerParams.kvPageIdxPtr + = reinterpret_cast(runnerParams.sparse_params.sparse_attn_indices); + } + else + { + tllmRunnerParams.kvPageIdxPtr + = reinterpret_cast(runnerParams.sparse_params.sparse_attn_ctx_indices); + } + tllmRunnerParams.kvPtr = runnerParams.sparse_params.sparse_kv_cache_pool; } mTllmGenFMHARunner->run(tllmRunnerParams); diff --git a/cpp/tensorrt_llm/kernels/sparseAttentionKernels.h b/cpp/tensorrt_llm/kernels/sparseAttentionKernels.h index 6c701a686157..73f940acc9d5 100644 --- a/cpp/tensorrt_llm/kernels/sparseAttentionKernels.h +++ b/cpp/tensorrt_llm/kernels/sparseAttentionKernels.h @@ -29,12 +29,13 @@ namespace kernels struct SparseAttentionParams { - int32_t* sparse_kv_indices{nullptr}; // [num_kv_heads, num_sparse_kv_indices] - int32_t* sparse_attn_indices{nullptr}; // [num_kv_heads, num_sparse_attn_indices] - int32_t* sparse_kv_offsets{nullptr}; // [num_contexts + 1] - int32_t* sparse_attn_offsets{nullptr}; // [num_generations + 1] - int32_t sparse_mla_topk{0}; // for DSA attention - void* sparse_mla_kv_cache_pool{nullptr}; // for DSA attention + int32_t* sparse_kv_indices{nullptr}; // [num_kv_heads, num_sparse_kv_indices] + int32_t* sparse_attn_indices{nullptr}; // [num_kv_heads, num_sparse_attn_indices] + int32_t* sparse_kv_offsets{nullptr}; // [num_contexts + 1] + int32_t* sparse_attn_offsets{nullptr}; // [num_generations + 1] + int32_t* sparse_attn_ctx_indices{nullptr}; // [num_kv_heads, num_tokens, sparse_topk] + int32_t sparse_topk{0}; + void* sparse_kv_cache_pool{nullptr}; int32_t sparse_attn_indices_block_size{1}; int32_t sparse_attn_indices_stride{0}; @@ -46,8 +47,9 @@ struct SparseAttentionParams << "sparse_attn_indices: " << this->sparse_attn_indices << std::endl << "sparse_kv_offsets: " << this->sparse_kv_offsets << std::endl << "sparse_attn_offsets: " << this->sparse_attn_offsets << std::endl - << "sparse_mla_topk: " << this->sparse_mla_topk << std::endl - << "sparse_mla_kv_cache_pool: " << this->sparse_mla_kv_cache_pool << std::endl + << "sparse_attn_ctx_indices: " << this->sparse_attn_ctx_indices << std::endl + << "sparse_topk: " << this->sparse_topk << std::endl + << "sparse_kv_cache_pool: " << this->sparse_kv_cache_pool << std::endl << "sparse_attn_indices_block_size: " << this->sparse_attn_indices_block_size << std::endl << "sparse_attn_indices_stride: " << this->sparse_attn_indices_stride << std::endl; return ss.str(); diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 6208767157d8..4adf6f5bc48a 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -116,7 +116,8 @@ class TllmGenFmhaKernel inline uint64_t hashID(int qkvLayout, int maskType, int kernelType, int scheduler, int multiCtasKvMode, int headDimPerCtaV, int headDimQk, int headDimV, int tileSizeKv, int numTokensPerPage, - int maxNumHeadsQPerKvInCta, bool reuseSmemKForV, bool uses2CtaMma, bool sparseMla, bool skipsSoftmax) const + int maxNumHeadsQPerKvInCta, bool reuseSmemKForV, bool uses2CtaMma, bool sparseAttention, + bool skipsSoftmax) const { TLLM_CHECK_WITH_INFO((headDimPerCtaV >= 32) && (headDimQk >= 32) && (headDimV >= 32) && (headDimPerCtaV <= 1024) && (headDimQk <= 1024) && (headDimV <= 1024), @@ -142,7 +143,7 @@ class TllmGenFmhaKernel // Bit 49 - 56: maxNumHeadsQPerKvInCta. // Bit 57 - 57: reuseSmemKForV. // Bit 58 - 58: uses2CtaMma. - // Bit 59 - 59: sparseMla. + // Bit 59 - 59: sparseAttention. // Bit 60 - 60: skipsSoftmax. return (static_cast(qkvLayout) << 0) | (static_cast(maskType) << 4) | (static_cast(kernelType) << 8) | (static_cast(scheduler) << 12) @@ -150,7 +151,7 @@ class TllmGenFmhaKernel | (static_cast(headDimQk >> 3) << 26) | (static_cast(headDimV >> 3) << 34) | (static_cast(tileSizeKv >> 6) << 42) | (static_cast(log2(numTokensPerPage)) << 44) | (static_cast(maxNumHeadsQPerKvInCta) << 49) | (static_cast(reuseSmemKForV) << 57) - | (static_cast(uses2CtaMma) << 58) | (static_cast(sparseMla) << 59) + | (static_cast(uses2CtaMma) << 58) | (static_cast(sparseAttention) << 59) | (static_cast(skipsSoftmax) << 60); } @@ -159,7 +160,7 @@ class TllmGenFmhaKernel return hashID(kernelMeta.mQkvLayout, kernelMeta.mMaskType, kernelMeta.mKernelType, kernelMeta.mTileScheduler, kernelMeta.mMultiCtasKvMode, kernelMeta.mHeadDimPerCtaV, kernelMeta.mHeadDimQk, kernelMeta.mHeadDimV, kernelMeta.mTileSizeKv, kernelMeta.mNumTokensPerPage, kernelMeta.mMaxNumHeadsQPerKvInCta, - kernelMeta.mReuseSmemKForV, kernelMeta.m2CtaMma, kernelMeta.mSparseMla, + kernelMeta.mReuseSmemKForV, kernelMeta.m2CtaMma, kernelMeta.mSparseAttn, kernelMeta.mSkipsSoftmaxWhenPossible); } @@ -354,10 +355,10 @@ class TllmGenFmhaKernel { // The maximum attention window (the maximum number of tokensKv that will be attended to). int maxAttentionWindow{params.mMaxSeqLenKv}; - // The sparseMla only selects topK tokensKv. - if (params.mSparseMla) + // The sparseAttention only selects topK tokensKv. + if (params.mSparseAttention) { - maxAttentionWindow = std::min(params.mMaxSeqLenKv, params.mSparseMlaTopK); + maxAttentionWindow = std::min(params.mMaxSeqLenKv, params.mSparseTopK); } // Some of the tilesKv will be skipped if the sliding window attention or chunked attention is used. if (isSlidingOrChunkedCausalMask(selectKernelParams.mMaskType)) @@ -392,7 +393,7 @@ class TllmGenFmhaKernel // Need to select a different kernel. selectKernelParams.mSelectNewKernel = true; } - else if (totalNumCtas < params.mMultiProcessorCount && isMlaGenKernel(params) && !params.mSparseMla + else if (totalNumCtas < params.mMultiProcessorCount && isMlaGenKernel(params) && !params.mSparseAttention && selectKernelParams.mTileSizeKv == 128 && tensorrt_llm::common::getEnvUseTileSizeKv64ForTrtllmGen()) { // Use smaller tileSizeKv to fully utilize the SMs. @@ -515,7 +516,7 @@ class TllmGenFmhaKernel // The sparseMla kernel will always use the 2CTA high-throughput kernel. // Check the conditions. - if (params.mNumHeadsQPerKv <= 32 || (params.mSparseMla && params.mNumHeadsQPerKv < 128) + if (params.mNumHeadsQPerKv <= 32 || (params.mSparseAttention && params.mNumHeadsQPerKv < 128) || useSwapsMmaAbMlaGenKernel(params)) { kernelType = FmhaKernelType::SwapsMmaAbForGeneration; @@ -530,7 +531,7 @@ class TllmGenFmhaKernel selectKernelParams.mMultiCtasKvMode = MultiCtasKvMode::GmemReductionWithSeparateKernel; } // The keepsMmaAbForGeneration sparseMla kernels only support numHeadsQPerKv = 128. - TLLM_CHECK_WITH_INFO(!params.mSparseMla || params.mNumHeadsQPerKv == 128, + TLLM_CHECK_WITH_INFO(!params.mSparseAttention || params.mNumHeadsQPerKv == 128, "The keepsMmaAbForGeneration sparseMla kernels only support numHeadsQPerKv = 128, got %d", params.mNumHeadsQPerKv); // The 2CTA keepsMmaAbForGeneration kernel is used when the numHeadsQPerKv is 128. @@ -588,6 +589,11 @@ class TllmGenFmhaKernel { maxNumHeadsQPerKvInCta = 128; } + // TODO(yuhang): check if this is correct. + if (params.mSparseAttention) + { + maxNumHeadsQPerKvInCta = 64; + } TLLM_CHECK_WITH_INFO((params.mNumHeadsQPerKv < maxNumHeadsQPerKvInCta || params.mNumHeadsQPerKv % maxNumHeadsQPerKvInCta == 0), "Not supported"); @@ -614,7 +620,7 @@ class TllmGenFmhaKernel // The number of tokens per page. int numTokensPerPage = params.mNumTokensPerPage; // SparseMla kernels use a fixed numTokensPerPage = 1. - if (params.mSparseMla) + if (params.mSparseAttention) { numTokensPerPage = 1; } @@ -639,8 +645,9 @@ class TllmGenFmhaKernel + ", headDimQk=" + std::to_string(params.mHeadDimQk) + ", headDimV=" + std::to_string(params.mHeadDimV) + ", tileSizeKv=" + std::to_string(selectKernelParams.mTileSizeKv) + ", numTokensPerPage=" + std::to_string(numTokensPerPage) + ", maxNumHeadsQPerKvInCta=" + std::to_string(maxNumHeadsQPerKvInCta) - + ", reuseSmemKForV=" + std::to_string(selectKernelParams.mReuseSmemKForV) + ", uses2CtaMma=" - + std::to_string(selectKernelParams.mUses2CtaMma) + ", sparseMla=" + std::to_string(params.mSparseMla) + + ", reuseSmemKForV=" + std::to_string(selectKernelParams.mReuseSmemKForV) + + ", uses2CtaMma=" + std::to_string(selectKernelParams.mUses2CtaMma) + + ", sparseAttention=" + std::to_string(params.mSparseAttention) + ", skipsSoftmax=" + std::to_string(selectKernelParams.mSkipsSoftmaxWhenPossible); TLLM_LOG_DEBUG("Searching for kernel traits: " + info); @@ -651,7 +658,7 @@ class TllmGenFmhaKernel static_cast(selectKernelParams.mMultiCtasKvMode), selectKernelParams.mHeadDimPerCtaV, params.mHeadDimQk, params.mHeadDimV, selectKernelParams.mTileSizeKv, numTokensPerPage, maxNumHeadsQPerKvInCta, selectKernelParams.mReuseSmemKForV, selectKernelParams.mUses2CtaMma, - params.mSparseMla, selectKernelParams.mSkipsSoftmaxWhenPossible), + params.mSparseAttention, selectKernelParams.mSkipsSoftmaxWhenPossible), info); } diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaReduction.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaReduction.cu index 1a0cca54dab9..c5e1f206306b 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaReduction.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaReduction.cu @@ -32,8 +32,8 @@ namespace kernels template -__global__ void __launch_bounds__(NumThreadsPerCta, 2) fmhaReductionKernel(KernelParams const params, bool sparseMla, - int32_t numCtasForReduction, int32_t numCtasForAllHeads, int32_t numHeadDimCtasV) +__global__ void __launch_bounds__(NumThreadsPerCta, 2) fmhaReductionKernel(KernelParams const params, + bool sparseAttention, int32_t numCtasForReduction, int32_t numCtasForAllHeads, int32_t numHeadDimCtasV) { // clang-format off @@ -80,10 +80,10 @@ __global__ void __launch_bounds__(NumThreadsPerCta, 2) fmhaReductionKernel(Kerne int32_t seqLenKv{params.ptrSeqLensKv[batchIdx]}; // Consider the causal-mask speculative decoding. seqLenKv = seqLenKv - ((params.mMaxSeqLenQ - 1) - ctaIdxQ); - // Consider sparseMlaTopK. - if (sparseMla) + // Consider sparseTopK. + if (sparseAttention) { - seqLenKv = min(seqLenKv, params.mSparseMlaTopK); + seqLenKv = min(seqLenKv, params.mNumSparseTopk); } // The actual number of CtasKv (TileSizeKv is always 128 for now). int32_t numCtasKv{min((seqLenKv + 127) / 128, params.mMaxNumCtasKv)}; @@ -388,7 +388,7 @@ void runFmhaReduction(TllmGenFmhaKernelMetaInfo const& kernelMeta, KernelParams // Launch the kernel. TLLM_CUDA_CHECK(cudaLaunchKernelEx( - &config, kernel, params, kernelMeta.mSparseMla, numCtasForReduction, numCtasForAllHeads, numHeadDimCtasV)); + &config, kernel, params, kernelMeta.mSparseAttn, numCtasForReduction, numCtasForAllHeads, numHeadDimCtasV)); } //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h index d8ba62edda38..7e805ee85241 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h @@ -286,9 +286,9 @@ struct TllmGenFmhaRunnerParams // Skip softmax threshold scale factor. float mSkipSoftmaxThresholdScaleFactor; // Whether to use sparse MLA. - bool mSparseMla; + bool mSparseAttention; // The top k value for sparse MLA. - int mSparseMlaTopK; + int mSparseTopK; // The cuda stream. cudaStream_t stream; // The layer index. diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h index 14b6c27e3944..186ff1e591fe 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/kernelParams.h @@ -148,6 +148,10 @@ struct KernelParams int64_t mNumHiddenEltsO; // The total number of pages in the paged-kv memory pool. int32_t mNumPagesInMemPool; + // The top k value for sparse MLA. + int32_t mNumSparseTopk; + // The number of tokensQ per CTA (used for groupsHeadsTokensQ generation kernel). + int32_t mNumTokensPerCtaQ; // The number of tokens per page (used if dynamic numTokensPerPage is enabled). int32_t mNumTokensPerPageLog2; // The output scale for FP8 quantization. @@ -165,10 +169,8 @@ struct KernelParams int32_t mStartTokenIdx; // The sum of sequence lengths for Q and K/V. int32_t mSumOfSeqLensQ, mSumOfSeqLensKv; - // The top k value for sparse MLA. - int32_t mSparseMlaTopK; - // The flag to use block sparse attention. - bool mUseBlockSparseAttention; + // Sparsity mode for sparse attention. + int32_t mSparsityMode; // Create the TMA shape/stride for Q. template @@ -713,12 +715,23 @@ struct KernelParams tileShapeKv[0] = numEltsInClampedHeadDimKv / numEltsDivisor; tileShapeKv[1] = numKeysPerTile; - // If sparse MLA is enabled, the shape and stride for K need to be updated for 2D layout (numTokensKvInPagedKv, - // headDimQk). - if (options.mSparseMla) + int32_t numInstsQ{kernelMeta.mStepQ / kernelMeta.mTileSizeQ}; + int32_t tileSizePerCtaQ{kernelMeta.mTileSizeQ * numInstsQ}; + // TODO(yuhangh): check the correctness + if (kernelMeta.mGroupsHeadsQ && !isSpecDecodingGenerationKernel(options.mKernelType)) + { + tileSizePerCtaQ = 1; + } + params.mNumTokensPerCtaQ = tileSizePerCtaQ; + + // If sparse attention is enabled, the shape and stride for KV need to be updated for 2D layout + // (numTokensKvInPagedKv, headDimQk). + if (options.mSparseAttention) { shapeK = std::vector{static_cast(options.mHeadDimQk), static_cast(INT_MAX)}; strideK = std::vector{1, static_cast(options.mHeadDimQk)}; + shapeV = std::vector{static_cast(options.mHeadDimV), static_cast(INT_MAX)}; + strideV = std::vector{1, static_cast(options.mHeadDimV)}; tileShapeKv[1] = 1; } @@ -844,10 +857,18 @@ struct KernelParams params.mStartTokenIdx = options.mSfStartTokenIdx; // The sparseMlaTopK needs to be a multiple of 4 as we use 16B cpAsync instructions for the indices. TLLM_CHECK_WITH_INFO( - !options.mSparseMla || (options.mSparseMlaTopK % 4) == 0, "SparseMlaTopK must be a multiple of 4"); - params.mSparseMlaTopK = options.mSparseMlaTopK; - params.mUseBlockSparseAttention = options.mUseBlockSparseAttention; + !options.mSparseAttention || (options.mSparseTopK % 4) == 0, "SparseTopK must be a multiple of 4"); + params.mNumSparseTopk = options.mSparseTopK; + if (options.mSparseAttention) + { + params.mSparsityMode = 0; + } + else if (options.mUseBlockSparseAttention) + { + params.mSparsityMode = 1; + } params.mSkipSoftmaxThresholdScaleFactor = options.mSkipSoftmaxThresholdScaleFactor; + return params; } }; diff --git a/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp b/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp index cd5ce16428a9..62c312bc07a9 100644 --- a/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp +++ b/cpp/tensorrt_llm/kernels/xqaDispatcher.cpp @@ -430,7 +430,7 @@ void XqaDispatcher::runImpl( tllmRunnerParams.mNumTokensPerPage = kv_cache_buffer.mTokensPerBlock; // Gather kv page offsets for sparse attention. - if (params.use_sparse_attention) + if (params.use_sparse_attention_gen_paged) { invokeGatherKvPageOffsets(reinterpret_cast(launchParams.sparse_kv_block_offsets), launchParams.sparse_seq_lengths, reinterpret_cast(kv_cache_buffer.data), @@ -445,7 +445,8 @@ void XqaDispatcher::runImpl( } else { - TLLM_CHECK_WITH_INFO(!params.use_sparse_attention, "Sparse attention is not supported for KVLinearBuffer."); + TLLM_CHECK_WITH_INFO( + !(params.use_sparse_attention_gen_paged), "Sparse attention is not supported for KVLinearBuffer."); static_assert(std::is_same_v); // Contiguous KV tllmRunnerParams.mQkvLayout = QkvLayout::ContiguousKv; diff --git a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp index 60e0d009394d..92466dfdd42b 100644 --- a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp @@ -64,8 +64,8 @@ void initBindings(nb::module_& m) nb::arg("softmax_stats_tensor") = std::nullopt, nb::arg("spec_decoding_bool_params"), nb::arg("spec_decoding_tensor_params"), nb::arg("sparse_kv_indices") = std::nullopt, nb::arg("sparse_kv_offsets") = std::nullopt, nb::arg("sparse_attn_indices") = std::nullopt, - nb::arg("sparse_attn_offsets") = std::nullopt, nb::arg("sparse_attn_indices_block_size"), - nb::arg("sparse_mla_topk") = std::nullopt, + nb::arg("sparse_attn_offsets") = std::nullopt, nb::arg("sparse_attn_ctx_indices") = std::nullopt, + nb::arg("sparse_attn_indices_block_size"), nb::arg("sparse_mla_topk") = std::nullopt, nb::arg("skip_softmax_threshold_scale_factor_prefill") = std::nullopt, nb::arg("skip_softmax_threshold_scale_factor_decode") = std::nullopt, nb::arg("skip_softmax_stat") = std::nullopt, nb::arg("cu_q_seqlens") = std::nullopt, diff --git a/cpp/tensorrt_llm/pybind/thop/bindings.cpp b/cpp/tensorrt_llm/pybind/thop/bindings.cpp index 7c017a10a6b4..4705f511d032 100644 --- a/cpp/tensorrt_llm/pybind/thop/bindings.cpp +++ b/cpp/tensorrt_llm/pybind/thop/bindings.cpp @@ -64,8 +64,8 @@ void initBindings(pybind11::module_& m) py::arg("softmax_stats_tensor") = std::nullopt, py::arg("spec_decoding_bool_params"), py::arg("spec_decoding_tensor_params"), py::arg("sparse_kv_indices") = std::nullopt, py::arg("sparse_kv_offsets") = std::nullopt, py::arg("sparse_attn_indices") = std::nullopt, - py::arg("sparse_attn_offsets") = std::nullopt, py::arg("sparse_attn_indices_block_size"), - py::arg("sparse_mla_topk") = std::nullopt, + py::arg("sparse_attn_offsets") = std::nullopt, py::arg("sparse_attn_ctx_indices") = std::nullopt, + py::arg("sparse_attn_indices_block_size"), py::arg("sparse_mla_topk") = std::nullopt, py::arg("skip_softmax_threshold_scale_factor_prefill") = std::nullopt, py::arg("skip_softmax_threshold_scale_factor_decode") = std::nullopt, py::arg("skip_softmax_stat") = std::nullopt, py::arg("cu_q_seqlens") = std::nullopt, diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index b0ee56a83f5d..d09098fa8eef 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -88,11 +88,11 @@ class RunnerBase c10::ArrayRef> spec_decoding_tensor_params, torch::optional attention_sinks, torch::optional sparse_kv_indices, torch::optional sparse_kv_offsets, torch::optional sparse_attn_indices, - torch::optional sparse_attn_offsets, int64_t const sparse_attn_indices_block_size, - int32_t const sparse_mla_topk, std::optional cu_q_seqlens, - std::optional cu_kv_seqlens, std::optional fmha_scheduler_counter, - std::optional mla_bmm1_scale, std::optional mla_bmm2_scale, - std::optional quant_q_buffer) const + torch::optional sparse_attn_offsets, torch::optional sparse_attn_ctx_indices, + int64_t const sparse_attn_indices_block_size, int32_t const sparse_mla_topk, + std::optional cu_q_seqlens, std::optional cu_kv_seqlens, + std::optional fmha_scheduler_counter, std::optional mla_bmm1_scale, + std::optional mla_bmm2_scale, std::optional quant_q_buffer) const = 0; }; @@ -149,11 +149,11 @@ class Runner : public RunnerBase c10::ArrayRef> spec_decoding_tensor_params, torch::optional attention_sinks, torch::optional sparse_kv_indices, torch::optional sparse_kv_offsets, torch::optional sparse_attn_indices, - torch::optional sparse_attn_offsets, int64_t const sparse_attn_indices_block_size, - int32_t const sparse_mla_topk, std::optional cu_q_seqlens, - std::optional cu_kv_seqlens, std::optional fmha_scheduler_counter, - std::optional mla_bmm1_scale, std::optional mla_bmm2_scale, - std::optional quant_q_buffer) const override + torch::optional sparse_attn_offsets, torch::optional sparse_attn_ctx_indices, + int64_t const sparse_attn_indices_block_size, int32_t const sparse_mla_topk, + std::optional cu_q_seqlens, std::optional cu_kv_seqlens, + std::optional fmha_scheduler_counter, std::optional mla_bmm1_scale, + std::optional mla_bmm2_scale, std::optional quant_q_buffer) const override { auto stream = at::cuda::getCurrentCUDAStream(qkv_or_q.get_device()); T* attention_input = static_cast(qkv_or_q.slice(0, token_offset).data_ptr()); @@ -404,15 +404,24 @@ class Runner : public RunnerBase = sparse_attn_indices.has_value() ? sparse_attn_indices.value().data_ptr() : nullptr; op.mRuntimeSparseAttentionParams.sparse_attn_offsets = sparse_attn_offsets.has_value() ? sparse_attn_offsets.value().data_ptr() : nullptr; + op.mRuntimeSparseAttentionParams.sparse_attn_ctx_indices + = sparse_attn_ctx_indices.has_value() ? sparse_attn_ctx_indices.value().data_ptr() : nullptr; op.mRuntimeSparseAttentionParams.sparse_attn_indices_block_size = sparse_attn_indices_block_size; op.mRuntimeSparseAttentionParams.sparse_attn_indices_stride = sparse_attn_indices.has_value() ? sparse_attn_indices.value().size(-1) : 0; - if (op.isMLAEnabled() && op.mUseSparseAttention) + if (op.useTllmGenSparseAttention()) { - op.mRuntimeSparseAttentionParams.sparse_mla_topk = sparse_mla_topk; + if (op.isMLAEnabled()) + { + op.mRuntimeSparseAttentionParams.sparse_topk = sparse_mla_topk; + } + else + { + op.mRuntimeSparseAttentionParams.sparse_topk = sparse_attn_ctx_indices.value().size(-1); + } if (op.useKVCache() && host_kv_cache_pool_pointers.has_value()) { - op.mRuntimeSparseAttentionParams.sparse_mla_kv_cache_pool = reinterpret_cast( + op.mRuntimeSparseAttentionParams.sparse_kv_cache_pool = reinterpret_cast( host_kv_cache_pool_pointers.value().index({pool_index, 0}).item()); } } @@ -631,8 +640,8 @@ void attention(torch::Tensor q, std::optional k, std::optional> spec_decoding_tensor_params, std::optional sparse_kv_indices, std::optional sparse_kv_offsets, std::optional sparse_attn_indices, std::optional sparse_attn_offsets, - int64_t const sparse_attn_indices_block_size, std::optional sparse_mla_topk, - std::optional skip_softmax_threshold_scale_factor_prefill, + std::optional sparse_attn_ctx_indices, int64_t const sparse_attn_indices_block_size, + std::optional sparse_mla_topk, std::optional skip_softmax_threshold_scale_factor_prefill, std::optional skip_softmax_threshold_scale_factor_decode, std::optional skip_softmax_stat, std::optional cu_q_seqlens, std::optional cu_kv_seqlens, std::optional fmha_scheduler_counter, std::optional mla_bmm1_scale, @@ -759,12 +768,18 @@ void attention(torch::Tensor q, std::optional k, std::optionalmIsSpecDecTree = spec_decoding_bool_params[2]; // is_spec_dec_tree op->mUseSparseAttention = false; + op->mUseTllmGenSparseAttentionPaged = false; op->mUseTllmGenSparseAttention = false; if ((sparse_kv_indices.has_value() && sparse_kv_indices.value().numel() > 0) - || (sparse_attn_indices.has_value() && sparse_attn_indices.value().numel() > 0)) + || (sparse_attn_indices.has_value() && sparse_attn_indices.value().numel() > 0) + || (sparse_attn_ctx_indices.has_value() && sparse_attn_ctx_indices.value().numel() > 0)) { op->mUseSparseAttention = true; if (sparse_attn_indices.has_value() && sparse_attn_indices.value().numel() > 0) + { + op->mUseTllmGenSparseAttentionPaged = true; + } + if (sparse_attn_ctx_indices.has_value() && sparse_attn_ctx_indices.value().numel() > 0) { op->mUseTllmGenSparseAttention = true; } @@ -898,8 +913,8 @@ void attention(torch::Tensor q, std::optional k, std::optional 0) && (attn_input_type != AttentionInputType::ContextOnly)) @@ -917,8 +932,8 @@ void attention(torch::Tensor q, std::optional k, std::optional k, std::optional> spec_decoding_tensor_params, std::optional sparse_kv_indices, std::optional sparse_kv_offsets, std::optional sparse_attn_indices, std::optional sparse_attn_offsets, - int64_t const sparse_attn_indices_block_size, std::optional sparse_mla_topk, - std::optional skip_softmax_threshold_scale_factor_prefill, + std::optional sparse_attn_ctx_indices, int64_t const sparse_attn_indices_block_size, + std::optional sparse_mla_topk, std::optional skip_softmax_threshold_scale_factor_prefill, std::optional skip_softmax_threshold_scale_factor_decode, std::optional skip_softmax_stat, std::optional cu_q_seqlens, std::optional cu_kv_seqlens, std::optional fmha_scheduler_counter, std::optional mla_bmm1_scale, diff --git a/tensorrt_llm/_torch/attention_backend/sparse/kernel.py b/tensorrt_llm/_torch/attention_backend/sparse/kernel.py index 3ca6f8ddd27c..8a469c4e9fe0 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/kernel.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/kernel.py @@ -1797,29 +1797,38 @@ def triton_rocket_reduce_scores( def _convert_req_index_to_global_index_kernel_with_stride_factor( req_id_ptr, # int32 [num_tokens] block_table_ptr, # int32 [num_requests, max_num_blocks_per_req] - token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] - out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS] + token_indices_ptr, # int32 [num_kv_heads, num_tokens, NUM_TOPK_TOKENS] + out_ptr, # int32 [num_kv_heads, num_tokens, kv_factor, NUM_TOPK_TOKENS] # shapes (compile-time where possible) max_num_blocks_per_req: tl.constexpr, BLOCK_SIZE: tl.constexpr, BLOCK_N: tl.constexpr, # tile width along columns # strides (in elements) stride_factor: tl.constexpr, # for strided memory layout adjustment layer_id: tl.constexpr, # for layer interleaving layout + num_kv_heads: tl.constexpr, + kv_factor: tl.constexpr, bt_stride0, bt_stride1, ti_stride0, ti_stride1, + ti_stride2, out_stride0, out_stride1, + out_stride2, + out_stride3, ): """ Triton kernel for converting request-local token indices to global KV cache pool indices. Derived from vllm's flashmla_sparse.py, with stride_factor fused in the kernel. + Now supports multiple KV heads and kv_factor dimensions. + Output shape: [num_kv_heads, num_tokens, kv_factor, NUM_TOPK_TOKENS] """ - # program_id(0) -> token_id (row) - # program_id(1) -> tile index along columns - token_id = tl.program_id(0) - tile_id = tl.program_id(1) + # program_id(0) -> kv_head_idx + # program_id(1) -> token_id (row) + # program_id(2) -> tile index along columns + kv_head_idx = tl.program_id(0) + token_id = tl.program_id(1) + tile_id = tl.program_id(2) # Each program covers BLOCK_N consecutive columns indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N) @@ -1828,7 +1837,7 @@ def _convert_req_index_to_global_index_kernel_with_stride_factor( req = tl.load(req_id_ptr + token_id) # Load token indices for this tile - ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1 + ti_ptr = token_indices_ptr + kv_head_idx * ti_stride0 + token_id * ti_stride1 + indice_id * ti_stride2 tok = tl.load(ti_ptr) # int32 # Only token == -1 should propagate as -1 @@ -1836,59 +1845,76 @@ def _convert_req_index_to_global_index_kernel_with_stride_factor( # Compute block id and in-block offset block_id = tok // BLOCK_SIZE - inblock_off = tok % BLOCK_SIZE + layer_id * BLOCK_SIZE + inblock_off = tok % BLOCK_SIZE # Guard block_table access valid_block = block_id < max_num_blocks_per_req bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1 base = tl.load(bt_ptr, mask=valid_block, other=0) - # If token == -1 OR block_id OOB, output -1 - # Otherwise: base * stride_factor + inblock_off + # KV cache pool: [num_blocks, num_layers, kv_factor, num_kv_heads, num_tokens_per_block] # (stride_factor accounts for layer interleaving in strided KV cache pools) - out_val = tl.where(is_invalid_tok | (~valid_block), -1, - base * stride_factor + inblock_off) + base_offset = base * stride_factor + \ + layer_id * kv_factor * num_kv_heads * BLOCK_SIZE + \ + kv_head_idx * BLOCK_SIZE + \ + inblock_off - # Store results - out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1 - tl.store(out_ptr_ij, out_val) + for kv_idx in range(kv_factor): + kv_factor_offset = kv_idx * num_kv_heads * BLOCK_SIZE + + out_val = tl.where(is_invalid_tok | (~valid_block), -1, + base_offset + kv_factor_offset) + + # Store results for this kv_factor index + # Output layout: [num_kv_heads, num_tokens, kv_factor, NUM_TOPK_TOKENS] + out_ptr_ij = (out_ptr + kv_head_idx * out_stride0 + + token_id * out_stride1 + kv_idx * out_stride2 + + indice_id * out_stride3) + tl.store(out_ptr_ij, out_val) def triton_convert_req_index_to_global_index( - req_id: torch.Tensor, # int32 [num_tokens] - block_table: torch. - Tensor, # int32 [num_requests, max_num_blocks_per_req] - token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS] - BLOCK_SIZE: int, - NUM_TOPK_TOKENS: int = 2048, - BLOCK_N: int = 128, # tile width along columns - stride_factor: + req_id: torch.Tensor, # int32 [num_tokens] + block_table: torch.Tensor, # int32 [num_requests, max_num_blocks_per_req] + token_indices: torch. + Tensor, # int32 [num_kv_heads, num_tokens, NUM_TOPK_TOKENS] + BLOCK_SIZE: int, + NUM_TOPK_TOKENS: int = 2048, + BLOCK_N: int = 128, # tile width along columns + stride_factor: int = None, # for strided memory layout (with layer interleaving), defaults to BLOCK_SIZE - layer_id: int = 0, # for layer interleaving layout + layer_id: int = 0, # for layer interleaving layout + num_kv_heads: int = 1, + kv_factor: int = 1, ): """ Convert request-local token indices to global KV cache pool indices. out[token_id, indice_id] = block_table[req_id[token_id], - token_indices[token_id, indice_id] // BLOCK_SIZE] * stride_factor - + token_indices[token_id, indice_id] % BLOCK_SIZE + token_indices[token_id, indice_id] // BLOCK_SIZE] * stride_factor + + layer_id * kv_factor * num_kv_heads * BLOCK_SIZE + + kv_factor_idx * num_kv_heads * BLOCK_SIZE + + kv_head_idx * BLOCK_SIZE + + token_indices[token_id, indice_id] % BLOCK_SIZE Args: - stride_factor: Memory stride between consecutive blocks (default: BLOCK_SIZE). + num_kv_heads: Number of KV heads (default: 1) + kv_factor: KV factor for additional dimension (default: 1) + stride_factor: Memory stride between consecutive blocks (default: BLOCK_SIZE * num_kv_heads * kv_factor). For non-contiguous pools with layer interleaving, use - (num_layers * BLOCK_SIZE) to account for memory gaps. + (num_layers * BLOCK_SIZE * num_kv_heads * kv_factor) to account for memory gaps. Only when token_indices[token_id, indice_id] == -1 do we output -1. For safety, we also output -1 if the derived block_id would be out-of-bounds. """ if stride_factor is None: - stride_factor = BLOCK_SIZE + stride_factor = BLOCK_SIZE * num_kv_heads * kv_factor assert req_id.dtype == torch.int32 assert block_table.dtype == torch.int32 assert token_indices.dtype == torch.int32 - assert token_indices.shape[1] == NUM_TOPK_TOKENS + assert token_indices.shape[-1] == NUM_TOPK_TOKENS assert NUM_TOPK_TOKENS % BLOCK_N == 0, \ f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by" \ f"BLOCK_N ({BLOCK_N})" @@ -1901,15 +1927,22 @@ def triton_convert_req_index_to_global_index( req_id_c = req_id.contiguous() block_table_c = block_table.contiguous() token_indices_c = token_indices.contiguous() - out = torch.empty_like(token_indices_c) + # Create output tensor with shape: [num_kv_heads, num_tokens, kv_factor, NUM_TOPK_TOKENS] + out = torch.empty((num_kv_heads, num_tokens, kv_factor, NUM_TOPK_TOKENS), + dtype=token_indices_c.dtype, + device=token_indices_c.device) # Strides in elements bt_stride0, bt_stride1 = block_table_c.stride() - ti_stride0, ti_stride1 = token_indices_c.stride() - out_stride0, out_stride1 = out.stride() + if token_indices_c.ndim == 2: # [num_tokens, NUM_TOPK_TOKENS] num_kv_heads=1 by default + ti_stride1, ti_stride2 = token_indices_c.stride() + ti_stride0 = 0 + else: # [num_kv_heads, num_tokens, NUM_TOPK_TOKENS] + ti_stride0, ti_stride1, ti_stride2 = token_indices_c.stride() + out_stride0, out_stride1, out_stride2, out_stride3 = out.stride() - # Exact 2D grid: tokens × column tiles - grid = (num_tokens, tiles_per_row) + # 3D grid: num_kv_heads × num_tokens × column tiles + grid = (num_kv_heads, num_tokens, tiles_per_row) _convert_req_index_to_global_index_kernel_with_stride_factor[grid]( req_id_c, @@ -1921,13 +1954,18 @@ def triton_convert_req_index_to_global_index( BLOCK_SIZE, BLOCK_N, stride_factor, - # strides layer_id, + num_kv_heads, + kv_factor, + # strides bt_stride0, bt_stride1, ti_stride0, ti_stride1, + ti_stride2, out_stride0, out_stride1, + out_stride2, + out_stride3, ) return out diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index e669925a3474..0577ba1b63e9 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -213,6 +213,7 @@ def plan( sparse_kv_offsets: Optional[torch.Tensor] = None, sparse_attn_indices: Optional[torch.Tensor] = None, sparse_attn_offsets: Optional[torch.Tensor] = None, + sparse_attn_ctx_indices: Optional[torch.Tensor] = None, sparse_attn_indices_block_size: int = 1, sparse_mla_topk: int = 0, skip_softmax_threshold_scale_factor_prefill: Optional[float] = None, @@ -261,6 +262,7 @@ def plan( sparse_kv_offsets (torch.Tensor): The batch offsets for the sparse KV indices, with shape of (num_contexts + 1) on GPU. sparse_attn_indices (torch.Tensor): The sparse indices for the attention layer, with shape of (num_heads_kv, num_sparse_tokens) on GPU. sparse_attn_offsets (torch.Tensor): The batch offsets for the sparse attention indices, with shape of (num_generations + 1) on GPU. + sparse_attn_ctx_indices (torch.Tensor): The sparse indices for the context attention layer, with shape of (num_heads_kv, num_ctx_tokens, num_sparse_tokens) on GPU. sparse_attn_indices_block_size (int): The granularity of the sparse attention indices, used by block sparse attention. sparse_mla_topk (int): The topk for the sparse MLA, used by DSA attention. skip_softmax_threshold_scale_factor_prefill (float): The scale factor for the skip softmax threshold in prefill phase. @@ -307,6 +309,7 @@ def plan( self.sparse_kv_offsets = sparse_kv_offsets self.sparse_attn_indices = sparse_attn_indices self.sparse_attn_offsets = sparse_attn_offsets + self.sparse_attn_ctx_indices = sparse_attn_ctx_indices self.sparse_attn_indices_block_size = sparse_attn_indices_block_size self.sparse_mla_topk = sparse_mla_topk self.helix_position_offsets = helix_position_offsets @@ -569,6 +572,7 @@ def run( self.sparse_kv_offsets, self.sparse_attn_indices, self.sparse_attn_offsets, + self.sparse_attn_ctx_indices, self.sparse_attn_indices_block_size, self.sparse_mla_topk, self.skip_softmax_threshold_scale_factor_prefill, @@ -1625,6 +1629,8 @@ def forward( q, k, metadata, **kwargs) sparse_attn_indices, sparse_attn_offsets = self.sparse_attn_predict( q, k, metadata, **kwargs) + sparse_attn_ctx_indices = self.sparse_attn_ctx_predict( + q, k, metadata, **kwargs) sparse_attn_indices_block_size = self.sparse_attention_config.get_indices_block_size( ) @@ -1685,6 +1691,7 @@ def forward( sparse_kv_offsets=sparse_kv_offsets, sparse_attn_indices=sparse_attn_indices, sparse_attn_offsets=sparse_attn_offsets, + sparse_attn_ctx_indices=sparse_attn_ctx_indices, sparse_attn_indices_block_size=sparse_attn_indices_block_size, sparse_mla_topk=metadata.sparse_mla_topk if hasattr( metadata, 'sparse_mla_topk') else 0, @@ -1940,6 +1947,19 @@ def sparse_attn_predict( """ raise NotImplementedError + def sparse_attn_ctx_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata: TrtllmAttentionMetadata, + **kwargs, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """ + Predict sparse attn ctx indices. It's implemented in the derived class. + """ + # TODO(yuhangh): complete this part logic + return None + def mla_rope_generation( self, fused_q: torch.Tensor, diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_attention.py b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py new file mode 100644 index 000000000000..6a5af5c6235c --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py @@ -0,0 +1,1022 @@ +""" +Unit tests for sparse attention with TrtllmAttention backend. +""" + +import math +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import pytest +import torch + +import tensorrt_llm +from tensorrt_llm._torch.attention_backend.sparse.kernel import ( + triton_convert_req_index_to_global_index, +) +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.bindings.executor import KvCacheConfig +from tensorrt_llm.mapping import Mapping + +ATOL = 1e-2 +RTOL = 1e-2 + + +@dataclass(kw_only=True, frozen=False) +class SparseScenario: + """Base configuration for sparse attention tests.""" + + dtype: torch.dtype = torch.float16 + kvcache_dtype: torch.dtype = torch.float16 + num_layers: int = 1 + num_heads: int = 32 + num_kv_heads: int = 8 + head_dim: int = 128 + page_size: int = 32 + num_pages: int = 16 + batch_size: int = 4 + + @property + def num_kv_groups(self) -> int: + return self.num_heads // self.num_kv_heads + + @property + def kv_cache_len(self) -> int: + return self.page_size * self.num_pages + + @property + def max_num_pages(self) -> int: + return self.batch_size * self.num_pages + + +@dataclass(kw_only=True, frozen=False) +class SparseContextScenario(SparseScenario): + """Configuration for context phase tests with sparse kv cache write.""" + + seq_lens: Tuple[int, ...] = (128,) + sparse_ratio: float = 0.5 + + def __post_init__(self): + if len(self.seq_lens) != self.batch_size: + raise ValueError( + f"seq_lens length {len(self.seq_lens)} must match batch_size {self.batch_size}" + ) + + @property + def max_seq_len(self) -> int: + return max(self.seq_lens) + + @property + def nnz_q(self) -> int: + return sum(self.seq_lens) + + +@dataclass(kw_only=True, frozen=False) +class SparseGenerationScenario(SparseScenario): + """Configuration for generation phase tests with sparse attention.""" + + past_kv_lens: Tuple[int, ...] = (256,) + num_contexts: int = 0 + sparse_ratio: float = 0.5 + + def __post_init__(self): + if len(self.past_kv_lens) != self.batch_size: + raise ValueError( + f"past_kv_lens length {len(self.past_kv_lens)} must match batch_size {self.batch_size}" + ) + + @property + def num_generations(self) -> int: + return self.batch_size - self.num_contexts + + @property + def max_past_kv_len(self) -> int: + return max(self.past_kv_lens) + + @property + def nnz_q(self) -> int: + return self.num_generations + + +class MockSparseAttentionConfig: + def get_indices_block_size(self) -> int: + return 1 + + +class TestSparseAttention(TrtllmAttention): + """TrtllmAttention subclass for testing with predetermined sparse indices.""" + + def __init__( + self, + *args, + sparse_kv_indices: Optional[torch.Tensor] = None, + sparse_kv_offsets: Optional[torch.Tensor] = None, + sparse_attn_indices: Optional[torch.Tensor] = None, + sparse_attn_offsets: Optional[torch.Tensor] = None, + sparse_attn_ctx_indices: Optional[torch.Tensor] = None, + **kwargs, + ): + kwargs["sparse_attention_config"] = MockSparseAttentionConfig() + kwargs["pos_embd_params"] = None + super().__init__(*args, **kwargs) + self._sparse_kv_indices = sparse_kv_indices + self._sparse_kv_offsets = sparse_kv_offsets + self._sparse_attn_indices = sparse_attn_indices + self._sparse_attn_offsets = sparse_attn_offsets + self._sparse_attn_ctx_indices = sparse_attn_ctx_indices + + def sparse_kv_predict(self, q, k, metadata, **kwargs): + return self._sparse_kv_indices, self._sparse_kv_offsets + + def sparse_attn_predict(self, q, k, metadata, **kwargs): + return self._sparse_attn_indices, self._sparse_attn_offsets + + def sparse_attn_ctx_predict(self, q, k, metadata, **kwargs): + return self._sparse_attn_ctx_indices + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """Repeat kv heads to match query heads.""" + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_key_value_heads, n_rep, slen, head_dim + ) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def _get_kv_cache_dtype(dtype: torch.dtype): + """Convert torch dtype to TensorRT-LLM dtype.""" + if dtype == torch.float16: + return tensorrt_llm.bindings.DataType.HALF + elif dtype == torch.bfloat16: + return tensorrt_llm.bindings.DataType.BF16 + else: + raise ValueError(f"Unsupported dtype: {dtype}") + + +def create_kv_cache_manager( + s: SparseScenario, kv_cache: Optional[torch.Tensor] = None +) -> KVCacheManager: + """Create kv cache manager for testing.""" + kv_cache_config = KvCacheConfig(max_tokens=s.max_num_pages * s.page_size) + mapping = Mapping(world_size=1, tp_size=1, rank=0) + + manager = KVCacheManager( + kv_cache_config, + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, + num_layers=s.num_layers, + num_kv_heads=s.num_kv_heads, + head_dim=s.head_dim, + tokens_per_block=s.page_size, + max_seq_len=s.max_num_pages * s.page_size, + max_batch_size=s.batch_size, + mapping=mapping, + dtype=_get_kv_cache_dtype(s.kvcache_dtype), + ) + + if kv_cache is not None: + for i in range(s.num_layers): + manager.get_buffers(i, kv_layout="HND").copy_(kv_cache[i]) + + return manager + + +def _generate_sparse_indices_for_batch( + seq_len: int, + sparse_ratio: float, + num_kv_heads: int, + device: torch.device, +) -> torch.Tensor: + """Generate sparse indices for a single batch with per-head patterns.""" + sparse_len = max(1, int(seq_len * sparse_ratio)) + batch_indices = [] + for _ in range(num_kv_heads): + indices = torch.randperm(seq_len, device=device)[:sparse_len].sort().values + batch_indices.append(indices) + return torch.stack(batch_indices, dim=0) + + +def generate_sparse_kv_indices( + s: SparseContextScenario, device: torch.device +) -> Tuple[torch.Tensor, torch.Tensor]: + """Generate sparse kv indices for context phase.""" + all_indices = [] + offsets = [0] + + for seq_len in s.seq_lens: + batch_indices = _generate_sparse_indices_for_batch( + seq_len, s.sparse_ratio, s.num_kv_heads, device + ) + all_indices.append(batch_indices) + offsets.append(offsets[-1] + batch_indices.shape[1]) + + indices = torch.cat(all_indices, dim=1).int() + offsets = torch.tensor(offsets, dtype=torch.int32, device=device) + return indices, offsets + + +def generate_sparse_attn_indices( + s: SparseGenerationScenario, device: torch.device +) -> Tuple[torch.Tensor, torch.Tensor]: + """Generate sparse attention indices for generation phase.""" + all_indices = [] + offsets = [0] + + for gen_idx in range(s.num_generations): + batch_idx = s.num_contexts + gen_idx + past_kv_len = s.past_kv_lens[batch_idx] + batch_indices = _generate_sparse_indices_for_batch( + past_kv_len, s.sparse_ratio, s.num_kv_heads, device + ) + all_indices.append(batch_indices) + offsets.append(offsets[-1] + batch_indices.shape[1]) + + indices = torch.cat(all_indices, dim=1).int() + offsets = torch.tensor(offsets, dtype=torch.int32, device=device) + return indices, offsets + + +def generate_sparse_attn_ctx_indices( + s: SparseContextScenario, device: torch.device +) -> torch.Tensor: + """ + Generate sparse attention context indices for context phase. + Returns: [num_kv_heads, num_tokens, num_sparse_tokens] with -1 padding. + """ + all_batch_indices = [] + token_offset = 0 + + for seq_len in s.seq_lens: + batch_indices = [] + for token_idx in range(seq_len): + # Each token can attend to all previous tokens including itself due to causal mask + available_kv_len = token_idx + 1 + sparse_len = max(1, int(available_kv_len * s.sparse_ratio)) + + per_head_indices = [] + for _ in range(s.num_kv_heads): + # TODO: check whether the indices are sorted + indices = torch.randperm(available_kv_len, device=device)[:sparse_len].sort().values + per_head_indices.append(indices) + + batch_indices.append(per_head_indices) + + all_batch_indices.append(batch_indices) + token_offset += seq_len + + # Flatten all batches together and find max sparse length + flattened_indices = [] + for batch in all_batch_indices: + flattened_indices.extend(batch) + + # Find the maximum sparse length across all tokens and heads + max_sparse_len = max( + max(indices.shape[0] for indices in token_indices) for token_indices in flattened_indices + ) + + # Build the tensor with -1 padding + total_tokens = sum(s.seq_lens) + result = torch.full( + (s.num_kv_heads, total_tokens, max_sparse_len), -1, dtype=torch.int32, device=device + ) + + for token_idx, token_indices in enumerate(flattened_indices): + for head_idx in range(s.num_kv_heads): + indices = token_indices[head_idx] + result[head_idx, token_idx, : len(indices)] = indices + + return result + + +def convert_sparse_attn_ctx_indices_to_global( + sparse_attn_ctx_indices: torch.Tensor, + metadata: TrtllmAttentionMetadata, + layer_idx: int = 0, + kv_factor: int = 2, +) -> torch.Tensor: + """ + Convert local sparse_attn_ctx_indices to global KV cache pool indices. + """ + num_kv_heads, num_tokens, num_sparse_tokens = sparse_attn_ctx_indices.shape + device = sparse_attn_ctx_indices.device + + tokens_per_block = metadata.kv_cache_manager.tokens_per_block + num_layers = metadata.kv_cache_manager.num_layers + stride_factor = num_layers * tokens_per_block * kv_factor * num_kv_heads + + # Build req_idx_per_token + num_contexts = metadata.num_contexts + seq_lens = ( + metadata.seq_lens[:num_contexts] + if hasattr(metadata.seq_lens, "__getitem__") + else metadata.seq_lens + ) + host_req_idx_per_token = torch.repeat_interleave( + torch.arange(num_contexts, dtype=torch.int32), seq_lens, dim=0 + ) + req_idx_per_token = host_req_idx_per_token.to(device) + + # Build block_table + block_ids_all = metadata.kv_cache_manager.get_batch_cache_indices( + metadata.request_ids[:num_contexts] + ) + + max_blocks_used = max(len(b) for b in block_ids_all) if block_ids_all else 1 + + host_block_table = torch.full((num_contexts, max_blocks_used), -1, dtype=torch.int32) + for i, blocks in enumerate(block_ids_all): + if len(blocks) > 0: + host_block_table[i, : len(blocks)] = torch.tensor(blocks, dtype=torch.int32) + + block_table = host_block_table.to(device) + + # Convert to global + global_indices = triton_convert_req_index_to_global_index( + req_idx_per_token, + block_table, + sparse_attn_ctx_indices, + BLOCK_SIZE=tokens_per_block, + NUM_TOPK_TOKENS=num_sparse_tokens, + BLOCK_N=64, + stride_factor=stride_factor, + layer_id=layer_idx, + num_kv_heads=num_kv_heads, + kv_factor=kv_factor, + ) + + return global_indices + + +def _extract_batch_tensors( + tensor: torch.Tensor, offset: int, length: int, shape_per_token: Tuple +) -> torch.Tensor: + """Extract and reshape tensors for a specific batch.""" + return tensor[offset : offset + length].view(length, *shape_per_token) + + +def build_expected_sparse_kv( + k: torch.Tensor, + v: torch.Tensor, + sparse_kv_indices: torch.Tensor, + sparse_kv_offsets: torch.Tensor, + s: SparseContextScenario, +) -> List[Tuple[torch.Tensor, torch.Tensor]]: + """Build expected sparse K and V values based on sparse indices.""" + expected_kvs = [] + token_offset = 0 + + for batch_idx, seq_len in enumerate(s.seq_lens): + sparse_len = max(1, int(seq_len * s.sparse_ratio)) + k_batch = _extract_batch_tensors(k, token_offset, seq_len, (s.num_kv_heads, s.head_dim)) + v_batch = _extract_batch_tensors(v, token_offset, seq_len, (s.num_kv_heads, s.head_dim)) + + expected_k = torch.zeros( + sparse_len, s.num_kv_heads, s.head_dim, device=k.device, dtype=k.dtype + ) + expected_v = torch.zeros_like(expected_k) + + start, end = sparse_kv_offsets[batch_idx].item(), sparse_kv_offsets[batch_idx + 1].item() + for head_idx in range(s.num_kv_heads): + indices = sparse_kv_indices[head_idx, start:end] + expected_k[:, head_idx] = k_batch[indices, head_idx] + expected_v[:, head_idx] = v_batch[indices, head_idx] + + expected_kvs.append((expected_k, expected_v)) + token_offset += seq_len + + return expected_kvs + + +def _extract_tokens_from_cache( + kv_buffer: torch.Tensor, + block_ids: List[int], + num_tokens: int, + num_kv_heads: int, + head_dim: int, + page_size: int, + dtype: torch.dtype, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Extract tokens from paged kv cache.""" + device = kv_buffer.device + k_cache = torch.zeros(num_tokens, num_kv_heads, head_dim, device=device, dtype=dtype) + v_cache = torch.zeros_like(k_cache) + + for token_idx in range(num_tokens): + block_idx = token_idx // page_size + offset_in_block = token_idx % page_size + block_id = block_ids[block_idx] + + for head_idx in range(num_kv_heads): + k_cache[token_idx, head_idx] = kv_buffer[block_id, 0, head_idx, offset_in_block, :].to( + dtype + ) + v_cache[token_idx, head_idx] = kv_buffer[block_id, 1, head_idx, offset_in_block, :].to( + dtype + ) + + return k_cache, v_cache + + +def extract_kv_from_paged_cache( + kv_cache_manager: KVCacheManager, + request_ids: List[int], + sparse_kv_offsets: torch.Tensor, + s: SparseContextScenario, + dtype: torch.dtype, +) -> List[Tuple[torch.Tensor, torch.Tensor]]: + """Extract K and V values from paged kv cache.""" + kv_buffer = kv_cache_manager.get_buffers(0, kv_layout="HND") + kv_caches = [] + + for batch_idx in range(s.batch_size): + num_sparse_tokens = ( + sparse_kv_offsets[batch_idx + 1].item() - sparse_kv_offsets[batch_idx].item() + ) + block_ids = kv_cache_manager.get_block_ids_per_seq([request_ids[batch_idx]])[0] + k_cache, v_cache = _extract_tokens_from_cache( + kv_buffer, block_ids, num_sparse_tokens, s.num_kv_heads, s.head_dim, s.page_size, dtype + ) + kv_caches.append((k_cache, v_cache)) + + return kv_caches + + +def _compute_causal_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + num_kv_groups: int, +) -> torch.Tensor: + """Compute causal attention for a single batch.""" + seq_len = q.shape[2] + head_dim = q.shape[3] + + k_expanded = repeat_kv(k, num_kv_groups) + v_expanded = repeat_kv(v, num_kv_groups) + + attn_weights = torch.matmul(q, k_expanded.transpose(-1, -2)) / math.sqrt(head_dim) + causal_mask = torch.triu( + torch.full((seq_len, seq_len), float("-inf"), device=q.device), diagonal=1 + ) + attn_weights = attn_weights + causal_mask + attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to( + q.dtype + ) + output = torch.matmul(attn_weights, v_expanded) + + return output + + +def reference_context_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: SparseContextScenario, +) -> torch.Tensor: + """Reference implementation for context phase.""" + outputs = [] + token_offset = 0 + + for seq_len in s.seq_lens: + q_batch = _extract_batch_tensors(q, token_offset, seq_len, (s.num_heads, s.head_dim)) + k_batch = _extract_batch_tensors(k, token_offset, seq_len, (s.num_kv_heads, s.head_dim)) + v_batch = _extract_batch_tensors(v, token_offset, seq_len, (s.num_kv_heads, s.head_dim)) + + q_batch = q_batch.view(1, seq_len, s.num_heads, s.head_dim).transpose(1, 2) + k_batch = k_batch.view(1, seq_len, s.num_kv_heads, s.head_dim).transpose(1, 2) + v_batch = v_batch.view(1, seq_len, s.num_kv_heads, s.head_dim).transpose(1, 2) + + output_batch = _compute_causal_attention(q_batch, k_batch, v_batch, s.num_kv_groups) + output_batch = output_batch.transpose(1, 2).reshape(seq_len, s.num_heads * s.head_dim) + outputs.append(output_batch) + token_offset += seq_len + + return torch.cat(outputs, dim=0) + + +def reference_context_sparse_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + sparse_attn_ctx_indices: torch.Tensor, + s: SparseContextScenario, +) -> torch.Tensor: + """ + Reference implementation for context phase with sparse attention. + Uses mask-based approach for each KV head. + """ + total_tokens = sum(s.seq_lens) + device = q.device + dtype = q.dtype + + # Reshape inputs: [num_tokens, num_heads, head_dim] + q_reshaped = q.view(total_tokens, s.num_heads, s.head_dim) + k_reshaped = k.view(total_tokens, s.num_kv_heads, s.head_dim) + v_reshaped = v.view(total_tokens, s.num_kv_heads, s.head_dim) + + outputs = [] + token_offset = 0 + + for seq_len in s.seq_lens: + q_batch = q_reshaped[ + token_offset : token_offset + seq_len + ] # [seq_len, num_heads, head_dim] + k_batch = k_reshaped[ + token_offset : token_offset + seq_len + ] # [seq_len, num_kv_heads, head_dim] + v_batch = v_reshaped[ + token_offset : token_offset + seq_len + ] # [seq_len, num_kv_heads, head_dim] + + batch_output = [] + + # Process each KV head + for kv_head_idx in range(s.num_kv_heads): + k_head = k_batch[:, kv_head_idx, :] + v_head = v_batch[:, kv_head_idx, :] + + # Build sparse mask for this head + sparse_mask = torch.full( + (seq_len, seq_len), float("-inf"), device=device, dtype=torch.float32 + ) + + for token_idx in range(seq_len): + global_token_idx = token_offset + token_idx + # Get sparse indices for this token: [num_sparse_tokens] + indices = sparse_attn_ctx_indices[kv_head_idx, global_token_idx] + # Filter out -1 padding + valid_indices = indices[indices >= 0] + # Set mask values to 0 for valid positions + sparse_mask[token_idx, valid_indices] = 0.0 + + # Apply causal mask on top of sparse mask + causal_mask = torch.triu( + torch.full((seq_len, seq_len), float("-inf"), device=device, dtype=torch.float32), + diagonal=1, + ) + combined_mask = sparse_mask + causal_mask + + # Process each query head in this KV group + for group_idx in range(s.num_kv_groups): + q_head_idx = kv_head_idx * s.num_kv_groups + group_idx + q_head = q_batch[:, q_head_idx, :] # [seq_len, head_dim] + + attn_scores = torch.matmul(q_head, k_head.T) / math.sqrt(s.head_dim) + attn_scores = attn_scores + combined_mask + attn_weights = torch.nn.functional.softmax( + attn_scores, dim=-1, dtype=torch.float32 + ).to(dtype) + + out_head = torch.matmul(attn_weights, v_head) + batch_output.append(out_head) + + # Concatenate all heads: [seq_len, num_heads, head_dim] -> [seq_len, num_heads * head_dim] + batch_output = torch.stack(batch_output, dim=1) + batch_output = batch_output.reshape(seq_len, s.num_heads * s.head_dim) + outputs.append(batch_output) + + token_offset += seq_len + + return torch.cat(outputs, dim=0) + + +def _get_selected_pages_tokens( + token_indices: torch.Tensor, + page_size: int, + kv_len: int, + device: torch.device, +) -> torch.Tensor: + """Convert token indices to page indices and gather all tokens from selected pages.""" + if len(token_indices) == 0: + return torch.tensor([], dtype=torch.long, device=device) + + page_indices = torch.unique((token_indices // page_size).sort().values) + selected_tokens = [] + + for page_idx in page_indices: + token_start = page_idx * page_size + token_end = min(token_start + page_size, kv_len) + selected_tokens.append(torch.arange(token_start, token_end, device=device)) + + return ( + torch.cat(selected_tokens) + if selected_tokens + else torch.tensor([], dtype=torch.long, device=device) + ) + + +def _compute_sparse_attention_per_head( + q_head: torch.Tensor, + k_sparse: torch.Tensor, + v_sparse: torch.Tensor, + head_dim: int, +) -> torch.Tensor: + """Compute attention for a single query head.""" + if len(k_sparse) == 0: + return torch.zeros(head_dim, device=q_head.device, dtype=q_head.dtype) + + attn_weights = torch.matmul(q_head, k_sparse.T) / math.sqrt(head_dim) + attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to( + q_head.dtype + ) + return torch.matmul(attn_weights, v_sparse) + + +def reference_generation_sparse_attention( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_new: torch.Tensor, + v_new: torch.Tensor, + sparse_attn_indices: torch.Tensor, + sparse_attn_offsets: torch.Tensor, + s: SparseGenerationScenario, +) -> torch.Tensor: + """Reference implementation for generation phase with sparse attention at page granularity.""" + outputs = [] + + for gen_idx in range(s.num_generations): + batch_idx = s.num_contexts + gen_idx + past_kv_len = s.past_kv_lens[batch_idx] + kv_len = past_kv_len + 1 + + k_full = k_cache[batch_idx, :kv_len].clone() + v_full = v_cache[batch_idx, :kv_len].clone() + k_full[past_kv_len] = k_new[gen_idx].view(s.num_kv_heads, s.head_dim) + v_full[past_kv_len] = v_new[gen_idx].view(s.num_kv_heads, s.head_dim) + + start, end = sparse_attn_offsets[gen_idx].item(), sparse_attn_offsets[gen_idx + 1].item() + q_batch = q[gen_idx].view(s.num_heads, s.head_dim) + head_outputs = [] + + for kv_head_idx in range(s.num_kv_heads): + kv_head_indices = sparse_attn_indices[kv_head_idx, start:end] + selected_token_indices = _get_selected_pages_tokens( + kv_head_indices, s.page_size, kv_len, q.device + ) + + if len(selected_token_indices) == 0: + head_outputs.extend( + [torch.zeros(s.head_dim, device=q.device, dtype=q.dtype)] * s.num_kv_groups + ) + continue + + k_sparse = k_full[selected_token_indices, kv_head_idx, :] + v_sparse = v_full[selected_token_indices, kv_head_idx, :] + + for group_idx in range(s.num_kv_groups): + q_head_idx = kv_head_idx * s.num_kv_groups + group_idx + out_head = _compute_sparse_attention_per_head( + q_batch[q_head_idx], k_sparse, v_sparse, s.head_dim + ) + head_outputs.append(out_head) + + outputs.append(torch.cat(head_outputs, dim=0)) + + return torch.stack(outputs, dim=0) + + +def _setup_context_test(s: SparseContextScenario): + """Setup common components for context test.""" + device = torch.device("cuda") + torch.manual_seed(42) + + q = torch.randn(s.nnz_q, s.num_heads * s.head_dim, device=device, dtype=s.dtype) + k = torch.randn(s.nnz_q, s.num_kv_heads * s.head_dim, device=device, dtype=s.dtype) + v = torch.randn(s.nnz_q, s.num_kv_heads * s.head_dim, device=device, dtype=s.dtype) + + sparse_kv_indices, sparse_kv_offsets = generate_sparse_kv_indices(s, device) + + kv_cache = torch.zeros( + s.num_layers, + s.max_num_pages, + 2, + s.num_kv_heads, + s.page_size, + s.head_dim, + device=device, + dtype=s.kvcache_dtype, + ) + kv_cache_manager = create_kv_cache_manager(s, kv_cache) + + request_ids = list(range(s.batch_size)) + kv_cache_manager.add_dummy_requests(request_ids, list(s.seq_lens)) + + metadata = TrtllmAttentionMetadata( + num_contexts=s.batch_size, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0] * s.batch_size), + seq_lens=torch.tensor(s.seq_lens, dtype=torch.int32), + max_num_requests=s.batch_size, + max_num_tokens=s.nnz_q, + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + prompt_lens=list(s.seq_lens), + ) + metadata.prepare() + + attention = TestSparseAttention( + layer_idx=0, + num_heads=s.num_heads, + head_dim=s.head_dim, + num_kv_heads=s.num_kv_heads, + sparse_kv_indices=sparse_kv_indices, + sparse_kv_offsets=sparse_kv_offsets, + ) + + return ( + device, + q, + k, + v, + sparse_kv_indices, + sparse_kv_offsets, + kv_cache_manager, + request_ids, + metadata, + attention, + ) + + +def _setup_generation_test(s: SparseGenerationScenario): + """Setup common components for generation test.""" + device = torch.device("cuda") + torch.manual_seed(42) + + token_nums = [past_len + 1 for past_len in s.past_kv_lens] + + q = torch.randn(s.num_generations, s.num_heads * s.head_dim, device=device, dtype=s.dtype) + k_new = torch.randn( + s.num_generations, s.num_kv_heads * s.head_dim, device=device, dtype=s.dtype + ) + v_new = torch.randn( + s.num_generations, s.num_kv_heads * s.head_dim, device=device, dtype=s.dtype + ) + + sparse_attn_indices, sparse_attn_offsets = generate_sparse_attn_indices(s, device) + + kv_cache = torch.randn( + s.num_layers, + s.max_num_pages, + 2, + s.num_kv_heads, + s.page_size, + s.head_dim, + device=device, + dtype=s.kvcache_dtype, + ) + kv_cache_manager = create_kv_cache_manager(s, kv_cache) + + request_ids = list(range(s.batch_size)) + kv_cache_manager.add_dummy_requests(request_ids, token_nums) + + metadata = TrtllmAttentionMetadata( + num_contexts=s.num_contexts, + kv_cache_params=KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=list(s.past_kv_lens) + ), + seq_lens=torch.tensor([1] * s.num_generations).int(), + max_num_requests=s.batch_size, + max_num_tokens=s.num_generations, + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + prompt_lens=list(s.past_kv_lens), + ) + metadata.prepare() + + attention = TestSparseAttention( + layer_idx=0, + num_heads=s.num_heads, + head_dim=s.head_dim, + num_kv_heads=s.num_kv_heads, + sparse_attn_indices=sparse_attn_indices, + sparse_attn_offsets=sparse_attn_offsets, + ) + + return ( + device, + q, + k_new, + v_new, + sparse_attn_indices, + sparse_attn_offsets, + kv_cache_manager, + request_ids, + metadata, + attention, + ) + + +def _build_reference_kv_cache( + kv_cache_manager, request_ids, s: SparseGenerationScenario, device, dtype +): + """Build reference K, V cache from paged format.""" + k_cache_ref = torch.zeros( + s.batch_size, s.kv_cache_len, s.num_kv_heads, s.head_dim, device=device, dtype=dtype + ) + v_cache_ref = torch.zeros_like(k_cache_ref) + + kv_buffer = kv_cache_manager.get_buffers(0, kv_layout="HND") + for batch_idx, past_kv_len in enumerate(s.past_kv_lens): + block_ids = kv_cache_manager.get_block_ids_per_seq([request_ids[batch_idx]])[0] + for block_local_idx, block_id in enumerate(block_ids): + token_start = block_local_idx * s.page_size + token_end = min(token_start + s.page_size, past_kv_len) + tokens_in_block = token_end - token_start + + for head_idx in range(s.num_kv_heads): + k_cache_ref[batch_idx, token_start:token_end, head_idx] = kv_buffer[ + block_id, 0, head_idx, :tokens_in_block, : + ].to(dtype) + v_cache_ref[batch_idx, token_start:token_end, head_idx] = kv_buffer[ + block_id, 1, head_idx, :tokens_in_block, : + ].to(dtype) + + return k_cache_ref, v_cache_ref + + +@pytest.mark.parametrize( + "s", + [ + SparseContextScenario(batch_size=2, seq_lens=(48, 64), sparse_ratio=0.5, num_pages=8), + SparseContextScenario( + batch_size=4, seq_lens=(96, 112, 128, 144), sparse_ratio=0.25, num_pages=16 + ), + SparseContextScenario(batch_size=1, seq_lens=(256,), sparse_ratio=0.75, num_pages=8), + SparseContextScenario(batch_size=3, seq_lens=(64, 96, 128), sparse_ratio=0.4, num_pages=12), + ], + ids=["batch2_var_seq", "batch4_var_seq", "batch1_seq256", "batch3_var_seq"], +) +def test_context_sparse_kv(s: SparseContextScenario): + """Test context phase with sparse kv cache write.""" + ( + device, + q, + k, + v, + sparse_kv_indices, + sparse_kv_offsets, + kv_cache_manager, + request_ids, + metadata, + attention, + ) = _setup_context_test(s) + + ref_output = reference_context_attention(q.clone(), k.clone(), v.clone(), s) + expected_kvs = build_expected_sparse_kv( + k.clone(), v.clone(), sparse_kv_indices, sparse_kv_offsets, s + ) + + qkv = torch.cat([q, k, v], dim=1) + output = attention.forward(qkv, None, None, metadata) + + assert output.shape == ref_output.shape, f"Shape mismatch: {output.shape} vs {ref_output.shape}" + torch.testing.assert_close(output, ref_output, atol=ATOL, rtol=RTOL) + print(f"Context sparse kv attention output test passed: {s}") + + actual_kvs = extract_kv_from_paged_cache( + kv_cache_manager, request_ids, sparse_kv_offsets, s, s.dtype + ) + + for batch_idx in range(s.batch_size): + actual_k, actual_v = actual_kvs[batch_idx] + expected_k, expected_v = expected_kvs[batch_idx] + torch.testing.assert_close( + actual_k, + expected_k, + atol=ATOL, + rtol=RTOL, + msg=f"K cache mismatch for batch {batch_idx} after sparse compaction", + ) + torch.testing.assert_close( + actual_v, + expected_v, + atol=ATOL, + rtol=RTOL, + msg=f"V cache mismatch for batch {batch_idx} after sparse compaction", + ) + + print(f"Context sparse kv cache content test passed: {s}") + kv_cache_manager.shutdown() + + +@pytest.mark.parametrize( + "s", + [ + SparseGenerationScenario( + batch_size=2, past_kv_lens=(96, 128), sparse_ratio=0.5, num_pages=16 + ), + SparseGenerationScenario( + batch_size=4, past_kv_lens=(192, 224, 256, 288), sparse_ratio=0.25, num_pages=32 + ), + SparseGenerationScenario(batch_size=1, past_kv_lens=(64,), sparse_ratio=0.75, num_pages=8), + SparseGenerationScenario( + batch_size=3, past_kv_lens=(128, 160, 192), sparse_ratio=0.4, num_pages=24 + ), + ], + ids=["batch2_var_kv", "batch4_var_kv", "batch1_kv64", "batch3_var_kv"], +) +def test_generation_sparse_attention(s: SparseGenerationScenario): + """Test generation phase with sparse attention computation.""" + ( + device, + q, + k_new, + v_new, + sparse_attn_indices, + sparse_attn_offsets, + kv_cache_manager, + request_ids, + metadata, + attention, + ) = _setup_generation_test(s) + + k_cache_ref, v_cache_ref = _build_reference_kv_cache( + kv_cache_manager, request_ids, s, device, s.dtype + ) + ref_sparse_output = reference_generation_sparse_attention( + q, k_cache_ref, v_cache_ref, k_new, v_new, sparse_attn_indices, sparse_attn_offsets, s + ) + + qkv = torch.cat([q, k_new, v_new], dim=1) + output = attention.forward(qkv, None, None, metadata) + + expected_shape = (s.num_generations, s.num_heads * s.head_dim) + assert output.shape == expected_shape, f"Shape mismatch: {output.shape} vs {expected_shape}" + assert torch.isfinite(output).all(), "Output contains non-finite values" + + torch.testing.assert_close(output, ref_sparse_output, atol=ATOL, rtol=RTOL) + print(f"Generation sparse attention test passed: {s}") + kv_cache_manager.shutdown() + + +@pytest.mark.parametrize( + "s", + [ + SparseContextScenario( + batch_size=2, + seq_lens=(128, 64), + sparse_ratio=0.5, + num_pages=8, + num_kv_heads=1, + num_heads=8, + head_dim=128, + ), + ], +) +def test_context_sparse_attention_mqa(s: SparseContextScenario): + """Test context phase with sparse attention using sparse_attn_ctx_indices (MQA setup).""" + device, q, k, v, _, _, kv_cache_manager, request_ids, metadata, _ = _setup_context_test(s) + + # Generate sparse attention context indices + sparse_attn_ctx_indices = generate_sparse_attn_ctx_indices(s, device) + + # Convert to global indices for attentionOp + global_sparse_attn_ctx_indices = convert_sparse_attn_ctx_indices_to_global( + sparse_attn_ctx_indices, metadata, layer_idx=0 + ) + + # Compute reference output using local indices + ref_output = reference_context_sparse_attention( + q.clone(), k.clone(), v.clone(), sparse_attn_ctx_indices, s + ) + + # Verify reference output shape + total_tokens = sum(s.seq_lens) + expected_shape = (total_tokens, s.num_heads * s.head_dim) + assert ref_output.shape == expected_shape, ( + f"Reference output shape mismatch: {ref_output.shape} vs {expected_shape}" + ) + assert torch.isfinite(ref_output).all(), "Reference output contains non-finite values" + + print(f"Context sparse attention MQA reference test passed: {s}") + + attention = TestSparseAttention( + layer_idx=0, + num_heads=s.num_heads, + head_dim=s.head_dim, + num_kv_heads=s.num_kv_heads, + sparse_attn_ctx_indices=global_sparse_attn_ctx_indices, # Use global indices here + ) + + qkv = torch.cat([q, k, v], dim=1) + output = attention.forward(qkv, None, None, metadata) + torch.testing.assert_close(output, ref_output, atol=ATOL, rtol=RTOL) + print(f"Context sparse attention MQA forward test passed: {s}") + + kv_cache_manager.shutdown() + + +if __name__ == "__main__": + s = SparseContextScenario( + batch_size=2, + seq_lens=(128, 64), + sparse_ratio=0.5, + num_pages=8, + num_kv_heads=1, + num_heads=8, + head_dim=128, + ) + test_context_sparse_attention_mqa(s)