[None][revert] Revert "[TRTLLM-11119][feat] Blackwell SageAttention, Integrate into AttentionOp API (NVIDIA#11718)" (NVIDIA#12679)

Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com>

Author: yunruis

cpp/tensorrt_llm/common/attentionOp.cpp

@@ -20,7 +20,6 @@
 #include "tensorrt_llm/common/envUtils.h"
 #include "tensorrt_llm/common/logger.h"
 #include "tensorrt_llm/common/memoryUtils.h"
-#include "tensorrt_llm/common/sageQuant.h"
 #include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention.h"
 #include "tensorrt_llm/kernels/flashMLA/flash_mla.h"
 #include "tensorrt_llm/kernels/gptKernels.h"
@@ -779,10 +778,6 @@ size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t
         = mNumAttnHeads * dim_k_per_head; // Assuming effective num_kv_heads = head_num for layout
     int const total_v_dim_all_heads
         = mNumAttnHeads * dim_v_per_head; // Assuming effective num_kv_heads = head_num for layout
-    bool const useSageAttnSeparateQkv = mEnableContextFMHA && !mIsMLAEnabled && mFmhaDispatcher->isSeparateQAndKvInput()
-        && (mSageAttnNumEltsPerBlkQ > 0 || mSageAttnNumEltsPerBlkK > 0 || mSageAttnNumEltsPerBlkV > 0)
-        && mFP8ContextFMHA;
-
     // Packed fp8 qkv buffer size for normal fp8 context FMHA
     size_t fp8_qkv_buffer_size = mFP8ContextFMHA && mEnableContextFMHA && !mFmhaDispatcher->isSeparateQAndKvInput()
         ? max_num_tokens * size_t(local_hidden_units_qo + 2 * local_hidden_units_kv)
@@ -808,22 +803,6 @@ size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t
             fp8_v_buf_size = mChunkPrefillBufferBatchSize * max_num_tokens * static_cast<size_t>(total_v_dim_all_heads);
         }
     }
-    else if (useSageAttnSeparateQkv)
-    {
-        fp8_q_buf_size = max_num_tokens * static_cast<size_t>(local_hidden_units_qo);
-        fp8_k_buf_size = max_num_tokens * static_cast<size_t>(local_hidden_units_kv);
-        fp8_v_buf_size = max_num_tokens * static_cast<size_t>(local_hidden_units_kv);
-    }
-
-    int32_t const q_max_n_blk = mSageAttnNumEltsPerBlkQ > 0 ? tc::divUp(input_seq_length, mSageAttnNumEltsPerBlkQ) : 0;
-    int32_t const k_max_n_blk = mSageAttnNumEltsPerBlkK > 0 ? tc::divUp(input_seq_length, mSageAttnNumEltsPerBlkK) : 0;
-    size_t const sage_q_sfs_buffer_size
-        = useSageAttnSeparateQkv ? sizeof(float) * mNumAttnHeads * batch_size * static_cast<size_t>(q_max_n_blk) : 0;
-    size_t const sage_k_sfs_buffer_size
-        = useSageAttnSeparateQkv ? sizeof(float) * mNumAttnKVHeads * batch_size * static_cast<size_t>(k_max_n_blk) : 0;
-    size_t const sage_v_sfs_buffer_size = useSageAttnSeparateQkv
-        ? sizeof(float) * tc::divUp(local_hidden_units_kv, std::max(1, mSageAttnNumEltsPerBlkV))
-        : 0;
 
     size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * max_num_tokens;
     size_t const encoder_padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * max_num_tokens;
@@ -839,7 +818,7 @@ size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t
         ? 0
         : (2 * size * cpMaxPaddedSequenceLength * getHeadSize() * (mNumHeads + 2 * mNumKVHeads) + cu_seqlens_size);
 
-    int const NUM_BUFFERS = 26;
+    int const NUM_BUFFERS = 23;
     size_t workspaces[NUM_BUFFERS];
     workspaces[0] = CUBLAS_WORKSPACE_SIZE;
     workspaces[1] = attention_mask_size;
@@ -863,10 +842,7 @@ size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t
     workspaces[19] = fmha_scheduler_counter;
     workspaces[20] = fmha_bmm1_scale_size;
     workspaces[21] = fmha_bmm2_scale_size;
-    workspaces[22] = sage_q_sfs_buffer_size;
-    workspaces[23] = sage_k_sfs_buffer_size;
-    workspaces[24] = sage_v_sfs_buffer_size;
-    workspaces[25] = cpWorkspaceSize;
+    workspaces[22] = cpWorkspaceSize;
     context_workspace_size = tc::calculateTotalWorkspaceSize(workspaces, NUM_BUFFERS);
 
     return context_workspace_size;
@@ -1442,10 +1418,6 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
         = mNumAttnHeads * dim_k_per_head; // Assuming effective num_kv_heads = head_num for layout
     int const total_v_dim_all_heads
         = mNumAttnHeads * dim_v_per_head; // Assuming effective num_kv_heads = head_num for layout
-    bool const useSageAttnSeparateQkv = mEnableContextFMHA && !mIsMLAEnabled && mFmhaDispatcher->isSeparateQAndKvInput()
-        && (mSageAttnNumEltsPerBlkQ > 0 || mSageAttnNumEltsPerBlkK > 0 || mSageAttnNumEltsPerBlkV > 0)
-        && mFP8ContextFMHA;
-
     // Packed fp8 qkv buffer size for normal fp8 context FMHA
     size_t fp8_qkv_buffer_size = mEnableContextFMHA && mFP8ContextFMHA && !mFmhaDispatcher->isSeparateQAndKvInput()
         ? params.num_tokens * (local_hidden_units_qo + 2 * local_hidden_units_kv)
@@ -1471,26 +1443,6 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
             fp8_v_buf_size = params.total_kv_len * static_cast<size_t>(total_v_dim_all_heads);
         }
     }
-    else if (useSageAttnSeparateQkv)
-    {
-        fp8_q_buf_size = params.num_tokens * static_cast<size_t>(local_hidden_units_qo);
-        fp8_k_buf_size = params.total_kv_len * static_cast<size_t>(local_hidden_units_kv);
-        fp8_v_buf_size = params.total_kv_len * static_cast<size_t>(local_hidden_units_kv);
-    }
-    int32_t const q_max_n_blk
-        = mSageAttnNumEltsPerBlkQ > 0 ? tc::divUp(params.input_seq_length, mSageAttnNumEltsPerBlkQ) : 0;
-    int32_t const k_max_n_blk
-        = mSageAttnNumEltsPerBlkK > 0 ? tc::divUp(params.input_seq_length, mSageAttnNumEltsPerBlkK) : 0;
-    // SageAttention V scales are shared across tokens and partitioned on the flattened hidden dimension (H * D).
-    int32_t const v_max_n_blk
-        = mSageAttnNumEltsPerBlkV > 0 ? tc::divUp(local_hidden_units_kv, mSageAttnNumEltsPerBlkV) : 0;
-    size_t const sage_q_sfs_buffer_size = useSageAttnSeparateQkv
-        ? sizeof(float) * mNumAttnHeads * params.batch_size * static_cast<size_t>(q_max_n_blk)
-        : 0;
-    size_t const sage_k_sfs_buffer_size = useSageAttnSeparateQkv
-        ? sizeof(float) * mNumAttnKVHeads * params.batch_size * static_cast<size_t>(k_max_n_blk)
-        : 0;
-    size_t const sage_v_sfs_buffer_size = useSageAttnSeparateQkv ? sizeof(float) * static_cast<size_t>(v_max_n_blk) : 0;
     size_t const padding_offset_size
         = mEnableContextFMHA ? 0 : sizeof(int) * params.batch_size * params.input_seq_length;
     size_t const encoder_padding_offset_size
@@ -1545,12 +1497,6 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
         = reinterpret_cast<float*>(nextWorkspacePtr(workspace_byte_ptr, offset, fmha_bmm1_scale_size));
     float* fmha_bmm2_scale_ptr
         = reinterpret_cast<float*>(nextWorkspacePtr(workspace_byte_ptr, offset, fmha_bmm2_scale_size));
-    float* sage_q_sfs_buf
-        = reinterpret_cast<float*>(nextWorkspacePtr(workspace_byte_ptr, offset, sage_q_sfs_buffer_size));
-    float* sage_k_sfs_buf
-        = reinterpret_cast<float*>(nextWorkspacePtr(workspace_byte_ptr, offset, sage_k_sfs_buffer_size));
-    float* sage_v_sfs_buf
-        = reinterpret_cast<float*>(nextWorkspacePtr(workspace_byte_ptr, offset, sage_v_sfs_buffer_size));
 
     T* gatherInBuffer = reinterpret_cast<T*>(nextWorkspacePtr(workspace_byte_ptr, offset, cpWorkspaceSize));
     T* gatherOutBuffer = gatherInBuffer + cpMaxPadedSequenceLength * getHeadSize() * (mNumHeads + 2 * mNumKVHeads);
@@ -1792,69 +1738,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
         }
         else
         {
-            if (useSageAttnSeparateQkv)
-            {
-                TLLM_CHECK_WITH_INFO(mFP8ContextFMHA, "SageAttention kernel runs under mFP8ContextFMHA option.");
-                TLLM_CHECK_WITH_INFO(
-                    mFmhaDispatcher->isSupported(), "SageAttention has no unfused fallback implemented.");
-                TLLM_CHECK_WITH_INFO(
-                    mSageAttnNumEltsPerBlkQ > 0 && mSageAttnNumEltsPerBlkK > 0 && mSageAttnNumEltsPerBlkV == 1,
-                    "SageQuant requires positive block sizes for Q and K while the block size for V must be 1.");
-                TLLM_CHECK_WITH_INFO(!params.kv_scale_quant_orig,
-                    "SageAttention disregards the configured params.kv_scale_quant_orig, invalidating the result.");
-                check_cuda_error(cudaMemsetAsync(sage_v_sfs_buf, 0, sage_v_sfs_buffer_size, stream));
-
-                tc::SageQuantParams qkParams{};
-                qkParams.headDim = getHeadSize();
-                qkParams.inputType = std::is_same_v<T, __nv_bfloat16> ? DATA_TYPE_BF16 : DATA_TYPE_FP16;
-                qkParams.quantType = mSageAttnQkInt8 ? DATA_TYPE_INT8 : DATA_TYPE_E4M3;
-                qkParams.vStage = 0;
-                qkParams.sumSeqLensV = params.total_kv_len;
-                qkParams.numHeadsV = mNumAttnKVHeads;
-                qkParams.ptrV = params.v_ptr;
-                qkParams.ptrVQuant = fp8_v_buf;
-                qkParams.ptrVScale = sage_v_sfs_buf;
-                qkParams.smCount = mMultiProcessorCount;
-                qkParams.stream = stream;
-
-                // Quantize into Fp8Q, SfsQ, SfsV
-                if (mSageAttnNumEltsPerBlkQ > 0)
-                {
-                    qkParams.sumSeqLensQk = params.num_tokens;
-                    qkParams.numHeads = mNumAttnHeads;
-                    qkParams.tokenBlockSize = mSageAttnNumEltsPerBlkQ;
-                    qkParams.ptrQk = attention_input;
-                    qkParams.ptrQkQuant = fp8_q_buf;
-                    qkParams.ptrQkScale = sage_q_sfs_buf;
-                    qkParams.vStage = 1;
-                    tc::invokeSageQuant(qkParams);
-                }
-                else
-                {
-                    invokeCudaCast(fp8_q_buf, attention_input, params.num_tokens * local_hidden_units_qo, stream);
-                }
-
-                // Quantize into Fp8K, SfsK, Fp8V
-                if (mSageAttnNumEltsPerBlkK > 0)
-                {
-                    qkParams.sumSeqLensQk = params.total_kv_len;
-                    qkParams.numHeads = mNumAttnKVHeads;
-                    qkParams.tokenBlockSize = mSageAttnNumEltsPerBlkK;
-                    qkParams.ptrQk = params.k_ptr;
-                    qkParams.ptrQkQuant = fp8_k_buf;
-                    qkParams.ptrQkScale = sage_k_sfs_buf;
-                    qkParams.vStage = 2;
-                    tc::invokeSageQuant(qkParams);
-                }
-                else
-                {
-                    invokeCudaCast(fp8_k_buf, params.k_ptr, params.total_kv_len * local_hidden_units_kv, stream);
-                }
-            }
-            else
-            {
-                invokeQKVPreprocessing(preprocessingParams, stream);
-            }
+            invokeQKVPreprocessing(preprocessingParams, stream);
         }
         sync_check_cuda_error(stream);
         {
@@ -1934,23 +1818,9 @@ int AttentionOp::enqueueContext(EnqueueContextParams<T> const& params, cudaStrea
         }
         else
         {
-            if (useSageAttnSeparateQkv)
-            {
-                fmhaParams.qkvPtr = nullptr;
-                fmhaParams.qPtr = reinterpret_cast<void const*>(fp8_q_buf);
-                fmhaParams.kPtr = reinterpret_cast<void const*>(fp8_k_buf);
-                fmhaParams.vPtr = reinterpret_cast<void const*>(fp8_v_buf);
-
-                fmhaParams.qScalePtr = sage_q_sfs_buf;
-                fmhaParams.kScalePtr = sage_k_sfs_buf;
-                fmhaParams.vScalePtr = sage_v_sfs_buf;
-            }
-            else
-            {
-                fmhaParams.qkvPtr = mFP8ContextFMHA ? reinterpret_cast<void const*>(fp8_qkv_buffer)
-                                                    : reinterpret_cast<void const*>(attention_input);
-                fmhaParams.qPtr = reinterpret_cast<void const*>(q_buf_2_);
-            }
+            fmhaParams.qkvPtr = mFP8ContextFMHA ? reinterpret_cast<void const*>(fp8_qkv_buffer)
+                                                : reinterpret_cast<void const*>(attention_input);
+            fmhaParams.qPtr = reinterpret_cast<void const*>(q_buf_2_);
         }
         // TODO: add contiguous kv buffer (cross-attention).
         fmhaParams.kvPtr = nullptr;
@@ -2877,22 +2747,13 @@ int AttentionOp::initialize() noexcept
             fmhaParams.attentionInputLayout = (mPagedKVCache && mPagedContextFMHA) ? AttentionInputLayout::Q_PAGED_KV
                                                                                    : AttentionInputLayout::PACKED_QKV;
         }
-        if (!mIsMLAEnabled && mFP8ContextFMHA
-            && (mSageAttnNumEltsPerBlkQ > 0 || mSageAttnNumEltsPerBlkK > 0 || mSageAttnNumEltsPerBlkV > 0))
-        {
-            fmhaParams.attentionInputLayout = AttentionInputLayout::SEPARATE_Q_K_V;
-        }
         fmhaParams.isSPadded = !mRemovePadding;
         fmhaParams.numQHeads = mNumAttnHeads;
         fmhaParams.numKvHeads = mNumAttnKVHeads;
         fmhaParams.numTokensPerBlock = mTokensPerBlock;
         fmhaParams.headSize = mHeadSize;
         fmhaParams.headSizeV = mHeadSize;
         fmhaParams.qScaling = mQScaling;
-        fmhaParams.sageBlockSizeQ = mSageAttnNumEltsPerBlkQ;
-        fmhaParams.sageBlockSizeK = mSageAttnNumEltsPerBlkK;
-        fmhaParams.sageBlockSizeV = mSageAttnNumEltsPerBlkV;
-        fmhaParams.dataTypeQkReinterpret = mSageAttnQkInt8 ? DATA_TYPE_INT8 : DATA_TYPE_E4M3;
 
         // mFmhaDispatcher is not used for generation MLA, but we still need to modify these values to avoid selecting
         // the wrong kernel, no matter mIsGenerationMLA is true or false
@@ -3199,10 +3060,6 @@ std::string AttentionOp::toString() const
     ss << "mPosShiftEnabled: " << std::boolalpha << mPosShiftEnabled << std::endl;
     ss << "mPagedContextFMHA: " << std::boolalpha << mPagedContextFMHA << std::endl;
     ss << "mFP8ContextFMHA: " << std::boolalpha << mFP8ContextFMHA << std::endl;
-    ss << "mSageAttnNumEltsPerBlkQ: " << mSageAttnNumEltsPerBlkQ << std::endl;
-    ss << "mSageAttnNumEltsPerBlkK: " << mSageAttnNumEltsPerBlkK << std::endl;
-    ss << "mSageAttnNumEltsPerBlkV: " << mSageAttnNumEltsPerBlkV << std::endl;
-    ss << "mSageAttnQkInt8: " << std::boolalpha << mSageAttnQkInt8 << std::endl;
     ss << "mFP8AttenOutput: " << std::boolalpha << mFP8AttenOutput << std::endl;
     ss << "mFP8ContextMLA: " << std::boolalpha << mFP8ContextMLA << std::endl;
     ss << "mDenseContextFMHA: " << std::boolalpha << mDenseContextFMHA << std::endl;

cpp/tensorrt_llm/common/attentionOp.h

@@ -119,10 +119,6 @@ class AttentionOp
         // this is a buffer of size [num_tokens, num_heads_q] with each element
         // representing the max and LSE/denominator of the softmax values
         float2* softmax_stats = nullptr;
-        // Optional SageAttention scaling factors.
-        float const* sage_attn_sfs_q = nullptr;
-        float const* sage_attn_sfs_k = nullptr;
-        float const* sage_attn_sfs_v = nullptr;
     };
 
     template <typename T>
@@ -523,12 +519,6 @@ class AttentionOp
     // Skip softmax threshold scale factor.
     float mSkipSoftmaxThresholdScaleFactorPrefill = 0;
     float mSkipSoftmaxThresholdScaleFactorDecode = 0;
-    // Optional SageAttention block sizes.
-    // Currently, these are only consumed by the TllmGen backend path.
-    int mSageAttnNumEltsPerBlkQ = 0;
-    int mSageAttnNumEltsPerBlkK = 0;
-    int mSageAttnNumEltsPerBlkV = 0;
-    bool mSageAttnQkInt8 = false;
 #ifdef SKIP_SOFTMAX_STAT
     uint32_t* mSkipSoftmaxTotalBlocks;
     uint32_t* mSkipSoftmaxSkippedBlocks;
@@ -551,8 +541,7 @@ class AttentionOp
             mAttnTpSize, mAttnTpRank, mAttnCpSize, mAttnCpRank, mUlyssesMQABroadcast, mEnableContextFMHA,
             mFMHAForceFP32Acc, mMultiBlockMode, mEnableXQA, mUseKVCache, mSkipAttn, mFuseFp4Quant,
             mNbMultiBlockSemaphores, mAttentionChunkSize.value_or(-1), mSkipSoftmaxThresholdScaleFactorPrefill,
-            mSkipSoftmaxThresholdScaleFactorDecode, mSageAttnNumEltsPerBlkQ, mSageAttnNumEltsPerBlkK,
-            mSageAttnNumEltsPerBlkV, mSageAttnQkInt8);
+            mSkipSoftmaxThresholdScaleFactorDecode);
     };
 
 private: