/*
 * Copyright (c) 2019-2022, NVIDIA CORPORATION.  All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "src/fastertransformer/layers/attention_layers/GlmDecoderSelfAttentionLayer.h"
#include "src/fastertransformer/kernels/decoder_masked_multihead_attention.h"
#include "src/fastertransformer/kernels/unfused_attention_kernels.h"
#include "src/fastertransformer/utils/logger.h"

namespace fastertransformer {

template<typename T>
struct TypeConverter {
    using Type = T;
};

template<>
struct TypeConverter<half> {
    using Type = uint16_t;
};

template<typename T>
void fusedQKV_masked_attention_dispatch(const T* q_buf,
                                        const T* k_buf,
                                        const T* v_buf,
                                        const T* relative_attention_bias,
                                        T* key_cache,
                                        T* value_cache,
                                        const int* cache_indir,
                                        T* context_buf,
                                        const bool* finished,
                                        const int* sequence_lengths,
                                        const int max_batch_size,
                                        const int inference_batch_size,
                                        const int beam_width,
                                        const int head_num,
                                        const int size_per_head,
                                        const int rotary_embedding_dim,
                                        const int layer_id,
                                        const int max_seq_len,
                                        const int max_input_len,
                                        const int* input_lengths,
                                        const int step,
                                        const float q_scaling,
                                        const int relative_attention_bias_stride,
                                        cudaStream_t stream)
{
    using DataType = typename TypeConverter<T>::Type;
    // Prepare the parameters.
    Masked_multihead_attention_params<DataType> params;
    memset(&params, 0, sizeof(params));
    int hidden_units = head_num * size_per_head;
    params.q_bias = nullptr;
    params.k_bias = nullptr;
    params.v_bias = nullptr;

    // Set the output buffer.
    params.out = reinterpret_cast<DataType*>(context_buf);

    // Set the input buffers.
    params.q = reinterpret_cast<const DataType*>(q_buf);
    params.k = reinterpret_cast<const DataType*>(k_buf);
    params.v = reinterpret_cast<const DataType*>(v_buf);
    params.stride = hidden_units;
    params.finished = const_cast<bool*>(finished);

    params.k_cache = reinterpret_cast<DataType*>(key_cache);
    params.v_cache = reinterpret_cast<DataType*>(value_cache);
    params.cache_indir = cache_indir;
    params.batch_size = inference_batch_size;
    params.beam_width = beam_width;
    params.seq_length = max_seq_len;
    params.length_per_sample = sequence_lengths;
    params.timestep = step - 1;
    params.num_heads = head_num;
    params.hidden_size_per_head = size_per_head;
    params.rotary_embedding_dim = rotary_embedding_dim;
    // Note: keep norm factor (sqrt(K_dim)) when adopting megatron T5 structure (may adjust)
    params.inv_sqrt_dh = layer_id + 1.0f;

    params.input_lengths = input_lengths;
    params.max_input_len = max_input_len;
    // TODO(bhsueh) Need better implementation
    if (relative_attention_bias != nullptr) {
        if (sizeof(T) == 4) {
            params.relative_attention_bias_float = reinterpret_cast<const float*>(relative_attention_bias);
        }
        else {
            params.relative_attention_bias_half = reinterpret_cast<const half*>(relative_attention_bias);
        }
    }
    params.relative_attention_bias_stride = relative_attention_bias_stride;

    masked_multihead_attention(params, stream);
}

template void fusedQKV_masked_attention_dispatch(const float* q_buf,
                                                 const float* k_buf,
                                                 const float* v_buf,
                                                 const float* relative_attention_bias,
                                                 float* key_cache,
                                                 float* value_cache,
                                                 const int* cache_indir,
                                                 float* context_buf,
                                                 const bool* finished,
                                                 const int* sequence_lengths,
                                                 const int max_batch_size,
                                                 const int inference_batch_size,
                                                 const int beam_width,
                                                 const int head_num,
                                                 const int size_per_head,
                                                 const int rotary_embedding_dim,
                                                 const int layer_id,
                                                 const int max_seq_len,
                                                 const int max_input_len,
                                                 const int* input_lengths,
                                                 const int step,
                                                 const float q_scaling,
                                                 const int relative_attention_bias_stride,
                                                 cudaStream_t stream);

template void fusedQKV_masked_attention_dispatch(const half* q_buf,
                                                 const half* k_buf,
                                                 const half* v_buf,
                                                 const half* relative_attention_bias,
                                                 half* key_cache,
                                                 half* value_cache,
                                                 const int* cache_indir,
                                                 half* context_buf,
                                                 const bool* finished,
                                                 const int* sequence_lengths,
                                                 const int max_batch_size,
                                                 const int inference_batch_size,
                                                 const int beam_width,
                                                 const int head_num,
                                                 const int size_per_head,
                                                 const int rotary_embedding_dim,
                                                 const int layer_id,
                                                 const int max_seq_len,
                                                 const int max_input_len,
                                                 const int* input_lengths,
                                                 const int step,
                                                 const float q_scaling,
                                                 const int relative_attention_bias_stride,
                                                 cudaStream_t stream);

template<typename T>
void GlmDecoderSelfAttentionLayer<T>::allocateBuffer()
{
    if (is_allocate_buffer_ == false) {
        qkv_buf_ =
            reinterpret_cast<T*>(allocator_->malloc(sizeof(T) * max_batch_size_ * 3 * local_hidden_units_, false));
        q_buf_ =
            reinterpret_cast<T*>(allocator_->malloc(sizeof(T) * max_batch_size_ * local_hidden_units_, false));
        k_buf_ =
            reinterpret_cast<T*>(allocator_->malloc(sizeof(T) * max_batch_size_ * local_hidden_units_, false));
        v_buf_ =
            reinterpret_cast<T*>(allocator_->malloc(sizeof(T) * max_batch_size_ * local_hidden_units_, false));
        context_buf_ =
            reinterpret_cast<T*>(allocator_->malloc(sizeof(T) * max_batch_size_ * local_hidden_units_, false));
        weights_buf_ = 
            reinterpret_cast<T*>(allocator_->malloc(sizeof(T) * d_model_ * 3 * local_hidden_units_, false));
        is_allocate_buffer_ = true;
    }
}

template<typename T>
void GlmDecoderSelfAttentionLayer<T>::allocateBuffer(size_t batch_size)
{
    FT_LOG_DEBUG(__PRETTY_FUNCTION__);
    qkv_buf_ =
        reinterpret_cast<T*>(allocator_->reMalloc(qkv_buf_, sizeof(T) * batch_size * 3 * local_hidden_units_, false));
    q_buf_ =
        reinterpret_cast<T*>(allocator_->reMalloc(q_buf_, sizeof(T) * batch_size * local_hidden_units_, false));
    k_buf_ =
        reinterpret_cast<T*>(allocator_->reMalloc(k_buf_, sizeof(T) * batch_size * local_hidden_units_, false));
    v_buf_ =
        reinterpret_cast<T*>(allocator_->reMalloc(v_buf_, sizeof(T) * batch_size * local_hidden_units_, false));
    context_buf_ =
        reinterpret_cast<T*>(allocator_->reMalloc(context_buf_, sizeof(T) * batch_size * local_hidden_units_, false));
    weights_buf_ = reinterpret_cast<T*>(allocator_->reMalloc(weights_buf_, sizeof(T) * d_model_ * 3 * local_hidden_units_, false));
    is_allocate_buffer_ = true;
}

template<typename T>
void GlmDecoderSelfAttentionLayer<T>::freeBuffer()
{
    if (is_allocate_buffer_) {
        allocator_->free(qkv_buf_);
        allocator_->free(q_buf_);
        allocator_->free(k_buf_);
        allocator_->free(v_buf_);
        allocator_->free(context_buf_);
        allocator_->free(weights_buf_);
        is_allocate_buffer_ = false;
    }
}

template<typename T>
bool GlmDecoderSelfAttentionLayer<T>::isValidBatchSize(size_t batch_size)
{
    if (batch_size <= max_batch_size_) {
        return true;
    }
    else {
        freeBuffer();
        max_batch_size_ = batch_size * 1.2;
        return true;
    }
}

template<typename T>
GlmDecoderSelfAttentionLayer<T>::GlmDecoderSelfAttentionLayer(size_t max_batch_size,
                                                        size_t head_num,
                                                        size_t size_per_head,
                                                        size_t local_head_num,
                                                        size_t rotary_embedding_dim,
                                                        size_t d_model,
                                                        const float q_scaling,
                                                        cudaStream_t stream,
                                                        cublasMMWrapper* cublas_wrapper,
                                                        IAllocator* allocator,
                                                        bool is_free_buffer_after_forward,
                                                        bool sparse,
                                                        int int8_mode):
    BaseAttentionLayer<T>(stream, cublas_wrapper, allocator, is_free_buffer_after_forward, sparse),
    max_batch_size_(max_batch_size),
    head_num_(head_num),
    size_per_head_(size_per_head),
    hidden_units_(head_num_ * size_per_head_),
    local_head_num_(local_head_num),
    local_hidden_units_(local_head_num_ * size_per_head_),
    rotary_embedding_dim_(rotary_embedding_dim),
    d_model_(d_model),
    q_scaling_(q_scaling),
    int8_mode_(int8_mode)
{
    FT_CHECK(size_per_head_ == 32 || size_per_head_ == 64 || size_per_head_ == 96 || size_per_head_ == 128
             || size_per_head_ == 160 || size_per_head_ == 192 || size_per_head_ == 224 || size_per_head_ == 256);
}

template<typename T>
GlmDecoderSelfAttentionLayer<T>::GlmDecoderSelfAttentionLayer(size_t max_batch_size,
                                                        size_t head_num,
                                                        size_t size_per_head,
                                                        cudaStream_t stream,
                                                        cublasMMWrapper* cublas_wrapper,
                                                        IAllocator* allocator,
                                                        bool is_free_buffer_after_forward,
                                                        bool sparse,
                                                        int int8_mode):
    GlmDecoderSelfAttentionLayer<T>(max_batch_size,
                                 head_num,
                                 size_per_head,
                                 head_num,
                                 0,
                                 head_num * size_per_head,
                                 1.0f,
                                 stream,
                                 cublas_wrapper,
                                 allocator,
                                 is_free_buffer_after_forward,
                                 sparse,
                                 int8_mode)
{
}

template<typename T>
GlmDecoderSelfAttentionLayer<T>::GlmDecoderSelfAttentionLayer(size_t max_batch_size,
                                                        size_t head_num,
                                                        size_t size_per_head,
                                                        const float q_scaling,
                                                        cudaStream_t stream,
                                                        cublasMMWrapper* cublas_wrapper,
                                                        IAllocator* allocator,
                                                        bool is_free_buffer_after_forward,
                                                        bool sparse,
                                                        int int8_mode):
    GlmDecoderSelfAttentionLayer<T>(max_batch_size,
                                 head_num,
                                 size_per_head,
                                 head_num,
                                 0,
                                 head_num * size_per_head,
                                 q_scaling,
                                 stream,
                                 cublas_wrapper,
                                 allocator,
                                 is_free_buffer_after_forward,
                                 sparse,
                                 int8_mode)
{
}

template<typename T>
GlmDecoderSelfAttentionLayer<T>::GlmDecoderSelfAttentionLayer(size_t max_batch_size,
                                                        size_t head_num,
                                                        size_t size_per_head,
                                                        size_t local_head_num,
                                                        cudaStream_t stream,
                                                        cublasMMWrapper* cublas_wrapper,
                                                        IAllocator* allocator,
                                                        bool is_free_buffer_after_forward,
                                                        bool sparse,
                                                        int int8_mode):
    GlmDecoderSelfAttentionLayer<T>(max_batch_size,
                                 head_num,
                                 size_per_head,
                                 local_head_num,
                                 0,
                                 head_num * size_per_head,
                                 1.0f,
                                 stream,
                                 cublas_wrapper,
                                 allocator,
                                 is_free_buffer_after_forward,
                                 sparse,
                                 int8_mode)
{
}

template<typename T>
GlmDecoderSelfAttentionLayer<T>::GlmDecoderSelfAttentionLayer(size_t max_batch_size,
                                                        size_t head_num,
                                                        size_t size_per_head,
                                                        size_t local_head_num,
                                                        size_t d_model,
                                                        const float q_scaling,
                                                        cudaStream_t stream,
                                                        cublasMMWrapper* cublas_wrapper,
                                                        IAllocator* allocator,
                                                        bool is_free_buffer_after_forward,
                                                        bool sparse,
                                                        int int8_mode):
    GlmDecoderSelfAttentionLayer<T>(max_batch_size,
                                 head_num,
                                 size_per_head,
                                 local_head_num,
                                 0,
                                 d_model,
                                 q_scaling,
                                 stream,
                                 cublas_wrapper,
                                 allocator,
                                 is_free_buffer_after_forward,
                                 sparse,
                                 int8_mode)
{
}

template<typename T>
GlmDecoderSelfAttentionLayer<T>::GlmDecoderSelfAttentionLayer(size_t max_batch_size,
                                                        size_t head_num,
                                                        size_t size_per_head,
                                                        size_t local_head_num,
                                                        size_t rotary_embedding_dim,
                                                        cudaStream_t stream,
                                                        cublasMMWrapper* cublas_wrapper,
                                                        IAllocator* allocator,
                                                        bool is_free_buffer_after_forward,
                                                        bool sparse,
                                                        int int8_mode):
    GlmDecoderSelfAttentionLayer<T>(max_batch_size,
                                 head_num,
                                 size_per_head,
                                 local_head_num,
                                 rotary_embedding_dim,
                                 head_num * size_per_head,
                                 1.0f,
                                 stream,
                                 cublas_wrapper,
                                 allocator,
                                 is_free_buffer_after_forward,
                                 sparse,
                                 int8_mode)
{
}

template<typename T>
GlmDecoderSelfAttentionLayer<T>::GlmDecoderSelfAttentionLayer(GlmDecoderSelfAttentionLayer<T> const& attention_layer):
    GlmDecoderSelfAttentionLayer<T>(attention_layer.max_batch_size_,
                                 attention_layer.head_num_,
                                 attention_layer.size_per_head_,
                                 attention_layer.local_head_num_,
                                 attention_layer.rotary_embedding_dim_,
                                 attention_layer.d_model_,
                                 attention_layer.q_scaling_,
                                 attention_layer.stream_,
                                 attention_layer.cublas_wrapper_,
                                 attention_layer.allocator_,
                                 attention_layer.is_free_buffer_after_forward_,
                                 attention_layer.sparse_,
                                 attention_layer.int8_mode_)
{
}

template<typename T>
GlmDecoderSelfAttentionLayer<T>::~GlmDecoderSelfAttentionLayer()
{
    cublas_wrapper_ = nullptr;
    freeBuffer();
}

template<typename T>
void GlmDecoderSelfAttentionLayer<T>::forward(std::vector<fastertransformer::Tensor>* output_tensors,
                                           const std::vector<fastertransformer::Tensor>* input_tensors,
                                           const AttentionWeight<T>* attention_weights)
{
    // input tensors:
    //      attention_input [batch_size, d_model_],
    //      finished [batch_size],
    //      sequence_lengths [batch_size]
    //      input_lengths [batch_size]
    //      max_input_length [1] on cpu
    //      step [1] on cpu
    //      cache_indirection [batch_size / beam_width, beam_width, max_seq_len]
    //      relative_attention_bias [1, head_num, step, step] or [1, head_num, max_seq_len, max_seq_len] (option)

    // output tensors:
    //      attention_output [batch_size, d_model_],
    //      key_cache [batch, local_head_num, size_per_head // x, max_seq_len, x]
    //      value_cache [batch, local_head_num, max_seq_len, size_per_head]

    FT_LOG_DEBUG(__PRETTY_FUNCTION__);
    FT_CHECK(input_tensors->size() == 9 || input_tensors->size() == 10);
    FT_CHECK(output_tensors->size() == 3);
    FT_CHECK(output_tensors->at(1).shape.size() == 5 || output_tensors->at(1).shape.size() == 3);
    FT_CHECK(output_tensors->at(2).shape.size() == 4 || output_tensors->at(2).shape.size() == 3);
    // FT_CHECK(isValidBatchSize(input_tensors->at(0).shape[0]));
    allocateBuffer(input_tensors->at(0).shape[0]);

    const T* attention_input = reinterpret_cast<const T*>(input_tensors->at(0).data);
    const bool* finished = reinterpret_cast<const bool*>(input_tensors->at(1).data);
    const int* sequence_lengths = reinterpret_cast<const int*>(input_tensors->at(2).data);
    const int* mask_positions = reinterpret_cast<const int*>(input_tensors->at(6).data);
    const int* cache_indir = reinterpret_cast<const int*>(input_tensors->at(7).data);
    const int layer_id = *((int*)(input_tensors->at(8).data));
    const T* relative_attention_bias =
        reinterpret_cast<const T*>(input_tensors->size() == 10 ? input_tensors->at(9).data : nullptr);
    const int relative_attention_bias_stride = input_tensors->size() == 10 ? input_tensors->at(9).shape[3] : 0;

    T* attention_out = (T*)(output_tensors->at(0).data);
    T* key_cache = (T*)(output_tensors->at(1).data);
    T* value_cache = (T*)(output_tensors->at(2).data);

    const int batch_size = input_tensors->at(0).shape[0];
    const int beam_width = input_tensors->at(7).shape[1];
    const int max_seq_len = output_tensors->at(1).shape[3];
    const int step = *(int*)(input_tensors->at(5).data);

    // printf("%d %d \n",batch_size, beam_width);

#ifdef SPARSITY_ENABLED
    const int m_padded = 8 * div_up(batch_size, 8);
    if (sparse_ && cublas_wrapper_->isUseSparse(1, 3 * local_hidden_units_, m_padded, d_model_)) {
        cublas_wrapper_->SpGemm(CUBLAS_OP_N,
                                CUBLAS_OP_N,
                                3 * local_hidden_units_,
                                m_padded,
                                d_model_,
                                attention_weights->query_weight.sp_kernel,
                                attention_input,
                                qkv_buf_);
    }
    else {
#endif
        if (int8_mode_ != 0 && batch_size <= 2) {
            FT_CHECK(attention_weights->query_weight.int8_kernel != NULL
                     && attention_weights->query_weight.quant_scale != NULL);
            int8WeightPerChannelLdkMultiplicationLauncher(attention_weights->query_weight.int8_kernel,
                                                          attention_input,
                                                          attention_weights->query_weight.quant_scale,
                                                          qkv_buf_,
                                                          batch_size,
                                                          3 * local_hidden_units_,
                                                          d_model_,
                                                          stream_);
        }
        else {
            if (int8_mode_ == 1) {
                FT_LOG_WARNING(
                    "[GlmDecoderSelfAttentionLayer<T>::forward] int8 glm doesn't support m > 2, run fp glm instead.\n");
            }
            if(attention_weights->query_weight.kernel) {
                cublas_wrapper_->Gemm(CUBLAS_OP_N,
                                  CUBLAS_OP_N,
                                  3 * local_hidden_units_,  // n
                                  batch_size,
                                  d_model_,  // k
                                  attention_weights->query_weight.kernel,
                                  3 * local_hidden_units_,  // n
                                  attention_input,
                                  d_model_,  // k
                                  qkv_buf_,
                                  3 * local_hidden_units_ /* n */);
            } else if(batch_size > MaxPerChannelLdkMultiplicationNum) {
                FT_CHECK(attention_weights->query_weight.int8_kernel != NULL || attention_weights->query_weight.int4_kernel != NULL);
                FT_CHECK(attention_weights->query_weight.quant_scale != NULL);
                if(attention_weights->query_weight.int8_kernel != NULL) {
                    invokeInt8WeightExtractionNoTrans(attention_weights->query_weight.int8_kernel,
                                                attention_weights->query_weight.quant_scale,
                                                weights_buf_,
                                                3 * local_hidden_units_,
                                                d_model_,
                                                stream_);
                } else {
                    invokeInt4WeightExtractionNoTrans(attention_weights->query_weight.int4_kernel,
                                                attention_weights->query_weight.quant_scale,
                                                weights_buf_,
                                                3 * local_hidden_units_,
                                                d_model_,
                                                stream_);
                }

                sync_check_cuda_error();
                cublas_wrapper_->Gemm(CUBLAS_OP_T,
                                  CUBLAS_OP_N,
                                  3 * local_hidden_units_,  // n
                                  batch_size,
                                  d_model_,  // k
                                  weights_buf_,
                                  d_model_,  // k
                                  attention_input,
                                  d_model_,  // k
                                  qkv_buf_,
                                  3 * local_hidden_units_ /* n */);
            } else {
                FT_CHECK(attention_weights->query_weight.int8_kernel != NULL || attention_weights->query_weight.int4_kernel != NULL);
                FT_CHECK(attention_weights->query_weight.quant_scale != NULL);
                if(attention_weights->query_weight.int8_kernel != NULL) {
                    int8WeightPerChannelLdkMultiplicationLauncher(attention_weights->query_weight.int8_kernel,
                                                          attention_input,
                                                          attention_weights->query_weight.quant_scale,
                                                          qkv_buf_,
                                                          batch_size,
                                                          3 * local_hidden_units_,
                                                          d_model_,
                                                          stream_);
                } else {
                    int4WeightPerChannelLdkMultiplicationLauncher(attention_weights->query_weight.int4_kernel,
                                                          attention_input,
                                                          attention_weights->query_weight.quant_scale,
                                                          qkv_buf_,
                                                          batch_size,
                                                          3 * local_hidden_units_,
                                                          d_model_,
                                                          stream_);
                }
                
            }
            
        }
#ifdef SPARSITY_ENABLED
    }
#endif
    sync_check_cuda_error();

    invokeGlmAddFusedQKVBiasTranspose(q_buf_,
                                        k_buf_,
                                        v_buf_,
                                        qkv_buf_,
                                        attention_weights->query_weight.bias,
                                        (int*)(input_tensors->at(3).data),
                                        mask_positions,
                                        layer_id,
                                        batch_size,
                                        1,
                                        step - 1,
                                        local_head_num_,
                                        size_per_head_,
                                        rotary_embedding_dim_,
                                        stream_);
    sync_check_cuda_error();

    fusedQKV_masked_attention_dispatch<T>(q_buf_,
                                          k_buf_,
                                          v_buf_,
                                          relative_attention_bias,
                                          key_cache,
                                          value_cache,
                                          cache_indir,
                                          context_buf_,
                                          finished,
                                          sequence_lengths,
                                          batch_size,
                                          batch_size,
                                          beam_width,
                                          local_head_num_,
                                          size_per_head_,
                                          0, //rotary_embedding_dim
                                          layer_id,
                                          max_seq_len,
                                          *(int*)(input_tensors->at(4).data),
                                          (int*)(input_tensors->at(3).data),
                                          step,
                                          q_scaling_,
                                          relative_attention_bias_stride,
                                          stream_);
    sync_check_cuda_error();

#ifdef SPARSITY_ENABLED
    if (sparse_ && cublas_wrapper_->isUseSparse(1, d_model_, m_padded, local_hidden_units_)) {
        cublas_wrapper_->SpGemm(CUBLAS_OP_N,
                                CUBLAS_OP_N,
                                d_model_,
                                m_padded,
                                local_hidden_units_,
                                attention_weights->attention_output_weight.sp_kernel,
                                context_buf_,
                                attention_out);
    }
    else {
#endif
        if (int8_mode_ != 0 && batch_size <= 2) {
            FT_CHECK(attention_weights->attention_output_weight.int8_kernel != NULL
                     && attention_weights->attention_output_weight.quant_scale != NULL);
            int8WeightPerChannelLdkMultiplicationLauncher(attention_weights->attention_output_weight.int8_kernel,
                                                          context_buf_,
                                                          attention_weights->attention_output_weight.quant_scale,
                                                          attention_out,
                                                          batch_size,
                                                          d_model_,
                                                          local_hidden_units_,
                                                          stream_);
        }
        else {
            if (int8_mode_ == 1) {
                FT_LOG_WARNING(
                    "[GlmDecoderSelfAttentionLayer<T>::forward] int8 glm doesn't support m > 2, run fp glm instead.\n");
            }
            if(attention_weights->attention_output_weight.kernel) {
                cublas_wrapper_->Gemm(CUBLAS_OP_N,
                                  CUBLAS_OP_N,
                                  d_model_,  // n
                                  batch_size,
                                  local_hidden_units_,  // k
                                  attention_weights->attention_output_weight.kernel,
                                  d_model_,  // n
                                  context_buf_,
                                  local_hidden_units_,  // k
                                  attention_out,
                                  d_model_ /* n */);
            }
            else if(batch_size > MaxPerChannelLdkMultiplicationNum) {
                FT_CHECK(attention_weights->attention_output_weight.int8_kernel != NULL || attention_weights->attention_output_weight.int4_kernel != NULL);
                FT_CHECK(attention_weights->attention_output_weight.quant_scale != NULL);
                if(attention_weights->attention_output_weight.int8_kernel != NULL) {
                    invokeInt8WeightExtractionNoTrans(attention_weights->attention_output_weight.int8_kernel,
                                                attention_weights->attention_output_weight.quant_scale,
                                                weights_buf_,
                                                d_model_,
                                                local_hidden_units_,
                                                stream_);
                } else {
                    invokeInt4WeightExtractionNoTrans(attention_weights->attention_output_weight.int4_kernel,
                                                attention_weights->attention_output_weight.quant_scale,
                                                weights_buf_,
                                                d_model_,
                                                local_hidden_units_,
                                                stream_);
                }

                sync_check_cuda_error();
                cublas_wrapper_->Gemm(CUBLAS_OP_T,
                                  CUBLAS_OP_N,
                                  d_model_,  // n
                                  batch_size,
                                  local_hidden_units_,  // k
                                  weights_buf_,
                                  local_hidden_units_,  // k
                                  context_buf_,
                                  local_hidden_units_,  // k
                                  attention_out,
                                  d_model_ /* n */);
            }
            else {
                FT_CHECK(attention_weights->attention_output_weight.int8_kernel != NULL || attention_weights->attention_output_weight.int4_kernel != NULL);
                FT_CHECK(attention_weights->attention_output_weight.quant_scale != NULL);
                if(attention_weights->attention_output_weight.int8_kernel != NULL) {
                    int8WeightPerChannelLdkMultiplicationLauncher(attention_weights->attention_output_weight.int8_kernel,
                                                          context_buf_,
                                                          attention_weights->attention_output_weight.quant_scale,
                                                          attention_out,
                                                          batch_size,
                                                          d_model_,
                                                          local_hidden_units_,
                                                          stream_);
                } else {
                    int4WeightPerChannelLdkMultiplicationLauncher(attention_weights->attention_output_weight.int4_kernel,
                                                          context_buf_,
                                                          attention_weights->attention_output_weight.quant_scale,
                                                          attention_out,
                                                          batch_size,
                                                          d_model_,
                                                          local_hidden_units_,
                                                          stream_);
                }
            }
        }
        sync_check_cuda_error();
#ifdef SPARSITY_ENABLED
    }
#endif

    if (is_free_buffer_after_forward_ == true) {
        freeBuffer();
    }
}

template class GlmDecoderSelfAttentionLayer<float>;
template class GlmDecoderSelfAttentionLayer<half>;
#ifdef ENABLE_BF16
template class GlmDecoderSelfAttentionLayer<__nv_bfloat16>;
#endif

}  // namespace fastertransformer
